ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

ULA的Wasserstein混合时间:理论解析与Python实践

ULA的Wasserstein混合时间:理论解析与Python实践 做采样或者说 MCMC 方向的人应该对这样一个趋势不陌生近两年关于 Unadjusted Langevin AlgorithmULA的收敛性论文标题里几乎都带着同一个固定组合——Wasserstein mixing time。乍看这像是一批纯理论工作和写代码跑实验的工程师没什么关系。但实际影响远比表面大贝叶斯推断里要回答“后验采样到底要跑多少步”生成模型里要回答“扩散模型反向过程离散化要用多小步长”背后都依赖同一套分析逻辑。这篇文章想把这根线完整串起来。我会先讲清楚 ULA 是什么、Wasserstein 距离为什么天然适合分析它、mixing time 到底刻画了什么然后从理论结果里提炼出一个对实践非常有用的结构ULA 的收敛误差可以分解为“初始距离的指数衰减”加上“离散化带来的偏差”步长选择就是在两者之间做权衡最后用 Python 写一组最小实验直观验证步长如何影响混合时间并给出工程上的建议和排错思路。读完这篇文章你能得到三样东西对 Wasserstein mixing time 这个概念的具体理解一套可以直接复制运行的 ULA 采样与距离度量代码以及在实际采样任务里选步长、判断收敛、避开常见坑的方法。文章不需要很强的数学背景高中水平的概率和矩阵运算就够了。1. 这篇文章真正要解决的问题先从一个最实际的问题出发。你在做贝叶斯推断后验分布没有解析形式只能用采样算法去逼近。你打开 Pyro 或者 Stan写了一个模型跑了几千步然后问自己这些样本够了吗要不要再多跑一万步传统 MCMC 的回答方式是看 trace plot、看 R-hat、看有效样本量。这些诊断方法有用但有一个根本缺陷它们只能告诉你“链可能已经混合了”不能告诉你“离真实目标分布还有多远”。而且诊断过程本身要依赖人的经验。ULA 的出现让“理论回答”成为可能。它是一种非常简单的采样算法从某个初始点出发每一步按照目标分布的梯度走一小步再加上一个高斯噪声。没有接受拒绝步骤没有复杂的 leapfrog 积分代码只要十行。但正因为少了 Metropolis 接受步骤它的链不再精确收敛到目标分布而是收敛到目标分布的一个“有偏近似”。于是问题来了这个偏差有多大需要跑多少步才能把偏差压到可接受范围答案恰好由 Wasserstein 距离下的混合时间给出。对 ULA 来说理论会告诉你给定目标分布的凸性和光滑性步长取多大、迭代多少步误差的期望上界是多少。这种“画出误差上界曲线”的能力是传统 MCMC 诊断很难提供的。什么人最应该读这篇文章三类。第一是正在研究采样算法、想读 ULA 收敛性论文但被数学符号劝退的人第二是在实际项目里用 Langevin 类算法做 Bayesian 推断、扩散模型生成需要理解步长和迭代数怎么权衡的人第三是单纯想搞懂“Wasserstein 距离为什么在深度学习时代这么火”的人。如果你不在这些人群里这篇文章也可以快速帮你建立采样与最优传输之间的桥梁。2. 先把三个概念对齐ULA、Wasserstein 距离、混合时间这一节不会堆公式只做一件事把题目里的三个关键词逐个讲清楚并说明它们之间的连接方式。2.1 ULA去掉接受步的 Langevin 动力学离散化先回忆连续时间 Langevin 动力学。给定目标分布 π(x) ∝ exp(−U(x))其中 U 称为势能函数那么随机微分方程dX_t −∇U(X_t) dt √2 dW_t的平稳分布就是 π。也就是说如果让粒子按照这个方程运动足够久它的位置分布会收敛到目标分布。这里 ∇U 等于 −∇log π(x)所以在贝叶斯场景里它就是后验密度的梯度。ULA 的做法非常直接用欧拉离散化去逼近这条连续轨迹X_{k1} X_k − h ∇U(X_k) √(2h) ξ_k其中 ξ_k 是标准高斯噪声h 是步长。这个迭代式的直观含义是粒子沿着目标分布的梯度方向爬升靠近高概率区域同时被噪声扰动保证探索性。这里的“Unadjusted”是关键。标准的 MALAMetropolis-adjusted Langevin Algorithm会在 ULA 的候选点之后加一个 Metropolis-Hastings 接受步骤保证链的平稳分布精确等于 π。而 ULA 直接跳过这一步。代价是链的平稳分布产生偏差收益是每一步的计算量大幅下降而且完全避免了在高维空间中接受率过低的问题。从实际项目角度ULA 的典型应用场景包括大规模贝叶斯推断、变分推断的补充采样、以及分数匹配生成模型里对逆向 SDE 的模拟。这些场景的共同特点是目标维数高、单步成本敏感无法承受 MALA 或 HMC 的额外计算开销。2.2 Wasserstein 距离把分布收敛看成运输问题Wasserstein 距离是度量两个概率分布之间差异的一种方式也叫推土机距离Earth Movers Distance。它的直观定义是把 μ 分布的质量“搬运”成 ν 分布的质量最少需要多少运输成本。数学上一阶和二阶 Wasserstein 距离分别定义为W₁(μ, ν) inf_{γ ∈ Π(μ,ν)} E_{(X,Y)∼γ} [ |X − Y| ]W₂(μ, ν) ( inf_{γ ∈ Π(μ,ν)} E_{(X,Y)∼γ} [ |X − Y|² ] )^{1/2}其中 Π(μ, ν) 是所有以 μ 和 ν 为边缘分布的联合分布也就是“搬运方案”的集合。W₂ 在深度学习和生成模型里非常常见WGAN 的判别器近似的是 W₁ 的变体扩散模型的收敛性分析则大量使用 W₂。为什么 Wasserstein 距离在采样算法的分析里如此受欢迎因为它同时兼顾了两件事分布的“位置差异”和“形状差异”。如果两个分布都集中在不同的地方Wasserstein 距离会直接反映簇中心的距离如果中心相同但方差不同它也能捕捉到方差的差异。这一点是 KL 散度和总变差距离都做不到的。一个很实用的性质是对于高斯分布W₂ 有精确的闭式解。这意味着当我们用 ULA 去采样一个高斯目标时可以精确计算采样分布与真实目标之间的距离非常适合做实验验证。2.3 混合时间从初值出发要走多少步才算“到站”混合时间来自马尔可夫链理论。给定一个马尔可夫链初始分布是 μ₀第 k 步的边际分布是 μ_k目标不变分布是 π。混合时间刻画的是从初始分布出发经过多少步之后μ_k 和 π 之间的距离缩小到某个阈值 ε 以内。在 Wasserstein 框架下可以这样定义τ_ε min { k ≥ 0 : W₂(μ_k, π) ≤ ε }注意这里的 μ_k 是“把链重复跑很多次得到的边际分布”不是某一条轨迹本身。理论家研究的是分布的收敛速度而实践中我们只有一条轨迹所以需要用遍历平均、时间平均去近似分布平均。对 ULA 要特别强调一点由于没有接受步骤ULA 的马尔可夫链并不以 π 为平稳分布。它的真实平稳分布是一个与步长 h 有关的分布 π_hh 越小越接近 π。因此“到站”需要区分两种含义一是链到达自己的平稳分布 π_h二是链到达真实目标 π。后者才是我们关心的混合时间而且它的极限不是 0而是 W₂(π_h, π)也就是算法的固有偏差。2.4 三个概念的连接方式把三者放在一起Wasserstein mixing time of ULA 的含义就是在 W₂ 度量下ULA 的迭代分布 μ_k 距离真实目标 π 到达精度 ε 所需的步数。这个问题可以拆成两步连续 Langevin 动力学本身的混合速度以及欧拉离散化对轨迹的扰动。下一节会展开讲为什么 W₂ 是分析这些问题最自然的选择。3. 为什么分析 ULA 要用 Wasserstein 距离有人会问分析采样算法的收敛性用 KL 散度或者总变差距离不行吗理论上当然可以但 Wasserstein 距离有四个不可替代的优势。第一Wasserstein 距离对空间位置敏感。总变差距离只关心“两个分布在同一集合上的概率差多大”完全不关心这些质量在空间上挪动了多远。如果目标分布有两个相距很远的模式链只停留在其中一个模式附近总变差距离可能已经很小但 W₂ 会明确告诉你还有一堆质量没有运到另一个模式距离还很大。对采样算法来说这种“空间上没探索到”的失败模式至关重要。第二Wasserstein 距离能捕捉矩的收敛。W₂ 收敛可以推出均值向量和协方差矩阵的收敛。在实践中用采样样本估计后验均值和后验方差是最高频的操作W₂ 收敛直接保证了这些矩估计的合理性。对比之下概率意义上的较弱收敛并不能给出同样的保证。第三Wasserstein 距离与梯度流理论深度关联。这一条对理论工作者更重要。著名的工作 Jordan-Kinderlehrer-OttoJKO证明Langevin 动力学对应的 Fokker-Planck 方程本质上是在 Wasserstein 空间里做 KL 散度的梯度下降。这个观点把概率分布的演化看成“在分布空间里的优化问题”W₂ 就是分布空间里的“距离度量”。采样算法于是可以被理解为对这条连续梯度流的离散化分析离散化误差也就是分析优化算法的收敛误差。这种几何视角让很多定理的证明变得自然。第四Wasserstein 距离与 log-Sobolev 不等式之间有优雅的桥梁。如果目标分布 π 满足对数 Sobolev 不等式LSI那么 Langevin 动力学的分布会在 KL 散度下指数收敛再通过 Talagrand 不等式 KL(μ|π) ≥ (λ/2) W₂²(μ, π) 推出 W₂ 也指数收敛。也就是说W₂ 的收敛速度可以直接从信息论不等式导出。这个链条也是当前扩散模型理论分析的核心工具。下表可以快速对比三种常用概率度量的特点度量直观含义对空间位移敏感捕捉高阶矩分析便利性总变差距离 TV两个分布在同一集合上的最大概率差否弱常用于离散 MCMC 分析KL 散度概率密度比的信息量弱通过 LSI 间接信息论工具丰富Wasserstein W₂质量搬运的最小平方成本是强与梯度流、最优输运理论联系紧密对 ULA 来说它的偏差恰恰体现为“分布整体被挪动了一点、被展宽了一点”这类误差用 W₂ 去度量最直观用 TV 反而可能看不出来。4. 理论主线混合时间的误差分解与步长指导现在进入本文最核心的理论部分。我不会完整复述定理证明但会给出一个对实践有指导意义的误差分解框架。理解这个框架之后你再看 ULA 相关的论文会发现它们大多在证明“某个具体假设下下面这个上界成立”。4.1 连续动力学的指数收缩先看连续时间的理想情况。如果势能函数 U 满足强凸条件即存在常数 m 0 使得对所有 x, y 都有⟨∇U(x) − ∇U(y), x − y⟩ ≥ m |x − y|²那么 Langevin 动力学的分布会以指数速度收敛到 π。用 W₂ 度量可以写成W₂(μ_t, π) ≤ e^{−mt} W₂(μ_0, π)这里的参数 m 就是强凸常数。直观上m 越大意味着目标分布在“谷底”附近越陡峭粒子越快被拉向高概率区域混合就越快。在更一般的假设下比如目标分布只满足 log-Sobolev 不等式而不是强凸收敛同样是指数型只是指数常数换成 LSI 常数。如果目标只是凸而不强凸收敛速度降为多项式阶如果目标非凸连续动力学本身也可能卡在局部模式里理论分析要困难得多。4.2 离散化偏差ULA 用的是欧拉离散化每一步都会引入离散化误差。这个误差有两个来源一是梯度在 [kh, (k1)h] 区间内被当作常数处理二是噪声的注入方式与连续布朗运动的积分不完全一致。离散化的净效应是链的平稳分布从 π 变成了 π_h而且 W₂(π_h, π) 的数量级是步长 h 的多项式。在强凸且光滑的假设下通常可以得到 O(h) 量级的偏差在更弱的光滑性假设下偏差可能退化为 O(√h) 量级。不同论文里的具体常数和指数会不一样但结构是一样的步长越小偏差越小。这里有一个容易踩的认知误区很多人以为 ULA 的误差主要来自随机性。实际上随机性部分会通过长时间平均互相抵消真正限制最终精度的往往是这个系统性偏差。即使你把链跑到无穷久W₂ 也不会变成 0而是停在 W₂(π_h, π) 附近。4.3 合起来两步误差上界把前面两项放在一起就可以写 ULA 的 Wasserstein mixing time 的主干结论。对第 k 步的边际分布 μ_k有类似形式的上界W₂(μ_k, π) ≤ (1 − mh)^k W₂(μ_0, π) B(h)其中第一项是初始距离的指数衰减衰减率由强凸常数 m 和步长 h 共同决定第二项 B(h) 是离散化带来的稳态偏差是步长的增函数。这个公式虽然是从论文里提炼出来的简化形式但它抓住了全部定量本质。第一项告诉你“从不好的初值出发需要跑多少步才能忘记初值”第二项告诉你“即使跑无限久误差的下限是什么”。两者之间天然存在一个权衡要减小 B(h)需要缩小步长但缩小步长会减慢 (1 − mh)^k 的衰减需要更多迭代步数才能达到同样的误差水平。4.4 对步长选择的指导根据上面的分解要达到 W₂ 精度 ε一般的策略是先根据偏差项确定步长上界要求 B(h) ≤ ε/2 左右于是步长需要取到与 ε 的多项式同阶。再根据收缩项确定迭代次数在步长确定后大约需要 k ≈ (1/(mh)) log(W₂(μ_0, π)/ε) 步把初始距离从 W₂(μ_0, π) 压缩到 ε/2。把两步合起来迭代次数的量级大概是 O(1/(mε^p))其中 p 由偏差项决定。这个次数还会通过初始距离 W₂(μ_0, π) 隐式地依赖维度 d一般来说维度越高需要的步数越多。这就是为什么论文里经常把结论写成“mixing time 是 O(d/ε²)”或者“O(d/ε² log(1/ε))”之类的形式。你不必背下具体数字但应该记住这个关系精度要求提高一个数量级迭代次数的需求会以多项式甚至更快的速度增长维度提高同样会推高步数需求。这个理论结构直接对应一个实践判断如果你发现 ULA 跑出来的样本和真实目标还有明显偏差先检查步长是否太大而不是盲目加迭代次数。步长不减跑再多步也只是在错误的平稳分布附近打转。5. Python 实现一维二维 ULA 与 W2 验证理论讲完开始动手。这一节会实现一个完整的 ULA 采样器并且用一维和二维两个目标分布验证 Wasserstein 距离下的收敛行为。全部代码都用 Python 和 NumPy 实现额外依赖 scipy 和 POT。5.1 环境准备与依赖建议使用 Python 3.9 及以上版本。需要安装以下包pip install numpy scipy matplotlib pot其中 POTPython Optimal Transport用于计算经验分布之间的 Wasserstein 距离scipy 用于生成理论分布的样本和分位数。如果只需要一维实验scipy 也够用二维经验 W₂ 计算推荐使用 POT因为自己实现最优传输求解器成本太高。5.2 实现 ULA 采样器首先写一个通用的 ULA 采样器。输入是目标分布的梯度函数、初始点、步长和迭代步数输出整条采样轨迹。# 文件名ula_core.py import numpy as np def ula_sampler(grad_log_pi, x0, step_size, n_steps, seed0): Unadjusted Langevin Algorithm (ULA) 采样器。 参数 ----- grad_log_pi : callable 对数目标密度梯度的函数 grad_log_pi(x) - ndarray x0 : ndarray 初始状态 step_size : float 步长 h对应公式 X h * grad sqrt(2h) * noise n_steps : int 迭代步数 seed : int 随机种子 返回 ----- samples : ndarray of shape (n_steps1, d) 包含初始点和每一步采样结果 rng np.random.default_rng(seed) x np.array(x0, dtypefloat) d np.size(x) samples np.zeros((n_steps 1, d)) samples[0] x.copy() for k in range(1, n_steps 1): noise rng.standard_normal(d) x x step_size * grad_log_pi(x) np.sqrt(2.0 * step_size) * noise samples[k] x.copy() return samples这个函数的实现要点有三个。第一噪声系数必须是 √(2h)这是 Langevin 动力学里扩散系数为 2 的离散化结果写错的话采样分布会完全错误。第二使用np.random.default_rng而不是全局的np.random方便复现。第三梯度方向是“往高概率爬升”的方向也就是 ∇log
返回列表