[Paper Review] A Simple Baseline for Bayesian Uncertainty in Deep Learning

Wesley J. Maddox, Timur Garipov+3 more|arXiv (Cornell University)|Feb 7, 2019

Gaussian Processes and Bayesian Inference63 references181 citations

TL;DR

SWAG introduces a scalable Gaussian posterior over neural network weights built from the SWA mean and a low-rank plus diagonal covariance estimated from SGD iterates, enabling Bayesian model averaging and improved uncertainty calibration across vision tasks.

ABSTRACT

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.

Motivation & Objective

• Motivate the need for reliable uncertainty representation in deep learning to aid decision-making in high-stakes domains.
• Propose a scalable Bayesian inference method that leverages SGD trajectories to approximate the posterior over network weights.
• Develop a practical algorithm (SWAG) that combines SWA with a low-rank plus diagonal covariance to form a Gaussian posterior.
• Demonstrate that SWAG yields well-calibrated predictions and competitive or superior uncertainty estimates across vision benchmarks.

Proposed method

• Build on Stochastic Weight Averaging (SWA) to use the SWA mean as the posterior mean.
• Estimate a diagonal covariance from running second moments of SGD iterates.
• Construct a low-rank covariance using the last K deviation vectors from SGD iterates.
• Form a Gaussian posterior N(theta_SWA, 1/2*(Sigma_diag + Sigma_low_rank)).
• Sample from the Gaussian to perform Bayesian model averaging for predictions.
• Provide an online procedure to update and store the necessary statistics with minimal overhead.

Experimental results

Research questions

• RQ1 Can SGD trajectories be used to approximate the local geometry of the posterior in deep networks?
• RQ2 Does a SWAG-based Gaussian posterior yield better uncertainty calibration than existing baselines across vision tasks?
• RQ3 Is SWAG effective for out-of-domain detection and transfer learning compared to alternatives like MC dropout and SGLD?
• RQ4 How does the low-rank plus diagonal approximation compare to diagonal-only covariances in practice?
• RQ5 Can SWAG improve calibration and predictive performance on language modeling and regression benchmarks as a broader baseline?

Key findings

• SWAG closely captures the local geometry of the posterior in the subspace spanned by SGD iterates.
• SWAG provides well-calibrated uncertainty estimates and higher test log-likelihood than several baselines on CIFAR-10/100 and ImageNet.
• SWAG outperforms many alternatives including MC dropout, SGLD, KFAC-Laplace, and SWA in uncertainty calibration.
• SWAG improves transfer learning performance and out-of-domain detection compared to several competitors.
• SWAG also yields improvements in language modeling perplexities and competitive results on regression tasks.