[论文解读] Neural Optimal Transport

Alexander Korotin, Daniil Selikhanovych+1 more|arXiv (Cornell University)|Jan 28, 2022

Medical Image Segmentation Techniques被引用 6

一句话总结

本文提出神经最优传输（NOT），一种基于深度学习的方法，用于在强成本和弱成本下计算确定性和随机最优传输映射与计划。通过证明神经网络可普遍逼近传输计划，作者构建了一个可微分、可扩展的框架，在无配对图像到图像翻译任务中表现优于以往方法，在CelebA、户外到教堂及手袋到鞋子基准测试中达到最先进水平。

ABSTRACT

We present a novel neural-networks-based algorithm to compute optimal transport maps and plans for strong and weak transport costs. To justify the usage of neural networks, we prove that they are universal approximators of transport plans between probability distributions. We evaluate the performance of our optimal transport algorithm on toy examples and on the unpaired image-to-image translation.

研究动机与目标

• 开发一种可扩展、可微分的深度学习框架，用于在强成本和弱成本下计算最优传输映射与计划。
• 证明深度神经网络是概率分布之间传输计划的通用逼近器，从而为在最优传输中使用神经网络提供理论依据。
• 解决现有方法的局限性，这些方法或无法强制执行边缘约束，或依赖于Langevin动力学等缓慢的采样过程。
• 通过学习保留结构和分布特性的传输映射，实现高质量的无配对图像到图像翻译。

提出的方法

• 提出一种基于弱最优传输对偶公式的新型优化框架，将其重新表述为极大极小鞍点问题。
• 使用由深度神经网络参数化的随机映射 $ T: \mathcal{X} \times \mathcal{Z} \to \mathcal{Y} $ 来建模条件传输计划。
• 实现一个可微分的目标函数，通过对偶变量 $ f $ 强制执行边缘约束，确保联合分布的前推匹配目标分布。
• 采用最小最大优化方案 $ \max_f \min_T $ 训练传输映射，支持端到端学习与反向传播。
• 将框架应用于 $ \mathbb{W}_2 $（确定性）和 $ \mathcal{W}_{2,\gamma} $（随机性）成本，支持一对一和一对多映射。
• 对弱成本函数 $ C(x, \mu) $ 使用变分近似，实现高效优化，无需依赖基于得分的模型或迭代采样。

实验结果

研究问题

• RQ1 深度神经网络是否能普遍逼近概率分布之间的最优传输计划？
• RQ2 一种可微分、端到端的深度学习框架是否能有效计算在强成本和弱成本下的确定性和随机最优传输映射？
• RQ3 所提出的神经最优传输方法是否在无配对图像到图像翻译任务中优于现有方法？
• RQ4 该方法在CelebA、户外到教堂及手袋到鞋子等多样化数据集上的泛化能力如何？
• RQ5 随机传输映射是否能够生成多样且逼真的翻译结果，同时保持结构一致性？

主要发现

• 所提出的神经最优传输方法在多个基准测试中，包括CelebA（女性到动漫）、户外到教堂及手袋到鞋子，均实现了最先进性能。
• 该方法成功计算了确定性（一对一）和随机性（一对多）传输映射，支持多样且高保真度的图像翻译。
• 消融研究显示，同步噪声向量不会在不同输入间产生共享的视觉伪影，表明随机映射能为每个输入生成多样化输出。
• 该框架在大规模视觉任务中表现出良好泛化能力，在 $ 128 \times 128 $ 和 $ 64 \times 64 $ 分辨率下均能生成高质量翻译结果。
• 理论证明表明深度神经网络是传输计划的通用逼近器，为在最优传输中使用神经网络提供了坚实基础。
• 该方法避免了计算成本高昂的采样过程（如Langevin动力学），相比基于得分的替代方法，训练和推理速度更快。