使用随机梯度马尔可夫链蒙特卡罗的大样本准确不确定性量化

# 利用随机梯度马尔可夫链蒙特卡洛进行大规模准确不确定性量化  
来源：https://arxiv.org/html/2606.00293  

###### 摘要  

诸如随机梯度下降（SGD）和随机梯度 Langevin 动力学（SGLD）等调优算法，在进行近似采样和不确定性量化时仍然具有挑战性，尤其是在批次大小较大或模型错误指定的实际相关设置中。现有提供调优指导的理论依赖于连续时间极限或较强的统计假设，在这些场景下会变得定量不准确。我们通过提出新的 SGD 和 SGLD（带或不带动量）离散时间近似来弥补这些不足，从而能够精确预测平稳协方差、迭代平均协方差以及积分自相关时间。此外，我们证明了定量的非渐近误差界，表明这些估计对于实际调优和不确定性量化足够准确。数值实验表明，我们的理论在现有方法失效的多种模型和数据生成分布下（包括使用 β-散度而非对数损失以获得统计鲁棒推断时）能够提供改进的调优指导。  

## 1 引言  

基于随机梯度的方法已成为机器学习中大规模优化的默认工具。诸如随机梯度下降（SGD）及其变体等算法主导了现代实践，因为子采样显著降低了每次迭代的计算成本，同时具有强大的经验性能和良好的泛化性质（Bottou, 2010；Hardt 等, 2016；Goodfellow 等, 2016）。从贝叶斯角度来看，基于子采样的马尔可夫链蒙特卡洛（MCMC）方法似乎提供了一条通向可扩展采样和不确定性量化（UQ）的类似路径。特别是，随机梯度 MCMC（SG-MCMC）算法，如随机梯度 Langevin 动力学（SGLD），用无偏小批量估计替换全数据似然梯度，期望以与 SGD 相当的计算成本实现后验采样（Welling & Teh, 2011；Li 等, 2016；Raginsky 等, 2017；Brosse 等, 2018；Nemeth & Fearnhead, 2021）。然而，在实践中，SG-MCMC 方法由于其步长、批次大小和温度参数必须精心选择以控制离散化偏差和混合行为，同时提供准确的 UQ，因此调优非常困难（Nemeth & Fearnhead, 2021；Coullon 等, 2023；Negrea 等, 2023；Rajpal 等, 2025；Kim 等, 2024；Mauri & Zanella, 2024；Alexos 等, 2022；Paulin 等, 2025；Akyildiz & Sabanis, 2024）。当统计模型错误指定时，这些挑战更加严峻，因为此时标准贝叶斯后验不再校准良好。同样的校准问题也出现在使用广义贝叶斯损失时，无论模型是否正确指定（Bissiri 等, 2016；Jewson 等, 2018）。  

近期工作已开始通过明确结合算法和统计的渐近视角来应对这些挑战。例如，Mandt 等（2017）采用了启发式视角，该视角受两条研究线索启发。第一条考虑随机逼近中的尺度极限，并表明在适当的空间和时间重缩放后，迭代轨迹联合收敛到连续时间 Ornstein–Uhlenbeck 过程（Kushner & Huang, 1981；Pflug, 1986；Walk, 1977；Kushner & Yang, 1993；Kushner & Yin, 2003）。第二条涉及贝叶斯后验的渐近行为，即 Bernstein–von Mises（或贝叶斯中心极限）定理（Kleijn & van der Vaart, 2012；Van der Vaart, 2000）。更近期的，Negrea 等（2023）和 Wang 等（2025）通过分析数据集大小和算法参数（如步长和批次大小）共同缩放的联合极限，形式化并扩展了 Mandt 等（2017）的启发式论证。此外，Wang & Huggins（2026）将 Negrea 等（2023）的结果扩展到具有局部隐变量的模型。这些结果刻画了迭代样本路径的极限随机过程，使得不仅能够确定极限平稳分布（这对 UQ 很重要），还能确定混合时间和迭代平均分布，后者决定了算法的计算效率和后验期望估计的准确性。因此，这些结果能够提供精确的调优建议，在最大化计算效率的同时，针对所需的不确定性量化形式（如频率覆盖（White, 1982）、贝叶斯模型不确定性（Kleijn & van der Vaart, 2012）或两者兼顾（Huggins & Miller, 2024））进行优化。  

然而，这些结果的一个主要局限是它们依赖于连续时间随机微分方程（SDE）极限，这仅在步长趋近于零时近似离散时间算法（Wang 等, 2025；Li 等, 2019）。这些极限近似在实践中最相关的大批次大小场景下会变得定量不准确。问题在于，使用大批次大小需要相对较大的步长（Goyal 等, 2017；Negrea 等, 2023），因此连续时间近似可能会严重误判平稳协方差结构，从而导致不准确的 UQ。图 1 展示了即使在简单的错误指定线性模型中，这些问题也可能出现。在该示例中，随着批次大小增加，由 SDE 极限导出的调优规则的准确性迅速下降，导致平稳协方差无法匹配三明治协方差 S⋆（White, 1982）。即使数据量增加，这种失败仍然存在，突显了连续时间近似在指导实际调优决策方面的根本局限（Wang 等, 2025）。  

近期工作使用了在大批次大小和/或大步长下仍然有效的离散时间近似（Dieuleveut 等, 2020；Liu 等, 2021；Ziyin 等, 2022）。尽管有前景，但现有结果要么假设噪声协方差恒定，要么仅适用于线性模型，或者不考虑模型错误指定。此外，大多数近似缺乏严格的非渐近误差保证；并且没有提供混合时间或迭代平均分布的估计。如图 1 所示，因此它们可能无法为不确定性量化提供可靠指导——在此例中是由于模型错误指定。  

图 1：带有异方差噪声的错误指定线性回归。数据根据 yn∼N(xn⊤θ⋆,1+‖xi‖22) 生成，其中 θ⋆∼N(0,ID) 固定，xn∼i.i.d.N(0,ID)。使用恒定步长 SGD 拟合线性模型。S⋆=J⋆−1I⋆J⋆−1：三明治协方差；Ŝ：基于不同理论得出的步长调优规则获得的协方差。  

表 1：用于调优 SG(L)D 进行采样的近似方法比较。参考文献指向与调优最直接相关的工作。大批次：该方法是否适用于大批次大小？非恒定噪声：该方法是否考虑了非恒定的随机梯度噪声？一般模型/损失：该方法是否考虑了模型错误指定或广义损失的使用？混合时间：该方法是否提供了混合时间和迭代平均协方差估计？误差界：是否提供了定量误差界？  

| 方法 | 大批次 | 非恒定噪声 | 一般模型/损失 | 混合时间 | 误差界 |
|------|--------|------------|----------------|----------|--------|
| 连续时间 (Mandt 等, 2017; Negrea 等, 2023; Wang 等, 2025) | ✗ | ✗ | ✓ | ✓ | ✓ |
| 离散二次 + 恒定噪声 (Dieuleveut 等, 2020; Liu 等, 2021) | ✓ | ✗ | ✓ | ✗ | ✓ |
| 线性回归 + 正确指定 (Ziyin 等, 2022) | ✓ | ✓ | ✗ | ✗ | ✗ |
| 离散二次 + 精确噪声 (本文) | ✓ | ✓ | ✓ | ✓ | ✓ |

在本文中，我们通过开发一个在批次大小较大且模型错误指定时仍然准确的离散时间理论框架来解决这些局限。表 1 比较了我们的方法与现有方法。我们的贡献如下：  

1. (次要) 我们引入了**代理算法框架**，阐明了现有方法的差异和局限，从而有助于确定需要进一步理论的方向（第 3 节）。  
2. (主要) 我们推导了**SGD 和 SGLD 的新离散时间近似**（带或不带动量），该近似在批次大小较大和模型错误指定时仍然准确（第 4 节）。  
3. (主要) 我们提供了**定量的非渐近误差分析**，证明得到的平稳协方差估计足够准确，可用于**采样和不确定性量化**的实际调优（第 4.2 节）。  
4. (主要) 我们利用结果提出了一种实用的、无需调优的可扩展不确定性量化流程（算法 1）。通过数值实验，我们展示了我们的理论在不同模型、批次大小范围和损失函数下提供了改进的调优指导（第 6 节）。  
5. (次要) 最后，尽管本文重点在于不确定性量化和采样，我们的结果也揭示了 SGD 的训练动力学和泛化行为，以及其在频率推断中的使用（Jantre 等, 2024; Hwang 等, 2022; Chang 等, 2017; Lyle 等, 2020; Mandt 等, 2017; Zhu 等, 2019; Lewkowycz 等, 2020; Keskar 等, 2017; Hoffer 等, 2017; Mori & Ueda, 2020）。为完整起见，我们通过一些初步实验说明了这些可能对更广泛机器学习社区感兴趣的方向（附录 E）。  

## 2 背景  

### 2.1 设定  

令 {xn}n=1N 表示观测数据，其中 xn∈X。对于参数 θ∈RD，假设存在观测级别的可微损失或负对数似然 ℓ:X×RD→R，以及正则化项 R:RD→R，在采样设置中应将其解释为负对数先验 −logπ0(θ)（相差一个加性常数）。这些共同决定了负势能（或损失）：  

L(θ):=N−1∑n=1Nℓ(xn,θ)+N−1R(θ). (2)  

定义随机梯度：  

Gt(θ):=B−1∑n∈St∇ℓ(xn,θ)+N−1∇R(θ), (3)  

其中 St={It1,It2,…,ItB} 是一个包含 B 个独立随机整数的集合，均匀地从 {1,…,N} 中有放回或无放回抽样。  

**随机梯度 Langevin 动力学**（SGLD；Welling & Teh, 2011）是一种马尔可夫链蒙特卡洛（MCMC）算法，单步更新方程为：  

θt = θt−1 − Λ Gt(θt−1) + √(2β−1Λ) ξt−1, (4)  

其中 Λ∈RD×D 是正定步长矩阵，β∈(0,∞] 是逆温度（通常设为 β=N），ξt−1∼i.i.d.N(0,I)。SGLD 是**子采样 MCMC** 算法的典型示例，其变体已应用于学习各种大规模模型（Ahn 等, 2012; Nemeth & Fearnhead, 2021; Aicher 等, 2025; Kim 等, 2024; Rajpal 等, 2025; Mauri & Zanella, 2024; Alexos 等, 2022; Paulin 等, 2025; Akyildiz & Sabanis, 2024）。