Markov Chain Monte Carlo without Evaluating the Target: An Auxiliary Variable Approach¶
会议: ICML 2026
arXiv: 2406.05242
代码: https://github.com/ywwes26/Auxiliary-MCMC
领域: 采样 / 贝叶斯推断 / MCMC
关键词: 辅助变量, 小批量 MCMC, 梯度型 proposal, doubly-intractable, Peskun 序
一句话总结¶
作者把 exchange、PoissonMH、TunaMH 三类"不算目标分布也能采样"的 MCMC 统一成一个用辅助变量的元算法,并在 proposal 与接受率两处同时引入辅助随机性,从而设计出小批量数据下仍保持精确平稳分布的梯度型 MCMC(Poisson–Barker、Poisson–MALA、Tuna–SGLD),实证显著超过 PoissonMH/TunaMH/SGLD 等基线。
研究背景与动机¶
领域现状:贝叶斯后验 \(\pi(\theta\mid x)\propto\pi(\theta)\prod_{i=1}^N \mathsf{p}_\theta(x_i)\) 的采样在两种情形下变得昂贵——一是 doubly-intractable,似然里有依赖 \(\theta\) 的归一化常数 \(Z(\theta)\);二是 tall data,\(N\) 极大、每步都要扫一遍数据。Exchange 算法 (Murray 2006) 处理前者,PoissonMH (Zhang & De Sa 2019) 与 TunaMH (Zhang et al. 2020) 用 Poisson 小批量处理后者。
现有痛点:这三类算法看起来各搞各的——一个生成合成数据来抵消 \(Z(\theta)\),一个把 likelihood 拆成 Poisson 因子,另一个用 minorization 技巧——但它们都只能用随机游走 proposal,在高维下混合极慢;而能利用梯度的 MALA / HMC 又必须扫全数据集,跟"少看数据"的初衷冲突。SGLD 试图绕过 MH 接受步、直接做带噪 SGD,但有持久的固定步长偏差,"基于小批量做 MH 修正"长期是开放问题。
核心矛盾:要保证精确平稳分布,传统 MH 必须计算 \(\pi(\theta'\mid x)/\pi(\theta\mid x)\) 这种昂贵比值;但要 scalable 又必须只看部分数据/合成样本。三种现有算法分别用了"换变量""泊松抽稀""无偏估计"等技巧来回避,却各自只解了一半问题(要么接受率不评目标、要么 proposal 不用梯度)。
本文目标:(i) 找出 exchange / PoissonMH / TunaMH 背后的公共结构;(ii) 把这个结构扩到 proposal 也能用辅助变量上,使梯度型 proposal 在小批量下也能保持精确平稳分布;(iii) 建立配套理论,说明新框架与"理想全数据链"的差距。
切入角度:把每一步 MH 的随机性显式拆成两路辅助变量 \(\omega_1,\omega_2\)——\(\omega_1\) 决定 proposal,\(\omega_2\) 估计目标比值——再用 involutive MCMC 的视角统一证明 detailed balance。
核心 idea:用"廉价估计"同时替代 proposal 设计与接受率计算里所有昂贵的项,只要联合分布 \(\mathbb{P}_{\theta,\theta'}(\omega_1,\omega_2)\) 在角标互换时进入接受率分子分母按规则匹配,就仍以 \(\pi\) 为不变分布。
方法详解¶
整体框架¶
作者先在 Section 2 给出一个公共子结构:任何"用辅助变量代替 \(\pi(\theta'\mid x)/\pi(\theta\mid x)\)"的 MH 步骤,都可以写成"采 \(\omega\sim P_{\theta\to\theta'}\) → 用 \(R_{\theta\to\theta'}(\omega)\) 当作比值估计 → 以 \(\min\{1,r\}\) 接受"。命题 1 给出充要条件:若 \(R_{\theta\to\theta'}(\omega)\pi(\theta\mid x)P_{\theta\to\theta'}(\omega)=\pi(\theta'\mid x)P_{\theta'\to\theta}(\omega)\),则 \(R\) 关于真比值无偏,并且链关于 \(\pi\) 可逆。Exchange、PoissonMH、TunaMH 都满足这个条件。
然后 Section 3 把这一架构扩成"双辅助变量元算法" (Algorithm 1):
- 从 \(\mathbb{P}_{\theta_t}(\cdot)\) 采 \(\omega_1\),决定 proposal 核 \(q_{\omega_1}(\theta_t,\cdot)\);
- 提议 \(\theta'\sim q_{\omega_1}(\theta_t,\cdot)\);
- 从 \(\mathbb{P}_{\theta_t,\theta'}(\cdot\mid\omega_1)\) 采 \(\omega_2\),用于估计目标比值;
- 用接受率 \(r=\dfrac{\pi(\theta'\mid x)\mathbb{P}_{\theta',\theta_t}(\omega_1,\omega_2)}{\pi(\theta_t\mid x)\mathbb{P}_{\theta_t,\theta'}(\omega_1,\omega_2)}\cdot\dfrac{q_{\omega_1}(\theta',\theta_t)}{q_{\omega_1}(\theta_t,\theta')}\) 决定是否接受。
设 \(\Omega_1\) 或 \(\Omega_2\) 为单点空间 \(\mathsf{NULL}\) 即可"关闭"对应辅助变量:(Null,Null) 是普通 MH;(Null,有) 退化为 Section 2 的旧框架(含 exchange/PoissonMH/TunaMH);\(\omega_1=\omega_2\) 是 Metropolis-within-Gibbs 视角的辅助 MH (Titsias & Papaspiliopoulos 2018);\(\omega_1\perp\omega_2\) 则让"设计 proposal"和"估计比值"用两组独立的小批量。命题 2 通过把 \((\theta,\omega_1,\theta',\omega_2)\) 视为 involutive MCMC 中的对合 \((\theta',\omega_1,\theta,\omega_2)\),给出统一的 detailed balance 证明。
关键设计¶
-
双辅助变量元算法:
- 功能:在 proposal 和接受率两处同时挂上辅助变量,让"梯度型 proposal + 小批量比值估计"在 MCMC 框架里第一次能共存且保持精确平稳分布。
- 核心思路:把每步 MH 的全部随机性写成 \((\omega_1,\theta',\omega_2)\),对合 \(f(\theta,\omega_1,\theta',\omega_2)=(\theta',\omega_1,\theta,\omega_2)\) 的雅可比为 1;接受率 \(r\) 中只要 \(\omega_1,\omega_2\) 的联合密度在 \((\theta,\theta')\) 互换下能成对抵消昂贵项,比如 PoissonMH 里 \(\mathbb{P}_\theta(\omega_1)\) 的 likelihood 部分被 \(\pi(\theta\mid x)\) 抵消、只剩小批量贡献,整个 \(r\) 就只依赖小批量数据。
- 设计动机:现有方法都是"两步走"——要么 proposal 用梯度但接受率扫全数据,要么接受率用小批量但 proposal 只能随机游走。双辅助变量统一了两件事,使梯度估计的开销跟比值估计的开销共享同一份小批量,每步成本不增反降。
-
Poisson–Barker / Poisson–MALA:locally balanced PoissonMH:
- 功能:把 PoissonMH 的随机游走 proposal 换成基于梯度的 locally balanced proposal(Barker 类或 MALA 类),同时复用 PoissonMH 的 Poisson 小批量采样来抵消归一化项。
- 核心思路:对应元算法中 \(\omega_1=\omega_2\) 的 Case 2。先按 PoissonMH 抽 \(\omega_1=(s_1,\dots,s_N)\sim\bigotimes_i \mathsf{Poi}(\lambda M_i/L+\phi_i(\theta;x))\),形成小批量 \(S=\{i\mid s_i>0\}\);proposal 用一维分解 \(q_{\omega_1}^{(g)}(\theta,\theta')=\prod_i q_{\omega_1,i}^{(g)}\),其中 \(q_{\omega_1,i}^{(g)}\propto g(e^{\partial_{\theta_i}\log(\pi(\theta\mid x)\mathbb{P}_\theta(\omega_1))(\theta_i'-\theta_i)})\mu_i(\theta_i'-\theta_i)\)。关键在于代理函数 \(\pi(\theta\mid x)\cdot\mathbb{P}_\theta(\omega_1)\) 只依赖小批量 \(S\),所以梯度算一次只扫几千个点;\(g(t)=t/(1+t)\) 得到 Poisson–Barker,\(g(t)=\sqrt{t}\) 得到 Poisson–MALA。
- 设计动机:locally balanced proposal (Zanella 2020; Livingstone & Zanella 2022) 在全数据下已经被证明比 MALA 更鲁棒;本文要做的是把那一套"用梯度信息塑形 proposal"的好处搬到小批量场景,并通过把 \(\omega_1\) 同时塞进 proposal 与接受率两边来抵消所有全数据计算。
-
Tuna–SGLD:让 SGLD 拥有精确 MH 修正:
- 功能:把 TunaMH 的随机游走 proposal 换成 SGLD 风格 proposal \(q_{\omega_1}(\theta,\cdot)\sim\mathcal{N}(\theta-\tfrac{\epsilon^2}{2}\tfrac{N}{K}\sum_{i\in B}\nabla_\theta U_i(\theta;x),\epsilon^2 I)\),再用 TunaMH 的 Poisson 小批量 \(\omega_2\) 估计目标比值,从而把 SGLD 变成关于 \(\pi\) 精确平稳的链。
- 核心思路:对应元算法 Case 3,\(\omega_1=B\) 是 size-\(K\) 的均匀小批量,独立于 \(\omega_2\);由于 \(\omega_1\) 的边缘分布不依赖 \(\theta\),接受率公式里 \(\omega_1\) 那一部分被消掉,最终 \(r=\dfrac{\pi(\theta'\mid x)\mathbb{P}_{\theta',\theta_t}(\omega_2)}{\pi(\theta_t\mid x)\mathbb{P}_{\theta_t,\theta'}(\omega_2)}\cdot\dfrac{q_{\omega_1}(\theta',\theta_t)}{q_{\omega_1}(\theta_t,\theta')}\),每步只看小批量。
- 设计动机:SGLD 的固定步长偏差长期没有干净的 MH 修正方案,Welling & Teh (2011) 把它列为开放问题;Tuna–SGLD 用 TunaMH 的辅助变量当"修正器",给出第一个仅用小批量数据就把 SGLD 变成精确采样器的方案。
损失函数 / 训练策略¶
无显式损失——所有算法都是迭代采样器。需调的超参是 PoissonMH 系列的 \(\lambda\)(控制小批量期望大小,文中常用 \(\lambda=0.0005L^2\) 到 \(0.01L^2\))、Tuna–SGLD 的批量大小 \(K\) 与步长 \(\epsilon\)、以及 locally balanced 类的 \(g\)。Pilot run 把步长调到目标接受率 0.25 / 0.4 / 0.55。
实验关键数据¶
主实验¶
实验包含三类任务:(i) 20 维异质截断高斯,\(N=10^5\)、\(\Sigma=\mathrm{diag}(1,0.95,\dots,0.05)\)、tempered posterior(\(\beta=10^{-5}\));(ii) 10 维 robust Student-\(t\) 线性回归,\(N=10^5\)、\(\nu=4\);(iii) MNIST 上的 Bayesian logistic regression。指标用 MSE 随时间曲线和按维度 ESS/s 的 min/median/max。
| 任务 | 方法 | Best ESS/s (Min, Med, Max) |
|---|---|---|
| 异质 Gaussian | MH | (0.05, 0.08, 0.47) |
| 异质 Gaussian | MALA | (0.10, 0.19, 2.77) |
| 异质 Gaussian | Barker | (0.12, 0.22, 1.53) |
| 异质 Gaussian | PoissonMH | (0.40, 0.66, 4.67) |
| 异质 Gaussian | Poisson–Barker | (0.91, 1.65, 12.16) |
| 异质 Gaussian | Poisson–MALA | (0.84, 1.65, 23.84) |
Poisson–Barker 在异质 Gaussian 上相对 PoissonMH 提升 1.37–7.12×,对 MALA 4.39–9.80×,对 Barker 6.58–15.58×,对随机游走 MH 13.62–70.12×;在 robust 线性回归任务上 Poisson–{MALA,Barker} 相对 PoissonMH 提升 1.80–8.79×,对全数据法接近或超过 100× 量级;SGLD 早期 MSE 下降快但很快进入由固定步长偏差导致的平台;MNIST Bayesian logistic regression 上 Tuna–SGLD 最快收敛且没有 SGLD 的偏差平台。
消融实验¶
| 配置 | 关键现象 | 说明 |
|---|---|---|
| Full Poisson–Barker | 全梯度型小批量 MH,最佳 ESS/s | 同时受益于梯度 proposal 与 Poisson 小批量比值估计 |
| 去掉梯度(= PoissonMH) | ESS/s 掉 1.4–7× | 验证 locally balanced proposal 的贡献 |
| 去掉 MH 修正(= SGLD) | MSE 早期降快但收敛到有偏分布 | 没有辅助变量纠偏就回到 SGLD 老问题 |
| 用 MALA 替换 Barker | 高接受率匹配最佳;低接受率(0.25)反而比 PoissonMH 差 | MALA 对步长更敏感,Barker 更鲁棒,与 Livingstone & Zanella 2022 一致 |
| 不同 \(\lambda\)(小批量大小) | 接受率/ESS 单调变化 | 偏小的 \(\lambda\) 噪声大但廉价,偏大则成本接近全数据 |
关键发现¶
- 关键贡献是"梯度 proposal"与"小批量接受率"的耦合;任何一边缺失都退化(缺梯度 = PoissonMH/TunaMH,缺修正 = SGLD)。
- Poisson–Barker 对接受率最不敏感,是默认推荐;Poisson–MALA 在高接受率下匹配但低接受率下变脆。
- Tuna–SGLD 首次给出"用小批量数据修正 SGLD"的可行方案,回答 Welling & Teh (2011) 留下的开放问题。
- Peskun 序 \(\mathbb{P}_{\mathsf{aux}}\prec\mathbb{P}_{\mathsf{MwG}}\prec\mathbb{P}_{\mathsf{ideal}}\) 保证:理想全数据链的渐近方差永远是新链的下界,没有"瞎赚"。
亮点与洞察¶
- 统一视角:把表面看上去毫不相干的三类算法(doubly-intractable 的 exchange、tall-data 的 PoissonMH 与 TunaMH)归约到同一个 detailed balance 等式 \(R\cdot\pi P=\pi' P'\),并指出它们对应 \(\omega_1=\mathsf{Null}\) 的特殊情形——这本身就是个干净的概念性贡献。
- involutive MCMC 视角:通过把 \((\omega_1,\theta',\omega_2)\) 一起对合,validity 证明压缩到一行,给后续扩展(如 \(a(t)\) 一般化接受函数、\(\varphi\) 非平凡对合)留了空间。
- 代理函数巧思:Poisson–Barker 里把梯度 \(\partial_\theta\log\pi\) 替换为 \(\partial_\theta\log(\pi\cdot\mathbb{P}_\theta(\omega_1))\),看上去只是加了一项 \(\log\mathbb{P}_\theta(\omega_1)\),但因为 \(\pi(\theta\mid x)\mathbb{P}_\theta(\omega_1)\) 只依赖小批量 \(S\),梯度计算成本和 PoissonMH 一致——这种"加一项让昂贵变廉价"的代理设计是可迁移到其他 minibatch 算法的 trick。
- 理论副产品:用统一框架重证 PoissonMH 与 TunaMH 的谱隙,PoissonMH 的界 \(\mathrm{Gap}(\mathbb{P}_{\mathsf{aux}})\ge \max\{\tfrac12 e^{-L^2/(\lambda+L)},e^{-1/e}e^{-L^2/(2\lambda)}\}\mathrm{Gap}(\mathbb{P}_{\mathsf{ideal}})\) 在 \(\lambda>L\) 时严格优于 Zhang & De Sa (2019),TunaMH 的 \(e^{-1/(2\chi)}\) 也严格优于 \(e^{-1/\chi-2\sqrt{\log 2/\chi}}\)。
局限与展望¶
- 依赖 PoissonMH/TunaMH 的技术假设:似然要可写成可被 Poisson 小批量抵消的形式(如 bounded Lipschitz \(U_i\)),并非任意 likelihood 都适用;非指数族或重尾 likelihood 上能否同样高效未验证。
- 实验规模偏小:最大数据集 MNIST + Bayesian logistic regression,没有评估 deep Bayesian neural network 或 LLM scale 的后验,距离实际大模型贝叶斯推断还有距离。
- 梯度估计方差:Tuna–SGLD 的 proposal 用的是普通 minibatch 梯度,方差大时接受率会下降;引入控制变量、SVRG 风格梯度估计、preconditioning 是自然延伸。
- 理论给的是相对界:所有谱隙界都是相对于 \(\mathbb{P}_{\mathsf{ideal}}\) 的乘性常数,理想链本身的混合速度仍要单独分析;与 PDMP 类样本器(Bouncy Particle、Zig-Zag)的实证比较缺失。
相关工作与启发¶
- vs Exchange / PoissonMH / TunaMH: 它们都是本框架在 \(\omega_1=\mathsf{Null}\) 时的特例;本文的扩展是允许 \(\omega_1\) 进入 proposal,从而首次让梯度信息进入小批量 MH。
- vs Pseudo-marginal MCMC (Andrieu & Roberts 2009): pseudo-marginal 把估计器值并入链状态,本文只保留 \(\theta\),每步重新生成辅助变量,因此是不同的 augmentation 构造。
- vs Involutive MCMC (Neklyudov et al. 2020): 本文给出 involutive MCMC 的一个具体且实用的实例族,并明确给出"双辅助变量"的具体应用(locally balanced + minibatch)。
- vs PDMP 子采样 (Bouchard-Côté et al. 2018; Bierkens et al. 2019): 它们用连续时间 + Poisson thinning 回避全数据,本文走离散 MH 路线,两条路在工程上互补。
- vs SGLD (Welling & Teh 2011): Tuna–SGLD 是 SGLD 的精确化版本,直接回答原文留下的"基于小批量数据做 MH 修正"开放问题。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 第一次把 doubly-intractable 与 tall-data 两条 MCMC 路线统一进同一个双辅助变量框架,并由此构造出小批量梯度型精确 MCMC。
- 实验充分度: ⭐⭐⭐⭐ 三个任务覆盖合成与真实数据,多基线对比充分;缺深度模型规模的后验实验。
- 写作质量: ⭐⭐⭐⭐⭐ 框架介绍清晰,命题 1–2 简洁有力,Case 1/2/3 与 Algorithm 2/3 的映射很容易跟。
- 价值: ⭐⭐⭐⭐⭐ 既是统一旧算法的概念性贡献,又给出可立刻替换 PoissonMH/TunaMH/SGLD 的实用算法,理论上还顺手收紧了两个已有谱隙界。