[论文解读] Neural Wasserstein Gradient Flows for Maximum Mean Discrepancies with Riesz Kernels

Fabian Altekrüger, Johannes Hertrich+1 more|arXiv (Cornell University)|Jan 27, 2023

Generative Adversarial Networks and Image Synthesis被引用 5

一句话总结

本论文提出基于生成神经网络近似传输与速度规划的神经Wasserstein梯度流，用于处理具有非光滑Riesz核的最大均值差异（MMD）泛函，避免依赖传统映射，从而能够处理奇异测度与绝对连续测度。该方法在MMD流与非光滑核下的barycenter计算中，实现了前后向格式向解析解的收敛。

ABSTRACT

Wasserstein gradient flows of maximum mean discrepancy (MMD) functionals with non-smooth Riesz kernels show a rich structure as singular measures can become absolutely continuous ones and conversely. In this paper we contribute to the understanding of such flows. We propose to approximate the backward scheme of Jordan, Kinderlehrer and Otto for computing such Wasserstein gradient flows as well as a forward scheme for so-called Wasserstein steepest descent flows by neural networks (NNs). Since we cannot restrict ourselves to absolutely continuous measures, we have to deal with transport plans and velocity plans instead of usual transport maps and velocity fields. Indeed, we approximate the disintegration of both plans by generative NNs which are learned with respect to appropriate loss functions. In order to evaluate the quality of both neural schemes, we benchmark them on the interaction energy. Here we provide analytic formulas for Wasserstein schemes starting at a Dirac measure and show their convergence as the time step size tends to zero. Finally, we illustrate our neural MMD flows by numerical examples.

研究动机与目标

• 开发一种基于神经网络的方法，用于计算具有非光滑Riesz核的MMD泛函的Wasserstein梯度流，能够实现奇异测度与绝对连续测度之间的过渡。
• 通过使用生成网络近似传输与速度规划，而非依赖传输映射，克服传统JKO格式与粒子流的局限性。
• 在目标测度未以解析形式已知且核函数为非凸非光滑的情况下，实现MMD流与barycenter的计算。
• 提供一种数值稳定且可扩展的框架，用于基于MMD的生成建模，突破绝对连续测度的限制。
• 通过收敛性分析与复杂流体的数值实验，验证神经格式的有效性，包括非连续测度演化。

提出的方法

• 利用生成神经网络近似Jordan-Kinderlehrer-Otto（JKO）后向格式，学习具有Riesz核的MMD泛函的传输计划分解。
• 采用基于Wasserstein最速下降流的前向格式，同样通过生成网络近似速度计划，实现从奇异初始测度的演化。
• 使用双重积分建模MMD中的相互作用能与势能项，其形式适合基于样本的近似，从而实现高效的梯度计算。
• 通过针对Wasserstein度量与Riesz核结构定制的损失函数训练神经网络，确保与底层最优传输动力学的一致性。
• 利用MMD与Riesz核具有切片表示的特性，通过一维排序实现快速梯度计算，提升训练效率。
• 将该方法应用于计算MMD barycenter与从复杂目标密度采样，包括非光滑与非凸能量景观。

实验结果

研究问题

• RQ1 神经网络能否有效近似具有非光滑Riesz核的MMD泛函的Wasserstein梯度流，特别是在流在奇异测度与绝对连续测度之间过渡时？
• RQ2 当初始测度为奇异时，神经后向与前向格式相较于经典粒子流在收敛性与稳定性方面表现如何？
• RQ3 生成网络在不假设测度绝对连续性的前提下，能在多大程度上建模传输与速度规划？
• RQ4 当时间步长趋于零时，所提方法是否收敛至MMD流的解析解？
• RQ5 当目标测度分别支持在不同子流形上时，神经格式能否成功计算MMD barycenter？

主要发现

• 随着时间步长趋于零，神经后向与前向格式均收敛至具有Riesz核的MMD流的解析解，验证了理论一致性。
• 该方法成功处理测度转换：Dirac测度可演化为绝对连续测度或奇异非Dirac测度，具体取决于维度，而传统粒子流无法建模此类行为。
• 数值实验表明，当从奇异初始测度开始时，两种神经格式在捕捉非光滑动力学方面均优于经典粒子流。
• Riesz核的使用使得通过切片实现高效梯度计算，从而加快训练并确保神经网络优化的稳定性。
• 神经格式成功计算了支持在不同子流形（如圆与正方形边界）上的测度之间的MMD barycenter，且收敛至理论barycenter。
• 在MNIST示例中，从均匀初始测度而非奇异测度开始，所有方法性能相近，表明神经格式对初始测度类型的鲁棒性。