Abstract
We propose SWA-Gaussian (SWAG), a simple, scalable, and general purposeapproach for uncertainty representation and calibration in deep learning.Stochastic Weight Averaging (SWA), which computes the first moment ofstochastic gradient descent (SGD) iterates with a modified learning rateschedule, 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 alow rank plus diagonal covariance also derived from the SGD iterates, formingan approximate posterior distribution over neural network weights; we thensample from this Gaussian distribution to perform Bayesian model averaging. Weempirically find that SWAG approximates the shape of the true posterior, inaccordance with results describing the stationary distribution of SGD iterates.Moreover, we demonstrate that SWAG performs well on a wide variety of computervision tasks, including out of sample detection, calibration, and transferlearning, in comparison to many popular alternatives including MC dropout, KFACLaplace, and temperature scaling.