Tamed Subgradient ULA：非光滑非凸目标分布的稳健采样方法

这次我们来看一个专门面向“非光滑 非凸”目标分布的采样算法The Tamed Subgradient Unadjusted Langevin Algorithm可以直接叫 Tamed Subgradient ULA。先讲清楚它在解决什么问题。贝叶斯推断、稀疏学习、生成模型里很多任务最终都要落到“从一个概率分布 π(x) ∝ exp(-f(x)) 中采样”。最经典的采样器叫 Unadjusted Langevin AlgorithmULA迭代上只做两件事朝势能下降方向走一步再加一个高斯噪声。但 ULA 有一个长期容易被忽略的边界条件f 必须足够光滑最好还满足凸性假设否则离散化之后的轨迹可能发散或者收敛结果偏离目标分布。Tamed Subgradient ULA 就是为这个边界条件设计的改进算法。它用次梯度替换梯度处理 f 不可导的问题用 taming 函数压缩漂移项处理梯度超线性增长导致的数值不稳定再把收敛分析放到非凸的耗散型假设上把“beyond convexity”拆成三个可以独立验证的工程点。本文会完整梳理它的设计动机、数学假设、Python 实现、实验验证和调参注意点。如果你正在做非凸贝叶斯后验采样、带 L1 正则项的稀疏模型或者任何“目标函数不那么规整”的采样任务这篇内容可以直接当作一个基础采样模板来用。1. 核心能力速览先给一张表快速定位这个算法。能力项说明算法类型基于 Langevin 动力学的离散时间马尔可夫链采样算法目标分布π(x) ∝ exp(-f(x))f 可以非光滑、非凸依赖接口次梯度 oracle输入 x返回 ∂f(x) 中任意一个次梯度核心技巧taming 函数压缩漂移 高斯噪声注入收敛性说明非渐近收敛分析以 Wasserstein 距离等概率距离衡量硬件要求纯 CPU 即可运行不依赖 GPU启动方式无固定启动脚本需自行实现迭代循环API 接口无内置 API可封装成 Python 函数供外部调用批量任务天然支持多链并行采样主要适用场景非凸贝叶斯推断、稀疏惩罚模型的采样、非光滑正则化问题关于“无固定启动脚本”需要说明一点它是一个数学算法不是附赠 WebUI 的开源软件包。要拿它做事需要自己写一个几十行的采样循环。这篇文章第 5 节会给出完整可运行的示例直接复制改改就能用。2. 背景Langevin 采样要解决什么问题2.1 从分布采样说起在机器学习里很多时候我们不是要“算一个函数的最小值”而是要“从某个分布里产生样本”。最典型的场景是贝叶斯后验推断给定数据 D后验密度通常写成π(x) ∝ exp(-f(x))其中 f(x) -log p(x) - log p(D|x)包含先验负对数密度和似然负对数。如果我们能从 π 中拿到近似独立的样本就可以估计后验均值、方差、分位数甚至做预测区间。高维分布直接采样是不现实的。主流路线有两条MCMC 和变分推断。Langevin 采样属于 MCMC 阵营它的特点是每一步都用到目标函数的梯度信息因此比随机游走的 Metropolis-Hastings 收敛更快在高维问题上优势明显。2.2 ULA 的基本迭代Unadjusted Langevin Algorithm 来自对连续时间 Langevin SDE 的 Euler-Maruyama 离散化dX_t -∇f(X_t) dt √2 dB_t对应的离散迭代是X_{k1} X_k - γ∇f(X_k) √(2γ) ξ_{k1}其中 ξ_{k1} ~ N(0, I_d) 是标准正态噪声γ 是步长。这个迭代的含义很直观沿负梯度方向走一步相当于往高概率区域移动加一个尺度为 √(2γ) 的高斯噪声相当于保留随机探索能力。注意“Unadjusted”这个词意思是它没有像 MALA 那样加 Metropolis-Hastings 接受拒绝步骤因此每一步的计算成本很低但会引入离散化偏差。在强凸 光滑假设下这个偏差可以被完整刻画算法的非渐近收敛性也研究得比较透彻。2.3 三个痛点把 ULA 推向真实任务时会遇到三个典型痛点。第一f 不光滑。很多真实模型的目标函数含有 L1 正则、Hinge Loss、分位数损失等不可导项。∇f 在不可导点不存在ULA 在形式上就不成立了。第二∇f 可能超线性增长。比如 f(x) (x² - 1)² 这种势能梯度是 4(x³ - x)随 |x| 增大会快速变大。Euler-Maruyama 离散化对超线性系数的 SDE 可能不稳定典型表现是模拟轨迹在几步之内直接发散到 NaN。第三f 非凸。凸性在大量 ULA 分析里是基本前提。但真实后验大多是多峰、非对称的非凸才是常态。丢掉凸性之后收敛分析必须换成耗散性等更弱的条件。Tamed Subgradient ULA 的三个设计正好一一对应这三个痛点。下面拆开讲。3. Tamed Subgradient ULA 的算法设计3.1 taming压住超线性梯度taming 的思想来自数值 SDE 领域的“tamed Euler scheme”。当漂移项系数 b(x) 非线性增长太快时不要直接使用 b(x)而是用T(b(x)) b(x) / (1 |b(x)|)。这个映射有三个关键性质。第一有界无论 |b(x)| 多大|T(b(x))| 始终不超过 1第二小梯度时保持近似当 |b(x)| 很小时分母接近 1T(b(x)) ≈ b(x)不会改变原动力学的主导行为第三具有更好的整体正则性这在理论分析里能让矩估计和耦合论证更容易建立。把它用到 ULA 里迭代变成X_{k1} X_k - γ · [g(X_k) / (1 |g(X_k)|)] √(2γ) ξ_{k1}其中 g(X_k) 是 f 在 X_k 处的次梯度。需要区分一下 taming 和普通的 gradient clipping。gradient clipping 通常写成 g · min(1, C/|g|)它同样能压住梯度的模但 clipping 在阈值点附近不光滑且阈值 C 需要手工设置。taming 的分式形式没有硬阈值在全空间都光滑还避免了“阈值到底设多少”的调参问题。你完全可以把 taming 理解成“带保险丝的梯度下降”正常梯度多大就走多远一旦梯度超过某个量级漂移项自动被压缩到不超过 1避免单步跳飞。3.2 subgradient处理不可导点当 f 不光滑时梯度不存在但次梯度通常存在。次梯度的定义是∂f(x) { g ∈ R^d : f(y) ≥ f(x) ⟨g, y-x⟩ o(|y-x|), ∀y }。对机器学习里常见的函数次梯度非常好算f(x) |x|x ≠ 0 时 g sign(x)x 0 时 g 可以取 [-1,1] 中任意值f(x) λ‖x‖₁ 平滑项次梯度等于平滑项的梯度加上 λ·sign(x)f(x) max(0, 1 - y·z)在可分界点选一个可行方向即可。Tamed Subgradient ULA 的核心假设是对任意 x都能拿到一个次梯度 g(x) ∈ ∂f(x)并且这个次梯度 oracle 满足后续分析需要的可测性和增长条件。注意一个关键点这里不要求次梯度有界。这正是 taming 存在的意义。非光滑部分交给次梯度超线性部分交给 taming两者分工明确互不干扰。3.3 完整的迭代流程综合起来算法单步流程如下输入步长 γ 0迭代次数 N初始点 X_0次梯度 oracle g(·)。计算次梯度 g_k g(X_k)tamingm_k g_k / (1 |g_k|)采样噪声 ξ_{k1} ~ N(0, I_d)更新 X_{k1} X_k - γ · m_k √(2γ) · ξ_{k1}丢弃前 B 步作为 burn-in保留后续样本。整个过程和 ULA 几乎一样唯一区别是第 2 步的 taming