[论文解读] Variational inference via Wasserstein gradient flows

Marc Lambert, Sinho Chewi+3 more|arXiv (Cornell University)|May 31, 2022

Markov Chains and Monte Carlo Methods被引用 12

一句话总结

本文提出了一种新颖的变分推断方法，利用高斯测度的Bures–Wasserstein空间上的Wasserstein梯度流，实现了对复杂后验分布的合理近似。该方法在对数凹目标下具备强大的理论保证，并通过基于ODE的优化有效逼近了多模态目标，使用高斯混合模型作为近似。

ABSTRACT

Along with Markov chain Monte Carlo (MCMC) methods, variational inference (VI) has emerged as a central computational approach to large-scale Bayesian inference. Rather than sampling from the true posterior $π$, VI aims at producing a simple but effective approximation $\hat π$ to $π$ for which summary statistics are easy to compute. However, unlike the well-studied MCMC methodology, algorithmic guarantees for VI are still relatively less well-understood. In this work, we propose principled methods for VI, in which $\hat π$ is taken to be a Gaussian or a mixture of Gaussians, which rest upon the theory of gradient flows on the Bures--Wasserstein space of Gaussian measures. Akin to MCMC, it comes with strong theoretical guarantees when $π$ is log-concave.

研究动机与目标

• 解决在目标后验为对数凹分布时，高斯及高斯混合近似在变分推断（VI）中缺乏理论保证的问题。
• 构建一个基于概率测度的Bures–Wasserstein空间上梯度流的系统性变分推断框架。
• 通过将变分卡尔曼滤波思想推广至连续时间动力系统，利用ODE建模时间演化后验近似。
• 提出一种可扩展的基于优化的算法，用于学习高斯分量的均值与协方差参数，无需计算Hessian矩阵。
• 通过使用高斯混合近似，超越平均场VI，以捕捉相关性和多模态结构。

提出的方法

• 将变分推断表述为在高斯测度的Bures–Wasserstein几何中，通过梯度流最小化提议分布与真实后验之间的KL散度。
• 通过分部积分和Wasserstein梯度流理论，推导出每个高斯分量的均值与协方差参数的连续时间ODE。
• 采用基于粒子的表示方法，使用K个高斯分量，每个分量通过耦合的ODE演化其均值和协方差矩阵的Cholesky分解因子。
• 使用四阶龙格-库塔格式对ODE进行数值积分，状态向量包含所有均值和向量化后的Cholesky参数。
• 通过直接演化下三角Cholesky因子而非协方差矩阵本身，来保证协方差矩阵的正定性。
• 在ODE中通过使用带sigma点的立方积分规则进行数值积分，以计算期望。

实验结果

研究问题

• RQ1 在Bures–Wasserstein空间上的Wasserstein梯度流能否为高斯及高斯混合近似提供一种理论基础坚实且可扩展的传统变分推断替代方法？
• RQ2 当目标后验为对数凹分布时，所提出的方法是否能获得类似Langevin扩散等MCMC方法的强收敛保证？
• RQ3 该方法在增加高斯分量数量时，对多模态后验的逼近能力如何？
• RQ4 该算法对初始化的敏感性如何？在实际中如何缓解不良初始化的影响？
• RQ5 该方法能否在不依赖Hessian信息的情况下，有效捕捉复杂的后验结构（如相关性和多模态）？

主要发现

• 所提方法在对数凹目标下实现了理论收敛保证，将已知于MCMC的强非渐近分析成果拓展至变分推断领域。
• 对于双模态目标，算法成功使用20个高斯分量捕捉了两个模式，等高线图显示与真实后验高度一致。
• KL散度在约30个积分步长（步长为0.1时为300步）后趋于稳定，表明优化过程已收敛。
• 增加高斯粒子数量可提升近似质量，对具有4–6个非等权重模态的目标，KL散度更小，密度估计更准确。
• 初始化显著影响收敛结果：初始位置靠近某模式的粒子会保持在该模式附近，而初始位置与各模式等距的粒子则会收敛至平均位置并伴随协方差增大，凸显了多组初始粒子的必要性。
• 该方法通过分部积分推导出无Hessian的更新方式，成功避免了Hessian计算，从而实现了高效且稳定的优化。