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-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.

翻译：最佳运输(OT)提供了一个多功能框架,用以以具有几何意义的方式比较复杂的数据分布。在概率计量之间计算瓦塞斯坦距离和大地测量的传统方法需要以网目为依存的域离散,并受到维度的诅咒。我们介绍了GeONet,一个网状的内分泌深神经操作网络,它从输入的初始分布和终端分布中学习非线性绘图,连接两个端点分布的瓦塞尔斯坦大地测量学。在离线培训阶段,GeONet学习了原始空间和双层空间的OT问题动态配制的支撑点最佳条件,其特征是结合的PDE系统。随后的推论阶段是瞬时的,可用于在线学习环境中的实时预测。我们证明,GeONet在模拟实例和CIFAR-10数据集中实现了与标准OT溶解器的类似测试精度,其推量按数量计算成本大大降低。