Stochastic Weight Averaging and EMA Weights

Training a neural network with stochastic gradient descent doesn't settle onto one clean minimum — it bounces around near the bottom of the loss surface for the rest of training, buffeted by the noise of different random mini-batches every step. Whichever point it happens to be at when you stop becomes "the model," even though that point is essentially a single noisy sample from a cloud of similarly-good weight settings the optimizer visited. A different random seed, or stopping one epoch earlier, would hand you a noticeably different final model. SWA and EMA both address this the same way: instead of trusting the last snapshot, average several snapshots together.

The analogy: averaging noisy GPS readings

A phone's GPS reading jitters around your true location moment to moment due to signal noise. Trust only the very last reading before your phone died, and you'd have a noisy estimate. But log readings for a minute and average them, and the jitter mostly cancels, leaving an estimate closer to your true position than any single reading. SWA and EMA do this to a network's weights: average several points visited near the end of training instead of trusting the one, individually-noisy point SGD happened to land on.

Two ways to average, same underlying idea

Stochastic Weight Averaging runs training normally for most of the schedule, then — typically once the learning rate has been lowered and the optimizer is circling a good region rather than making large jumps — starts a separate running average of the weights, updated at fixed intervals (say, once per epoch):

$$\bar{w} \leftarrow \frac{\bar{w} \cdot n + w_{\text{current}}}{n + 1}$$

In words: the new average is the old average, weighted by how many snapshots contributed to it so far, blended with the current weights, giving every collected snapshot equal say.

EMA weights instead keep a running average updated at every step with a decay factor $\beta$ close to 1 (e.g., $0.999$):

$$\bar{w} \leftarrow \beta\, \bar{w} + (1-\beta)\, w_{\text{current}}$$

In words: the average leans heavily on its own past value and only slightly toward the freshest weights, changing slowly and smoothly, giving recent snapshots more weight than distant ones (unlike SWA's flat average). Both produce a weight vector no single SGD step ever visited, but which tends to sit in a flatter region of the loss surface — and flatter regions are empirically associated with better generalization, because small mismatches between training and live data move you less far up the loss surface than a sharp, narrow minimum would.

Drag the settings above and watch SGD-style steps overshoot and oscillate near a minimum rather than settling exactly on it — that oscillation is exactly the noise SWA and EMA are built to average away.

Worked example 1: EMA by hand over four steps

Suppose a single weight takes the noisy values $2.0, 2.4, 1.8, 2.2$ over four steps, tracked with $\beta = 0.5$ (exaggerated for a hand-computable example), starting the average at the first value:

$$\bar w_1 = 2.0,\quad \bar w_2 = 0.5(2.0)+0.5(2.4)=2.2,\quad \bar w_3 = 0.5(2.2)+0.5(1.8)=2.0,\quad \bar w_4 = 0.5(2.0)+0.5(2.2)=2.1$$

The raw values swing from 1.8 to 2.4 (a range of 0.6); the EMA only swings from 2.0 to 2.2 (a range of 0.2) — visibly smoother, and closer to the values' true center of roughly 2.1.

Worked example 2: SWA's flat average versus the last snapshot

Suppose the last five epochs produced weight snapshots (simplified to one number) of $3.1, 2.9, 3.3, 2.8, 3.0$. The final snapshot alone is $3.0$. The SWA average is $(3.1+2.9+3.3+2.8+3.0)/5 = 3.02$ — close in this case. But if the final snapshot had instead been an unlucky outlier, say $2.5$, from a mini-batch that pushed weights the wrong way right before stopping, trusting it alone would give a noticeably worse model than the $3.02$ average, since the average dilutes any one unlucky step across all five.

What this means in practice

SWA is typically combined with a cyclical or constant learning rate near the end of training, so the averaged snapshots are meaningfully different points circling one basin, not points still descending; EMA needs no separate phase — it runs alongside normal training the whole time, and the EMA copy is what gets deployed. Both add negligible compute cost: an extra weight buffer and a cheap update, no extra forward or backward passes.

Practice

1. With EMA decay $\beta = 0.9$ and raw values $10, 12, 8$ (starting the average at $10$), compute $\bar w$ after each step.
2. Why does SWA typically wait until late in training, after the learning rate has been lowered, to start averaging?
3. In your own words, explain why a flatter region of the loss surface should be more robust to the mismatch between training data and live data than a sharp, narrow minimum.