©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 的表达式。
们此处不具体讨论 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 一个距离函数，并通过它定义流形上任何两点
的测地线
