GeONet: a neural operator for learning the Wasserstein geodesicOptimal 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-dependent 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 a
simulation example and the CIFAR-10 dataset with considerably reduced
inference-stage computational cost by orders of magnitude.
arxiv.org