这次我们来看一个偏理论、但和现在像 SDE 扩散模型、采样算法、MCMC 收敛性分析都强相关的主题Wasserstein mixing time of the unadjusted Langevin algorithm也就是“未调整 Langevin 算法在 Wasserstein 距离下的混合时间”。如果你平时主要做应用落地可能更熟悉“用 Langevin 采样器生成样本”或者“用扩散模型逐步去噪”。但真正影响采样质量、步长选择、计算效率的是这一类算法在分布层面的收敛速度。而 Wasserstein 距离下的 mixing time 分析恰好回答了一个非常关键的问题离散化后的 Langevin 算法要走到与目标分布足够接近最短需要多少轮迭代。这篇文章我会先讲清楚这个论题到底在分析什么然后拆解其中涉及的核心概念接着给出一个可以在本地用 Python 复现的数值实验框架。也就是说我们不只停留在公式层面还会用手上的普通 CPU 电脑跑起来验证不同步长、不同维数、不同目标分布下算法混合时间的变化趋势。如果你是研究采样理论、MCMC 加速、扩散模型扩散过程设计或者想深入理解 SGLD、RMSprop 类随机优化方法背后概率视角的读者这篇内容值得收藏。1. 核心能力速览能力项说明研究对象Unadjusted Langevin AlgorithmULA在 Wasserstein 距离下的有限时间混合行为核心概念Wasserstein 距离、mixing time、离散化误差、收敛速率、强凸/非强凸目标分布理论工具梯度流、耦合方法、同步耦合、收缩性分析、非渐近界主要输出理解 ULA 的迭代步数与 Wasserstein 距离之间关系的可复现实验框架计算门槛普通 CPU 即可完成 2D/5D/10D 的小规模验证实验显存需求不需要 GPU显存占用为 0支持平台Windows / Linux / macOS需要 Python 3.8 以上启动方式Jupyter Notebook 或命令行脚本是否支持 API不涉及这是理论研究与数值验证项目是否支持批量任务支持可以通过配置批量扫描步长、维数、目标分布参数适合读者研究采样理论、MCMC、扩散模型、随机优化的算法工程师与研究生需要提前说明这不是一个开箱即用的工程工具而是一套“理论分析 数值验证”思路。它不能直接帮你生成图片或者做语言模型推理但它能帮你判断当你给 Langevin 算法设置步长 0.01 时跑一万步到底够不够。从材料看这个主题的重点是“mixing time”也就是达到混合的时间。研究的是“未调整”版本这意味着我们不做 Metropolis 接受拒绝修正而是直接用离散化后的随机微分方程迭代。这样做的好处是分析方便、每一步的计算量低但代价是存在离散化偏差。Wasserstein 距离则是比较两个概率分布之间“搬土成本”的一种度量方式它比 KL 散度或总变差距离更能体现样本在真实几何空间中的移动成本。2. 适用场景与使用边界这个论题适合下面几类场景理论分析你需要估计采样算法的迭代复杂度例如从初始分布出发经过多少步之后样本逼近目标分布。算法对比比较 ULA、未调整的 Milstein 方法、随机梯度版本 SGLD 在相同 Wasserstein 精度下的迭代次数。步长调度研究研究固定步长、衰减步长、自适应步长对混合时间的影响。扩散模型理解扩散模型的采样过程本质上接近一个离散化的反向 SDE与 ULA 的收敛分析存在深层联系。不太适合的场景也很明确如果你的目标是快速生成大量高维样本Unadjusted Langevin 算法通常不如现代的 HMC 或 SGLD 系列实用。另外这个主题本身不涉及神经网络训练或大规模数据处理不会直接给你一个高精度的生成模型。使用边界上要注意Wasserstein 混合时间的理论结果往往依赖强凸条件、对数 Sobolev 不等式或 Poincare 不等式等结构假设。如果目标分布是多模态的、高度非凸的固定步长 ULA 可能在有限时间内根本无法混合。此时不要期望数值实验会得到整齐的线性下降曲线实际观测到的很可能是 Wasserstein 距离长期停滞。还要考虑学术规范和使用合规。如果你要在一篇论文里引用这个主题下的结果必须准确交代实验的目标分布、初始分布、样本量、Wasserstein 估计方式和随机种子。不要只贴一张漂亮的收敛曲线而忽略对步长、维度、迭代次数的完整描述。若实验涉及真实业务数据或敏感信息建议先用合成数据完成方法验证再在受控环境中使用真实数据。3. 前置知识与环境准备3.1 理解 ULA 的迭代公式Unadjusted Langevin Algorithm 针对一个目标分布而通常最简单的一阶近似是 Euler-Maruyama 离散化这个迭代可以拆成两个部分第一项是梯度下降项推动样本往概率密度高的方向移动第二项是高斯噪声项保证采样过程不会坍缩到单一最大值点。从 Wasserstein 混合时间的视角我们关心的是从某个初始分布 出发迭代 k 步之后样本经验分布与真实分布 之间的 2-Wasserstein 距离有多小。这个“距离对迭代次数”的关系就是混合时间曲线。一般来说强凸分布下会观测到明显的指数型收缩而在非强凸分布下收敛会变得缓慢甚至停滞。3.2 Wasserstein 距离的估计方式2-Wasserstein 距离的定义是其中 是联合分布约束条件是两个边缘分布分别为 和 。实际数值实验中如果已知两个分布都是高斯分布那么存在解析式Wasserstein 距离可以直接计算。但更通用的做法是从两个分布中分别采样一批样本然后计算两个经验分布之间的 Wasserstein 距离。Python 中 SciPy 提供了scipy.stats.wasserstein_distance但这个方法只支持一维。高维问题中可以使用 POT 库Python Optimal Transport。这里我们主要关心二维和十维以内的实验所以可以用 POT 的高维接口。3.3 环境准备清单实际运行时只需要一个 Python 环境和几个科学计算库。推荐使用 Python 3.9 或更高版本。python -m venv langvin_mixing_env source langvin_mixing_env/bin/activate # Windows 下执行 langvin_mixing_env\Scripts\activate pip install numpy scipy matplotlib pot如果只是想快速体验也可以直接用 Jupyter Notebook但为了观察批量扫描结果我更推荐以脚本方式组织代码。环境版本建议如下依赖库版本建议用途Python3.9 或以上运行环境NumPy1.24 或以上矩阵运算与随机数生成SciPy1.10 或以上线性代数、科学计算POT0.9.0 或以上高维 Wasserstein 距离计算Matplotlib3.7 或以上绘制收敛曲线建议先把固定随机种子否则不同批次实验的可比性会变差。这也是后面做批量任务时最容易忽略的细节。4. 从理论到代码搭建 ULA 混合时间验证框架4.1 代码结构按功能拆成三个模块targets.py定义目标分布包括强凸高斯分布、高斯混合分布。ula.py实现 ULA 采样器记录每步的样本快照。mixing_experiment.py负责计算 Wasserstein 距离、组织批量实验。先把采样器写好。import numpy as np class ULA: def __init__(self, grad_log_target, step_size, seedNone): self.grad_log_target grad_log_target self.step_size step_size self.rng np.random.default_rng(seed) def sample(self, x0, num_steps): 从初始点 x0 出发迭代 num_steps 步返回每一步的样本快照。 x np.array(x0, dtypefloat) dim x.shape[0] trajectory np.zeros((num_steps 1, dim)) trajectory[0] x for k in range(1, num_steps 1): # 梯度下降项 grad self.grad_log_target(x) # 噪声项 noise self.rng.normal(0.0, np.sqrt(2.0 * self.step_size), sizedim) x x self.step_size * grad noise trajectory[k] x return trajectory这个采样器会把每一步的样本全部保存下来。后面我们计算混合时间时需要观察每步经验分布与目标分布之间的差距所以不能只记录最终样本。4.2 目标分布设置先测试最简单的强凸高斯分布这样我们可以用高斯分布的解析 Wasserstein 距离来验证估计方法的正确性。import numpy as np def target_gaussian(dim): mean np.zeros(dim) cov np.eye(dim) * 2.0 precision np.linalg.inv(cov) def grad_log_target(x): return -precision (x - mean) return grad_log_target, mean, cov接着设置高斯混合分布作为一个非强凸的测试用例def target_gaussian_mixture(dim, mu1, mu2, sigma): inv_sigma 1.0 / sigma def grad_log_target(x): p1 np.exp(-0.5 * np.sum((x - mu1) ** 2) * inv_sigma) p2 np.exp(-0.5 * np.sum((x - mu2) ** 2) * inv_sigma) total p1 p2 grad1 -(x - mu1) * inv_sigma * p1 grad2 -(x - mu2) * inv_sigma * p2 return (grad1 grad2) / total return grad_log_target4.3 Wasserstein 距离计算二维/低维情况下直接用 POT 的经验分布 Wasserstein 距离。import ot def wasserstein2_from_samples(samples1, samples2): 计算两组样本之间的 2-Wasserstein 距离。 a np.ones(len(samples1)) / len(samples1) b np.ones(len(samples2)) / len(samples2) M ot.dist(samples1, samples2, metricsqeuclidean) W2_sq ot.emd2(a, b, M) return np.sqrt(W2_sq)注意这里估算的是经验分布之间的 Wasserstein 距离当样本数量少时估计误差会偏大。因此混合时间曲线需要有一定的样本量支撑一般每个时间点取 2000 到 5000 个样本比较稳妥。如果目标分布是高斯的还可以用解析公式来对照def gaussian_w2(mean1, cov1, mean2, cov2): 两个高斯分布之间的 2-Wasserstein 距离解析式。 diff_mean np.linalg.norm(mean1 - mean2) cov_product np.linalg.eigvalsh(cov1).dot(np.linalg.eigvalsh(cov2)) if cov1.shape[0] 1 else None cov_sqrt np.linalg.sqrtm(np.linalg.multi_dot([cov1, cov2])) trace_term np.trace(cov1 cov2 - 2 * np.linalg.sqrtm(cov_sqrt)) return np.sqrt(diff_mean ** 2 trace_term)更稳妥的做法是先用高斯目标分布跑通整个代码再用 POT 估计经验 Wasserstein 距离并与解析距离对照这样能验证整套实验流程没有写错。5. 功能测试与效果验证5.1 测试 1固定步长下强凸高斯分布的混合时间实验设计目标分布二维高斯分布均值为零协方差矩阵为 2 倍单位阵。初始分布远离目标分布的高斯分布均值偏移 10 个单位。步长从 0.1、0.05、0.02、0.01 中选择一组。迭代步数2000 步。样本数每步同时跟踪 500 个独立采样路径用于估计经验分布。核心代码如下import numpy as np import matplotlib.pyplot as plt from scipy.stats import norm from ula import ULA from targets import target_gaussian def run_experiment(dim2, step_size0.05, num_steps1000, num_chains500, seed42): grad_log_target, true_mean, true_cov target_gaussian(dim) ula ULA(grad_log_target, step_size, seedseed) # 初始分布均值偏移 x0 np.full(dim, 10.0) all_trajectories [] for i in range(num_chains): initial_noise np.random.normal(0, 0.5, sizedim) traj ula.sample(x0 initial_noise, num_steps) all_trajectories.append(traj) all_trajectories np.array(all_trajectories) # shape: (num_chains, num_steps1, dim) # 计算每个迭代步的经验分布与目标分布之间的 Wasserstein 距离 w2_distances [] true_samples np.random.multivariate_normal(true_mean, true_cov, size2000) for k in range(0, num_steps 1, 5): samples_at_k all_trajectories[:, k, :] # 这里使用 POT 估算 # 为了演示这里用解析解当目标分布是高斯的经过线性迭代后样本分布也近似高斯 # 但更通用做法是直接调用 POT mean_empirical samples_at_k.mean(axis0) cov_empirical np.cov(samples_at_k.T) w2_est gaussian_w2(true_mean, true_cov, mean_empirical, cov_empirical) w2_distances.append((k, w2_est)) return w2_distances运行该实验后预期可以看到一条随着迭代步数增加而平坦下降的曲线。步长越小最终能达到的 Wasserstein 距离越低但需要更多迭代步数才能降低到相同水平。5.2 测试 2高维扩展把dim从 2 改为 10观察曲线变化。一般来说维度增加会带来两个问题第一初始距离增加需要更多步数才能把均值项收缩第二协方差结构的收敛速度变慢。这里不用 GPUCPU 也能跑完。主要计算开销来自 Wasserstein 距离估计如果使用 POT样本量较大时耗时明显。实际测试建议先跑一个链看整体轨迹是否稳定。再跑多链计算经验分布距离。最后再考虑大样本量和批量实验。5.3 测试 3高斯混合分布的非强凸现象将目标分布换成两个相距较远的峰观察 ULA 的混合行为。def run_mixture_experiment(): dim 2 mu1 np.array([-5.0, -5.0]) mu2 np.array([5.0, 5.0]) grad_log_target target_gaussian_mixture(dim, mu1, mu2, sigma1.0) ula ULA(grad_log_target, step_size0.02, seed42) x0 np.array([0.0, 0.0]) traj ula.sample(x0, 5000) # 打印最后一步样本位置 print(Final sample:, traj[-1])如果步长太小样本可能长时间停留在某一个峰附近无法在另一个峰之间切换。这个现象在绘制的样本轨迹图上会很明显。注意这里不要直接输出“混合时间已失败”这样的结论而是观察Wasserstein 距离曲线是否出现平台期。出现平台期就意味着步长需要进一步缩小或者需要引入 tempering / 模拟退火策略。5.4 判断实验结果是否成功的标准强凸高斯情形Wasserstein 距离随迭代步数单调下降并且最终稳定在一个小值附近。混料情形如果目标分布是多模态需要看到两个峰之间都有足够的样本被探索到否则视为混合不完全。所有实验都建议保存w2_distances到 CSV方便后续画多条曲线对比。6. 批量实验设计与自动化运行6.1 批量参数扫描mixing time 分析最怕的就是只跑一组参数就下结论。步长、维数、初始分布偏移量都会改变结论。建议把实验组织成套娃配置import itertools import pandas as pd configs [ {dim: 2, step_size: 0.05, num_steps: 1000, num_chains: 200}, {dim: 2, step_size: 0.02, num_steps: 2000, num_chains: 200}, {dim: 2, step_size: 0.01, num_steps: 5000, num_chains: 200}, {dim: 10, step_size: 0.02, num_steps: 5000, num_chains: 100}, {dim: 10, step_size: 0.01, num_steps: 10000, num_chains: 100}, ] results [] for cfg in configs: w2 run_experiment(**cfg) for k, value in w2: results.append({dim: cfg[dim], step_size: cfg[step_size], iter: k, w2: value}) df pd.DataFrame(results) df.to_csv(mixing_time_results.csv, indexFalse)6.2 批量任务队列设计如果实验规模继续扩大可以做成一个小型任务队列而不是在单次脚本中循环执行。这里不引入 redis 或 celery直接用 Python 的concurrent.futures做进程级并行即可。每个进程执行一组参数组合结果写入独立文件。from concurrent.futures import ProcessPoolExecutor def process_one_config(cfg): w2 run_experiment(**cfg) cfg_id fdim_{cfg[dim]}_step_{cfg[step_size]} with open(fresult_{cfg_id}.csv, w) as f: f.write(iter,w2\n) for k, value in w2: f.write(f{k},{value}\n) return cfg_id if __name__ __main__: with ProcessPoolExecutor(max_workers2) as executor: futures {executor.submit(process_one_config, cfg): cfg for cfg in configs} for f in futures: print(Done:, f.result())这样有两个好处第一单个实验崩溃不会影响其他实验第二每个实验结果独立落地方便增量分析。6.3 记录元信息批量实验最关键的一步是记录每个实验的元信息。建议在结果目录下生成一个config.json{ target: gaussian_mixture, dim: 2, step_size: 0.02, num_steps: 5000, num_chains: 200, initial_mean_offset: 10.0, seed: 42, random_state: 20240501 }后面做对比时会发现没有元信息的实验结果基本就是废数据。7. 资源占用与性能观察这个主题的数值实验不需要 GPU也不需要关心显存占用。主要开销集中在两个方面采样器生成样本的计算量。Wasserstein 距离估计的计算量。在 2 维情形下500 个链、2000 步迭代每个链需要计算 2000 次梯度。耗时很小最多几秒钟。但 POT 的ot.emd2需要求解最优传输问题在样本量为 5000 时每次调用可能耗时 0.1 到 0.5 秒。如果每 5 步计算一次总共 400 次 Wasserstein 估计需要几十秒到几分钟不等。10 维情形下耗时会明显增大。建议每 20 步或 50 步再计算一次 Wasserstein 距离不用每一步都计算。要观察内存占用可以使用 Python 的tracemalloc或直接看系统任务管理器。采样轨迹数据如果全部保存在内存中带来的内存压力较大。例如 500 个链、10000 步、10 维数据float64 类型占用约 4 亿字节约 400MB。如果继续加大规模需要把轨迹数据分块保存或只保存每隔若干步的样本。性能优化建议使用向量化批量采样同时更新多个链而不是循环更新。减少 Wasserstein 距离评估频率。对高维问题避免使用经验协方差的解析 Wasserstein 公式直接用 POT但只对降维后样本或使用切片 Wasserstein 距离做近似。下面是一个向量化 ULA 批量采样的简化写法def batched_ula_sampling(grad_log_target, x_init, step_size, num_steps): x_init: (num_chains, dim) x x_init.copy() num_chains, dim x.shape trajectory np.zeros((num_steps 1, num_chains, dim)) trajectory[0] x noise_std np.sqrt(2.0 * step_size) grad_scale step_size for k in range(1, num_steps 1): grad grad_log_target(x) noise np.random.normal(0.0, noise_std, size(num_chains, dim)) x x grad_scale * grad noise trajectory[k] x return trajectory由于grad_log_target对批量输入的操作方式必须支持二维数组广播所以需要把之前的高斯梯度函数改为向量化版本def grad_log_target_vectorized(x): # x: (num_chains, dim) return -x / sigma_sq这里sigma_sq是方差参数。向量化后采样速度可以提升几十倍。8. 常见问题与排查方法问题现象可能原因排查方式解决方案样本剧烈发散数值变为 NaN 或 Inf步长过大离散化不稳定打印前几步的梯度范数和样本范数调小步长一般从 0.01 开始测试Wasserstein 距离曲线无法下降初始分布离目标太远或目标分布非凸检查初始样本均值与目标均值的距离增加迭代步数或缩小初始偏移量多模态目标中样本只停留在一个峰内步长太小无法跨越势垒绘制样本轨迹图或二维直方图调大步长或改为 simulated temperingPOT 计算 Wasserstein 距离过慢样本量过大、维度较高测试不同样本量下的单次耗时只对部分时间点计算或使用切片 Wasserstein批量任务中途崩溃单个配置导致内存过大或数值发散查看对应配置文件为每个任务增加 try/except 和错误日志结果无法复现未固定随机种子检查全局随机种子和采样器内种子每次实验固定 seed并记录配置强烈建议在实验目录下加一个debug_log.txt把每次运行的配置、数值范围、错误信息写进去否则多组实验后很难定位问题。9. 最佳实践与使用建议先把最简单的二维强凸高斯实验跑通。只有当你能看到一条不断下降的 Wasserstein 距离曲线再开始研究更复杂的问题。这是整个项目最稳妥的起点。实验设计上固定步长扫描比单次长时间运行更有价值。步长过大的发散现象、步长过小的低速混合现象在曲线上非常直观。建议把步长从小到大排列至少跑 5 组画出多条 Wasserstein 收敛曲线叠加图。训练和采样过程中注意随机种子分离采样器内部使用独立np.random.default_rng(seed)不要直接依赖全局随机状态。尤其在批量任务中不同进程使用同一个全局随机种子会产生完全相关的模拟结果。文件组织上推荐这样的目录结构mixing_time_project/ ├── ula.py ├── targets.py ├── mixing_experiment.py ├── configs/ │ └── exp_gaussian_2d.json ├── results/ │ ├── raw/ │ └── figures/ └── logs/输入素材、配置文件、实验结果原始数据、最终图标分开存放。分析时永远保留原始轨迹数据而不是只保存一个均值 Wasserstein 距离序列因为后续可能还要复核误差棒或不同单条链的方差。如果这个分析要用于课程报告或论文必须在方法部分说明目标分布是什么、初始分布是什么、采样器是否离散化、是否做投影或裁剪、Wasserstein 距离估计方法是什么、样本量是多少、是否固定随机种子。这比任何理论推导都更容易被读者挑刺。涉及真实数据或隐私场景时建议使用差分隐私框架或把数据先做脱敏再放到采样器里。属于学术引用时也需要准确引用相关定理来源不夸大自创结论。10. 总结与下一步这个主题最值得花时间的地方不是背诵某个混合时间上界而是亲自观察“步长、维数、目标几何形状”三个因素如何同时影响收敛曲线。一旦跑通了二维强凸案例再切换到高斯混合分布你会直观理解为什么理论分析都要求强凸或对数 Sobolev 条件。最先建议验证的功能是固定步长下二维高斯的混合时间曲线这是整套实验的基准。最容易踩的坑是步长设置太大导致样本轨迹发散你会误以为 ULA 完全不能用实际上只是需要把步长降到 0.01 甚至更小。如果后续要深入可以从三个方向扩展第一把 ULA 换成 SGLD引入随机梯度噪声观察额外噪声对 mixing time 的影响第二加入外推项或重启动机制研究加速效果第三把采样换成投影梯度、镜像梯度、近端梯度形式分析限制在约束集中时的 Wasserstein 混合行为。这三个方向每一步都能和现在的代码框架衔接上。