A Simple Baseline for Bayesian Uncertainty in Deep Learning ( NeurIPS 2019 )Permalink

AbstractPermalink

SWAG ( =SWA-Gaussian ) : simple, scalable, general purpose approach for…

• (1) uncertainty representation
• (2) calibration in deep learning

1. IntroductionPermalink

Representing Uncertainty is crucial.

This paper…

• use the information contained in the SGD trajectory to efficiently approximate the posterior distn over the weights of NN
• find that “Gaussian distn” fitted to the first 2 moments of SGD captures the local geometry of posterior!

2. Related WorkPermalink

2-1. Bayesian MethodsPermalink

MCMCPermalink

• HMC (Hamiltonian Monte Carlo)
• SGHMC (stochastic gradient HMC)
  • allows for stochastic gradients to be used in Bayesian Inference
  • (crucial for both scalability & exploring a space of solutions to provide good generalization)
• SGLD (stochastic gradient Langevin dynamics)

Variational InferencePermalink

• Reparameterization Trick

Dropout Variational InferencePermalink

• spike and slab variational distribution
• optimize dropout probabilities as well

Laplace ApproximationPermalink

• assume a Gaussian posterior, $\mathcal{N}\left(\theta^{*}, \mathcal{I}\left(\theta^{*}\right)^{-1}\right)$.
  • $\mathcal{I}\left(\theta^{*}\right)^{-1}$ : inverse of the Fisher information matrix

2-2. SGD based approximationPermalink

• averaged SGD as an MCMC sampler

2-3. Methods for Calibration of DNNsPermalink

(SSDE) Ensembles of several networks

3. SWA-Gaussian for Bayesian Deep LearningPermalink

propose SWAG for Bayesian model averaging & uncertainty estimation

3-1. Stochastic Gradient Descnet (SGD)Permalink

standard training of DNNs

$\Delta \theta_{t}=-\eta_{t}\left(\frac{1}{B} \sum_{i=1}^{B} \nabla_{\theta} \log p\left(y_{i} \mid f_{\theta}\left(x_{i}\right)\right)-\frac{\nabla_{\theta} \log p(\theta)}{N}\right)$.

• loss function : NLL & regularizer
  • NLL : $-\sum_{i} \log p\left(y_{i} \mid f_{\theta}\left(x_{i}\right)\right)$
  • regularizer : $\log p(\theta)$

3-2. Stochastic Weight Averaging (SWA)Permalink

main idea of SWA

• run SGD with a constant learning rate schedule,

  starting from a pre-trained solution & average the weights

• $\theta_{\text {SWA }}=\frac{1}{T} \sum_{i=1}^{T} \theta_{i}$.

3-3. SWAG-DiagonalPermalink

simple diagonal format for the covariance matrix.

maintain a running average of the 2nd uncentered moment for each weight,

then compute the covariance using the following standard identity at the end of training:

• $\overline{\theta^{2}}=\frac{1}{T} \sum_{i=1}^{T} \theta_{i}^{2}$ .
• $\Sigma_{\text {diag }}=\operatorname{diag}\left(\overline{\theta^{2}}-\theta_{\text {SWA }}^{2}\right)$.
• Approximate posterior distribution : $\mathcal{N}\left(\theta_{\mathrm{SWA}}, \Sigma_{\text {Diag }}\right)$

3-4. SWAG : Low Rank plus Diagonal Covariance StructurePermalink

full SWAG algorithm

• diagonal covariance approximation : TOO restrictive
• more flexible low-rank plus diagonal posterior approximation

sample covariance matrix of SGD :

• sum of outer products

• $\Sigma=\frac{1}{T-1} \sum_{i=1}^{T}\left(\theta_{i}-\theta_{\mathrm{SWA}}\right)\left(\theta_{i}-\theta_{\mathrm{SWA}}\right)^{\top}$ ( rank = $T$ )

  • but don’t know $\theta_{SWA}$ during training

• $\Sigma \approx \frac{1}{T-1} \sum_{i=1}^{T}\left(\theta_{i}-\bar{\theta}_{i}\right)\left(\theta_{i}-\right.\left.\bar{\theta}_{i}\right)^{\top}=\frac{1}{T-1} D D^{\top}$.

  • where $D$ is the deviation matrix comprised of columns $D_{i}=\left(\theta_{i}-\bar{\theta}_{i}\right)$

    and $\bar{\theta}_{i}$ is the running estimate of the parameters’ mean obtained from the first $i$ samples

Combine (1) & (2)

• (1) low-rank approximation : $\Sigma_{\text {low-rank }}=\frac{1}{K-1} \cdot \widehat{D} \widehat{D}^{\top}$
• (2) diagonal approximation : $\Sigma_{\text {diag }}=\operatorname{diag}\left(\overline{\theta^{2}}-\theta_{\text {SWA }}^{2}\right)$

$\rightarrow$ $\mathcal{N}\left(\theta_{\text {SWA }}, \frac{1}{2} \cdot\left(\Sigma_{\text {diag }}+\Sigma_{\text {low-rank }}\right)\right)$.

To sample from SWAG, we use….

• $\tilde{\theta}=\theta_{\mathrm{SWA}}+\frac{1}{\sqrt{2}} \cdot \Sigma_{\mathrm{diag}}^{\frac{1}{2}} z_{1}+\frac{1}{\sqrt{2(K-1)}} \widehat{D} z_{2}$.

  where $z_{1} \sim \mathcal{N}\left(0, I_{d}\right), z_{2} \sim \mathcal{N}\left(0, I_{K}\right)$.

Full AlgorithmPermalink

3-5. Bayesian Model Averaging with SWAGPermalink

MAP

• posterior $\log p(\theta \mid \mathcal{D})=\log p(\mathcal{D} \mid \theta)+\log p(\theta)$.
• prior $p(\theta)$ : regularization in optimization

Bayesian procedure “marginalizes” the posterior over $\theta$

• $p\left(y_{*} \mid \mathcal{D}, x_{*}\right)=\int p\left(y_{*} \mid \theta, x_{*}\right) p(\theta \mid \mathcal{D}) d \theta$.

• using MC sampling…

  $p\left(y_{*} \mid \mathcal{D}, x_{*}\right) \approx \frac{1}{T} \sum_{t=1}^{T} p\left(y_{*} \mid \theta_{t}, x_{*}\right), \quad \theta_{t} \sim p(\theta \mid \mathcal{D})$.

Prior Choice

• (typically) weight decay is used to regularize DNN

• when SGD is used with momentum $\rightarrow$ implicit regularization