[论文解读] Stochastic natural gradient descent draws posterior samples in function space

Samuel Smith, Daniel Duckworth+3 more|arXiv (Cornell University)|Jun 25, 2018

Machine Learning and Algorithms被引用 4

一句话总结

本文证明了当模型预测趋近于真实条件分布时，小批量自然梯度下降（NGD）在函数空间中渐近地从贝叶斯后验分布中采样，温度为 $ T \approx \epsilon N / (2B) $。本文提出了随机NGD，一种新型优化器，通过引入Jeffreys先验来校正参数化依赖性，从而在整个参数空间中实现有效的后验采样。

ABSTRACT

Recent work has argued that stochastic gradient descent can approximate the Bayesian uncertainty in model parameters near local minima. In this work we develop a similar correspondence for minibatch natural gradient descent (NGD). We prove that for sufficiently small learning rates, if the model predictions on the training set approach the true conditional distribution of labels given inputs, the stationary distribution of minibatch NGD approaches a Bayesian posterior near local minima. The temperature $T = \\epsilon N / (2B)$ is controlled by the learning rate $\\epsilon$, training set size $N$ and batch size $B$. However minibatch NGD is not parameterisation invariant and it does not sample a valid posterior away from local minima. We therefore propose a novel optimiser, "stochastic NGD", which introduces the additional correction terms required to preserve both properties.

研究动机与目标

• 建立小批量自然梯度下降与函数空间中贝叶斯后验采样之间的理论对应关系。
• 确定NGD的平稳分布何时在局部极小值附近近似为贝叶斯后验的条件。
• 通过推导校正更新规则，解决标准小批量NGD中缺乏参数化不变性的问题。
• 提出并验证一种新优化器——随机NGD，其同时保持参数化不变性和有效的后验采样。
• 通过实证验证，NGD最优泛化发生在 $ B \approx \epsilon N / 2 $ 时，与贝叶斯预测一致。

提出的方法

• 在 $ \epsilon \to 0 $ 极限下推导小批量NGD的平稳分布，表明当模型预测趋近真实条件分布时，其收敛至温度为 $ T = \epsilon N / (2B) $ 的贝叶斯后验。
• 识别出标准小批量NGD由于Fisher信息矩阵的参数依赖性而缺乏参数化不变性。
• 提出随机NGD（SNGD），通过引入源自Fisher信息矩阵导数的校正项以恢复不变性。
• 引入等价于Jeffreys先验的乘法偏差项，确保平稳分布在整个参数空间中保持为有效的后验分布。
• 使用自适应平滑的Fisher信息矩阵移动平均，以在训练过程中降低方差。
• 通过最终梯度更新的集成来估计后验预测分布，并评估泛化性能。

实验结果

研究问题

• RQ1 在何种条件下，小批量NGD渐近地从函数空间的贝叶斯后验分布中采样？
• RQ2 温度 $ T \approx \epsilon N / (2B) $ 如何控制NGD的平稳分布？
• RQ3 为何标准小批量NGD无法实现参数化不变性，以及如何加以修正？
• RQ4 能否通过修改NGD更新规则，同时保持参数化不变性和有效的后验采样？
• RQ5 NGD最优泛化所需的批量大小是否与贝叶斯预测一致，即 $ B \approx \epsilon N / 2 $？

主要发现

• 当 $ f_{\omega}(x_i) \to P(Y|x_i) $ 时，小批量NGD在局部极小值附近的平稳分布趋近于温度为 $ T \approx \epsilon N / (2B) $ 的贝叶斯后验。
• 实证结果表明，NGD集成的测试交叉熵在 $ B \approx \epsilon N / 2 $ 时最小化，证实了预测的最优比例。
• 对于 $ N=4096 $，测试交叉熵的最小值出现在 $ B=256 $，与贝叶斯预测 $ B \propto N $ 一致。
• 当 $ N $ 增加至 16384 时，测试交叉熵的最小值区域变得更为平坦，与理论预期一致。
• 在MNIST上的CNN实验中，当温度 $ T $ 保持恒定时，测试交叉熵基本不受学习率影响，证实了线性缩放规则 $ B \propto \epsilon $。
• 当 $ T \gtrsim 1 $ 时，测试交叉熵迅速上升，表明泛化性能在此温度以上显著下降，与贝叶斯理论一致。