[Paper Review] 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 Algorithms4 citations

TL;DR

This paper proves that minibatch natural gradient descent (NGD) asymptotically samples from a Bayesian posterior in function space when the model predictions approach the true conditional distribution, with temperature $ T \approx \epsilon N / (2B) $. It introduces stochastic NGD, a novel optimizer that corrects for parameterization dependence by incorporating the Jeffreys prior, enabling valid posterior sampling throughout parameter space.

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.

Motivation & Objective

• To establish a theoretical correspondence between minibatch natural gradient descent and Bayesian posterior sampling in function space.
• To identify the conditions under which NGD's stationary distribution approximates a Bayesian posterior near local minima.
• To resolve the lack of parameterization invariance in standard minibatch NGD by deriving a corrected update rule.
• To propose and validate a new optimizer, stochastic NGD, that preserves both parameterization invariance and valid posterior sampling.
• To empirically verify that optimal generalization in NGD occurs at $ B \approx \epsilon N / 2 $, aligning with Bayesian predictions.

Proposed method

• Derives the stationary distribution of minibatch NGD in the $ \epsilon \to 0 $ limit, showing it converges to a Bayesian posterior at temperature $ T = \epsilon N / (2B) $ when model predictions approach the true conditional distribution.
• Identifies that standard minibatch NGD lacks parameterization invariance due to the parameter dependence of the Fisher information matrix.
• Proposes stochastic NGD (SNGD), which adds correction terms derived from the derivative of the Fisher information matrix to restore invariance.
• Introduces a multiplicative bias term equivalent to the Jeffreys prior, ensuring the stationary distribution remains a valid posterior across parameter space.
• Uses a moving average of the Fisher information with adaptive smoothing to reduce variance during training.
• Employs ensembling over final gradient updates to estimate posterior predictive distributions and evaluate generalization performance.

Experimental results

Research questions

• RQ1 Under what conditions does minibatch NGD asymptotically sample from a Bayesian posterior in function space?
• RQ2 How does the temperature $ T \approx \epsilon N / (2B) $ control the stationary distribution of NGD?
• RQ3 Why does standard minibatch NGD fail to be parameterization invariant, and how can this be corrected?
• RQ4 Can a modified NGD update rule preserve both parameterization invariance and valid posterior sampling with stochastic gradients?
• RQ5 Does the optimal batch size for NGD generalization align with Bayesian predictions at $ B \approx \epsilon N / 2 $?

Key findings

• When $ f_{\omega}(x_i) \to P(Y|x_i) $, the stationary distribution of minibatch NGD approaches a Bayesian posterior at temperature $ T \approx \epsilon N / (2B) $ near local minima.
• Empirical results show that the test cross entropy of NGD ensembles is minimized when $ B \approx \epsilon N / 2 $, confirming the predicted optimal ratio.
• For $ N=4096 $, the minimum test cross entropy occurs at $ B=256 $, matching the Bayesian prediction of $ B \propto N $.
• As $ N $ increases to 16384, the minimum in test cross entropy becomes flatter, consistent with theoretical expectations.
• In CNN experiments on MNIST, test cross entropy is largely independent of learning rate at constant $ T $, confirming the linear scaling rule $ B \propto \epsilon $.
• The test cross entropy rises rapidly when $ T \gtrsim 1 $, indicating that generalization degrades beyond this temperature, consistent with Bayesian theory.