Batch and match: black-box variational inference with a score-based divergence
Diana Cai, Chirag Modi, Loucas Pillaud-Vivien, Charles C. Margossian, Robert M. Gower, David M. Blei, Lawrence K. Saul
ICML)
stat.ML, cs.AI, cs.LG, stat.CO
2024-02-23
提出 BaM:用 score 散度取代 ELBO,proximal 闭式更新做黑盒变分推断,梯度评估比 ADVI 少一个数量级以上,且对学习率稳健。
黑盒变分推断(BBVI)的标配做法是最大化 ELBO:选一个简单分布族近似算不出归一化常数的后验,重参数化后交给 SGD。两个老毛病随之而来:梯度估计方差大、学习率难调;维度一高,或把变分族从对角高斯换成能建模相关性的全协方差高斯,病得更重。Stan、PyMC、Pyro 里的 BBVI 基本都走这条路。此前的 GSM(2023)绕开了 SGD,对高斯变分族解非线性方程更新,但不最小化任何显式散度,目标非高斯时要靠启发式才能收敛。这篇问的问题很具体:换一个散度、保留全协方差高斯族,每步优化能不能直接解出来?
BaM(batch and match)把目标换成作者定义的 score-based divergence:在变分分布 q 下,对 target 与 q 的 score(log 密度的梯度)之差做二次加权平均,权重取 q 的协方差的逆。只需未归一化目标的 score;散度非负、两分布相等时为零、对仿射变换不变。最后这条性质来自 Cov(q) 加权,普通 Fisher 散度没有。
每轮迭代两步:
闭式解是关键设计。它不经过任何线性化,属于 stochastic proximal point 一族,这类方法对步长类超参的稳健性是公认的。三个极限便于把握:λ→0 不动;B=1 且 λ→∞ 退回 GSM;B→∞ 且 λ→∞ 对高斯目标一步到位。理论上,目标为高斯且批量无穷时,均值与协方差的归一化误差以 (1−δ)^t 指数衰减(δ=λβ/(1+λ),β 刻画初始化质量),且对任意 λ>0、任意初始化成立,不必像梯度法那样把步长绑在未知的 Hessian 最大特征值上。单步代价 O(D³),批量远小于维度时可降到 O(D²B+B³)。
基线是 ADVI(Stan 默认的 ELBO 实现)和 GSM。
| 设置 | 基线 | 结果 |
| 高斯目标,D=4/16/64/256,看 forward KL | ADVI | 梯度评估少多个数量级 |
| sinh-arcsinh 构造的非高斯目标(变偏度、变尾重) | ADVI、GSM | 收敛快于 ADVI;偏度 1.8 时 GSM 与 Score 变体发散,BaM 稳定 |
| posteriordb 三个贝叶斯模型(D=7/10/13,以 HMC 采样为参考) | ADVI、GSM | 相对均值误差更低;8-schools 上相对标准差误差反而更高 |
| CIFAR-10 深度生成模型,潜变量 D=256 | ADVI、GSM | B=300 时快一个数量级以上 |
具体误差值都在对数坐标的收敛曲线里,论文没有给数值表。CIFAR-10 实验有一个可复述的硬数字:预算 3000 次梯度评估,ADVI 的最优配置是 B=10、T=300,BaM 用 B=300、T=10 拿到相近的重构 MSE。差别落在并行性上:BaM 单轮的 B 次梯度互相独立、可并行,ADVI 的梯度基本只能串行;论文的 wallclock 曲线显示 BaM 在该实验确实更快。
对在 Stan、PyMC、Pyro 里跑 ADVI 的人,这是同一种黑盒设定(只需 log 未归一化密度及其梯度)下的替换候选:更少的梯度评估、闭式更新、免掉 SGD 的学习率试错,批内梯度天然适合 GPU 并行。这个散度还有独立用途:归一化常数未知时也可计算,能当 goodness-of-fit 诊断用,KL 做不到。边界也要说清:这是高斯变分族内部的算法升级,近似天花板仍由「目标能不能被高斯逼近」决定。
作者自己列的:收敛证明只覆盖高斯目标加无限批量,有限批量与非高斯目标都没有理论;变分族目前只有高斯;没做数据子采样,大数据量的贝叶斯推断还没碰。读下来另有多处要留意。λ 在实践里没那么省心:高斯目标用常数 λ=BD 就够,非高斯目标必须换成衰减的 BD/(t+1),等于还是要调调度。非高斯目标上 BaM 收到的 forward KL 比 ADVI 差(reverse KL 相当),偏态重的目标本来就超出高斯族的表达范围。批量要和维度相当才发力,B=10 配 D=256 表现差,单轮的内存与矩阵开方开销随之上涨。8-schools 上相对标准差误差收得更高,说明全协方差高斯不是处处占便宜。