[论文解读] Wasserstein variational gradient descent: From semi-discrete optimal transport to ensemble variational inference

Luca Ambrogioni, Umut Güçlü+1 more|arXiv (Cornell University)|Nov 7, 2018

Generative Adversarial Networks and Image Synthesis参考文献 26被引用 5

一句话总结

本文提出 Wasserstein 变分梯度下降（WVGD），一种基于粒子的变分推断方法，利用半离散最优传输来最小化 Wasserstein 散度而非 KL 散度。通过利用最优传输映射，WVGD 避免了模式崩溃，且无需依赖基于核函数的排斥力，从而实现更精确的后验近似，并提供可解释的、与每个粒子相关的局部变分近似。

ABSTRACT

Particle-based variational inference offers a flexible way of approximating complex posterior distributions with a set of particles. In this paper we introduce a new particle-based variational inference method based on the theory of semi-discrete optimal transport. Instead of minimizing the KL divergence between the posterior and the variational approximation, we minimize a semi-discrete optimal transport divergence. The solution of the resulting optimal transport problem provides both a particle approximation and a set of optimal transportation densities that map each particle to a segment of the posterior distribution. We approximate these transportation densities by minimizing the KL divergence between a truncated distribution and the optimal transport solution. The resulting algorithm can be interpreted as a form of ensemble variational inference where each particle is associated with a local variational approximation.

研究动机与目标

• 解决现有基于粒子的变分推断方法的局限性，特别是 SVGD 中对任意核函数用于粒子排斥的依赖。
• 通过利用最优传输理论，克服有限粒子下基于粒子的推断中的模式崩溃问题。
• 通过最优传输映射，为每个粒子提供一种原则化的关联方式，实现局部变分近似。
• 实现集成变分推断，使每个粒子能够建模后验分布中的不同模式或区域。
• 开发一种方法，可应用于非摊销推断（amortized inference）之外的场景，与以往的 Wasserstein VI 方法不同，支持对后验分布的直接优化。

提出的方法

• 将变分推断表述为最小化粒子近似与真实后验之间的半离散最优传输散度。
• 利用最优传输问题的解，推导出将每个粒子映射到后验连续区域的传输密度。
• 使用参数族（如高斯混合模型）近似最优传输映射，并最小化近似与真实传输映射之间的 KL 散度。
• 基于 Wasserstein 梯度流推导出粒子动力学更新规则，实现粒子位置与局部变分参数的迭代优化。
• 将粒子动力学与兼容深度学习库的随机优化框架集成。
• 利用所得的传输密度，为每个粒子形成局部变分近似，从而实现集成 VI 的解释。

实验结果

研究问题

• RQ1 能否利用最优传输理论构建一种基于粒子的变分推断方法，以避免人为引入的排斥力？
• RQ2 粒子与后验之间的最优传输映射能否提供一种自然且可解释的后验空间分区？
• RQ3 在有限粒子设置下，最小化半离散 Wasserstein 散度是否能比最小化 KL 散度带来更优的后验近似？
• RQ4 该方法能否在无需摊销推断的情况下应用，与以往的 Wasserstein VI 方法不同？
• RQ5 在高维设置下，WVGD 与 SVGD 在近似精度与收敛性方面的表现如何比较？

主要发现

• 在高斯混合模型基准测试中，WVGD 在平方误差方面始终优于 SVGD，即使 SVGD 使用三倍于 WVGD 的粒子数。
• 在贝叶斯逻辑回归任务中，WVGD 使用 5 个粒子所获得的 ELBO 高于标准 VI（1 个粒子），在波士顿数据集中提升最高达 160，在癌症数据集中提升最高达 320。
• 该方法成功捕捉了后验分布中的多个模式，如通过将参数空间划分为与不同粒子相关联的独立区域所展示的。
• WVGD 学习到的传输密度对应于数据的语义上有意义的解释，例如在高斯过程协方差函数中表现出的振荡型与噪声型行为。
• WVGD 中的粒子动力学由于最优传输中固有的解释性排除效应（explaining-away effect），自然避免了模式崩溃，且无需依赖基于核函数的排斥力。
• 该方法在多种数据集（Iris、Boston、Diabetes、Cancer）上表现稳健，随着粒子数量的增加，ELBO 持续提升。