PaperlayerView exact revision on arXiv

Source-linked AI summary

Stochastic Gradient Descent as Approximate Bayesian Inference

Stephan Mandt, Matthew D. Hoffman, David M. Blei

arXiv:1704.04289v2stat.MLcs.LG

TL;DR

The paper asks how SGD can support efficient approximate Bayesian sampling and inference. It analyzes constant SGD as a stationary stochastic process, tunes its parameters to match posteriors, and extends the framework to variational EM, momentum, MCMC, and averaging. The resulting analysis includes an optimal averaged stochastic-gradient sampler under stated assumptions and characterizes approximation errors for related MCMC methods.

  • Problem

    The paper examines how SGD and stochastic-gradient MCMC relate, including how to obtain efficient approximate Bayesian sampling from SGD variations.

  • Method

    The authors model constant SGD through its stationary distribution and tune learning rate, minibatch size, preconditioning, momentum, and averaging to approximate Bayesian posteriors.

  • Results

    The paper derives KL-optimal tuning, a variational EM algorithm, momentum sampling adjustments, MCMC approximation analyses, and an optimal averaged sampler yielding one effectively independent sample per dataset pass under assumptions.

  • Takeaways & Limitations

    Constant SGD and its extensions can serve as approximate Bayesian inference and sampling methods by controlling the stationary distribution through optimization-algorithm parameters.

  • Takeaways & Limitations

    The finite-window iterate-averaging sampler requires assumptions that can conflict: quadratic costs imply discretization error when the step size exceeds 2/λmax, and ill-conditioned Hessians may require preconditioning.

Abstract

from arXiv · show

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.

1. Introduction

The paper asks how SGD relates to scalable Bayesian sampling and develops a stochastic-process view in which constant SGD approximates posterior inference. It extends this view to tuning, variational EM, momentum, MCMC, and iterate averaging.

  • Motivation: Constant SGD approaches an optimum and then samples around it through a stationary distribution that can approximate a Bayesian posterior.Its stationary distribution is controlled by algorithm parameters and typically mixes faster than asymptotically exact stochastic-gradient MCMC.
  • Core framework: The paper models constant SGD as a multivariate Ornstein-Uhlenbeck process with a Gaussian stationary distribution parameterized by learning rate, minibatch size, and preconditioning.This provides the basis for analyzing a broad class of SGD algorithms.
  • Approximate inference: The authors minimize KL divergence analytically between the stationary distribution and the posterior, relating optimal step sizes and preconditioners to local Hessians and noise covariances.The resulting preconditioners connect to AdaGrad, RMSProp, and Fisher scoring.
  • Variational EM: Constant SGD yields a variational EM algorithm that uses SGD to optimize hyperparameters while performing approximate Bayesian inference.The paper demonstrates the approach on Bayesian multinomial regression.
  • Momentum: Momentum changes the stationary covariance scale but not its shape, allowing momentum SGD to support approximate Bayesian inference through damping adjustment.The scaling factor is a simple function of the damping coefficient.
  • Iterate averaging: Polyak iterate averaging can produce an optimal stochastic-gradient MCMC sampler with one effectively independent sample per dataset pass under certain assumptions.The result also implies that useful-sample costs for SG-MCMC scale linearly with dataset size.

2. Related Work

The paper situates its contribution among stochastic differential equations, variational inference, scalable MCMC, and stochastic optimization. Its distinguishing focus is stationary distributions and KL matching between multivariate sampling distributions and Bayesian posteriors.

  • Scalable MCMC: SGLD adds artificial noise whose dominance increases as the step size decays, whereas this paper studies constant-SGD stationary sampling.The paper analyzes approximation behavior for stochastic-gradient Langevin dynamics and SGFS.
  • Scalable MCMC: SGFS preconditions gradients with inverse gradient-noise covariance to improve mixing, and this paper extends its analysis to finite learning rates and practical diagonal preconditioners.The paper also derives a criterion for SGFS's free noise parameter.
  • Stochastic optimization: The paper contrasts its use of constant SGD for approximate sampling with prior work that anneals the posterior over time for global optimization.Both approaches draw analogies between SGD and scalable MCMC but pursue opposite objectives.
  • Stochastic differential equations: Unlike prior SGD diffusion work focused on dynamical properties, this paper focuses on approximately converged stationary distributions.It introduces KL minimization between multivariate sampling distributions and the posterior.
  • Variational inference: The approach differs from black-box variational inference by controlling approximate-posterior shape through optimization parameters rather than parameters of a simple variational distribution.The paper reports similar practical performance to Gaussian reparameterization-gradient BBVI.

3. Continuous-Time Limit Revisited

The paper approximates constant-step SGD by a continuous-time stochastic process to characterize its behavior near a local optimum. Under Gaussian, locally constant noise and quadratic-loss assumptions, this process has an analytically tractable stationary distribution.

  • Continuous-time approximation: The analysis begins with a constant-step SGD process and replaces its discrete-time dynamics with a continuous-time stochastic differential equation.The approximation is justified when gradients or learning rates are small enough for discretization error to be negligible.
  • Assumptions: The framework assumes stochastic gradients are Gaussian and their covariance is approximately constant near a local optimum.The latter is motivated by iterates remaining in a sufficiently small region around the optimum.
  • Assumptions: The local loss is assumed to be well approximated by a quadratic with positive-definite Hessian near the optimum.This assumption is intended for a low-variance quasi-stationary distribution around a deep local minimum.
  • Ornstein-Uhlenbeck process: These assumptions yield a multivariate Ornstein-Uhlenbeck process whose stationary distribution is analytically Gaussian.The Gaussian stationary distribution is the paper's core tool for analyzing constant SGD.
  • Stationary distribution: The stationary covariance is proportional to learning rate and inversely proportional to Hessian magnitude and minibatch size.This characterizes the stationary distribution generated by constant-step SGD.
  • Regimes: Constant SGD has an early search phase, a directed approach phase, and a final phase in which sampling noise dominates near the local optimum.The formalism is expected to be most valid in the latter two regimes, especially near stationarity.
  • Caveat: If gradient-noise covariance is rank-deficient, the stationary covariance is also rank-deficient, complicating the subsequent analysis.The paper assumes full-rank noise covariance and notes that small isotropic Gaussian noise could enforce it.

4. SGD as Approximate Inference

The paper treats constant SGD as approximate Bayesian inference by tuning its stationary distribution to minimize KL divergence from a posterior. This perspective yields preconditioned and diagonal variants, hyperparameter optimization, momentum sampling, and conditions under which sampling can be exact.

  • Empirical checks: The predicted covariance of SGD iterates closely matches empirical covariance in the wine-quality regression experiment and agrees well with Figure 1’s theoretical prediction.Figure 1 compares posterior and stationary distributions for linear and logistic regression using preconditioned constant SGD, constant SGD, and BBVI.
  • Variational interpretation: Constant SGD is tuned by minimizing the KL divergence between its stationary distribution and the posterior.The loss is set to the negative log-joint distribution, and the stationary distribution depends on the learning rate, minibatch size, and preconditioning matrix.
  • Variational interpretation: Three constant-SGD variants use a scalar learning rate, a full preconditioning matrix, or a diagonal preconditioner to trade off efficiency and accuracy.Each variant generates samples from an approximate posterior.
  • Optimal tuning: Under the paper’s assumptions, the optimal learning rate is inversely proportional to average noise covariance and proportional to minibatch-to-dataset size.The learning rate is obtained by analytically minimizing the KL divergence for Gaussian posterior and stationary distributions.
  • Optimal tuning: The optimal full preconditioner is derived analytically, while restricting it to be diagonal gives an optimal diagonal preconditioner.The full inverse noise covariance yields exact posterior samples under assumptions A1-A4.
  • Extensions: Constant SGD also supports variational EM for hyperparameter optimization and an approximate-posterior sampler based on momentum.For momentum, the effective combination ϵµS determines the KL divergence, while parameter choices still affect Markov-chain dynamics.

5. Analyzing Stochastic Gradient MCMC Algorithms

The paper quantifies how finite-step stochastic-gradient MCMC methods deviate from the posterior and derives optimal preconditioning and noise choices. For SGFS, the variational analysis identifies conditions under which the stationary distribution exactly matches the posterior.

  • Stochastic Gradient Langevin Dynamics: Finite learning rates make Langevin dynamics overestimate the posterior covariance.The discrepancy vanishes as ϵ→0, when the stationary distribution becomes identical to the posterior.
  • Stochastic Gradient Fisher Scoring: The analysis rederives the Fisher-scoring preconditioner as KL-optimal for stochastic-gradient Fisher scoring.The result applies to the scalable MCMC algorithm and its practical diagonal-preconditioning variant.
  • Stochastic Gradient Fisher Scoring: The optimal SGFS preconditioner can make the KL divergence zero, yielding posterior samples without decreasing the learning rate to zero.This exact result assumes the sampling-noise distribution is estimated perfectly.
  • Practical approximations: Diagonal preconditioning makes SGFS only approximately Bayesian, but its deviation from the exact posterior can be quantified and optimized.The paper derives the optimal diagonal preconditioner from the KL divergence.
  • Practical approximations: Injecting Gaussian noise can improve stability by compensating for small gradient-noise variances and preventing degenerate behavior when gradient covariance is rank-deficient.This guidance can require adding noise in dimensions where gradient variance is too small.

6. A Bayesian View on Iterate Averaging

The paper analyzes iterate averaging as both an optimization estimator and a posterior sampler. It shows that averaging achieves preconditioner-independent covariance, can produce posterior samples under optimistic assumptions, and has important efficiency and conditioning limits.

  • Iterate Averaging for Optimization: Iterate averaging achieves the best possible stochastic-gradient convergence rate, even relative to methods using Hessian information.Consequently, constant preconditioning cannot improve its convergence speed under the stated setting.
  • Iterate Averaging for Optimization: The averaged estimator is unbiased for the objective minimum when initialization is drawn from constant SGD’s stationary distribution.Its accuracy after T iterations is characterized by a covariance computed from the OU process autocorrelation.
  • Iterate Averaging for Optimization: The estimator covariance depends on total processed examples TS, not minibatch size or step size.This establishes efficiency independence from those two tuning parameters for the iterate-averaging scheme.
  • Preconditioning: Any constant positive-definite preconditioner leaves the stationary covariance of iterate averaging unchanged.Changes to the stationary distribution and autocorrelation matrix exactly cancel.
  • Posterior Sampling: Under stronger-than-usual assumptions, iterate averaging can generate one exact posterior sample per pass through the dataset.The claimed cost is exactly N gradient calls, where N is the number of observations.
  • Posterior Sampling: Good posterior samples require dataset size to be sufficiently large relative to the Hessian condition number.When the condition number is large relative to N, the paper suggests using a preconditioner approximating A^-1.
  • Efficiency and limits: The sublinear-time lower-bound argument depends on assumptions 1–5, which can be violated despite being optimistic.The paper therefore presents the no-sublinear-sampling conclusion within that assumption set.

7. Experiments

Experiments support the paper’s stochastic-process analysis and show that constant SGD and iterate averaging can approximate posterior inference. IASG converges quickly, while constant SGD performs competitively with scalable MCMC and supports inexpensive hyperparameter selection.

  • Experiments found empirical support for the predicted stationary distributions of constant SGD with KL-optimal learning rates.
  • Constant SGD competed with the evaluated MCMC methods for posterior approximation on real-world datasets.The results suggest that approximating the preconditioning matrix matters more than artificial noise in scalable MCMC.
  • 7.1 Confirming the Stationary Distribution’s Covariance: Iterate averaging with a constant learning rate and theory-selected averaging window produced samples closely resembling the exact posterior in synthetic linear regression.The experiment used ε = 0.005, minibatch size S = 1, and an averaging window of N/S = 10^4.
  • 7.2 Iterate Averaging as Approximate MCMC: IASG found fast approximate posterior covariance estimates under both MAP and random initialization.The comparisons considered marginal variances against ground truth for IASG, SGLD, and NUTS.
  • 7.2 Iterate Averaging as Approximate MCMC: With MAP initialization, IASG converged faster than SGLD or NUTS but slightly biased the smallest variance.SGLD estimates remained highly variable even after 10,000 dataset sweeps.
  • 7.3 Optimizing Hyperparameters: Constant SGD and BBVI yielded similar validation results, making constant SGD an inexpensive alternative to cross-validation or other VEM methods for hyperparameter selection.

8. Conclusions

The conclusions recast constant-rate SGD and its extensions as approximate inference methods through a stochastic-process perspective. The paper derives KL-based tuning, a variational EM procedure, and sampling results, while identifying further exploration of algorithmic similarities as future work.

  • Under specified assumptions, SGD is approximated by a multivariate Ornstein-Uhlenbeck process with an analytically computable stationary distribution.
  • Learning rate, preconditioning, and averaging period control the stationary distribution and can be tuned to minimize its KL divergence from a Bayesian posterior.
  • The Bayesian view yields a double-SGD variational EM scheme that uses decreasing hyperparameter learning rates and constant learning rates for model parameters.The paper reports that the scheme is easy to implement and prevents degenerate solutions.
  • Future work includes comparing noise characteristics across black-box variational inference and Langevin-type MCMC and further studying iterate averaging as a Bayesian algorithm.
  • In linear regression, the SGD sampling distribution differs from the objective’s distribution because it is also affected by gradient-noise covariance.

Appendix B. Stationary Covariance

The appendix uses the analytic solution of the Ornstein-Uhlenbeck process to derive an algebraic relation for its stationary covariance. Positivity of the Hessian’s eigenvalues ensures the relevant lower-limit contribution vanishes.

  • The Ornstein-Uhlenbeck process has an analytic solution expressed through a stochastic integral.
  • Defining Σ as E[θ(t)θ(t)⊤], the appendix derives an algebraic relation for the stationary covariance of the multivariate process.
  • The lower limit of the stochastic integral vanishes because the eigenvalues of A are positive.

Appendix C. Reverse KL Divergence Setup

The appendix contrasts forward and reverse KL objectives for matching stationary and posterior distributions. Reverse KL tuning is generally impractical unless the Hessian and stationary covariance commute, whereas the momentum derivation uses stationary second moments.

  • Reverse KL divergence cannot generally be re-expressed in terms of BB⊤ through the stationary-covariance relation.
  • When the Hessian A and stationary covariance Σ commute, the reverse-KL derivation can proceed to an optimal learning-rate result.
  • Reverse KL replaces the inverse trace of noise covariance with the trace of its inverse and emphasizes different noise eigenvalues than forward KL.
  • Because A and Σ rarely commute in practice, the paper identifies KL(q||f) as the only viable option.
  • The momentum stationary distribution is derived by setting the second-moment equations to zero and accounting for the noise contribution to E[dv dv⊤].
  • The stationary position–momentum cross-correlation is antisymmetric, enabling the derivation of the momentum covariance result.

Appendix E. Stochastic Gradient Fisher Scoring

This appendix derives the KL-divergence expression for stochastic gradient Fisher scoring by eliminating the unknown stationary covariance using Lyapunov identities and covariance reparameterization.

  • The derivation starts from an Ornstein-Uhlenbeck process with minibatch size S = 1.
  • The unknown stationary covariance Σ is removed from the variational bound using the stationary covariance identity A′Σ + ΣA′⊤= B′B′⊤.The resulting criterion is written as HAΣ + ΣAH = ϵHBB⊤H + HEE⊤H.
  • Tr(AΣ) = 1 2Tr(H(ϵBB⊤+ EE⊤)) provides the key trace simplification used in the KL calculation.
  • Reparameterizing Σ = TH separates T from H, enabling simplification of the covariance-dependent criterion.
  • The resulting KL divergence is expressed in terms of trace and log-determinant terms involving H, A, T, and the noise covariance components.

Appendix F. Square Root Preconditioning

This appendix analyzes square-root preconditioning, using a matrix based on the diagonal noise covariance and choosing a scalar learning rate to minimize posterior mismatch.

  • The preconditioner is proportional to the square root of the diagonal entries of the noise covariance.
  • An additional scalar learning rate ϵ is introduced alongside the diagonal preconditioning matrix.
  • Theorem 7 gives the constant learning rate that minimizes KL divergence between the process stationary distribution and the posterior.
  • Minimizing the KL divergence over ϵ yields Eq. 47.

Appendix G. Iterage Averaging

This appendix derives the covariance of averaged iterates for an Ornstein-Uhlenbeck process by combining non-equal-time covariance identities with Hessian eigendecomposition, then simplifies the result for isotropic covariance.

  • The non-equal-time stationary covariance is derived from the formal Ornstein-Uhlenbeck solution together with Eq. 29.At equal times, the covariance reduces to E[θ(t)θ(t)⊤] = Σ.
  • The averaged-iterate covariance is defined over T time steps as the equal-time covariance of the time-averaged iterates.
  • The calculation splits into two contributions, corresponding to the two terms in the averaged covariance expression.
  • Using A = UΛU⊤, the first contribution becomes ϵT A−1Σ + 1 ϵ2T 2 UΛ−2(e−ϵTΛ −I)U⊤Σ.
  • The second contribution is the analogous expression with Σ placed on the left, and combining both terms yields the full covariance formula.
  • For Σ = ϵ 2S I, the result simplifies to D = 1 ST A−1 + 1 ϵST 2 (UΛ−2(e−ϵTΛ −I)U⊤), matching Eq. 33.
Loading 1704.04289v2…