[论文解读] Scalable Unbalanced Optimal Transport using Generative Adversarial Networks

Karren Yang, Caroline Uhler|arXiv (Cornell University)|Oct 26, 2018

Model Reduction and Neural Networks被引用 15

一句话总结

该论文提出了一种基于生成对抗网络（GAN）的可扩展框架，用于处理不平衡最优传输（OT），通过联合学习传输映射和缩放因子来建模源分布与目标分布之间的质量变化。通过将不平衡OT表述为类似Monge的问题，并利用凸共轭对偶性进行松弛，该方法实现了类似GAN的高效、随机、交替梯度更新，从而在MNIST、USPS、CelebA以及斑马鱼单细胞RNA-seq数据的群体建模任务中取得了最先进性能。

ABSTRACT

Generative adversarial networks (GANs) are an expressive class of neural generative models with tremendous success in modeling high-dimensional continuous measures. In this paper, we present a scalable method for unbalanced optimal transport (OT) based on the generative-adversarial framework. We formulate unbalanced OT as a problem of simultaneously learning a transport map and a scaling factor that push a source measure to a target measure in a cost-optimal manner. In addition, we propose an algorithm for solving this problem based on stochastic alternating gradient updates, similar in practice to GANs. We also provide theoretical justification for this formulation, showing that it is closely related to an existing static formulation by Liero et al. (2018), and perform numerical experiments demonstrating how this methodology can be applied to population modeling.

研究动机与目标

• 为解决经典最优传输在处理源分布与目标分布之间质量变化方面的局限性。
• 开发一种可扩展、连续且可微的不平衡OT框架，支持传输与质量变化的联合建模。
• 通过引入对偶变量与散度正则化，将类似GAN的训练动态扩展至不平衡OT，形式化为松弛的Monge问题。
• 实现在高维、连续设置（如生物群体建模与图像生成）中不平衡OT的实际应用。
• 为连续测度提供现有熵-传输求解器的可扩展替代方案，提升训练稳定性和收敛性。

提出的方法

• 将不平衡OT表述为包含随机传输映射 $ T $ 和缩放因子 $ \xi $ 的Monge型问题，最小化包含传输成本与散度惩罚的代价函数。
• 利用 $ \psi $-散度的凸共轭对偶性松弛该问题，将其转化为适合随机梯度下降的对偶优化问题。
• 采用受GAN启发的交替梯度下降算法，其中生成器学习 $ T $ 和 $ \xi $，判别器学习对偶变量 $ f $。
• 对 $ T $ 和 $ \xi $ 使用ReLU与Sigmoid/Softplus输出层的神经网络，并应用梯度惩罚以稳定训练。
• 在训练中引入重缩放技巧：通过用 $ 1/\xi $ 缩放 $ T $ 的损失，当 $ \xi $ 较小时加速收敛。
• 提出一种可扩展算法（附录中算法2），用于在连续设置下求解Liero等人（2018）提出的最优熵-传输问题，扩展了先前方法的适用范围。

实验结果

研究问题

• RQ1 能否将类似GAN的框架适配于不平衡最优传输问题，并显式建模质量变化？
• RQ2 如何将不平衡OT的Monge表述松弛为适合基于深度学习的随机优化的对偶优化问题？
• RQ3 哪些训练动态与正则化策略（如梯度惩罚）能提升连续测度下不平衡OT的稳定性和收敛性？
• RQ4 所提出的方法能否在高维、连续数据集（如图像与单细胞组学数据）上实现可扩展性？
• RQ5 在准确性和效率方面，联合学习传输与缩放因子的方法相较于现有熵-传输求解器有何表现？

主要发现

• 所提方法在MNIST、USPS和CelebA数据集上的群体转换任务中成功建模了质量变化与传输过程。
• 在斑马鱼单细胞RNA-seq数据上，该方法准确捕捉了发育过程中的动态转变，呈现逼真的细胞群体演化特征。
• 对 $ \xi $ 和 $ f $ 应用梯度惩罚显著提升了训练稳定性，尤其在训练初期效果明显。
• 对 $ T $ 的损失采用重缩放技巧，当 $ \xi $ 值较小时可加快收敛速度，避免对匹配度差的样本更新过慢。
• 与现有熵-传输求解器相比，该方法在性能上具有竞争力，且在连续和高维测度上展现出更优的可扩展性。
• 所提出的算法2为连续不平衡OT提供了Chizat等人（2016）方法的可扩展替代方案，扩展了熵-传输方法的应用范围。