[论文解读] Variational Wasserstein gradient flow

Jiaojiao Fan, Qinsheng Zhang+2 more|arXiv (Cornell University)|Dec 4, 2021

Advanced Neuroimaging Techniques and Applications被引用 8

一句话总结

本文提出了一种变分Wasserstein梯度流方法，将依赖于密度的优化目标替换为参数化函数上的变分公式，从而实现在高维经验分布上的高效、可扩展计算。通过利用f-散度的变分形式和输入凸神经网络（ICNNs），该方法避免了昂贵的密度估计和Hessian行列式近似，实现了在CIFAR10数据集上SOTA的FID分数（23.1），并展现出比GAN更稳定训练和更优的生成样本质量。

ABSTRACT

Wasserstein gradient flow has emerged as a promising approach to solve optimization problems over the space of probability distributions. A recent trend is to use the well-known JKO scheme in combination with input convex neural networks to numerically implement the proximal step. The most challenging step, in this setup, is to evaluate functions involving density explicitly, such as entropy, in terms of samples. This paper builds on the recent works with a slight but crucial difference: we propose to utilize a variational formulation of the objective function formulated as maximization over a parametric class of functions. Theoretically, the proposed variational formulation allows the construction of gradient flows directly for empirical distributions with a well-defined and meaningful objective function. Computationally, this approach replaces the computationally expensive step in existing methods, to handle objective functions involving density, with inner loop updates that only require a small batch of samples and scale well with the dimension. The performance and scalability of the proposed method are illustrated with the aid of several numerical experiments involving high-dimensional synthetic and real datasets.

研究动机与目标

• 解决由于密度估计和Hessian矩阵计算导致的高维空间中Wasserstein梯度流计算不可行的问题。
• 开发一种在经验分布上计算Wasserstein梯度流的数值稳定且可扩展的方法。
• 将依赖于密度的优化目标函数替换为一种允许基于样本评估的变分公式。
• 相比标准GAN，提升训练稳定性并避免模式崩溃。
• 在高维合成数据和真实世界数据集（包括MNIST和CIFAR10）上展示该方法的有效性。

提出的方法

• 基于f-散度的变分表示形式构建目标函数，通过在参数化函数类上进行最大化。
• 使用输入凸神经网络（ICNNs）将最优传输映射参数化为凸函数的梯度。
• 通过在小批量样本上的内层更新，替代显式密度评估或Hessian行列式近似的需求。
• 通过随机优化实现JKO格式，其中每个时间步长最小化正则化的Wasserstein距离加上变分目标。
• 通过一系列传输映射，使初始分布（如高斯分布）沿梯度流动力学演化至目标分布。
• 通过传播样本经过学习到的传输映射，实现沿流的采样和密度评估。

实验结果

研究问题

• RQ1 在高维设置下，Wasserstein梯度流目标函数的变分公式是否能消除对显式密度估计的需求？
• RQ2 与现有基于JKO的方法相比，该方法在计算效率和可扩展性方面表现如何？
• RQ3 该变分公式是否能保持诸如矩匹配和嵌入不等式等几何与统计性质？
• RQ4 该方法是否能在无需模式崩溃的情况下，实现在CIFAR10等图像生成基准上的SOTA样本质量？
• RQ5 该方法的性能如何随不同步长和流步数变化？

主要发现

• 所提方法在CIFAR10上实现了23.1的Fréchet Inception Distance（FID）分数，优于GAN及其他基于流的生成模型。
• 采用步长a=5.0时训练收敛稳定且未出现模式崩溃，而更大的值会导致不稳定和样本质量下降。
• 与需要Hessian行列式近似的先前JKO方法相比，该方法展现出更优的可扩展性和计算效率。
• 变分目标满足关于积分概率度量的矩匹配性质和嵌入不等式，支持其理论合理性。
• FID分数随时间步数收敛，表明即使在增加流步数时，学习动态仍保持稳定。
• 该方法在CIFAR10上实现了具有竞争力的Inception Score（7.48 ± 0.12），与AE-OT-GAN和OTM等SOTA模型相当。