A Simple Baseline for Bayesian Uncertainty in Deep Learning

By Wesley J. Maddox, Timur Garipov, Pavel Izmailov, Dmitry P. Vetrov, Andrew Gordon Wilson. 2019

1,004 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Safety, Robustness & Interpretability, Probabilistic & Bayesian Methods

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Safety, Robustness & Interpretability

In Probabilistic & Bayesian Methods

Citations in the corpus

Efficient and Scalable Bayesian Neural Nets with Rank-1 Factors (2020, background): “Con-clusions compared to SWAG and Subspace Inference are consistent with CIFAR-10’s.”
Practical Deep Learning with Bayesian Principles (2019, background): “Many other approaches have recently been proposed to compute posterior approximations by training deterministic networks [46, 37, 38].”
Deep Ensembles: A Loss Landscape Perspective (2019, methodology): “This is similar to SWAG-diagonal [Maddox et al., 2019].”
Bayesian Deep Learning and a Probabilistic Perspective of Generalization (2020, methodology): “MultiSWAG combines multiple independently trained SWAG approximations [29], to create a mixture of Gaussians approximation to the posterior, with each Gaussian centred on a different basin.”
A Review and Comparative Study on Probabilistic Object Detection in Autonomous Driving (2020, background): “Still, it is reported that approximation techniques in BNNs are highly sensitive to hyper-parameters, and hard to scale to large datasets and network architectures [33].”
Questions for Flat-Minima Optimization of Modern Neural Networks (2022, methodology): “Maddox et al. (2019) approximate the posterior distribution over neural network weights by parameterizing a Gaussian distribution with the SWA solution as the ﬁrst moment.”
Git Re-Basin: Merging Models modulo Permutation Symmetries (2022, methodology): “Therefore, we recommend the recalculation of batch statistics after merging models (Izmailov et al., 2018; Wortsman et al., 2021; Maddox et al., 2019; Wang et al., 2021).”
TRADI: Tracking deep neural network weight distributions (2019, methodology): “Alternatives to ensembles use a network with multiple prediction heads [35], collect weight checkpoints from local minima and average them [23] or fit a distribution over them and sample networks [37].”
Beyond temperature scaling: Obtaining well-calibrated multiclass probabilities with Dirichlet calibration (2019, methodology): “An alternative to post-hoc calibration is to modify the classiﬁer learning algorithm itself: MMCE [17] trains neural networks by optimising the combination of log-loss with a kernel-based measure of calibration loss; SWAG [19] models the posterior distribution over the weights of the neural network ”
In Defense of Pseudo-Labeling: An Uncertainty-Aware Pseudo-label Selection Framework for Semi-Supervised Learning (2021, background): “Uncertainty and Calibration Estimating network prediction uncertainty has been a deeply studied topic (Graves, 2011; Blundell et al., 2015; Louizos & Welling, 2016; Lakshminarayanan et al., 2017; Malinin & Gales, 2018; Maddox et al., 2019; Welling & Teh, 2011).”
Uncertainty as a Form of Transparency: Measuring, Communicating, and Using Uncertainty (2020, methodology): “These methods can be broadly classified into two categories: Bayesian approaches [5, 20, 32, 51, 57, 63, 69, 92, 118, 184, 194] and Non-Bayesian, or frequentist, approaches [2, 98, 102, 107, 111, 177].”
A Survey on Bayesian Deep Learning (2016, background): “Due to BNN’s long history, the term ‘Bayesian deep learning’ sometimes specifically refers to ‘Bayesian neural networks’ [73, 128].”
What Are Bayesian Neural Network Posteriors Really Like? (2021, methodology): “…et al., 2020), expectation propagation (Hernández-Lobato & Adams, 2015), and leveraging the stochastic gradient descent (SGD) trajectory, either for a deterministic approximation, or sampling as in SGLD (Mandt et al., 2017; Maddox et al., 2019; Izmailov et al., 2018; Wilson & Izmailov, 2020).”
Generalized Out-of-Distribution Detection: A Survey (2021, positive): “Recent works attempt several less principled approximations including MC-dropout (Gal & Ghahramani, 2016) and deep ensembles (Dietterich, 2000; Lakshminarayanan et al., 2017; Maddox et al., 2019) for faster and better estimates of uncertainty.”