以Langevin Diffusion为例，探索如何将采样视为优化问题

最新推荐文章于 2024-11-04 10:51:20 发布

PaperWeekly

最新推荐文章于 2024-11-04 10:51:20 发布

阅读量145

点赞数

文章标签：人工智能深度学习

原文链接：https://mp.weixin.qq.com/s?__biz=MzIwMTc4ODE0Mw==&mid=2247679596&idx=1&sn=fcf9a1519d8ef438941da826ba98ef60&chksm=97b8c1f36ec21fb5a9b6f80269165b9e2a53d045b5850aa22ffbf987de6d4fec91c20cf78aff&scene=126&sessionid=0

版权

©PaperWeekly 原创 · 作者 | Rui-Yang Zhang

单位 | Lancaster University UK

研究方向 | Computational Statistics

采样（sampling）与优化（optimization）的联系是错综复杂的。近年来，Applied Mathematics、Computational Statistics、Machine Learning 等领域的学者开始逐步探索采样与优化这两个课题之间的联系。本文以 Langevin Diffusion 为例，粗略刻画一些比较成熟的、有启发性的脉络。

更多这个方向的内容和背景可以参考 Sinho Chewi, Jonathan Niles-Weed, Philippe Rigollet 前段时间挂在 arxiv 上的书 Statistical Optimal Transport（https://arxiv.org/abs/2407.18163）。

视采样为优化问题

采样与优化最直接的联系是通过 Variational Inference（VI），又名 Variational Bayes，是机器学习中十分常见的推断手段。Variational Inference 的核心思想如下：

假定我们有目标概率分布，我们试图用来逼近它，其中是一族以为参数的概率分布，例如多元高斯。我们通过优化来找到一个合适的以作为目标函数的 approximation。之后，一切需要使用的工作我们都用来代替。

理论上，如果足够大（最好涵盖），那么我们的逼近的质量就会很好，后续的任务的结果也是。Variational Inference 在计算统计中的直接竞争对手是 MCMC。MCMC 帮助我们从目标分布的采样，并用采样做 empirical distribution 以逼近。一般来说，MCMC 相较于 VI 速度慢但质量高。更多可见 arXiv:1601.00670v9。

回到 VI。我们可以将其视为用优化（KL divergence）为途径作采样的一种方法。我们会在下文中进一步探索我们具体如何用优化的眼光去诠释采样。

什么是 Langevin Diffusion？

Langevin Diffusion，又叫做 Overdamped Langevin Equation，在一维中的表达式是

其中，是随机过程，是布朗运动，而的 equilibrium distribution 的 density 是（注：equilibrium 的存在要求 V 满足一些 regularity conditions，详见 Roberts and Tweedie（1996）Theorem 2.1。）。

如果我们对 Langevin Diffusion 作 Euler-Maruyama Discretisation，我们就可以得到 Unadjusted Langevin Algorithm：

这个算法又叫做 Langevin Monte Carlo（LMC），是 MCMC 算法 Metropolis-Adjusted Langevin Algorithm（MALA）的变种，在计算统计和机器学习中十分常见。LMC 的输出是一系列，而这些样本可以用作的（有偏 biased）采样。如果我们在每一步中做 Metropolis Adjustment，则输出的是无偏采样。

Langevin Diffusion的Fokker-Planck Equation

Fokker-Planck Equation（FPE）帮助我们将 SDE 从 particle 的表达形式转换成 distribution 的表达形式。这里不去过多解释 FPE，只套用公式。我们有

其中是的 probability distribution。如果我们的 Langevin Diffusion 有 equilibrium，那么我们的 p 就不依赖于 t，所以上面的式子等于 0。这也可以帮助我们发现我们的 equilibrium 是。我们会在之后再一次回到这个 Fokker Planck 的表达式。

Wasserstein 空间的定义

们此处不具体讨论 Optimal Transport，以及太多的关于 Wasserstein 空间的细节。我们将 Wasserstein 空间简单的想象成一种可以让我们在其中定义一些几何的概率分布空间。我们用指代我们的 Wasserstein space。我们有 metric space ，其中是定义在上的所有 finite variance 的概率分布，而是 2-Wasserstein metric，定义为

其中是概率分布的所有 coupling。更多的可以参考：

https://zhuanlan.zhihu.com/p/663591093

Wasserstein Gradient Flow

在 Euclidean 空间中的 gradient flow 比较容易定义。如果我们有 differentiable 目标函数，并且我们想要对其做 minimisation，那么我们可以设定以下的 gradient flow：

由起点出发，连续地随着目标函数 F 下降，直到 flow 收敛。显然，这里的目标函数如果是非凸且有多个 local minima，那么我们可能会陷入 local optimal 中。这个是没有办法避免的。

上面提到的 gradient flow 是一个 ODE，是连续时间的。除非我们有 ODE 的解析解，我们都没有办法完美、连续地沿着这个 ODE 运动。所以我们需要离散化，去 approxiamte。如果我们用（explicit）Euler discretisation，我们就有

而这个是我们熟知的 gradient descent。不同的对于 gradient flow 的离散化可以帮助我们获得不同的离散时间算法。这里不去做过多的介绍。

在 Euclidean 空间中的 gradient flow 固然有趣，但我们这里更关系在概率分布中建立 gradient flow。我们这里将在 Wasserstein 空间中搭建 gradient flow（又称 Wasserstein gradient flow），我们试图对于目标函数进行优化，建立类似于以下的 ODE：

定义

在场论中我们有 Eulerian 和 Lagrangian 两种诠释法。对于一个位于 time-varying vector field 中的粒子，它的运动轨迹通过 Lagrangian specification 是；而它的 Eulerian specification 是，其中是的 distribution function。这个 ODE 也叫做 continuity equation。

我们不难发现，对于 Eulerian specification，vector field 不是 unique 的。但我们对于 vector field 的选择（尽管满足我们的 continuity equation）会影响我们在 Lagrangian specification 中的理解。

我们有以下最优的选择使得我们的 kinetic energy 是最小的：，其中是从到的 optimal transport。这个的证明可以参考（https://chewisinho.github.io/main.pdf）的 Theorem 1.3.19。

定义

我们需要在 Wasserstein 空间中求导。这里，我们借助黎曼几何的手段，赋予 Wasserstein 空间一个 Riemannian metric，并由此定义导数。

对于任意一个黎曼流形和定义在其上的 Riemannian metric 和 tangent space ，我们可以 induce 一个距离函数，并通过它定义流形上任何两点的测地线

其中 argmin 的 curve 就是测地线。其次，如果我们有定义在流形上的函数，我们可以定义它在点的导数为唯一的使得所有以速度 v 在 t=0 时穿过 p 点的 curves 满足。

回到 Wasserstein space。我们可以有 Wasserstein geodesic，定义为

我们不去解释为什么这个定义是合理的。需要注意的是，这里因为我们在 Wasserstein space，所以我们的点都是 probability distributions，且我们需要让我们的 vector field 满足 continuity equation。更重要的是，我们可以推导出，Wasserstein geodesic 可以被建造，并其可以被表达为，，其中。

之后，对于某个定义在 Wasserstein space 上的函数 F，我们有它的 first variation 满足，则我们有并且右手边的是在 Euclidean space 中的 gradient。

拼接起来

如果我们把我们现有的组件都拼接起来，我们有 continuity equation 并且我们的 vector field 因为是 gradient flow，所以选为，于是我们得到 Wasserstein gradient flow：

如果我们选择特定的 objective function ，我们可以推导出，所以我们有：对于 KL divergence 的 Wasserstein gradient flow 为

于之前的 Langevin Diffusion 的 Fokker Planck Equation 吻合！这个发现最早出现于【JKO1998】（Richard Jordan, David Kinderlehrer, and Felix Otto. The variational formulation of the fokker– planck equation. SIAM journal on mathematical analysis, 29(1):1–17, 1998）。

接着，如果我们对于 Wasserstein gradient flow 做 Forward-Flow splitting discretisation scheme，我们就可以得到 LMC 的 distribution 的诠释。这个最早是在这篇文章中发现的：Andre Wibisono. Sampling as optimization in the space of measures: The langevin dynamics as a composite optimization problem. In Conference on Learning Theory, pages 2093–3027. PMLR, 2018。

然后呢？

发现了 Langevin diffusion 可以被理解成一种 Wasserstein gradient flow 的特殊离散化后，我们能够做什么呢？

Durmus et al 在他们 2019 年的 PMLR 上通过这个发现得到了 LMC 的新的 convergence rate，见：Alain Durmus, Szymon Majewski, and B la˙zej Miasojedow. Analysis of langevin monte carlo via convex optimization. The Journal of Machine Learning Research, 20(1):2666–2711, 2019。

除此之外，我们可以将这个想法（即，将采样算法用某种 gradient flow 来诠释）扩展到别的算法上。

例如 Stein Variational Gradient Descent （见 Liu Qiang 的 Stein Variational Gradient Descent as Gradient Flow（https://arxiv.org/abs/1704.07520）与 Duncan et al 的 On the geometry of Stein variational gradient descent）。例如 tempering （见 Chopin et al 的 A connection between tempering and entropic mirror descent）。

还有，我们可以通过优化和采样的联系，将优化中的工具引入采样中。一个例子是 Chris Nemeth 组的工作，将优化中的 learning rate free 的技术引入 sampling 中，见 Coin Sampling: Gradient-Based Bayesian Inference without Learning Rates 和 Learning Rate Free Sampling in Constrained Domains。

其次，我们可以考虑用不同目标函数的 Wasserstein gradient flow。比较出名的有用 Maximum Mean Discrepancy，可以参考 Arthur Gretton 组的一系列工作，例如 Maximum Mean Discrepancy Gradient Flow， Deep MMD Gradient Flow without adversarial training。