[Paper Review] Stochastic Gradient Descent as Approximate Bayesian Inference

Stephan Mandt, Matthew D. Hoffman+1 more|arXiv (Cornell University)|Apr 13, 2017

Stochastic Gradient Optimization Techniques39 references108 citations

TL;DR

The paper reframes constant-step SGD as a stochastic process whose stationary distribution can approximate a Bayesian posterior, derives optimal SGD hyperparameters for this purpose, and extends the view to momentum, preconditioning, and SGD-based MCMC variants.

ABSTRACT

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.

Motivation & Objective

• Provide a probabilistic interpretation of constant SGD as an approximate posterior sampler.
• Derive optimal SGD hyperparameters (learning rate and preconditioning) to minimize KL divergence to the posterior.
• Show how momentum and preconditioning affect the stationary distribution for approximate inference.
• Develop variational EM and scalable MCMC viewpoints using the OU-process framework.
• Analyze implications for iterate averaging and stochastic-gradient MCMC algorithms.

Proposed method

• Model SGD with a constant learning rate as a multivariate Ornstein-Uhlenbeck process around a local optimum.
• Assume Gaussian gradient noise and a quadratic local loss to obtain analytic stationary distributions.
• Minimize KL divergence between the stationary distribution and a Gaussian posterior to derive optimal SGD settings.
• Extend to preconditioning matrices and diagonal variants to achieve better posterior matching.
• Treat SGD with momentum as a scaled covariance transformation in the OU framework for approximate inference.
• Compare constant-SGD-based posteriors to BBVI and analyze hyperparameter optimization via a variational EM perspective.

Experimental results

Research questions

• RQ1 Can constant SGD be tuned to yield an approximate Bayesian posterior distribution over parameters?
• RQ2 How should learning rate and preconditioning be chosen to minimize KL divergence to the posterior?
• RQ3 What is the effect of momentum on the stationary distribution and its use for approximate sampling?
• RQ4 How do stochastic-gradient MCMC methods (SGLD, SGFS) relate to SGD under the OU-process framework and what are their approximation errors?
• RQ5 Can iterate averaging provide optimal sampling properties within this framework?

Key findings

• The stationary distribution of constant-SGD is Gaussian and can approximate the posterior; KL divergence guides optimal hyperparameters.
• Theorem 1 gives the optimal scalar learning rate for KL-optimality: epsilon* = 2S/N · D / Tr(BB^T).
• Theorem 2 shows the optimal full preconditioner H* = (2S/N)(BB^T)^{-1} for matching the stationary distribution to the posterior; diagonal variants are also characterized.
• Momentum scales the stationary covariance but preserves its shape, enabling approximate sampling.
• For SG-MCMC methods, the OU-process view justifies preconditioning as optimal and clarifies finite-learning-rate errors; iterate averaging can yield a near-optimal sampler but implies linear cost in data passes.
• Iterate averaging can yield exactly one effectively independent sample per data pass under certain assumptions.