Stochastic natural gradient descent draws posterior samples in function space

By Samuel L. Smith, Daniel Duckworth, Quoc V. Le, Jascha Narain Sohl-Dickstein. 2018

8 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Statistics and Probability, Probabilistic & Bayesian Methods

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Yasaman Bahri
2. Ethan Dyer
3. Guy Gur-Ari
4. Jascha Narain Sohl-Dickstein
5. Greg Yang
6. Laurence Aitchison
7. Aitor Lewkowycz
8. Jaehoon Lee

In Statistics and Probability

1. Yasaman Bahri
2. Jascha Narain Sohl-Dickstein
3. Laurence Aitchison
4. Ethan Dyer
5. Guy Gur-Ari
6. Boris Hanin
7. Mihai Nica
8. Jaehoon Lee

In Probabilistic & Bayesian Methods

1. Ethan Dyer
2. Guy Gur-Ari
3. Laurence Aitchison
4. Jascha Narain Sohl-Dickstein
5. Yasaman Bahri
6. Greg Yang
7. Stefano Ermon
8. Yang Song

Citations in the corpus

• The large learning rate phase of deep learning: the catapult mechanism (2020, background): “According to this suggestion, training with a large learning rate (or with a small batch size) can improve performance because it leads to more stochasticity during training (Mandt et al., 2017; Smith et al., 2017; Smith & Le, 2018; Smith et al., 2018).”
• A unified theory of adaptive stochastic gradient descent as Bayesian filtering (2018, methodology): “…approach that works well in practice in natural-gradient methods (Amari, 1998; Martens & Grosse, 2015; Grosse & Martens, 2016; Song & Ermon, 2018; Smith et al., 2018) is to use a Fisher-Information based estimate, where Importantly, this Fisher-Information based estimate does not depend on…”