Stochastic Gradient Descent as Approximate Bayesian Inference
This work provides a theoretical foundation for using SGD in Bayesian inference, potentially benefiting researchers and practitioners in machine learning who need scalable approximate inference methods, though it is incremental in extending existing SGD perspectives.
The paper tackles the problem of interpreting Stochastic Gradient Descent (SGD) as approximate Bayesian inference by showing that constant SGD simulates a Markov chain with a stationary distribution, and it derives methods to adjust tuning parameters to match this distribution to a posterior, minimizing Kullback-Leibler divergence, and proposes scalable algorithms like the Averaged Stochastic Gradient Sampler.
Stochastic Gradient Descent with a constant learning rate (constant SGD) simulates a Markov chain with a stationary distribution. With this perspective, we derive several new results. (1) We show that constant SGD can be used as an approximate Bayesian posterior inference algorithm. Specifically, we show how to adjust the tuning parameters of constant SGD to best match the stationary distribution to a posterior, minimizing the Kullback-Leibler divergence between these two distributions. (2) We demonstrate that constant SGD gives rise to a new variational EM algorithm that optimizes hyperparameters in complex probabilistic models. (3) We also propose SGD with momentum for sampling and show how to adjust the damping coefficient accordingly. (4) We analyze MCMC algorithms. For Langevin Dynamics and Stochastic Gradient Fisher Scoring, we quantify the approximation errors due to finite learning rates. Finally (5), we use the stochastic process perspective to give a short proof of why Polyak averaging is optimal. Based on this idea, we propose a scalable approximate MCMC algorithm, the Averaged Stochastic Gradient Sampler.