在非凸、非光滑势能面上做采样是贝叶斯推断、逆问题、生成模型和统计物理里绕不开的问题目标分布通常写成 (\pi(x)\propto e^{-V(x)})而 (V(x)) 往往是不可导的比如带稀疏约束的 (L_1) 正则项。很多人第一反应是直接用梯度型采样器但梯度不存在怎么办更麻烦的是当损失函数非凸时我们还能不能给出可控的复杂度估计最近在优化与采样交叉方向研究中一个比较重要的思路是把 Moreau–Yosida 正则化和未调整 Langevin 采样结合起来同时用 Active-Trace 技巧给出更精细的非渐近复杂度上界。这篇文章会从背景、公式推导到代码验证把这套思路拆开讲清楚。我会先介绍 Unadjusted Langevin SamplingULA的基本框架再解释为什么需要 Moreau–Yosida 正则化然后说明 Active-Trace 复杂度界限到底刻画了什么最后用一个小实验演示这套方法在数值上如何表现。适合有一定优化或采样基础、想读理论论文但又被抽象符号劝退的读者如果你想直接看结论也可以跳到第七节的总结和工程建议部分。1. 背景为什么采样问题如此依赖梯度1.1 从概率分布到积分计算在很多机器学习问题中我们真正想要计算的并不是某个函数的最小值而是一个概率分布下的期望[ \mathbb{E}_{x\sim \pi}[f(x)] \frac{\int f(x)e^{-V(x)}\mathrm{d}x}{\int e^{-V(x)}\mathrm{d}x}. ]如果变量维度很低可以用求积公式硬算但在高维空间里归一化常数 (\int e^{-V(x)}\mathrm{d}x) 几乎无法直接计算。马尔可夫链蒙特卡洛MCMC提供了一条绕开归一化常数的路径构造一个以 (\pi) 为不变分布的马尔可夫链迭代替换得到近似样本然后用这些样本做蒙特卡洛估计。1.2 Langevin 动力学的核心地位Langevin 动力学起初来自统计物理它描述的粒子在势能场中同时受到确定性漂移和随机噪声的作用。写成随机微分方程SDE就是[ \mathrm{d}x_t -\nabla V(x_t)\mathrm{d}t \sqrt{2}\mathrm{d}B_t. ]在很温和的条件下这个连续时间过程的稳态分布恰好是[ \pi(x) \propto e^{-V(x)}. ]这个结论极具吸引力因为它意味着我们只需要不停迭代就能得到近似服从 (\pi) 的样本。然而计算机无法真正连续模拟 SDE必须做时间离散化。最简单的离散格式是 Euler–Maruyama把它应用到 Langevin SDE 上就得到 ULA 或者更广义的 Langevin 采样器。2. 从连续动力学到离散采样ULA 的完整过程2.1 ULA 的递推公式假设步长为 (h)在第 (k) 步迭代中ULA 按下面的规则更新[ x_{k1} x_k - h\nabla V(x_k) \sqrt{2h}\xi_k, \qquad \xi_k\sim \mathcal{N}(0,I). ]这里有几个关键细节需要解释。第一(-\nabla V(x_k)) 是确定性漂移它让样本朝势能下降方向移动。第二(\sqrt{2h}\xi_k) 是随机扰动它保证链不会坍缩到某个局部极小值上而是能保持与 (\pi) 匹配的“扩散强度”。第三步长 (h) 同时控制漂移和扩散的幅度噪声的标准差是 (\sqrt{2h})不是常数。这一点容易被写错尤其是在把 ULA 和随机梯度下降做类比时。如果目标势函数是强凸且光滑的可以证明 ULA 的迭代分布 (\mu_k) 会以指数速度收敛到目标分布 (\pi)。也就是说Wasserstein-2 距离满足类似下面的递归上界[ W_2^2(\mu_{k1}, \pi) \le (1 - mh)W_2^2(\mu_k, \pi) C d h^2. ]其中 (m) 是凸性参数(d) 是维度。这一项 (C d h^2) 来自离散化误差正是“未调整”这个关键词的来源。2.2 “未调整”是什么意思标准的 Metropolis-adjusted Langevin AlgorithmMALA会在 ULA 提议后增加一步接受/拒绝判断从而确保链关于 (\pi) 可逆消除离散化的渐近偏差。但引入接受步骤会带来额外的计算开销并且在高维问题中接受率可能很低。无修正的 ULA 省去了接受判断所以迭代更快但也付出了代价离散化引入的偏差始终存在只能用更小的步长来抑制。因此 ULA 的非渐近分析关注的是“给定精度 (\varepsilon)最多需要多少步能使迭代分布靠近 (\pi)”。2.3 非光滑势函数的直接障碍现在问题来了。如果 (V(x)) 包含不可导项例如稀疏贝叶斯模型里常见的 (V(x)g(x)\lambda\lVert x\rVert_1)那么 (\nabla V(x_k)) 在某个坐标等于 0 的时刻根本不存在。此时 ULA 无法直接执行。一种朴素做法是把不可导项近似成“非常陡峭”的可导函数比如把 (|x|) 改写成 (\sqrt{x^2\epsilon^2})。这种方式实现简单但会引入额外的近似误差而且对 (\epsilon) 的选择非常敏感数值上容易出现震荡。更优雅的方案是用 Moreau–Yosida 正则化把不可导势函数处理成一个全局保持形状但整体光滑的函数。3. Moreau–Yosida 正则化把不可导变成可导3.1 Moreau–Yosida 包络的定义对于任意适当的凸函数或者更一般的函数 (V)Moreau–Yosida 正则化定义如下[ V_\lambda(x)\min_{z} \left{ V(z) \frac{1}{2\lambda}\lVert x-z\rVert^2 \right}. ]其中 (\lambda0) 是正则化参数。这个公式的含义是在 (x) 点取值时不光看 (V(x)) 本身而是允许 (z) 稍微偏离 (x)换取 (V(z)) 的下降偏离越多二次惩罚越大。与之相伴的是 Moreau 包络的近似算子也叫近端算子[ \operatorname{prox}_{\lambda V}(x)\arg\min_{z} \left{ V(z) \frac{1}{2\lambda}\lVert x-z\rVert^2 \right}. ]如果 (V) 是凸函数那么近似算子的解存在且唯一。很多问题中近端算子有显式解这极大方便了算法实现。3.2 核心性质包络的梯度等于残差Moreau–Yosida 正则化有一个非常漂亮的结论[ \nabla V_\lambda(x)\frac{1}{\lambda} \left( x - \operatorname{prox}_{\lambda V}(x) \right). ]也就是说我们不需要对不可导的 (V) 求梯度只需要计算一个近端算子的残差。只要近端算子可以高效求解我们就得到了一个全局定义的“光滑梯度”。更进一步即便 (V) 本身非凸在某些条件下 (V_\lambda) 也会比 (V) 更光滑。如果 (V) 是凸函数(V_\lambda) 是光滑凸函数且梯度是 (\frac{1}{\lambda})-Lipschitz 的。这为后续的非渐近收敛分析铺平了道路。3.3 (L_1) 范数的具体例子以 (V(x)\lambda_0|x|) 为例。它的近端算子就是著名的软阈值算子[ \operatorname{prox}_{\lambda\lambda_0|\cdot|}(x)\begin{cases} x - \lambda\lambda_0, x \lambda\lambda_0,\ 0, |x|\le \lambda\lambda_0,\ x \lambda\lambda_0, x -\lambda\lambda_0. \end{cases} ]把该式代入残差公式会发现 Moreau–Yosida 包络的梯度等价于对“超过阈值 (\lambda\lambda_0) 的部分”做线性收缩。这个函数本质上就是 Huber 损失函数的光滑形式在 (|x|) 很小时原来 (L_1) 的尖点被替换成一段二次曲线而远离原点的地方仍然近似保持线性增长。因此基于 Moreau–Yosida 的 ULA 更新可以改写成[ x_{k1}x_kh \cdot \frac{1}{\lambda} \left( x_k - \operatorname{prox}_{\lambda V}(x_k) \right) \sqrt{2h}\xi_k. ]注意这里的计算过程完全避开了“(V) 是否可导”的问题同时保留了目标分布的稀疏结构当某分量落入软阈值区间时梯度残差为 0但噪声依然注入所以链不会“死锁”在尖点上。4. Active-Trace 复杂度界限如何刻画采样收敛4.1 为什么传统遍历性分析不够用对于强凸且光滑的势函数已知的 ULA 复杂度很成熟要达到 Wasserstein-2 精度 (\varepsilon)迭代次数一般可以写成维度、条件数和精度参数的多项式。但在非凸、非光滑问题里问题变得困难得多。一方面目标分布可能有多个局部极值区域。链可能在某一个盆地附近停留很长时间全局遍历性很难快速实现。另一方面非光滑点构成的活动集通常测度为 0随机链几乎不会精确落在非光滑点上但它会反复逼近这些点。如果我们不仔细刻画这类“活性轨迹”的影响复杂度分析会过度乐观。4.2 Active-Trace沿着链追踪“临界活动”Active-Trace 的思想脱胎于非凸优化中对“活性集”的分析。传统优化理论经常假设迭代点落在某个光滑区域内部从而可以局部地用强凸/强光滑性质推导线性收敛。但实际算法在接近最优解时经常使某些约束或正则项处于临界状态这些“临界活动”才决定了最终收敛速度。放到采样问题上Active-Trace 复杂度考虑的不是“单一步长是否把分布推向目标”而是把整条随机链产生的状态序列看作一条轨迹统计那些靠近非光滑点或者活动集的迭代次数和距离。若这条轨迹长期停留在临界区域之外复杂度分析退化成普通光滑情形若轨迹频繁与临界区域交互那就要把 Moreau–Yosida 正则化带来的动态平衡计入代价。这种刻画方式的一个直接好处是它不再要求势函数全局凸或者全局光滑而是鼓励把采样过程分成“光滑地形”和“临界过渡”两种状态来分别定价最终得到精度 (\varepsilon) 和参数 (\lambda, h) 之间的精细折中关系。4.3 复杂度上界的主要结构在 Active-Trace 分析框架里最核心的不是“最终一步的分布误差”而是“沿轨迹累加的某种势能误差”。若定义某个辅助函数 (\Phi(x))那么分析会证明类似下面的形式[ \sum_{k0}^{T-1} \mathbb{E}\left[\Phi(x_{k1}) - \Phi(x_k)\right] \le A(\lambda) B(h,\lambda,d). ]左边是沿着链步的代际差之和经过望远镜求和可以变成首末两端的势能差右边第一项来自 Moreau–Yosida 包络与原势函数之间的偏差第二项来自离散随机梯度带来的噪声。在这个框架下要达到采样精度 (\varepsilon)给定最优的 (\lambda) 和步长 (h)可以从上述不等式中解出总迭代次数的上界。因为随机游走本身带有维数相关的扩散噪声所以最终复杂度通常表现为维度项 光滑参数项 非光滑活动项 精度项。需要强调的是这里并不会出现比纯光滑采样低一阶的情况因为 ULA 的离散噪声始终存在但通过 Active-Trace可以把非光滑部分对常数的影响刻画得更紧凑不会因为个别不可导点就悲观地把整个势函数视为“完全非光滑”。4.4 与已知经典结论的对照经典强凸光滑情形中ULA 的 Wasserstein 距离上界通常为[ W_2^2(\mu_T,\pi) \lesssim e^{-mhT} W_2^2(\mu_0,\pi) O(h). ]这里第二个项和维度相关。如果我们把精度要求设为 (W_2(\mu_T,\pi)\le \varepsilon)一个常规结论是迭代次数量级大致为[ O\left(\frac{d}{\varepsilon^2}\right) ] 乘上一些条件数因子。当引入 Moreau–Yosida 正则化且势函数非光滑时参数 (\lambda) 会额外控制平滑性与近似误差之间的 trade-off。(\lambda) 如果取得太大包络 (V_\lambda) 和原 (V) 可能相差太远导致采样到的分布并非真正来自 (\pi)而是来自某个被过度平滑的分布(\lambda) 取得太小近似算子区域非常窄数值上接近不光滑算法又会退化。Active-Trace 复杂度可以帮助我们理解这类平衡它给出了使总误差不超过 (\varepsilon) 时(\lambda) 和 (h) 应该随 (\varepsilon) 缩放的合理速率从而避免“步长已经很小但非光滑偏差依然很大”的尴尬局面。5. 数值实验从简单代码理解这套机制的直观含义理论分析很抽象最好还是用一个小实验来看这些步骤到底发生了什么。我们用 Python 写一个 1 维问题目标势函数是非凸且部分不可导的。5.1 实验环境本文代码在以下环境中验证过。Python 3.10 NumPy 1.24 Matplotlib 3.7可选用于画图这里的版本不是关键只要 NumPy 可用代码逻辑都能运行。下面不依赖 PyTorch也不需要 GPU。5.2 基础 ULA 采样器先来看一个双峰光滑势[ V_1(x)\frac{(x^2 - 4)^2}{8}. ]它有两个势阱分别位于 (x-2) 和 (x2)。对应梯度为[ \nabla V_1(x)\frac{x^3}{2} - 2x. ]基础 ULA 采样代码如下import numpy as np def grad_V1(x): return 0.5 * x**3 - 2.0 * x def ula_sample(x0, h, steps, grad_fn, burn_in500): rng np.random.default_rng(42) x x0 samples [] for k in range(steps burn_in): x x - h * grad_fn(x) np.sqrt(2.0 * h) * rng.standard_normal() if k burn_in: samples.append(x) return np.array(samples, dtypefloat) x0 0.1 h 0.05 samples ula_sample(x0, h, 20000, grad_V1) print(样本均值, samples.mean()) print(样本标准差, samples.std())运行后可以得到一个围绕双峰分布的样本集。因为目标分布对称均值应接近 0标准差接近目标分布的离散程度。如果步长 (h) 过大比如取h 0.5会发生什么样本方差会迅速爆炸因为离散 Euler 格式在强梯度区域不再稳定。我们把这个作为对照组。5.3 Moreau–Yosida 处理不可导势函数现在构造一个更贴近主题的目标[ V_2(x)\frac{(x^2 - 4)^2}{8} \gamma |x|. ]后半部分在 (x0) 处不可导。真实目标分布既保留双峰结构又在原点附近因为 (|x|) 的惩罚而压低概率。按前面所说如果把整个势函数代入 Moreau–Yosida 包络梯度的计算可以改写成一个近端步骤。但在数值上为了演示我们手动实现其梯度替代形式不可导项 (|x|) 被替换成软阈值残差。def prox_l1(x, lam, gamma): # prox_{lam * gamma * |x|}(x) thresh lam * gamma if x thresh: return x - thresh elif x -thresh: return x thresh else: return 0.0 def grad_smooth_part(x): return 0.5 * x**3 - 2.0 * x def grad_my_total(x, lam, gamma): # smooth part gradient Moreau-Yosida gradient of gamma*|x| prox_val prox_l1(x, lam, gamma) return grad_smooth_part(x) (x - prox_val) / lam def my_ula_sample(x0, h, lam, gamma, steps, burn_in500): rng np.random.default_rng(7) x x0 samples [] for k in range(steps burn_in): g grad_my_total(x, lam, gamma) x x - h * g np.sqrt(2.0 * h) * rng.standard_normal() if k burn_in: samples.append(x) return np.array(samples, dtypefloat) samples_my my_ula_sample(0.5, h0.02, lam0.1, gamma1.0, steps30000) print(Moreau-Yosida ULA 样本均值, samples_my.mean()) print(小于 0 的比例, (samples_my 0).mean())这里最重要的不是代码本身而是理解grad_my_total做了什么原本 (L_1) 尖点处梯度不存在但通过 Moreau 包络梯度变成了近端残差。只要 (|\text{prox}_{\lambda V}(x)-x|) 不为 0就会产生一个指向原点的力如果 (x) 已经在可接受区间内残差为 0此时仅靠噪声维持随机探索。小技巧可以把lam分别设为0.01, 0.1, 1.0观察样本直方图变化。lam很小时近似更精确但步长必须设置更小否则数值不稳定lam很大时平滑后的目标可能失去尖点惩罚的锐利感样本中心区域概率偏高。5.4 观察收敛用滑动平均判断链是否稳定直接看采样结果可能不够直观。我们可以观察链的滑动均值来判断是否已经遍历两个峰def running_average(samples, window500): out np.zeros_like(samples, dtypefloat) for i in range(samples.shape[0]): left max(0, i - window 1) out[i] samples[left:i1].mean() return out samples_raw ula_sample(0.1, 0.02, 10000, grad_V1) avg running_average(samples_raw) print(前 200 步均值, avg[:200].mean()) print(后 200 步均值, avg[-200:].mean())如果前 200 步均值和末尾 200 步均值差别很大说明链可能还没有很好地穿越两个峰。此时需要增大步长或者在初始化时多次随机起点。Active-Trace 复杂度说明中“轨迹长期停留在某一活动区”的问题在实验中就表现为这种滑动均值长期偏移。6. 常见误区与排查思路按照理论推导和上面实验我整理几个容易被忽视的坑点。问题现象常见原因解决思路采样结果发散数值变成 NaN步长 (h) 过大离散化不稳定缩小步长检查梯度范数是否过大样本集中在某个单峰另一个峰很少访问链没有充分混合或步长太小适当增大噪声/步长或用并行链、插值策略用 Moreau–Yosida 后样本分布和理论不一致(\lambda) 太大过度平滑原势函数减小 (\lambda)同时匹配更小步长势函数不可导项是 (L_1)但直接调用自动微分失败自动微分框架无法对不可导点给出统一梯度用近端算子残差替换梯度把噪声方差写成 (\sqrt{h}) 而不是 (\sqrt{2h})随机微分方程归一化约定不一致回到 SDE 检查扩散系数样本后处理没有做 burn-in 丢弃初始位置离高概率区域太远丢弃前 10%~20% 的样本或从众数附近初始化不同随机种子结果差异大链未收敛不能归因于随机性增加迭代数比较多次独立链更具体地说当你在自定义采样器里复现某个理论结论时最容易出错的通常是步长 (h) 与势函数尺度不匹配。如果势函数梯度量级是 (L)那么步长应明显小于 (2/L) 这类稳定性阈值在强非凸区域这个限制可能更严格。Moreau–Yosida 参数 (\lambda) 的选择也很微妙。它在目标函数上引入了一个二次惩罚型平滑估计因此最终采样器实际上是在对目标分布的光滑代理分布采样。为了控制偏差(\lambda) 应当与目标精度一起缩小而不是固定为一个很大的值。Active-Trace 的复杂度的价值正在于它把 (\lambda)、步长 (h)、精度 (\varepsilon) 和维度 (d) 放在同一个不等式框架里告诉我们这些量应该如何联合缩放而不是分别调节导致浪费算力。7. 最佳实践与工程建议7.1 什么时候应该使用 ULA 而不是 MALAMALA 有接受/拒绝步骤能保证渐近无偏但每步都要计算一次额外的接受率且在高维空间中接受率可能迅速下降。对于离线大规模贝叶斯推断如果目标是快速得到近似样本ULA 更合适如果应用要求严格的无偏性且维度不太高可以考虑 MALA、HMC 或 SMC 这类方法。Moreau–Yosida 类型的采样方法适合势函数包含近端可分解项的场合。常见信号处理模型写出来常有如下形式[ V(x) \frac{1}{2\sigma^2}\lVert Ax-y\rVert^2 \lambda_0 \lVert x\rVert_1\mathcal{I}_{C}(x). ]其中 (\mathcal{I}_C) 是指示函数用于表达约束。这类问题可以写成三项之和。只要每项的近端算子都能高效计算即便整体非凸也可以用类似策略实现端到端采样。7.2 参数调节建议在正式跑大规模数据前建议先在减小版数据或低维数据上做一次快速实验确定三件事目标分布的典型尺度样本是否存在多模态梯度的最大范数量级用于反推合理步长不可导项的影响区间例如 (L_1) 的软阈值宽度。之后在一个小时间预算内做步长扫描比较相邻步长下的经验分布差异。如果差异很小说明当前步长已经进入“小偏差区”如果差异显著说明还需要缩小步长。与此同时不要忽视随机种子数。最好运行 3 到 5 条独立链利用 Gelman-Rubin 诊断这类工具判断链是否收敛。7.3 关于复杂度结论的实际价值很多读者会觉得“复杂度上界反正只是理论上界工程里没什么用”。但不完全是这样。Active-Trace 复杂度这样的结论能提供几个非常实际的指导它表明在某些非光滑条件下ULA 不会因为局部不可导点而灾难性地退化它说明复杂度依赖维度主要以噪声项的方差形式出现而不是以某些难以估计的全局常数形式出现它强调了 Moreau–Yosida 参数与步长需要按同步比例缩小这能防止“调了半天步长最后发现偏差点来自平滑偏差”这类盲区。所以工程上不一定要逐字推导证明但至少要学到“参数不能独立地调到很小”“误差主要来自两个方向离散化偏差和正则化偏差”这两个直觉。7.4 数值稳定性建议实现过程中如果遇到梯度范数过大可以选择对梯度做裁剪或者使用自适应步长。但自适应步长在采样器里会引入额外的非平稳性需要谨慎验证。一个更稳妥的方法是使用重启动x x0 for k in range(total_steps): g grad(x) if not np.isfinite(g).all(): x x0 # 重新初始化 continue x x - h * g np.sqrt(2.0 * h) * rng.standard_normal()不过这种重启动会破坏链的平稳性所以统计样本前要再次丢弃 burn-in。对于近端算子的计算建议单独封装并做单元测试。比如用 (L_1) 近端验证软阈值公式def test_prox_l1(): assert abs(prox_l1(3.0, 0.5, 1.0) - 2.5) 1e-12 assert abs(prox_l1(-3.0, 0.5, 1.0) 2.5) 1e-12 assert abs(prox_l1(0.1, 0.5, 1.0)) 1e-12这个测试在重构代码时非常有用可以防止把梯度和近端算子符号写反。8. 总结与延伸这篇文章从目标分布采样问题出发梳理了 Unadjusted Langevin Sampling 的基础解释了 ULA 为什么会因为势函数不可导而失效然后引入 Moreau–Yosida 正则化说明如何用近端算子把不可导问题转换成可计算的光滑梯度。最后介绍了 Active-Trace 复杂度界限的思想它不是只看最后一步的分布误差而是沿着马尔可夫链的轨迹追踪临界活动从而得到更贴合非凸非光滑场景的复杂度控制方式。如果你下一步想深入研究建议重点阅读以下方向近端算子在不同正则函数下的闭式解包括 (L_1)、总变差、指示函数Wasserstein 距离和 Fisher–Rao 度量在采样收敛分析中的作用将 Moreau–Yosida 正则化和随机梯度估计相结合的随机化版本对应大规模数据场景与 score-based diffusion 模型的关系因为扩散模型的采样器本质上也是 Langevin 型动力学。代码实验永远是检验理论直觉的最好方式。你可以把目标势函数换成二维高斯混合并加入 (L_1) 约束观察不同 (\lambda) 和步长对样本质量的影响再用链与链之间的分布距离作为收敛判断。这样既能加深对本篇结论的理解也能为将来阅读原始论文的定理证明提供直观支撑。