How is SWAG different from SWA?

SWA returns an averaged model; SWAG adds a parameter distribution and prediction averaging.

Does increasing samples make the posterior approximation exact?

No. It improves numerical averaging under the chosen Gaussian.

Why keep a low-rank term?

It captures selected correlations among parameter movements that a diagonal approximation misses.

AI Glossary / Probabilistic machine learning / Advanced

Stochastic weight averaging Gaussian approximation

Definition

Stochastic weight averaging Gaussian approximation, usually called SWA-Gaussian or SWAG, fits a Gaussian distribution over neural-network parameters using statistics collected along a training trajectory. It extends stochastic weight averaging by estimating parameter variation as well as the mean, then averages predictions from sampled parameter settings to approximate Bayesian prediction.

What it means

Ordinary SWA produces one averaged parameter vector. SWAG treats that vector as the center of a distribution and uses the trajectory's variation to construct its covariance. The full method combines a diagonal variance estimate with a low-rank component that captures some coordinated movement among parameters.

This is an approximate posterior representation. Training iterates are not automatically exact posterior draws, and a Gaussian fitted around one explored region cannot represent every possible mode of a neural-network posterior. Its usefulness must be evaluated through predictive behavior, calibration, and relevant robustness tests rather than inferred solely from the word Bayesian.

How it works

Start from a trained model and continue stochastic-gradient training with a suitable exploration schedule. Accumulate running averages of weights and their elementwise squares. Their difference yields a diagonal variance estimate. Also retain a limited collection of deviations from running weight means to describe correlated directions without storing a full covariance matrix.

In the original full SWAG construction, the covariance is one half the sum of the diagonal and low-rank covariance terms. The low-rank term uses the retained deviation matrix and a normalization based on its number of columns. This differs from the diagonal-only variant.

Sample new parameter vectors from the fitted Gaussian. For networks with batch normalization, recompute the required activation statistics for each sampled model. Evaluate each model and average predictive probabilities or predictive distributions. Averaging logits or weights is generally a different operation. More samples reduce Monte Carlo error around the fitted approximation; they do not remove errors in that approximation itself.

Three practical examples

1. Classification confidence: average class probabilities from SWAG parameter samples and evaluate reliability on held-out examples. Compare uncertainty estimates with actual error frequencies rather than assuming all low-confidence cases are difficult for the same reason.

2. Regression prediction: combine sampled models' predictive distributions, including their observation-noise models where applicable. Separate variation in predicted means from total predictive uncertainty.

3. Review prioritization: investigate whether predictive uncertainty helps identify records needing manual inspection. Test that proposed workflow on relevant shifts and rare cases instead of treating parameter spread as a universal anomaly detector.

Think of it this way

Think of keeping not just the average location of several nearby designs, but also the directions in which their components tend to move together. You generate further plausible nearby designs and compare their outputs. The quality of that exploration still depends on whether the original neighborhood represented the important possibilities.

Worked example

Consider an original two-parameter covariance illustration using stipulated summary statistics. Let the SWA mean be (1, 2), the diagonal covariance estimate have diagonal entries 4 and 1, and the retained deviation matrix have two columns: (1, 1) and (minus 1, minus 1). This example starts from those summaries rather than reconstructing a complete training run.

With K = 2 retained columns, the low-rank covariance is D times D transpose divided by K minus 1. It is the matrix with every entry equal to 2. Combining it with the diagonal estimate and multiplying by one half gives covariance rows (3, 1) and (1, 1.5).

The off-diagonal covariance is 1, so the two parameter fluctuations are positively related. The determinant is 3 times 1.5 minus 1 = 3.5, consistent with a positive definite covariance.

For the simple linear output Y = w1 + w2, the fitted Gaussian gives mean 3 and variance 3 + 1.5 + 2 times 1 = 6.5. This calculation includes covariance; summing only the two diagonal entries would incorrectly give 4.5. It demonstrates the meaning of the covariance structure, not measured uncertainty calibration for a real network.

Illustrative case study

Fictional case: Cedar Imaging wants to improve the usefulness of confidence estimates for document classification. Its analysts fit SWAG around a trained network, record the collection schedule and retained covariance rank, and compare predictive averages with the single SWA model.

The team uses validation data to select settings and keeps a separate test set for the final comparison. It measures predictive likelihood and calibration alongside accuracy and records the extra time required for sampling and normalization refresh. Uncertainty on unfamiliar document formats receives a separate check. This fictional case specifies an experiment without promising that SWAG will detect every unfamiliar input or improve every metric.

Common misconceptions and mistakes

SWAG is not simply Gaussian noise added with an arbitrary variance. Its covariance is derived from collected trajectory statistics, with explicit scaling choices. It is not a full-covariance posterior fit and does not guarantee coverage of separated modes. Confusing variance with standard deviation changes sampling amplitudes. Forgetting normalization updates or averaging raw logits instead of predictive probabilities can also change the intended procedure.

Business use

The business attraction is uncertainty estimation using information already present in continued training, without fitting many completely independent networks. The tradeoffs include storing deviation vectors, additional forward passes, and evaluating whether the approximation remains useful outside the training distribution. A deployment assessment should compare quality against total inference cost and report the sampling count and covariance settings needed to reproduce the result.

Check your understanding

Answer: 1.

What is the mean of the example output?

Answer: It is 1 plus 2 = 3.

Answer: 2.

Why is its variance 6.5 rather than 4.5?

Answer: The output variance includes twice the covariance between the two weights.

Answer: 3.

Can one Gaussian necessarily describe multiple distant posterior modes?

Answer: No; that is a limitation of the representation.

Related terms

Stochastic weight averaging
Calibration
Batch normalization
Bayesian model averaging
Gaussian approximation
low-rank covariance
epistemic uncertainty

Learn next

Understand SWA's running mean first. Then derive diagonal and correlated Gaussian sampling, and compare averaging model probabilities with predicting at the average weights. Finally, evaluate calibration independently of the training objective.

Calibration

Sources

Maddox and colleagues, A Simple Baseline for Bayesian Uncertainty in Deep Learning, Sections 3.3–3.5, introduces SWAG covariance construction and predictive averaging: https://proceedings.neurips.cc/paper/2019/file/118921efba23fc329e6560b27861f0c2-Paper.pdf

The covariance calculation and Cedar Imaging scenario are original teaching examples; the scenario is fictional.

Independent review consulted Sections 3.3–3.5 of the original SWAG paper.

Previous: Stochastic weight averaging Next: Stock sampling