[论文解读] Learning High Dimensional Wasserstein Geodesics

Liu Shu, Shaojun Ma+3 more|arXiv (Cornell University)|Feb 5, 2021

Seismic Imaging and Inversion Techniques参考文献 42被引用 5

一句话总结

本文提出一种基于深度学习的方法，通过拉格朗日乘子和KKT条件将最优传输问题转化为极大极小鞍点问题，从而计算高维Wasserstein测地线。该方法使用深度神经网络参数化传输映射和速度场，实现基于样本的双向训练，无需显式正则化或降维，即可同时估计Wasserstein距离、最优映射和测地线路径。

ABSTRACT

We propose a new formulation and learning strategy for computing the Wasserstein geodesic between two probability distributions in high dimensions. By applying the method of Lagrange multipliers to the dynamic formulation of the optimal transport (OT) problem, we derive a minimax problem whose saddle point is the Wasserstein geodesic. We then parametrize the functions by deep neural networks and design a sample based bidirectional learning algorithm for training. The trained networks enable sampling from the Wasserstein geodesic. As by-products, the algorithm also computes the Wasserstein distance and OT map between the marginal distributions. We demonstrate the performance of our algorithms through a series of experiments with both synthetic and realistic data.

研究动机与目标

• 解决高维空间中计算Wasserstein测地线的挑战，其中传统基于PDE的方法受维数灾难影响。
• 开发一种可扩展的、基于样本的算法，避免空间离散化和正则化项，使其适用于高维场景。
• 通过统一框架同时计算Wasserstein距离、最优传输映射和分布之间的测地线路径。
• 通过在测地线上提供自然采样机制，推动其在机器学习、机器人学和控制领域的实际应用。
• 克服现有OT方法仅关注距离或映射估计的局限，实现对完整测地线轨迹的建模。

提出的方法

• 利用拉格朗日乘子将动态最优传输问题表述为极大极小问题，转化为鞍点问题。
• 应用KKT条件以缩小搜索空间并简化优化，获得更易处理的公式化形式。
• 使用深度神经网络参数化速度场和对偶势函数，实现在高维空间中的灵活函数逼近。
• 设计一种双向、基于样本的训练算法，交替优化前向和反向传输路径。
• 引入预条件处理和双向动力学，提升高维空间中训练的稳定性和收敛性。
• 采用小批量采样进行端到端网络训练，实现可扩展的推理与测地线采样，无需显式正则化约束。

实验结果

研究问题

• RQ1 基于深度学习的方法是否能在不依赖空间离散化或显式正则化的情况下计算高维Wasserstein测地线？
• RQ2 在测地线计算过程中，如何同时估计最优传输映射和Wasserstein距离？
• RQ3 所提出的双向学习策略在高维OT问题中在多大程度上提升了稳定性和收敛性？
• RQ4 该方法能否生成测地线路径上的真实中间样本，从而支持图像到图像翻译和数据生成等应用？
• RQ5 在具有复杂高维分布的真实数据（如图像和手写数字）上，该方法表现如何？

主要发现

• 所提方法成功在高维空间（如28×28 MNIST数字和256×256图像）中计算出Wasserstein测地线，无需降维。
• 算法在源分布与目标分布之间生成了视觉上逼真的中间样本，如在图像颜色迁移和数字变换任务中所展示。
• 在图像翻译任务（如夏季到秋季森林景观）中，生成图像与真实图像高度相似，色彩分布匹配度高。
• 通过与POT库在合成数据和真实数据集上的对比验证，该方法在Wasserstein距离和最优传输映射的计算中表现出高精度。
• 双向训练方案提升了稳定性和收敛性，即使在复杂高维场景中也能实现可靠的测地线采样。
• 该方法在基于样本的测地线计算中达到最先进性能，适用于一般凸成本函数（包括Lp范数），并支持前向与反向传输。