[论文解读] Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark

Alexander Korotin, Lingxiao Li+4 more|arXiv (Cornell University)|Jun 3, 2021

Advanced Neuroimaging Techniques and Applications参考文献 38被引用 10

一句话总结

本文提出一种基于输入凸神经网络（ICNNs）的连续Wasserstein-2基准，用于生成具有解析已知最优传输映射的连续概率测度对，从而实现对神经最优传输求解器的严格评估。尽管在图像生成任务中表现优异，许多求解器仍无法准确恢复真实的传输映射，暴露出最优传输精度与生成质量之间的关键脱节。

ABSTRACT

Despite the recent popularity of neural network-based solvers for optimal transport (OT), there is no standard quantitative way to evaluate their performance. In this paper, we address this issue for quadratic-cost transport -- specifically, computation of the Wasserstein-2 distance, a commonly-used formulation of optimal transport in machine learning. To overcome the challenge of computing ground truth transport maps between continuous measures needed to assess these solvers, we use input-convex neural networks (ICNN) to construct pairs of measures whose ground truth OT maps can be obtained analytically. This strategy yields pairs of continuous benchmark measures in high-dimensional spaces such as spaces of images. We thoroughly evaluate existing optimal transport solvers using these benchmark measures. Even though these solvers perform well in downstream tasks, many do not faithfully recover optimal transport maps. To investigate the cause of this discrepancy, we further test the solvers in a setting of image generation. Our study reveals crucial limitations of existing solvers and shows that increased OT accuracy does not necessarily correlate to better results downstream.

研究动机与目标

• 为连续、高维设置下的神经最优传输求解器缺乏标准化定量评估提供解决方案。
• 构建具有解析已知最优传输映射的连续概率测度对，用于二次成本最优传输（Wasserstein-2）。
• 评估现有OT求解器在高维空间（包括图像分布）中恢复真实传输映射的保真度。
• 探究OT求解器精度与生成建模下游性能之间的关系。
• 揭示尽管在GAN及其他应用中表现出色，当前求解器仍存在局限性。

提出的方法

• 使用输入凸神经网络（ICNNs）参数化传输映射，其梯度可生成给定测度的最优传输映射。
• 通过ICNN梯度对基测度进行前向推导，构建连续概率测度，确保最优传输映射的解析可计算性。
• 利用ICNN梯度可精确生成对应二次成本OT问题的最优传输映射的特性，实现真实值的精确计算。
• 基于此构造方法，在高维空间（包括64×64 CelebA人脸图像）生成基准数据集。
• 在传输映射恢复任务和基于GAN的图像生成设置下评估求解器，以比较OT精度与下游性能。
• 采用三重循环随机极小-极大-极小优化框架训练GAN，以适应MM和MM:R等具有复杂优化结构的非对称求解器。

实验结果

研究问题

• RQ1 神经最优传输求解器是否能在连续、高维设置下准确恢复真实最优传输映射？
• RQ2 Wasserstein-2距离估计的改进在多大程度上与更好的生成建模性能相关？
• RQ3 为何一些OT映射恢复性能较差的求解器在GAN中仍能取得强大的生成结果？
• RQ4 梯度偏差与支撑集不匹配如何影响OT求解器在图像生成中的性能？
• RQ5 所提出的基于ICNN的基准能否推广至其他代价函数或测度类型？

主要发现

• 许多流行的神经OT求解器即使在低维设置下也表现出显著的映射恢复误差，其性能与简单基线相当。
• 基于ICNN的参数化方法产生最精确的求解器，其中ICNN-based W2求解器对真实传输映射的保真度最高。
• 尽管映射恢复准确，W2求解器并未实现最先进的生成性能，表明OT精度与下游质量之间存在脱节。
• 采用极大极小公式化的MM和MM:R求解器因梯度偏差和训练不稳定而产生模糊图像，尽管MM:R通过反向架构设计有所改善。
• QC求解器中拟合的传输映射显著提升了图像质量，生成结果与W2求解器相当，表明即使求解器并非最精确，映射保真度仍可增强生成效果。
• 该基准揭示，更高的OT精度并不必然带来更好的生成性能，挑战了OT在GAN中作用的传统假设。