PaperlayerView exact revision on arXiv

Source-linked AI summary

A Simple Baseline for Bayesian Uncertainty in Deep Learning

Wesley Maddox, Timur Garipov, Pavel Izmailov, Dmitry Vetrov, Andrew Gordon Wilson

arXiv:1902.02476v2cs.LGcs.AIcs.CVstat.ML

TL;DR

Deep learning models often lack calibrated uncertainty, while Bayesian methods can be difficult to scale to modern architectures. SWAG uses SGD trajectories to fit a low-rank-plus-diagonal Gaussian posterior around the SWA solution and samples it for Bayesian model averaging. The method captures posterior geometry approximately and performs well across calibration, out-of-domain detection, transfer learning, and language modeling tasks.

  • Problem

    Deep learning models can produce overconfident, miscalibrated predictions, while existing Bayesian approaches are often hyperparameter-sensitive and difficult to scale.

  • Method

    SWAG fits a Gaussian weight-posterior approximation using the SWA mean and a low-rank-plus-diagonal covariance derived from SGD iterates.

  • Results

    SWAG captures posterior geometry approximately and provides well-calibrated uncertainty across computer-vision settings, out-of-domain detection, transfer learning, and language modeling.

  • Takeaways & Limitations

    SWAG is presented as a practical, scalable step toward accurate Bayesian deep learning for large modern neural networks.

  • Takeaways & Limitations

    SWAG requires SGD with weight decay or explicit L2 regularization, and may underestimate variance along random directions outside the SGD trajectory subspace.

Abstract

from arXiv · show

We propose SWA-Gaussian (SWAG), a simple, scalable, and general purpose approach for uncertainty representation and calibration in deep learning. Stochastic Weight Averaging (SWA), which computes the first moment of stochastic gradient descent (SGD) iterates with a modified learning rate schedule, has recently been shown to improve generalization in deep learning. With SWAG, we fit a Gaussian using the SWA solution as the first moment and a low rank plus diagonal covariance also derived from the SGD iterates, forming an approximate posterior distribution over neural network weights; we then sample from this Gaussian distribution to perform Bayesian model averaging. We empirically find that SWAG approximates the shape of the true posterior, in accordance with results describing the stationary distribution of SGD iterates. Moreover, we demonstrate that SWAG performs well on a wide variety of tasks, including out of sample detection, calibration, and transfer learning, in comparison to many popular alternatives including MC dropout, KFAC Laplace, SGLD, and temperature scaling.

1 Introduction

Deep learning models often produce overconfident, miscalibrated predictions, while existing Bayesian approaches can be difficult to scale. SWAG uses SGD trajectories to construct an efficient Gaussian posterior approximation and achieves well-calibrated predictions across diverse tasks.

  • Uncertainty representation matters because machine-learning models support decisions where rare mistakes can be costly.The paper highlights medical diagnosis and autonomous vehicles as examples.
  • Existing Bayesian deep-learning methods provide probabilistic uncertainty but are often sensitive to hyperparameters and difficult to scale.
  • SWAG approximates the neural-network weight posterior using information contained in the SGD trajectory.
  • SWAG combines the SWA mean with a low-rank-plus-diagonal covariance of SGD iterates to define a Gaussian posterior over weights.SWA averages SGD iterates under a high constant learning-rate schedule.
  • The SGD-iterate Gaussian captures local posterior geometry approximately well within the low-dimensional subspace spanned by the trajectory.The paper motivates this view through analyses of the stationary distribution of SGD iterates, while noting that their assumptions do not hold for deep neural networks.
  • SWAG provides well-calibrated uncertainty across computer-vision settings and achieves higher test likelihood than several alternatives across CIFAR-10, CIFAR-100, and ImageNet.Compared methods include MC-Dropout, temperature scaling, SGLD, KFAC-Laplace, and SWA.

2 Related Work

Bayesian deep learning represents uncertainty by averaging predictions over distributions of neural-network parameters, but modern networks make exact inference difficult. Related approaches include MCMC, variational, dropout, Laplace, SGD-based, and calibration methods with different scalability and approximation strategies.

  • Bayesian model averaging places a distribution over model parameters and marginalizes them to form a predictive distribution.
  • Modern neural networks make Bayesian inference difficult because they have millions of parameters, highly non-convex posteriors, and often require mini-batch optimization.
  • MCMC: HMC provides a gold-standard MCMC approach but requires full gradients, making it computationally intractable for modern neural networks.SGHMC extends the framework to stochastic gradients for scalability and exploration.
  • Variational Inference: Variational inference fits Gaussian posterior approximations, while recent methods often focus on smaller-scale datasets and architectures.Noisy optimization methods such as noisy Adam and noisy KFAC provide another variational-inference line of work.
  • Dropout Variational Inference: Dropout variational inference interprets test-time dropout predictions as approximate Bayesian inference and requires ensembling dropout predictions at test time.
  • Laplace Approximations: Laplace approximations assume a Gaussian posterior around a MAP estimate, with practical variants using diagonal or KFAC curvature approximations.
  • SGD-Based Approximations: SGD-based approximations use SGD iterates as samplers or estimate their covariance to construct Gaussian approximations with calibrated parameter confidence intervals under stated conditions.
  • Calibration: Temperature scaling rescales neural-network logits using a validation set and one hyperparameter to improve calibration.Related work also includes ensembles and calibrated regression based on rescaling.

3 SWA-Gaussian for Bayesian Deep Learning

SWAG approximates the posterior over neural-network weights by combining SWA’s mean with diagonal and low-rank covariance estimates derived from SGD iterates. It samples this Gaussian approximation for Bayesian model averaging, while remaining efficient and applicable under standard SGD-based training.

  • Posterior construction: SWAG uses the SWA solution as the mean of a Gaussian posterior approximation over neural-network weights.SWA averages weights traversed by SGD under a constant learning-rate schedule.
  • Posterior construction: The diagonal approximation maintains running averages of each weight’s second uncentered moment and forms ΣDiag using a standard covariance identity.The resulting distribution is N(θSWA, ΣDiag), called SWAG-Diagonal.
  • Low-rank covariance: SWAG adds a low-rank covariance estimate formed from recent SGD-iterate deviations to make the posterior approximation more flexible than a diagonal covariance.The covariance estimate uses a deviation matrix and retains the last K deviation vectors, producing a rank-K approximation.
  • Low-rank covariance: The final SWAG posterior is Gaussian with mean θSWA and covariance equal to the sum of diagonal and low-rank components.Sampling uses stored deviation vectors and can be performed without storing those models on a GPU.
  • Bayesian model averaging: Bayesian model averaging samples weights from the SWAG Gaussian, updates batch-normalization statistics when needed, and averages predictive distributions across samples.This contrasts with MAP prediction, which uses only one parameter setting rather than marginalizing over the posterior.
  • Scope and assumptions: SWAG is applicable to most deep neural networks without modifying training when SGD uses weight decay or explicit L2 regularization.The paper notes that using Adam and other stochastic first-order methods remains a possible direction for future work.

4 Does the SGD Trajectory Capture Loss Geometry?

The SGD trajectory captures posterior geometry within the low-dimensional subspace spanned by its iterates, but SWAG variance estimates are not guaranteed to be adequate across the full weight space.

  • SWAG analyzes posterior geometry along covariance eigenvectors using the ℓ2-regularized cross-entropy loss around the SWA solution.The analysis evaluates loss as a function of distance along directions defined by the SWAG covariance.
  • SWAG captures posterior geometry in the subspace spanned by SGD iterates.
  • The SGD-iterate subspace is very low-dimensional relative to the full neural-network weight space.
  • SWAG variance estimates are not guaranteed to be adequate along all directions in weight space.Because SGD explores a low-dimensional subspace, variance may be underestimated along random directions.

5 Experiments

Experiments evaluate SWAG uncertainty, calibration, accuracy, transfer learning, out-of-domain detection, language modeling, and regression against multiple approximate Bayesian and calibration baselines. Across these settings, SWAG generally provides strong likelihood, calibration, and perplexity results, with particular advantages when test distributions differ from training distributions.

  • Experiments compare SWAG and SWAG-Diagonal with MC dropout, temperature scaling, SGLD, Laplace approximations, deep ensembles, SGD-iterate ensembles, and SWA.
  • Negative log-likelihood evaluates both predictive accuracy and uncertainty quality, while reliability diagrams compare confidence with accuracy.Reliability diagrams bin test examples by maximum predicted probability and compare mean confidence with accuracy.
  • 5.1 Calibration and Uncertainty Estimation on Image Classification Tasks: SWAG and SWAG-Diagonal perform comparably or better than alternatives across considered tasks, with SWAG best overall.SWA with temperature scaling is competitive but requires a validation set and is ineffective under distribution shift in CIFAR-10-to-STL-10 transfer.
  • 5.1 Calibration and Uncertainty Estimation on Image Classification Tasks: SWAG and SWAG-Diagonal achieve good calibration across evaluated image-classification settings, while low-rank plus diagonal SWAG is generally better calibrated than SWAG-Diagonal.SWAG is comparable with SWA in predictive accuracy and generally outperforms other approaches.
  • 5.2 Comparison to ensembling SGD solutions: An ensemble of 5 SGD solutions reaches NLL 0.6478, competitive with a single SWAG solution requiring 5× less computation to train.An ensemble of 3 SWAG models achieves NLL 0.6178.
  • 5.4 Language Modeling with LSTMs: SWAG improves out-of-domain detection, transfer learning, and language-modeling performance, including substantially better test perplexities than SWA on Penn Treebank and WikiText-2.The experiments also include UCI regression comparisons with additional approximate Bayesian inference methods.

6 Discussion

The discussion presents SWAG as a practical step toward scalable Bayesian deep learning, grounded in the observation that neural-network posteriors are close to Gaussian within the SGD-trajectory subspace. Bayesian model averaging in that subspace can improve predictions over SGD or SWA solutions.

  • SWAG is presented as a step toward practical, scalable, and accurate Bayesian deep learning for large modern neural networks.
  • The posterior over neural-network parameters is close to Gaussian in the subspace spanned by the SGD trajectory.
  • Bayesian model averaging within the SGD-trajectory subspace can improve predictions over SGD or SWA solutions.

A Asymptotic Normality of SGD

SGD can have an approximately Gaussian stationary distribution under restrictive assumptions, but several assumptions fail in deep neural networks. Despite these failures, SGD trajectories can still capture useful posterior geometry in relevant subspaces.

  • Asymptotic theory: Under decaying learning rates and regularity conditions, SGD iterates converge asymptotically to a Gaussian distribution.The limiting covariance is expressed using the inverse Hessian and gradient-log-likelihood covariance.
  • SGD dynamics: On a two-dimensional quadratic, isotropic gradient noise produces the correct Gaussian stationary distribution, while momentum changes its scale but not its shape.These conclusions hold when the learning rate is not too large.
  • Scope: For non-isotropic gradient noise, the stationary distribution of SGD differs from the exact posterior distribution.The quadratic experiments are directly predicted by the analyzed theory.
  • Deep-network implications: For deep neural networks, non-convexity and over-parameterization invalidate key theoretical assumptions, although SGD iterates can remain approximately Gaussian within an attraction basin.The paper connects this local behavior to posterior geometry in the low-dimensional subspace spanned by SGD iterates.
  • Assumptions: The analysis assumes Gaussian, parameter-independent full-rank gradient noise, small learning rates, and a locally quadratic loss with positive-definite Hessian.These assumptions support the continuous-time approximation and the derivation of an optimal constant learning rate.

B.1 Assumption 2: Gradient Covariance Noise.

The gradient-noise covariance is approximately stable during the stationary phase, but the theory-derived learning rates are too large for practical deep-network training. Deep-network Hessians also violate the assumed positive-definiteness condition.

  • Covariance estimation: Estimating gradient-noise covariance requires dividing stochastic-gradient variance by the batch size, because V(ĝ(θ)) = C(θ)/B.Failing to account for this scaling inflates the covariance estimate by a factor of B.
  • Momentum: Momentum changes the theoretically optimal learning rate by a factor involving the momentum term.The experiments use μ = 0.1, corresponding to ρ = 0.9 in PyTorch’s SGD implementation.
  • Gradient covariance: During the stationary phase, the trace of the gradient-noise covariance is essentially constant, unlike during training from scratch.The discrepancy from early training is attributed to large initial gradients before the stationary distribution is reached.
  • Optimal learning rates: The computed optimal learning rates are too large, ranging from 1–3 at minimum versus standard rates of 0.1 or 0.01.Figure 5 reports the same qualitative mismatch as an approximately tenfold overestimate.
  • Hessian assumption: The Hessian of a trained PreResNet-164 on CIFAR-100 has minimum eigenvalue −272, contradicting the positive-definiteness assumption.The maximum eigenvalue is 3580 for comparison.

C Further Geometric Experiments

Across architecture–dataset pairs, SWAG captures posterior geometry in the subspace explored by SGD and achieves useful predictive behavior with relatively low-rank covariance estimates.

  • Posterior geometry: Across PreResNet-110 and VGG-16 on CIFAR-10 and CIFAR-100, SWAG captures posterior geometry in the subspace spanned by the SGD trajectory.Figure 6 examines cross-sections and two-dimensional surfaces along leading SWAG covariance eigenvectors.
  • Covariance rank: Thirty SWAG samples with rank-20 covariance match SWA accuracy over training for PreResNet56 on CIFAR-100.Rank-20 and rank-140 SWAG, as well as SWA, outperform ensembles of SGD iterates.
  • Sampling and scale: For WideResNet on CIFAR-100, about 30 samples provide stable accuracy and negative log-likelihood, while scales below 1 perform better and 0.5 and 0.25 are best.The figure compares SWAG, SWAG-Diag, and SWA across sample counts and covariance scales.
  • Experimental scope: The section presents these experiments as tests of SWAG hyperparameters and current theoretical limitations.The experiments connect geometric fidelity with choices of covariance rank and sampling scale.

D.2 Number of Samples in the Forwards Pass

SWAG uses Monte Carlo samples for Bayesian model averaging, with few samples often sufficient for likelihood estimates but more needed for stable accuracy. Its calibration advantages persist across shifted and transfer-learning settings, subject to practical trade-offs.

  • Sampling theory: Monte Carlo integration converges at rate 1/√K, where K is the number of samples, although good results may require few samples.The predictive integral generally lacks a closed-form expression.
  • Number of samples: On WideResNet28x10 and CIFAR-100, about 3 samples match SWA negative log-likelihood, while about 30 samples are needed for stable accuracy.The authors used 30 samples in most experiments and recommend tuning the count using validation performance and compute resources.
  • Covariance scale: SWAG covariance depends on the SGD learning rate and momentum, so the learning rate can be selected using validation negative log-likelihood and covariance rescaled during sampling.For WideResNet on CIFAR-100, the covariance scale is varied to reduce this dependence.
  • Batch normalization: Updating batch-normalization parameters after sampling significantly improves negative log-likelihood, though it can slow inference.The requirement is demonstrated for WideResNet on CIFAR-100.
  • Calibration: SWAG typically outperforms alternatives for uncertainty calibration and improves calibration over SWA, including under distribution shift and transfer learning.It performs on par or better than temperature scaling, while temperature scaling fails to achieve good calibration in transfer learning.
  • Practical trade-offs: Deep ensembles often calibrate somewhat better but require several times longer training than SWAG.This trade-off makes SWAG particularly useful when training time is limited but inference time can be increased.

E.2 Out-of-Domain Image Detection

The experiments evaluate SWAG’s uncertainty quality across out-of-domain detection, calibration, likelihood, accuracy, transfer, and regression. SWAG generally provides strong calibration and likelihood performance, while its low-rank covariance improves over diagonal approximations.

  • Out-of-Domain Detection: SWAG and SWAG-Diagonal produce more distinct in-domain and out-of-domain entropy distributions than the other methods.The comparison uses symmetrized, discretized KL divergence between predictive-entropy distributions.
  • Likelihood and Accuracy: SWAG typically outperforms the baselines in negative log-likelihood while remaining competitive with SWA and temperature-scaled SWA in test accuracy.SWAG-Diagonal is generally inferior to SWAG for log-likelihood but outperforms SWA.
  • Calibration: SWAG is competitive with temperature scaling for expected calibration error, while its low-rank plus diagonal covariance substantially improves calibration over the diagonal variant.SWAG-Diagonal also achieves better calibration than SWA.
  • Regression: On small UCI regression datasets, SWAG outperforms DVI and other methods on three of six datasets and remains competitive on the other three.Its computational time is the same as SGD, whereas DVI is reported to be 300x slower.
  • Regression: SWAG is considerably better calibrated than SGD on four of six UCI datasets, while neither method is ever far beneath the target 95% coverage.The comparison uses coverage of 95% credible sets for SWAG and 95% confidence regions for SGD.

H Classification Experimental Details and Parameters

The classification experiments use standard architectures, shared SWA-style training settings, and fixed sampling choices for SWAG and competing approximate-inference methods. Hyperparameters are tuned or specified separately for each baseline and dataset setting.

  • Architectures: The experiments use Preactivation-ResNet-164 and WideResNet28x10 implementations for CIFAR classification.The architecture implementations are identified by their corresponding repositories.
  • Shared Settings: Across datasets and architectures, methods use the SWA learning-rate schedule and weight decay, with Pre-ResNet trained for 300 epochs and averaging beginning after epoch 160.The paper uses its own PyTorch implementations for all methods.
  • SWAG: SWAG uses rank K = 20, 30 weight samples for Bayesian model averaging, and the remaining hyperparameters from SWA.The sampling and covariance-rank choices are fixed across the experiments described here.
  • Baselines: KFAC-Laplace is tuned on validation data for every model and dataset, while temperature scaling selects its temperature by minimizing validation negative log-likelihood.KFAC uses a KFAC Fisher approximation and custom sampling implementation.
  • Baselines: MC-Dropout averages predictions over 30 dropout masks, with dropout rates selected from {0.1, 0.05, 0.01} using validation performance.The selected rate was 0.05 for VGG-16 and WideResNet28x10.
  • Baselines: SGLD uses pretrained SGD checkpoints, a decaying learning-rate schedule, and gradient-noise rescaling because unrescaled settings caused divergence.For WideResNet and PreResNet, the runs last 100 and 150 epochs respectively.
Loading 1902.02476v2…