Natural Gradient Descent for Improving Variational Inference Based Classification of Radio Galaxies¶

会议: NeurIPS 2025
arXiv: 2511.13224
代码: 有
领域: 贝叶斯深度学习 / 射电天文
关键词: 自然梯度下降, 变分推断, 贝叶斯神经网络, 不确定性校准, 射电星系分类

一句话总结¶

研究使用自然梯度下降优化器 iVON 替代标准 SGD 来优化变分推断中的 BNN 参数，在射电星系分类中获得更好的不确定性校准，同时保持与 HMC 和 BBB-VI 相当的预测性能。

未来射电天文巡天预计产出 EB 级数据，需要统计上鲁棒的 ML 模型。贝叶斯神经网络（BNN） 提供原则性的不确定性建模方式。

前期基准工作发现： - HMC（哈密顿蒙特卡洛）在预测精度、校准和分布偏移检测上整体最佳，但计算成本极高（7天） - BBB-VI（Bayes by Backprop）性能尚可但面临初始化敏感、收敛慢、冷后验效应等问题

核心问题：标准 VI 使用 SGD 优化变分参数，但 VI 的参数空间是统计流形（黎曼流形），每个点对应一个概率分布。SGD 假设欧几里得空间，可能不是在分布空间中最高效的优化方向。

自然梯度下降（NGD） 用 Fisher 信息矩阵（FIM）的逆对梯度进行预条件化，账户了统计流形的几何结构，提供更直接的分布空间路径。

使用 iVON（Improved Variational Online Newton） 算法替代 SGD 优化 BNN 的变分参数。iVON 基于贝叶斯学习规则（BLR）框架，将变分推断与自然梯度下降统一。

BLR 框架下的自然梯度更新：
- 变分分布选择多元高斯 \(q(\boldsymbol{\theta}) = \mathcal{N}(\boldsymbol{\theta}|\mathbf{m}, \mathbf{S}^{-1})\)
- 自然梯度更新自然参数：\(\boldsymbol{\lambda}_{t+1} \leftarrow \boldsymbol{\lambda}_t - \alpha \nabla_{\boldsymbol{\mu}}\{\mathbb{E}_{q}[l(\boldsymbol{\theta})] - \mathcal{H}(q)\}\)
- 等价于类牛顿法更新，需要梯度和 Hessian 信息
iVON 的可扩展近似：
- 对角 Hessian 近似：降低计算从 \(O(d^2)\) 到 \(O(d)\)
- 重参数化技巧估计二阶信息：\(\hat{\mathbf{h}} = \hat{\mathbf{g}} \cdot (\boldsymbol{\theta} - \mathbf{m}) / \boldsymbol{\sigma}^2\)——通过测量梯度对参数随机扰动的响应来近似曲率
- 几何修正项：保证精度矩阵正定性，确保变分分布在训练全程有效
- 均值和标准差更新：\(\mathbf{m} \leftarrow \mathbf{m} - \alpha \frac{(\hat{\mathbf{g}} + \delta\mathbf{m})}{(\mathbf{h} + \delta)}\)，\(\boldsymbol{\sigma} \leftarrow \frac{1}{\sqrt{\text{ess}(\mathbf{h} + \delta)}}\)
与 BBB-VI 的关键区别：
- BBB-VI 在欧几里得空间分别对均值和方差计算梯度：\(\mathbf{m} \leftarrow \mathbf{m} - \alpha \nabla_\mathbf{m} \mathcal{L}\)
- iVON 使用自然梯度，分母中的曲率估计 \((\mathbf{h} + \delta)\) 提供自适应步长
- iVON 使用1个 MC 样本即可，与 BBB 一致

推断方法	测试错误率 ↓	UCE ↓	训练时间
HMC	4.16 ± 0.45	14.76 ± 0.95	7 天
BBB-VI	3.94 ± 0.01	12.77 ± 6.11	40 分钟
iVON (ess=10N)	3.07 ± 1.47	8.37 ± 4.12	25 分钟
iVON (ess=100N)	3.36 ± 1.23	12.19 ± 6.57	25 分钟

数据集	HMC	BBB-VI	iVON
MiraBest (iD)	低能量 ✓	低能量 ✓	低能量 ✓
GalaxyMNIST (远 OoD, 光学)	高能量 ✓	高能量 ✓	高能量 ✓
MIGHTEE (近 OoD, 不同射电望远镜)	高能量 ✓	中等区分 ✓	无法可靠区分 ✗

优化器作为归纳偏置：本文揭示了一个重要观察——优化器的选择不仅影响收敛速度，还决定了模型学到的表示类型（分布式/冗余 vs 压缩/局部），从而影响不同下游任务的表现
自然梯度下降利用参数空间的黎曼几何，为 VI 优化提供了更「自然」的方向
实验虽小但分析深入，每个发现都有具体的物理应用意义