使用子采样马尔可夫链蒙特卡罗的潜变量模型大规模不确定性量化

# 大规模潜变量模型的子采样马尔可夫链蒙特卡洛不确定性量化
来源：https://arxiv.org/html/2606.00309

###### 摘要

随机梯度 Langevin 动力学结合 Gibbs 更新（SGLD–Gibbs）为潜变量模型中的近似贝叶斯推断提供了一种高度可扩展的方法。然而，如何以原则性的方式调整算法的超参数以确保不确定性估计具有统计意义仍不清楚。在这项工作中，我们通过为 SGLD–Gibbs 开发统计缩放极限理论来填补这一调参指导上的空白。我们在适当的时间-空间重缩放下推导了全局参数和潜变量的联合渐近极限。我们证明全局参数收敛到扩散型极限，而每个潜变量收敛到一个跳跃过程，这反映了间歇性 Gibbs 更新的使用。这种联合跳跃-扩散结构揭示了潜变量随机性如何贡献于全局参数的平稳分布。我们利用我们的结果提出了 SGLD–Gibbs 超参数调整的明确指导，以确保有意义的不确定性量化。数值实验表明，根据我们的指导调整的 SGLD–Gibbs 在参数估计、不确定性量化和预测性能方面优于随机变分推断。

## 1 引言

随机梯度方法，如随机梯度下降（SGD）和随机梯度 Langevin 动力学（SGLD），已成为大规模优化和近似贝叶斯推断的核心工具（Nemirovski 等，2009 (https://arxiv.org/html/2606.00309#bib.bib36)；Moulines & Bach，2011 (https://arxiv.org/html/2606.00309#bib.bib32)；Bottou 等，2018 (https://arxiv.org/html/2606.00309#bib.bib4)；Welling & Teh，2011 (https://arxiv.org/html/2606.00309#bib.bib48)；Nemeth & Fearnhead，2021 (https://arxiv.org/html/2606.00309#bib.bib35)）。在近似采样方面，潜变量模型（LVM）是 SG(L)D 最常被引用的应用之一。例子包括高斯混合模型、混合成员随机块模型（Li 等，2016 (https://arxiv.org/html/2606.00309#bib.bib24)）、潜在狄利克雷分配（Patterson & Teh，2013 (https://arxiv.org/html/2606.00309#bib.bib37)）、贝叶斯矩阵分解（Ahn 等，2015 (https://arxiv.org/html/2606.00309#bib.bib1)）、混合效应模型（Danaher，2023 (https://arxiv.org/html/2606.00309#bib.bib9)）和离散选择模型（Loaiza-Maya & Nibbering，2023 (https://arxiv.org/html/2606.00309#bib.bib25)；Loaiza-Maya 等，2024 (https://arxiv.org/html/2606.00309#bib.bib26)）。在这些应用中，经常使用 *SGLD–Gibbs* 方案，其中 SGLD 更新步骤通过一个或多个潜变量的条件抽取来构建，从而使得每次迭代的计算成本随数据规模呈有利增长。

然而，关于如何调整 SGLD–Gibbs 的算法超参数（如步长、小批量大小和逆温度）几乎没有严格的指导。此外，如何使用 SGLD–Gibbs 获得有意义的不确定性量化仍不清楚。大量工作（Walk，1977 (https://arxiv.org/html/2606.00309#bib.bib45)；Pflug，1986 (https://arxiv.org/html/2606.00309#bib.bib39)；Kushner & Yin，2003 (https://arxiv.org/html/2606.00309#bib.bib21)；Negrea 等，2023 (https://arxiv.org/html/2606.00309#bib.bib34)；Wang 等，2025 (https://arxiv.org/html/2606.00309#bib.bib46)）已采用缩放极限分析来研究标准 SG(L)D。这些方法将 SG(L)D 样本路径与连续时间随机过程联系起来，并刻画出优化精度、渐近行为和数值效率的特征。此类分析已被证明对理解超参数调整和不确定性量化特别有用（Mandt 等，2017 (https://arxiv.org/html/2606.00309#bib.bib27)；Negrea 等，2023 (https://arxiv.org/html/2606.00309#bib.bib34)）。然而，现有的缩放极限结果并不直接适用于潜变量模型。

在这项工作中，我们通过联合分析在适当时空重缩放下的全局参数和潜变量来填补这一空白，从而提供 SGLD–Gibbs 动力学的统一渐近刻画。我们证明全局参数收敛到扩散型极限，而每个潜变量收敛到一个独立的跳跃过程。我们进一步证明，全局扩散与潜变量跳跃之间的相互作用从根本上改变了全局参数的噪声结构。特别地，潜变量贡献了一个额外的变异性来源，该变异性取决于每次迭代使用的 Gibbs 样本数量。我们利用这些结果来推导关于不确定性量化和超参数调整的具体指导。实验发现，在混合建模和主题建模的应用中，与随机变分推断相比，SGLD–Gibbs 能提供更高的准确性和更可靠的不确定性量化。

### 1.1 相关工作与替代方法

鉴于其广泛使用，SG(L)D 方法已从多个角度得到研究，包括有限样本误差界、收敛速率和平稳分布（例如，Mcleish，1976 (https://arxiv.org/html/2606.00309#bib.bib29)；Ruppert，1988 (https://arxiv.org/html/2606.00309#bib.bib43)；Polyak & Juditsky，1992 (https://arxiv.org/html/2606.00309#bib.bib40)；Kushner & Yin，2003 (https://arxiv.org/html/2606.00309#bib.bib21)；Negrea 等，2023 (https://arxiv.org/html/2606.00309#bib.bib34)；Rakhlin 等，2011 (https://arxiv.org/html/2606.00309#bib.bib42)；Dieuleveut 等，2020 (https://arxiv.org/html/2606.00309#bib.bib10)；Mou 等，2020 (https://arxiv.org/html/2606.00309#bib.bib31)；Cheng 等，2020 (https://arxiv.org/html/2606.00309#bib.bib5)；Srikant，2024 (https://arxiv.org/html/2606.00309#bib.bib44)；Anastasiou 等，2019 (https://arxiv.org/html/2606.00309#bib.bib2)；Ge 等，2015 (https://arxiv.org/html/2606.00309#bib.bib12)；Jin 等，2017 (https://arxiv.org/html/2606.00309#bib.bib18)）。与我们工作最相关的是随机逼近算法的缩放极限理论，该理论表明，在适当的重缩放下，SGD 和 SGLD 轨迹收敛到 Ornstein-Uhlenbeck 扩散（Kushner & Huang，1981 (https://arxiv.org/html/2606.00309#bib.bib22)；Kushner & Yang，1993 (https://arxiv.org/html/2606.00309#bib.bib23)；Kushner & Yin，2003 (https://arxiv.org/html/2606.00309#bib.bib21)；Negrea 等，2023 (https://arxiv.org/html/2606.00309#bib.bib34)）。进一步的结果刻画了混合时间、平稳协方差和平均迭代的行为（Mandt 等，2017 (https://arxiv.org/html/2606.00309#bib.bib27)；Negrea 等，2023 (https://arxiv.org/html/2606.00309#bib.bib34)；Collins-Woodfin 等，2024 (https://arxiv.org/html/2606.00309#bib.bib7)；Qian 等，2024 (https://arxiv.org/html/2606.00309#bib.bib41)；Kushner & Yang，1993 (https://arxiv.org/html/2606.00309#bib.bib23)）。

最近几项工作扩展了这一理论，以改进随机梯度算法的不确定性量化。Wang 等（2025 (https://arxiv.org/html/2606.00309#bib.bib46)）为缩放极限的扩散近似开发了非渐近函数误差界。Wang 等（2026 (https://arxiv.org/html/2606.00309#bib.bib47)）开发了随机梯度算法的离散时间代理理论，阐明了在大批量或非渐近机制下基于扩散的不确定性量化何时仍然有效。另一条研究路线研究 SGD 在高维机制（参数维度d→∞d\\to\infty）下的缩放极限，得出低维汇总统计量的平均场或动力学方程（Arous 等，2022 (https://arxiv.org/html/2606.00309#bib.bib3)；Collins-Woodfin 等，2023 (https://arxiv.org/html/2606.00309#bib.bib6)；Mignacco 等，2021 (https://arxiv.org/html/2606.00309#bib.bib30)）。

变分贝叶斯方法，包括平均场变分贝叶斯、在线变分贝叶斯、随机变分推断（SVI）及相关变分近似，也可以为潜变量模型推断提供可扩展的方案（Hoffman 等，2013 (https://arxiv.org/html/2606.00309#bib.bib16)，2010 (https://arxiv.org/html/2606.00309#bib.bib15)；Kucukelbir 等，2017 (https://arxiv.org/html/2606.00309#bib.bib20)）。然而，它们量化后验不确定性的能力可能受到根本性限制（Gelman 等，2013 (https://arxiv.org/html/2606.00309#bib.bib13)；Margossian 等，2025 (https://arxiv.org/html/2606.00309#bib.bib28)；Giordano 等，2018 (https://arxiv.org/html/2606.00309#bib.bib14)）。例如，Margossian 等（2025 (https://arxiv.org/html/2606.00309#bib.bib28)）证明，当真后验分布存在依赖结构时，基于因子分解的变分近似通常无法正确估计后验不确定性。根据所最小化的散度，变分贝叶斯产生的不确定性估计即使模型设定正确，也常常校准不良。

## 2 预备知识与问题设定

本节介绍本文所考虑的一类潜变量模型，描述 SGLD–Gibbs 算法，并回顾关于随机梯度方法缩放极限的已有结果。

### 2.1 潜变量模型

我们考虑一类一般的潜变量模型，其中每个观测值都与一个未观测到的潜变量相关联。令 \{\(X_i, z_i\)\}_{i=1}^n 表示独立的数据对，包括观测数据 \(X_i \in \mathcal{X}\) 和潜变量 \(z_i \in \mathcal{Z}\)。联合分布由全局参数 \(\theta \in \Theta \subset \mathbb{R}^d\) 参数化，并具有以下分解形式：

\[
p(X_i, z_i \mid \theta) = p(z_i \mid \theta) \, p(X_i \mid z_i, \theta),
\]
(2)

其中 \(\pi_0(\theta)\) 是 \(\theta\) 上的先验分布。观测的边缘似然为 \(p(X_i \mid \theta) = \int p(X_i, z_i \mid \theta) \, dz_i\)，对应的对数似然为 \(\ell(\theta; X_i) := \log p(X_i \mid \theta)\)。这一公式涵盖了许多常用模型，包括混合模型、混合成员随机块模型、主题模型和贝叶斯矩阵分解。有关此类模型的学习和近似贝叶斯推断的系统性讨论，请参见 Murphy（2023 (https://arxiv.org/html/2606.00309#bib.bib33)）。

### 2.2 带 Gibbs 更新的 SGLD

我们研究结合潜变量 Gibbs 更新的随机梯度 Langevin 动力学。令 \(b \in \{1,\dots,n\}\) 表示小批量大小。在第 \(k\) 次迭代时，从 \(\{1,\dots,n\}\) 中有放回地均匀采样一个小批量索引集合 \(I_k = \{I_k(1), \dots, I_k(b)\}\)。这一约定主要是为了理论上的方便。Negrea 等（2023 (https://arxiv.org/html/2606.00309#bib.bib34)）已建立了无放回抽样的类似缩放极限结果，我们预期本文的结论在无放回抽样下保持不变，仅当 \(b\) 与 \(n\) 同阶时存在常数水平的偏差。给定当前全局参数 \(\theta_k\)，算法分两步进行：

#### (i) 潜变量的 Gibbs 更新。

对于每个 \(i \in I_k\)，潜变量从其条件后验中重新采样：

\[
z_{i,k+1} \sim p(z_i \mid X_i, \theta_k),
\]
(3)

而索引不在小批量中的潜变量保持不变。

#### (ii) 全局参数的 SGLD 更新。

使用带有刷新后潜变量的随机梯度估计器：

\[
G_k(\theta) := \frac{1}{n} \nabla \log \pi_0(\theta) + \frac{1}{b} \sum_{i \in I_k} \nabla_\theta \log p\left(X_i, z_{i,k+1} \mid \theta\right),
\]

全局参数更新如下：

\[
\theta_{k+1} = \theta_k + \frac{h}{2} \Gamma G_k(\theta_k) + \sqrt{\frac{h}{\beta}} \Gamma^{1/2} \xi_k,
\]
(4)

其中 \(h > 0\) 是步长，\(\beta \in (0,\infty]\) 是逆温度，\(\Gamma \in \mathbb{R}^{d \times d}\) 是正定预处理矩阵，\(\xi_k \sim \mathcal{N}(0, I_d)\)。完整流程总结在算法 1 中。

**算法 1** 用于潜变量模型的 SGLD–Gibbs  
1：输入：步长 \(h\)，批量大小 \(b\)，逆温度 \(\beta\)，预处理器 \(\Gamma\)，初始值 \((\theta_0, \{z_{i,0}\}_{i=1}^n)\)  
2：对于 \(k = 0,1,2,\ldots\) 执行  
3：     采样小批量 \(I_k \subset \{1,\ldots,n\}\)，满足 \(|I_k| = b\)  
4：     对于每个 \(i \in I_k\) 执行  
5：         采样 \(z_{i,k+1} \sim p(z_i \mid X_i, \theta_k)\)  
6：     结束循环  
7：     使用更新后的 \(\{z_{i,k+1}\}_{i \in I_k}\) 通过 SGLD 步更新 \(\theta_{k+1}\)  
8：结束循环

### 2.3 随机梯度方法的缩放极限

在不涉及潜变量的一般设定中，假设观测值来自未知分布 \(P_\star\) 的独立同分布样本。模型由全局参数 \(\theta \in \Theta \subset \mathbb{R}^d\) 参数化，每个观测值 \(X_i\) 具有形如 \(p(X_i \mid \theta)\) 的似然，以及 \(\theta\) 上的先验分布 \(\pi_0(\theta)\)。最优参数定义为 \(\theta_\star := \operatorname{arg\,min}_\theta \mathbb{E}\left[\ell(X,\theta)\right]\)，其中 \(X \sim P_\star\)。

回顾 \(I_k \subset \{1,\dots,n\}\) 表示第 \(k\) 个小批量。SGLD 使用方程 4 中给出的单步更新，但此时的随机梯度估计器为：

\[
G_k(\theta) := \frac{1}{n} \nabla \log \pi_0(\theta) + \frac{1}{b} \sum_{i \in I_k} \nabla_\theta \log p\left(X_i \mid \theta\right).
\]
(5)

缩放极限理论将离散时间的随机梯度算法与连续时间的随机过程在适当的时空重缩放下联系起来。令 \(\theta_k^{(n)} \in \mathbb{R}^d\) 表示第 \(k\) 次迭代的全局参数，\(\hat{\theta}^{(n)}\) 表示满足一阶条件 \(\sum_{i=1}^n \nabla \ell(\hat{\theta}^{(n)}; X_i) = 0\) 的临界点。定义重缩放后的连续时间过程：

\[
\vartheta_t^{(n)} = n^{\mathfrak{w}} \left( \theta_{\lfloor n^{\mathfrak{a}} t \rfloor}^{(n)} - \hat{\theta}^{(n)} \right),
\]
(6)

其中 \(\mathfrak{w} > 0\) 和 \(\mathfrak{a} > 0\) 分别表示空间和时间缩放指数。

然后，在适当的缩放机制下，该过程在分布上收敛到一个 Ornstein-Uhlenbeck 过程，其漂移和扩散系数取决于预处理器 \(\Gamma\) 以及一阶和二阶 Fisher 信息矩阵：

\[
I_\star := \mathbb{E}\left[ \left[ \nabla_\theta \ell(\theta_\star; X) \right]^{\otimes 2} \right]
\]