Saved in:
Bibliographic Details
Main Authors: Gracyk, Andrew, Chen, Xiaohui
Format: Preprint
Published: 2022
Subjects:
Online Access:https://arxiv.org/abs/2209.14440
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866929352721039360
author Gracyk, Andrew
Chen, Xiaohui
author_facet Gracyk, Andrew
Chen, Xiaohui
contents Optimal transport (OT) offers a versatile framework to compare complex data distributions in a geometrically meaningful way. Traditional methods for computing the Wasserstein distance and geodesic between probability measures require mesh-specific domain discretization and suffer from the curse-of-dimensionality. We present GeONet, a mesh-invariant deep neural operator network that learns the non-linear mapping from the input pair of initial and terminal distributions to the Wasserstein geodesic connecting the two endpoint distributions. In the offline training stage, GeONet learns the saddle point optimality conditions for the dynamic formulation of the OT problem in the primal and dual spaces that are characterized by a coupled PDE system. The subsequent inference stage is instantaneous and can be deployed for real-time predictions in the online learning setting. We demonstrate that GeONet achieves comparable testing accuracy to the standard OT solvers on simulation examples and the MNIST dataset with considerably reduced inference-stage computational cost by orders of magnitude.
format Preprint
id arxiv_https___arxiv_org_abs_2209_14440
institution arXiv
publishDate 2022
record_format arxiv
spellingShingle GeONet: a neural operator for learning the Wasserstein geodesic
Gracyk, Andrew
Chen, Xiaohui
Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
Optimal transport (OT) offers a versatile framework to compare complex data distributions in a geometrically meaningful way. Traditional methods for computing the Wasserstein distance and geodesic between probability measures require mesh-specific domain discretization and suffer from the curse-of-dimensionality. We present GeONet, a mesh-invariant deep neural operator network that learns the non-linear mapping from the input pair of initial and terminal distributions to the Wasserstein geodesic connecting the two endpoint distributions. In the offline training stage, GeONet learns the saddle point optimality conditions for the dynamic formulation of the OT problem in the primal and dual spaces that are characterized by a coupled PDE system. The subsequent inference stage is instantaneous and can be deployed for real-time predictions in the online learning setting. We demonstrate that GeONet achieves comparable testing accuracy to the standard OT solvers on simulation examples and the MNIST dataset with considerably reduced inference-stage computational cost by orders of magnitude.
title GeONet: a neural operator for learning the Wasserstein geodesic
topic Machine Learning
Artificial Intelligence
Computer Vision and Pattern Recognition
url https://arxiv.org/abs/2209.14440