A neural network rarely finishes training at the only useful point in parameter space. Late in training, stochastic gradient descent (SGD) can visit several nearby parameter settings that all perform reasonably well, while the final checkpoint represents only one of them.
Stochastic weight averaging (SWA) turns that observation into a simple training technique: collect model parameters from multiple points late in an SGD trajectory and compute their arithmetic mean. The result is one model with averaged weights, so inference does not require running an ensemble of all collected checkpoints.
The idea sounds almost too simple, but the details matter. Averaging checkpoints from unrelated training runs can fail, averaging too early can mix poor solutions into the result, and models with batch normalization need special handling after their parameters are averaged.
This article builds a practical mental model for SWA. You will learn what is actually averaged, why the learning-rate schedule matters, how SWA differs from checkpoint ensembling and exponential moving averages, how to evaluate it fairly, and when a simpler final checkpoint is the better choice.
Start with the smallest useful example
Suppose a tiny model has only two trainable parameters. Late in training, SGD produces these checkpoints:
checkpoint A: [2.0, 4.0]
checkpoint B: [2.4, 3.6]
checkpoint C: [2.2, 4.2]Their equal-weight average is:
w_swa = ([2.0, 4.0] + [2.4, 3.6] + [2.2, 4.2]) / 3
= [2.2, 3.9333...]SWA uses [2.2, 3.9333...] as the parameters of a single model.
For a real network, the same operation is applied element by element across every averaged parameter tensor. You do not average predictions at inference time, and you do not keep three active models just to make one prediction.
This distinction is important. An ensemble might evaluate A, B, and C separately and then combine their predictions. SWA instead creates a new parameter vector from A, B, and C and evaluates that one network.
The useful mental model is a trajectory, not three independent models
Neural-network optimization is non-convex, so arbitrary parameter averaging is unsafe. Two independently trained networks can implement similar functions while arranging hidden units differently. A coordinate-wise average between those models can land at a poor solution.
SWA is more constrained. The original method averages points sampled along one SGD trajectory while using a learning-rate schedule that keeps exploring a region of parameter space. Because those points come from the same continuing optimization path, parameter correspondence is preserved in a way that is not guaranteed across unrelated runs.
If the collected parameters are w_1, w_2, ..., w_n, the SWA parameters are
w_swa = (1 / n) * sum(w_i for i = 1..n)You do not need to store every checkpoint to compute this. Maintain a running average:
average_1 = w_1
average_n = average_(n-1) + (w_n - average_(n-1)) / nAfter each update, the running value equals the arithmetic mean of all parameters collected so far. Memory therefore needs roughly one additional model-sized parameter copy rather than a growing archive of checkpoints.
Why SWA changes the usual end of training
A conventional schedule often reduces the learning rate toward a small value so optimization settles near one point. SWA needs a somewhat different late-training behavior because averaging several nearly identical checkpoints provides little benefit.
The SWA procedure proposed by Izmailov and colleagues uses a constant or cyclical learning rate during the averaging phase. The purpose is to let SGD continue moving through a useful region rather than collapse immediately onto one final point. Parameters sampled along that trajectory are then averaged.
A simplified training plan looks like this:
1. Train normally for most of the budget.
2. Enter an SWA phase with an appropriate non-vanishing learning rate.
3. Periodically add the current parameters to the running average.
4. Finish training and prepare the averaged model for evaluation.
5. Recompute batch-normalization statistics when required.
6. Compare the SWA model with the normal checkpoint on held-out data.The start time, learning rate, and collection frequency are hyperparameters, not universal constants. A schedule that works for one architecture and dataset does not become a guarantee for another.
Why averaging can help
Imagine that the low-loss region around a solution is shaped like a broad valley. SGD can move around that valley and visit several points with good training loss. The arithmetic mean of those points may lie more centrally in the region than an individual point near one side.
This is the geometric intuition behind SWA. The original SWA work reported that averaging SGD iterates with its modified learning-rate schedule could find solutions associated with wider optima and improve generalization on the evaluated image-classification settings.
That intuition should not be turned into a universal rule that flatter-looking parameters must generalize better. Neural-network loss geometry depends on parameterization, and SWA’s empirical benefit must still be measured on the task that matters.
The practical claim is narrower: late-training weight averaging is a low-complexity experiment that can produce a useful single model when the optimizer trajectory supplies compatible, good-quality parameter samples.
Do not confuse SWA with checkpoint ensembling
Suppose three checkpoints each produce a probability vector for an input. A prediction ensemble combines those output vectors:
p_ensemble = (p_A + p_B + p_C) / 3SWA instead averages parameters first:
w_swa = (w_A + w_B + w_C) / 3
p_swa = model(x; w_swa)In a nonlinear neural network, these expressions are generally not equal:
model(x; average(weights)) != average(model(x; weights))An ensemble can gain robustness from multiple independently evaluated models, but inference cost usually grows because several forward passes are required. SWA keeps the inference structure of one model. Its extra cost is mainly during training and model preparation, where the running parameter copy is maintained.
If inference latency permits an ensemble and maximum predictive quality is the objective, compare the methods empirically rather than assuming SWA replaces ensembling.
SWA is also different from an exponential moving average
An exponential moving average (EMA) also maintains a second set of model parameters, but recent checkpoints receive more weight. A common update has the form
w_ema <- decay * w_ema + (1 - decay) * w_currentWith a decay such as 0.999, older information fades gradually as new parameters arrive. SWA, by contrast, gives equal weight to the parameter samples included in its arithmetic mean.
The training intent can differ too. SWA is commonly paired with a late learning-rate schedule designed to keep sampling a region of solutions. EMA is often updated throughout training as a smoothed version of the evolving model.
Neither rule is inherently superior. If your training recipe already has a well-tested EMA, replacing it with SWA adds another experimental variable. Compare them under the same data split, training budget, and evaluation protocol.
Batch normalization needs explicit attention
Averaging trainable parameters is not the whole story for networks that use batch normalization.
Batch-normalization layers maintain running activation statistics such as means and variances. Those statistics were observed while the ordinary training model moved through its parameter trajectory. After SWA constructs a new averaged set of weights, the old running statistics may no longer match the activations produced by that averaged model.
A common solution is to make a pass over training data with the final SWA model to recompute batch-normalization statistics before evaluation. PyTorch, for example, provides torch.optim.swa_utils.update_bn() for this purpose.
This is an implementation detail with model-quality consequences. Forgetting it can make a correct parameter average look worse simply because normalization statistics are stale.
Architectures without batch normalization do not require this particular recalculation. Other stateful layers or framework-specific buffers should still be checked according to their documented semantics rather than assumed to behave like ordinary trainable parameters.
A practical implementation pattern
The following pseudo-code shows the important state transitions without tying the idea to one framework:
model = initialize_model()
average = empty_average()
count = 0
for epoch in training_epochs:
train_one_epoch(model)
if epoch >= swa_start:
set_or_step_swa_learning_rate()
average = running_parameter_average(average, model, count)
count += 1
else:
step_normal_learning_rate_schedule()
swa_model = copy_model_with_parameters(average)
recompute_normalization_statistics_if_needed(swa_model)
evaluate(swa_model)In production code, prefer the averaging utilities supplied and documented by your framework when they match your training setup. For example, current PyTorch documentation provides AveragedModel for maintaining averaged parameters and SWALR for annealing toward an SWA learning rate.
The pseudo-code is intentionally simplified. Distributed training, mixed precision, optimizer state, checkpoint resumption, parameter freezing, and custom model buffers can all affect where and how the average should be maintained.
Evaluate the right comparison
An SWA experiment is useful only if its baseline is fair. At minimum, record the normal model and the SWA model produced by the same training run or by carefully matched runs.
Compare the metric that corresponds to the deployment task, such as validation loss, accuracy, F1, ranking quality, or another domain-specific measure. Also measure operational properties that matter to the system.
SWA usually keeps the same architecture and parameter count as the base model, so it should not be treated as model compression. It also does not inherently make a forward pass cheaper. If latency changes, investigate the actual runtime path instead of attributing the difference to averaging itself.
Training cost can increase because the SWA phase may extend training or preserve a learning rate that continues exploration. Memory also increases while training because an averaged parameter copy is stored. These costs may be modest relative to the full training job, but they are not zero.
Common mistakes
Averaging unrelated checkpoints
Do not assume that independently trained models can be averaged just because they share an architecture. Neural networks can represent similar functions with different parameter arrangements. SWA is based on averaging compatible points from an optimization trajectory, not arbitrary model files.
Starting before the model reaches a useful region
If early, poor-quality parameters enter an equal-weight average, their influence does not decay away as it would under an EMA. Choose the averaging phase deliberately and validate the start point.
Letting the learning rate vanish during collection
If collected checkpoints are nearly identical, SWA has little trajectory to average. The learning-rate strategy is part of the method, not an unrelated tuning detail.
Forgetting normalization statistics
For batch-normalized models, stale running statistics can invalidate the comparison. Recompute or otherwise handle them using the framework’s supported procedure.
Reporting only the averaged model
Keep the ordinary checkpoint as a baseline. If SWA does not improve the held-out metric enough to justify added training complexity, the normal checkpoint is the simpler deployment artifact.
When SWA is a good experiment
SWA is worth testing when you already have a stable SGD-style training pipeline, can afford a late averaging phase, and want to explore whether a single averaged solution generalizes better than the final trajectory point. It is especially convenient when deployment should still use one model rather than a multi-checkpoint ensemble.
It is less attractive when training is already extremely expensive and cannot be extended, when the optimizer recipe does not provide useful compatible samples to average, or when a validated EMA or other checkpoint-selection method already meets the requirement. It is also not a substitute for fixing data leakage, poor labels, an unsuitable objective, or a weak evaluation design.
Conclusion
Stochastic weight averaging is best understood as a change to the end of optimization: keep SGD exploring a useful region, collect compatible parameter states, and replace the final point with their equal-weight average.
The arithmetic is simple; the training context is what makes it meaningful. Use a deliberate averaging phase, preserve a normal-checkpoint baseline, refresh batch-normalization statistics when necessary, and evaluate the resulting single model on the metrics that matter. If those conditions are handled carefully, SWA becomes a practical experiment rather than just another checkpoint trick.