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.

翻译：最优传输（OT）提供了一个多功能框架，用于以几何意义明确的方式比较复杂的数据分布。计算概率测度间Wasserstein距离与测地线的传统方法需要针对网格的域离散化，并受维度灾难困扰。本文提出GeONet，一种与网格无关的深度神经算子网络，它学习从初始与终端分布的输入对到连接两个端点分布的Wasserstein测地线的非线性映射。在离线训练阶段，GeONet学习原始空间与对偶空间中OT问题动态表述的鞍点最优性条件，该条件由耦合偏微分方程组刻画。随后的推理阶段可瞬时完成，并能部署于在线学习环境中的实时预测。我们通过仿真示例与MNIST数据集证明，GeONet在测试精度上达到与标准OT求解器相当的水平，同时推理阶段计算成本降低了数个数量级。