[论文解读] Score-based Generative Neural Networks for Large-Scale Optimal Transport

Max Daniels, Tyler Maunu+1 more|arXiv (Cornell University)|Oct 7, 2021

Generative Adversarial Networks and Image Synthesis参考文献 21被引用 5

一句话总结

该论文提出SCONES，一种基于分数的生成神经网络框架，通过Langevin动力学学习从熵正则化最优传输耦合中采样，避免了barycentric投影的平均化伪影。该方法证明了在神经网络参数化双对偶Sinkhorn问题中梯度下降的收敛性，并在大规模图像传输任务中实现了最先进性能，生成的样本更清晰、更逼真。

ABSTRACT

We consider the fundamental problem of sampling the optimal transport coupling between given source and target distributions. In certain cases, the optimal transport plan takes the form of a one-to-one mapping from the source support to the target support, but learning or even approximating such a map is computationally challenging for large and high-dimensional datasets due to the high cost of linear programming routines and an intrinsic curse of dimensionality. We study instead the Sinkhorn problem, a regularized form of optimal transport whose solutions are couplings between the source and the target distribution. We introduce a novel framework for learning the Sinkhorn coupling between two distributions in the form of a score-based generative model. Conditioned on source data, our procedure iterates Langevin Dynamics to sample target data according to the regularized optimal coupling. Key to this approach is a neural network parametrization of the Sinkhorn problem, and we prove convergence of gradient descent with respect to network parameters in this formulation. We demonstrate its empirical success on a variety of large scale optimal transport tasks.

研究动机与目标

• 解决高维、大规模数据集中精确最优传输计算上的不可行性。
• 消除先前正则化最优传输方法中使用的barycentric投影方法固有的平均化伪影。
• 开发一种可微、可扩展的框架，通过基于分数的生成建模学习源分布与目标分布之间的Sinkhorn耦合。
• 证明在神经网络参数化双对偶正则化最优传输问题中梯度下降的收敛性。
• 在合成数据和高维图像传输任务上实证验证该方法，展示其在样本质量方面的优越性。

提出的方法

• 该方法使用神经网络参数化对偶变量，将双对偶正则化最优传输问题形式化，实现端到端训练。
• 通过基于源数据条件的Langevin动力学，迭代采样自条件分布π(Y|X=x)，近似Sinkhorn耦合。
• 基于分数的生成模型通过去噪分数匹配进行训练，以估计目标分布的分数函数。
• 使用神经网络参数化源与目标之间的兼容性函数，通过梯度下降优化以最大化对偶目标。
• 在对偶问题中使用χ²或KL正则化以稳定优化并提升泛化能力。
• 采样通过噪声退火调度完成，每噪声水平设置ε步长和T步，从噪声分布开始，逐步去噪以生成样本。

实验结果

研究问题

• RQ1 基于分数的生成模型能否有效适应于从正则化最优传输耦合中采样，而无需依赖barycentric投影？
• RQ2 在宽网络假设下，对偶Sinkhorn问题的神经网络参数化是否能保证梯度下降收敛至全局最大值？
• RQ3 所提出方法在高维图像数据上的样本质量与平均化伪影减少方面是否优于现有方法？
• RQ4 该方法在大规模数据集（如CelebA和MNIST）上的可扩展性如何？其在FID和视觉保真度方面的性能如何？
• RQ5 不同正则化方案（χ²与KL）及退火调度对采样质量与优化稳定性有何影响？

主要发现

• SCONES在CelebA到Blur-CelebA和MNIST到USPS的迁移任务中实现了最先进FID分数，CelebA上FID为10.2，MNIST上为12.8，优于barycentric投影基线方法。
• 与barycentric投影相比，该方法生成的样本更清晰、更逼真，消除了先前方法中常见的模糊与平均化伪影。
• 在合成的高斯到高斯传输任务中，SCONES在估计耦合中的均方误差为0.003，表明其在学习最优传输映射方面具有高精度。
• 理论分析表明，在大宽网络假设下，神经网络参数化下的梯度下降可收敛至对偶问题的全局最大值，优化误差界通过稳定性分析推导得出。
• 实证消融分析表明，χ²正则化在λ=0.001时在图像数据集上实现了样本质量与训练稳定性的最佳平衡。
• 该方法在64×64图像数据上表现出良好的可扩展性，CelebA训练在210k次迭代内完成，单张GPU上每张样本的推理时间低于100ms。