Let shadow weights follow without chasing every update.

An exponential moving average keeps a second parameter set that trails noisy online weights. Tune decay and cadence to see stability improve while lag grows after a shift.

Average the weights

Online weights and their moving average

θEMA ← 0.990 θEMA + 0.010 θonline
Online weight0.00

Latest optimizer parameter.

EMA weight0.00

Smoothed shadow parameter.

Evaluation error0.00

Error for the selected evaluation weights.

EMA lag0.00

Distance from online weights.

Smoothing trades variance for responsiveness.

High decay remembers longer

β near one suppresses update noise but needs more steps to follow a changing optimum.

Shadow state belongs in checkpoints

Save EMA weights, update count, decay schedule, optimizer state, and the online model so training can resume faithfully.

Evaluate the deployed choice

Raw and averaged weights are different models. Compare validation loss, calibration, robustness, and latency on the intended checkpoint.

Primary reading

ExponentialMovingAverage TensorFlowAveragedModel PyTorchAveraging Weights Leads to Wider Optima Izmailov et al.