Does SWA require an ensemble at prediction time?

No. The averaged weights define one model.

Is it identical to exponential moving averaging?

No. Their weighting of selected states differs.

Can I average every checkpoint from initialization onward?

The arithmetic is possible, but early poor states can damage the result; the collection schedule is a design choice to validate.

AI Glossary / Probabilistic machine learning / Intermediate

Stochastic weight averaging

Definition

Stochastic weight averaging, or SWA, forms a single neural-network model by taking an equal average of selected parameter vectors along a training trajectory. It commonly uses a constant or cyclical learning rate during a later training phase so that the collected weights explore a useful region instead of simply converging to one final point.

What it means

Training produces many parameter settings, but deployment often uses only the last or best validation checkpoint. SWA uses several selected settings to construct a new parameter vector. The resulting network still makes predictions as one model.

This is weight averaging, not prediction averaging. For a nonlinear network, predicting with averaged parameters generally differs from averaging the predictions of the original networks. SWA is also different from an exponential moving average: ordinary SWA gives equal weight to its selected checkpoints, whereas an exponential average emphasizes recent ones. Neither name alone specifies a complete training and evaluation procedure.

How it works

Begin with a standard training phase. At an appropriate point, collect parameter snapshots under a learning-rate schedule that allows continued exploration. Maintain their running arithmetic mean; it is not necessary to store every snapshot just to calculate that mean.

If n snapshots have already been averaged, incorporate a new snapshot by multiplying the existing average by n, adding the new weights, and dividing by n plus one. Parameters must correspond to the same architecture and aligned parameterization.

For networks with batch normalization, refresh activation statistics for the averaged network using suitable training data before evaluation. Statistics associated with earlier individual weights may not describe the averaged model correctly. Compare the completed SWA model with the baseline using held-out data. The original work's favorable empirical findings motivate the method, but they are not a guarantee for every architecture, dataset, or averaging schedule.

Three practical examples

1. Image classification: compare a late-training SWA model with a final checkpoint under the same evaluation protocol. Include the batch-normalization refresh when relevant and measure both accuracy and calibration.

2. Tabular neural prediction: investigate whether averaging selected weights reduces sensitivity to a single late-training checkpoint. Keep data splits fixed so a changing test set does not obscure the comparison.

3. Deployment simplification: test a weight-averaged model when serving a full prediction ensemble is too expensive. Measure how closely its predictions meet the application's requirements rather than assuming it exactly reproduces an ensemble.

Think of it this way

Think of several nearby drafts of the same mechanical design, all using the same component coordinates. Averaging those coordinates produces one new design to test. It is different from operating all the original designs and averaging their outputs. Mixing unrelated coordinate systems could produce nonsense even if each design works individually.

Worked example

Consider an original scalar illustration. Three selected checkpoints have parameter values 0, 2, and 4. Their equal weight average is (0 + 2 + 4) divided by 3 = 2. If a fourth checkpoint has value 6, the running update gives (3 times 2 + 6) divided by 4 = 3, matching the direct average of all four values.

To distinguish weight and prediction averaging, consider a one-parameter logistic model that predicts sigmoid(w times x), with no bias and x = 1. For the first three checkpoints, predictions are approximately 0.5000, 0.8808, and 0.9820. Their prediction average is approximately 0.7876.

The weight-averaged model uses w = 2 and predicts approximately 0.8808 instead. The two operations give different answers because the sigmoid is nonlinear. This calculation demonstrates arithmetic and an important distinction; the chosen values are not claimed to be a realistic trained SWA trajectory or evidence of improved generalization.

For a linear prediction w times x, the two averaging operations would agree. That special equality should not be transferred to a general multilayer network.

Illustrative case study

Fictional case: Juniper Vision trains a classifier for an internal document-routing task. The team saves the normal final-checkpoint baseline and develops an SWA comparison from selected later training states. It records the averaging start, collection frequency, learning-rate schedule, and any batch-normalization refresh.

Analysts evaluate both models on the same untouched test set after selecting settings with validation data. They examine performance on uncommon document types as well as the overall average. If the SWA model misses an important subgroup more often, they investigate rather than declaring success from the method's reputation. This fictional workflow makes no claim of measured gains.

Common misconceptions and mistakes

SWA does not mean averaging arbitrary independently trained networks without checking parameter alignment. Neural-network symmetries can make direct averaging of unrelated weights unreliable. It is not a full Bayesian posterior approximation, and it does not automatically supply calibrated uncertainty. A flatter-looking loss plot in one parameterization is also not a universal certificate of better generalization. Validate the actual resulting model.

Business use

The attraction is a potentially useful training modification that still serves one network at inference time. Costs include maintaining an additional parameter average, collecting suitable states, and sometimes refreshing normalization statistics. Teams should compare total training and validation effort, not only inference latency. A deployment decision should use the same acceptance criteria as any other model change, including relevant subgroup performance and reproducibility.

Check your understanding

Answer: 1.

What is the average after adding the fourth example checkpoint?

Answer: It is 3.

Answer: 2.

Why did the logistic predictions not average to the prediction at the average weight?

Answer: The prediction function is nonlinear in the weight.

Answer: 3.

Why refresh batch-normalization statistics?

Answer: The averaged weights can change activation distributions relative to those of individual checkpoints.

Related terms

Checkpoint
Ensemble
Batch normalization
Weight averaging
exponential moving average
learning-rate schedule
SWAG

Learn next

Implement an online arithmetic mean of matching parameter vectors. Then compare weight averaging with prediction averaging on a nonlinear model and study the effect of collection schedules and normalization statistics.

Sources

Izmailov and colleagues, Averaging Weights Leads to Wider Optima and Better Generalization, introduces SWA: https://proceedings.mlr.press/r16/izmailov18a.html

PyTorch's official optimization documentation describes equal averaging and batch-normalization refresh: https://docs.pytorch.org/docs/2.14/optim.html

The scalar calculations and Juniper Vision case are original teaching examples; the case is fictional.

Independent review consulted the original author manuscript at https://arxiv.org/pdf/1803.05407 and the official PyTorch weight-averaging and batch-normalization sections. The PMLR landing page opened; its linked PDF returned an internal retrieval error.

Previous: Stochastic variational inference Next: Stochastic weight averaging Gaussian approximation