[论文解读] Large-Scale Wasserstein Gradient Flows

Petr Mokrov, Alexander Korotin+4 more|arXiv (Cornell University)|Jun 1, 2021

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

一句话总结

该论文提出了一种可扩展的方法，利用输入凸神经网络（ICNN）参数化JKO格式中的最优传输映射，实现大规模Wasserstein梯度流的高效采样与密度估计，无需域离散化或粒子模拟。该方法通过随机梯度下降实现精确的扩散建模与未归一化密度采样，已在Fokker-Planck动力学、贝叶斯推断和非线性滤波任务中取得成功。

ABSTRACT

Wasserstein gradient flows provide a powerful means of understanding and solving many diffusion equations. Specifically, Fokker-Planck equations, which model the diffusion of probability measures, can be understood as gradient descent over entropy functionals in Wasserstein space. This equivalence, introduced by Jordan, Kinderlehrer and Otto, inspired the so-called JKO scheme to approximate these diffusion processes via an implicit discretization of the gradient flow in Wasserstein space. Solving the optimization problem associated to each JKO step, however, presents serious computational challenges. We introduce a scalable method to approximate Wasserstein gradient flows, targeted to machine learning applications. Our approach relies on input-convex neural networks (ICNNs) to discretize the JKO steps, which can be optimized by stochastic gradient descent. Unlike previous work, our method does not require domain discretization or particle simulation. As a result, we can sample from the measure at each time step of the diffusion and compute its probability density. We demonstrate our algorithm's performance by computing diffusions following the Fokker-Planck equation and apply it to unnormalized density sampling as well as nonlinear filtering.

研究动机与目标

• 解决高维扩散过程中JKO步骤在Wasserstein梯度流中的计算不可行性问题。
• 通过同时访问样本和密度，实现对时变概率测度的连续、可微分近似。
• 通过ICNN参数化传输映射，克服域离散化和基于粒子方法的局限性。
• 支持机器学习应用，如未归一化后验采样和非线性滤波，实现可扩展、可微分的动力学。
• 提供一种通用框架，通过Brenier定理与ICNN避免昂贵的最优传输计算，以求解Wasserstein梯度流。

提出的方法

• 利用Brenier定理将最优传输映射表示为凸势函数的梯度，并由输入凸神经网络（ICNN）参数化。
• 通过最小化Wasserstein空间中的正则化自由能泛函，使用随机梯度下降（SGD）优化每个JKO步骤。
• 通过连续的ICNN-based传输映射参数化时间离散化的Fokker-Planck动力学，实现对测度演化的顺序学习。
• 通过ICNN-based映射的可逆结构，实现对演化测度 $\rho_t$ 的采样与密度估计。
• 避免显式计算Wasserstein距离，依赖ICNN的梯度结构与凸性，确保优化问题的良好设定。
• 在JKO步骤中引入熵正则化，以确保严格凸性，并通过对偶公式实现高效优化。

实验结果

研究问题

• RQ1 输入凸神经网络能否在高维空间中高效且可扩展地近似Wasserstein梯度流的JKO步骤？
• RQ2 如何在不进行域离散化或粒子模拟的情况下，同时实现对演化测度 $\rho_t$ 的采样与密度估计？
• RQ3 ICNN-based参数化在多大程度上保持了Wasserstein梯度流的几何结构，同时仍适用于随机优化？
• RQ4 该方法能否在机器学习任务中有效应用于未归一化后验采样与非线性滤波？
• RQ5 与传统最优传输求解器或粒子方法相比，使用ICNN进行JKO步骤在计算与统计权衡方面有何表现？

主要发现

• 所提出的ICNN-based JKO方法可在无需域离散化或粒子模拟的情况下，实现对演化测度 $\rho_t$ 的高效采样与密度估计。
• 通过SGD顺序学习传输映射，该方法实现了对Fokker-Planck动力学的精确近似，且无需显式计算Wasserstein距离。
• 该方法支持未归一化后验采样与非线性滤波任务，在贝叶斯机器学习中展现出实际应用价值。
• ICNN的使用使得测度演化的连续、可微分建模成为可能，从而在每个时间步均可实现生成与密度评估。
• 该方法可扩展至高维空间，避免了基于网格离散化的指数级计算成本，适用于复杂的真实世界扩散过程。
• 在分子发现与种群动力学的初步实验表明其具有更广泛的应用潜力，在MOSES分子数据集上表现优异。