Source-linked AI summary
A Variational Analysis of Stochastic Gradient Algorithms
Stephan Mandt, Matthew D. Hoffman, David M. Blei
TL;DR
The paper asks whether constant-rate SGD can provide useful approximate posterior inference instead of merely optimization. It models SGD as a multivariate Ornstein-Uhlenbeck process, tunes its parameters by minimizing KL divergence to the posterior, and extends the analysis to scalable Fisher scoring and hyperparameter optimization. The resulting framework gives constant-rate SGD a variational Bayesian interpretation and supports posterior approximation alongside hyperparameter learning.
Problem
The paper addresses whether the simplest constant-rate SGD can serve as a sensible approximate Bayesian inference method rather than only an optimization algorithm.
Method
The paper models constant SGD as a multivariate Ornstein-Uhlenbeck process and tunes learning rates and preconditioners by minimizing KL divergence to the posterior.
Results
The analysis derives optimal constant-SGD parameters, shows that SGFS’s preconditioner is optimal, and analyzes diagonal preconditioning and numerical stability.
Takeaways & Limitations
Constant-rate SGD can provide approximate posterior samples while also supporting Bayesian hyperparameter optimization.
Takeaways & Limitations
The theoretical approximation assumes locally stable, low-variance behavior and can become only approximately Bayesian when SGFS uses a diagonal preconditioner.
Abstract
from arXiv · showhide
Stochastic Gradient Descent (SGD) is an important algorithm in machine learning. With constant learning rates, it is a stochastic process that, after an initial phase of convergence, generates samples from a stationary distribution. We show that SGD with constant rates can be effectively used as an approximate posterior inference algorithm for probabilistic modeling. Specifically, we show how to adjust the tuning parameters of SGD such as to match the resulting stationary distribution to the posterior. This analysis rests on interpreting SGD as a continuous-time stochastic process and then minimizing the Kullback-Leibler divergence between its stationary distribution and the target posterior. (This is in the spirit of variational inference.) In more detail, we model SGD as a multivariate Ornstein-Uhlenbeck process and then use properties of this process to derive the optimal parameters. This theoretical framework also connects SGD to modern scalable inference algorithms; we analyze the recently proposed stochastic gradient Fisher scoring under this perspective. We demonstrate that SGD with properly chosen constant rates gives a new way to optimize hyperparameters in probabilistic models.
1. Introduction
The paper interprets constant-step-size SGD as a stochastic process whose stationary distribution can approximate a posterior. It tunes learning rates, preconditioners, and minibatch-related parameters by minimizing KL divergence, extending the perspective to Fisher scoring and hyperparameter optimization.
- Motivation: Constant SGD converges toward an optimum and then samples a stationary distribution around it because of persistent gradient noise.Unlike traditional SGD, it retains a constant learning rate rather than decreasing it.
- Variational view: The stationary distribution depends on the learning rate, preconditioning matrix, and minibatch size, allowing these parameters to be tuned toward the posterior.The paper minimizes the Kullback-Leibler divergence between the stationary distribution and the exact posterior.
- Variational view: A multivariate Ornstein-Uhlenbeck interpretation yields Gaussian stationary distributions for a broad class of SGD algorithms.The stationary covariance is non-trivial and is parameterized by algorithmic tuning choices.
- Scalable inference: The analysis covers stochastic gradient Fisher scoring, showing that its suggested preconditioning matrix is optimal and deriving a stability criterion for its free noise parameter.It also analyzes the effect of replacing the preconditioner with a diagonal approximation.
- Hyperparameter optimization: Constant-rate SGD can simultaneously support posterior inference and optimization of model hyperparameters in Bayesian multinomial logistic regression.The demonstration uses normal priors and compares this approach with stochastic variational methods.
2. Continuous-Time Limit Revisited
The paper approximates discrete SGD with a continuous-time stochastic process under assumptions about gradient noise, local curvature, and discretization. This produces a multivariate Ornstein-Uhlenbeck process with an analytically tractable Gaussian stationary distribution.
- Problem setup: The analysis begins by expressing minibatch-based stochastic gradients as unbiased estimates of the full gradient.A minibatch consists of uniformly sampled loss contributions.
- Continuous-time approximation: The discrete SGD updates are approximated by a continuous-time stochastic differential equation that is easier to analyze.The approximation neglects discretization errors under an explicit assumption.
- Assumptions: The Ornstein-Uhlenbeck approximation assumes Gaussian gradient noise, approximately constant noise covariance, and locally quadratic loss behavior.The noise-covariance assumption is justified when iterates remain near a local optimum.
- Ornstein-Uhlenbeck process: The resulting process has an analytic Gaussian stationary distribution that serves as the paper’s core tool for analyzing SGD.This makes stationary distributions and KL divergences tractable.
- Stationary behavior: The stationary covariance increases with learning rate and decreases with curvature magnitude and minibatch size.The relationship characterizes the stationary distribution generated by constant-step-size SGD.
3. SGD as Approximate Inference
The paper tunes constant SGD so its stationary distribution approximates a Bayesian posterior by minimizing KL divergence, yielding several approximate-inference variants. The same variational perspective identifies optimal preconditioners for stochastic gradient Fisher scoring, quantifies diagonal-approximation bias, and supports gradient-based hyperparameter learning.
- Empirical support: In illustrative linear- and logistic-regression examples, empirical and predicted iterate covariances closely match, supporting the modeling assumptions.The figures use parameters selected by minimizing KL divergence between the OU stationary distribution and the posterior.
- Constant SGD: SGD parameters can be tuned to minimize KL divergence between its stationary distribution and the posterior, turning constant SGD into a posterior inference algorithm.The tunable parameters include learning rate, minibatch size, and optionally a preconditioning matrix.
- Constant SGD: The paper proposes constant-step, full-preconditioned, and diagonal-preconditioned SGD variants that trade efficiency and accuracy while generating approximate-posterior samples.The diagonal restriction is motivated by the impracticality of large dense matrices in high-dimensional applications.
- Optimal tuning: The optimal constant learning rate and preconditioners are derived analytically from the KL objective under assumptions including an approximately Gaussian posterior.The results connect tuning to Hessian and noise-covariance structure, while the diagonal preconditioner is the inverse of the noise matrix’s diagonal part.
- Stochastic Gradient Fisher Scoring: Stochastic gradient Fisher scoring is recovered as a KL-optimal preconditioning choice, and its exact idealized form has zero KL divergence from the posterior.The paper also analyzes diagonal or scalar approximations and quantifies the bias caused by truncating the preconditioner.
- Hyperparameter optimization: The stationary-distribution interpretation justifies jointly optimizing model parameters and hyperparameters as a variational expectation-maximization procedure.The paper presents gradient-based hyperparameter learning in constant SGD as a cheaper alternative to cross-validation.
4. Experiments
The experiments test constant SGD’s stationary-distribution predictions, compare its posterior approximation with alternative inference methods, and evaluate gradient-based hyperparameter learning.
- 4.1. Stationary distributions: The experiments test whether constant SGD’s stationary distributions match the Ornstein-Uhlenbeck predictions.The study compares empirical and predicted covariances and reports good evidence supporting the theoretical assumptions.
- 4.1. Stationary distributions: The evaluation uses linear and logistic regression across wine quality, protein structure, and skin segmentation data sets.The data sets contain 4,898, 45,730, and 245,057 instances, respectively.
- 4.1. Stationary distributions: Constant SGD’s posterior approximation is not worse than BBVI’s in the reported comparisons.Posterior and stationary-distribution samples are compared through projections and covariance assessments.
- 4.1. Comparison with inference methods: Constant SGD can compete with the evaluated stochastic-gradient MCMC methods in approximating the posterior.The comparison includes stochastic gradient Langevin dynamics and stochastic gradient Fisher scoring, with and without diagonal approximation.
- 4.2. Hyperparameter optimization: The results suggest constant SGD can provide an inexpensive alternative to cross-validation or other VEM methods for hyperparameter selection.The experiment uses Bayesian multinomial logistic regression with normal priors and evaluates validation loss across L2 regularization values.
5. Related Work
The related work situates this paper between scalable MCMC, stochastic optimization, and stochastic-differential-equation analyses. Its distinct focus is stationary distributions and KL-based tuning of SGD to approximate posteriors.
- Scalable MCMC: Unlike scalable MCMC methods that target asymptotically exact posterior samples, this work tunes constant SGD to approximate the posterior with its stationary distribution.The comparison emphasizes approximate inference rather than exact long-run sampling.
- Variational inference: The paper focuses on stationary distributions from constant learning rates, whereas related variational work mainly studies nonstationary sampling distributions.The related work also distinguishes the paper’s model hyperparameters from SGD’s optimization hyperparameters.
- Stochastic optimization: The paper derives optimal learning-rate and preconditioning criteria through a different argument from earlier optimal-control results.It also introduces a scalar prefactor for the optimal learning rate.
- Scalable MCMC: Where some scalable-MCMC work anneals posteriors for global optimization, this paper uses constant SGD as an approximate sampler.The distinction is the direction of the optimization–sampling connection.
- Stochastic differential equations: The paper applies stochastic differential equations to stationary-distribution analysis rather than primarily to SGD’s dynamical or convergence properties.The related literature includes analyses of momentum, convergence, and stochastic-gradient MCMC.
6. Conclusions
The paper reframes constant-rate SGD as an approximate Bayesian inference method by optimizing the divergence between its stationary distribution and the posterior. It derives tuning criteria and reports empirical evidence that the approach supports posterior approximation and hyperparameter optimization.
- Conclusions: The paper chooses constant learning rates to minimize KL divergence between SGD’s stationary distribution and the posterior.This replaces the conventional goal of decreasing the learning rate to zero.
- Conclusions: The analysis yields criteria for selecting learning rates, minibatch sizes, and preconditioning matrices.These criteria arise from approximating SGD with a continuous-time Ornstein-Uhlenbeck process.
- Conclusions: The paper presents an analysis of stochastic gradient Fisher scoring and finds that simple SGD can compete with stochastic variational methods for empirical Bayesian hyperparameter optimization.The conclusion frames this as an empirical result of the demonstrated approach.
A. Stationary Covariance
The appendix derives the stationary covariance of the multivariate Ornstein-Uhlenbeck process from its formal solution. The derivation relies on the vanishing lower integral limit under a positivity condition on the drift matrix.
- Stationary covariance: The stationary covariance is defined as Σ = E[θ(t)θ(t)⊤] and derived from the formal Ornstein-Uhlenbeck solution.The derivation produces an algebraic relation for the covariance matrix.
- Stationary covariance: The lower limit of the stochastic integral vanishes because the eigenvalues of A are positive.This positivity condition is used to obtain the stationary covariance relation.
B. Stochastic Gradient Fisher Scoring
The section formulates the stochastic-gradient analysis through an Ornstein–Uhlenbeck process and derives a variational criterion involving its stationary covariance. The covariance is reparameterized and the resulting expression is identified as the main-paper result.
- The analysis starts from an Ornstein–Uhlenbeck process.
- The variational bound is expressed using redefined matrices A′ and B′.
- The unknown stationary covariance Σ is eliminated using the identity A′Σ + ΣA′⊤ = B′B′⊤.
- The covariance is reparameterized as Σ = TH before evaluating the KL divergence.
- The resulting expression is reported as the result given in the main paper.
C. Square root preconditioning
This section analyzes square-root preconditioning, using a diagonal noise-covariance transformation and an additional scalar learning rate. It states that the KL-minimizing constant learning rate is obtained from the corresponding divergence.
- The method preconditions with a matrix proportional to the square root of the noise covariance’s diagonal entries.
- The preconditioner is diagonal, containing square roots of the noise covariance’s diagonal elements, and uses an additional scalar learning rate ϵ.
- The theorem gives the constant learning rate that minimizes KL divergence between the process’s stationary distribution and the posterior.
- The proof obtains the relevant KL divergence from Theorem 2 by setting G−1 ≡ H.
- Minimizing the KL divergence over ϵ yields Eq. 34.