What is stochastic weight averaging (SWA)?

**Stochastic Weight Averaging (SWA) creates a new set of neural-network weights by taking the simple arithmetic mean of several weight snapshots from the *same training run*. Instead of deploying only the latest optimized weights, SWA combines selected points along the optimizer's later trajectory into one model**. It is not an ensemble that averages multiple models' predictions.

How does SWA work?

1. Train the main model normally: An optimizer continues updating the network using training batches.
2. Start averaging later in training: Once the model has learned useful features, take weight snapshots at selected updates or epochs—not necessarily at every update.
3. Average each parameter separately: The snapshots must belong to the same architecture and compatible parameter positions within the same training trajectory.
4. Maintain an SWA copy: The optimizer still trains the current model; the SWA copy is updated by averaging, without its own gradient steps.
5. Evaluate the averaged model: After handling normalization statistics, compare validation performance against the latest training-model checkpoint.

For n selected snapshots θ₁, θ₂, …, θₙ:

θ_SWA = (θ₁ + θ₂ + … + θₙ) / n

This calculation applies element-wise to every trainable weight tensor. You can compute it without storing every checkpoint using the running average:

θ_SWA,new = (n × θ_SWA,old + θ_next) / (n + 1)

Here n counts snapshots already averaged, not all optimizer steps since training began.

Beginner-friendly numerical example

Consider a toy model with two weights. Suppose we save three checkpoints late in the same training run:

• Checkpoint 1: θ₁ = [2, 4]
• Checkpoint 2: θ₂ = [4, 8]
• Checkpoint 3: θ₃ = [6, 6]

Step 1 — Average the first two snapshots:

θ_SWA,2 = ([2, 4] + [4, 8]) / 2 = [3, 6]

Step 2 — Include the third snapshot:

θ_SWA,3 = (2 × [3, 6] + [6, 6]) / 3

= ([6, 12] + [6, 6]) / 3 = [4, 6]

Check using the direct arithmetic mean:

θ_SWA = ([2, 4] + [4, 8] + [6, 6]) / 3 = [4, 6]

The latest model weights are [6, 6], while the SWA weights are [4, 6]. Every snapshot contributes one-third to the final average. These numbers illustrate the arithmetic, not a guarantee that [4, 6] gives better predictions.

When are weights collected, and what about learning rate?

• Typically start late: Begin collecting snapshots after ordinary training has made substantial progress.
• Use a suitable learning-rate schedule: The original SWA work used a constant or cyclical learning rate during the averaging phase so the optimizer explores multiple nearby solutions instead of settling at one point.
• Collect repeatedly: For example, sample the model after each epoch during an SWA phase. A different sampling cadence is possible; choose it deliberately.
• Keep training the original model: Averaging is an additional process alongside optimizer updates, not a replacement for the optimizer.

In PyTorch, torch.optim.swa_utils.AveragedModel(model) maintains an equal-weight running average by default. After a selected optimizer update, call swa_model.update_parameters(model). torch.optim.swa_utils.SWALR can move the learning rate toward a chosen fixed SWA rate.

Why does BatchNorm need special attention?

BatchNorm stores running means and variances of activations in buffers. These statistics depend on the model's weights and data distribution; simply averaging parameters does not automatically make the existing BatchNorm statistics correct for the new SWA weights.

• Common approach: After the final SWA update, run torch.optim.swa_utils.update_bn(train_loader, swa_model) to recompute BatchNorm statistics by passing training inputs through the SWA model.
• Alternative: Some implementations average buffers using AveragedModel(..., use_buffers=True); check whether that approach gives appropriate validation results.
• No BatchNorm: Skip this step when the model has no BatchNorm-style running statistics.

Benefits, limitations, and applications

• Benefits: The original work found that averaging points visited by SGD can land in a broader, well-performing region of weight space and sometimes improve generalization, robustness of the selected checkpoint, or model calibration.
• Applications: Image classifiers and other deep networks trained along a compatible optimization path; it is easy to test as a late-training extension.
• Limitations: Requires additional training time for collecting snapshots and a separate parameter copy. Averages of very different or incompatible solutions may perform poorly; SWA does not guarantee better accuracy. Recomputing BatchNorm statistics adds a data pass.
• Training versus inference: Average parameters during training; evaluate and deploy the single averaged model, normally with no further snapshot averaging at inference.

Brief comparison with EMA: SWA normally gives equal weight to chosen snapshots; an exponential moving average gives more weight to recent parameters. The detailed EMA update and decay-factor explanation is covered in “What is an exponential moving average (EMA) of neural network weights?” (Question 254).

Key takeaway: SWA averages several compatible neural-network weight snapshots gathered during training to obtain one alternative model, which must be validated before use.

Research: Izmailov et al. (2018), *Averaging Weights Leads to Wider Optima and Better Generalization*. Implementation: PyTorch — Weight Averaging (SWA and EMA) and PyTorch AveragedModel.