Stochastic Weight Averaging

Training with stochastic gradient descent doesn't settle smoothly into one exact point — the weights keep bouncing around near the bottom of the loss surface from batch to batch, even late in training, because each mini-batch's gradient is a noisy estimate. If you just grab whatever weights happen to be current when you decide to stop, you've picked one particular point out of that bouncing cloud, somewhat arbitrarily. Stochastic Weight Averaging (SWA) asks a different question: what if, instead of picking one point from that late-training cloud, you averaged several of them together?

The analogy: many blurry photos versus one sharp guess

If you're trying to pin down exactly where a hovering drone is, and your camera is a little shaky, one single photo might catch it slightly to the left or slightly to the right of its true position — that's the noise. But if you take ten photos over a few seconds and average their positions, the shake tends to cancel out because it wobbles in different directions each time, and the averaged position is closer to the drone's true, steady hovering point than any single shaky photo. SWA treats the late-training weight snapshots the same way: each one is a noisy read of "where the good region of the loss surface is," and averaging several of them cancels out some of that noise.

Averaging, and why it lands somewhere flatter

Concretely, SWA runs ordinary training as normal, then in a final phase — often with a constant or cyclical learning rate rather than one that's still decaying to near zero — it periodically snapshots the current weights and maintains a running average:

$$\bar\theta_n = \frac{n \cdot \bar\theta_{n-1} + \theta_n}{n+1}$$

In words: the new running average ($\bar\theta_n$) is the old running average ($\bar\theta_{n-1}$) weighted by how many snapshots have gone into it so far ($n$), blended with the newest snapshot ($\theta_n$), then renormalized. Each new snapshot pulls the average only slightly, so it settles toward a stable central point among the recent snapshots rather than any one of them individually.

Because the loss surface near a good minimum tends to be curved in some directions and much flatter in others, and different noisy SGD steps wander in different directions within that flat region, averaging several nearby points tends to land nearer the center of that flat region rather than out toward one of its steeper edges — which is exactly the kind of flat minimum associated with better generalization.

Worked example 1: averaging three snapshots

Suppose late in training, three snapshots of a single weight are recorded: $\theta_1 = 2.10$, $\theta_2 = 1.95$, $\theta_3 = 2.05$. The running average updates as: $\bar\theta_1 = 2.10$; $\bar\theta_2 = \frac{1(2.10)+1.95}{2} = 2.025$; $\bar\theta_3 = \frac{2(2.025)+2.05}{3} = \frac{4.05+2.05}{3} = 2.033$. The final averaged weight, $2.033$, sits centrally among the three noisy snapshots rather than matching any single one — and none of the three individual snapshots is guaranteed to be the "best" one on its own.

Worked example 2: why the average can beat every individual snapshot

Suppose test accuracy for snapshots 1, 2, and 3 individually were $90.1\%$, $89.7\%$, and $90.3\%$ — each a noisy read of the model's true quality. Averaging the weights (not the accuracies) into $\bar\theta$ and then evaluating that single averaged model on the test set often produces something like $91.0\%$ — better than any individual snapshot. This isn't guaranteed by the arithmetic (averaging weights is not the same as averaging predictions), but empirically it happens often enough that SWA is a reliable, nearly-free accuracy boost, because it exploits the specific way SGD's noise wanders around a flat basin rather than any single noisy readout of it.

What this means in practice

SWA is nearly free: it costs one extra set of weights to store and average, and no extra gradient computations, since it just re-uses snapshots from training you were already doing. It's commonly paired with a cyclical or constant learning rate at the end of training specifically so the snapshots explore the basin rather than converge to a single point too tightly to have anything interesting to average. It's distinct from SWA-Gaussian (SWAG), a related technique that also models the spread of those snapshots for uncertainty estimation, not just their average.

Practice

1. Four snapshots of a weight are $1.0, 1.2, 0.9, 1.1$. Compute the running average after all four.
2. Why does SWA usually pair with a constant or cyclical learning rate rather than one that's still decaying toward zero?
3. Why is averaging weights different from averaging the predictions of several models (an ensemble)?