A neural network’s final training step is not necessarily its most useful checkpoint. Stochastic optimization keeps moving the parameters as it follows noisy mini-batch gradients, so two nearby checkpoints can behave slightly differently even when training is otherwise healthy.
An exponential moving average (EMA) of model weights gives you a second set of parameters that changes more smoothly. Instead of evaluating only the latest training weights, you maintain a weighted history in which recent weights matter most and older weights gradually fade away.
EMA is simple, but several details determine whether it helps: when to start averaging, how to choose the decay, which state to average, and how to evaluate without accidentally continuing training from the wrong parameters. This article builds the method from one scalar example and turns it into a practical training pattern.
Keep two versions of the weights
Suppose the optimizer owns the ordinary training weights. Call them theta.
EMA maintains a separate copy, theta_ema. After an optimizer update, refresh that copy with:
theta_ema = decay * theta_ema + (1 - decay) * thetawhere decay is between 0 and 1.
The optimizer still updates theta. Gradients are not normally computed through theta_ema, and the EMA copy does not replace the optimizer’s working parameters during training.
A useful mental model is:
mini-batch
-> gradients
-> optimizer updates theta
-> EMA observes new theta and updates theta_ema
training continues with theta
evaluation can use theta_emaThis separation is the core implementation rule. EMA is a shadow copy of the parameters, not a second optimizer.
Work through the smallest example
Use one scalar parameter so the arithmetic is visible. Let the EMA decay be 0.9, and initialize the EMA from the initial training value:
theta = 10.0
theta_ema = 10.0After the first optimizer step, suppose theta becomes 12.0:
theta_ema = 0.9 * 10.0 + 0.1 * 12.0
= 10.2After the next step, suppose theta becomes 11.0:
theta_ema = 0.9 * 10.2 + 0.1 * 11.0
= 10.28The training parameter moved from 10 to 12 to 11. The EMA moved from 10 to 10.2 to 10.28.
That slower movement is intentional. The EMA suppresses some short-term variation in the optimization path while still following sustained changes.
Understand what the decay controls
Repeatedly expanding the update shows that EMA is a weighted combination of parameter values from different training steps. Ignoring initialization for a moment, a weight from k updates ago is multiplied by approximately:
decay^kWith a larger decay, old parameter values fade more slowly. With a smaller decay, the EMA follows the current model more closely.
For intuition, consider the weight assigned to a value 100 updates in the past:
decay = 0.90 -> 0.90^100 is very small
decay = 0.99 -> 0.99^100 is about 0.366
decay = 0.999 -> 0.999^100 is about 0.905These numbers do not mean that one decay is universally better. They show that the same numeric decay has a very different time scale depending on how often the EMA is updated.
If you update EMA after every optimizer step, its effective history is measured in optimizer steps. If you update it only every ten optimizer steps, the same decay spans a much longer portion of training.
This is why copying a decay value without copying the update schedule can produce very different behavior.
Think in terms of an effective averaging window
A rough way to reason about EMA memory is:
effective window ~ 1 / (1 - decay)This is an intuition, not a hard cutoff. EMA has an exponentially decaying tail rather than a fixed-size window.
The approximation gives useful scale:
decay 0.9 -> roughly 10 updates
decay 0.99 -> roughly 100 updates
decay 0.999 -> roughly 1000 updatesThe right scale depends on training dynamics. If parameters are changing rapidly because the model is early in training or the learning rate is large, a very slow EMA can lag far behind. Later, when updates are smaller, a longer average may provide a useful stable checkpoint.
Choose the decay by validation behavior and training scale rather than by the number of nines alone.
Update EMA after the optimizer step
The ordering of operations matters. A typical training iteration is:
1. compute loss with theta
2. compute gradients
3. optimizer updates theta
4. update theta_ema from the new thetaUpdating EMA before the optimizer step means it observes the previous parameter state instead. That is not mathematically invalid, but it shifts the sequence being averaged and can make an implementation disagree with the intended algorithm or another training run.
Be explicit about the order, especially when reproducing a training recipe.
Gradient accumulation adds another detail. If several micro-batches contribute gradients to one optimizer step, EMA usually follows optimizer updates, not micro-batches. Updating EMA after every micro-batch would repeatedly average the same unchanged parameters while gradients are merely accumulating.
Initialize the shadow weights deliberately
The simplest initialization is:
theta_ema = copy(theta)Then every EMA value begins on the same scale as the trained parameter.
Another family of implementations starts an accumulator near zero and applies bias correction, similar in spirit to correcting an exponential average that has not accumulated much history. That can be valid too, but the formulas and checkpoint semantics differ.
Do not mix the two approaches accidentally. If you initialize EMA from the model weights, applying a bias correction designed for a zero-initialized accumulator changes the intended estimate.
Some training recipes also delay EMA until a warm-up point or vary the decay early in training. The motivation is straightforward: a long-memory average can retain too much influence from poorly trained early weights. Whether a warm-up helps depends on the optimization trajectory, so treat it as a tunable training choice rather than a requirement of EMA itself.
Decide which model state belongs in the average
For ordinary neural-network parameters, averaging floating-point trainable weights is straightforward. Real model checkpoints can contain more state than those parameters.
For example, batch normalization commonly keeps running statistics in addition to learned scale and bias parameters. Those running statistics are buffers, not ordinary gradient-updated weights. Optimizer state such as momentum estimates is separate again.
A robust implementation must define what EMA model means:
trainable parameters -> usually averaged
optimizer state -> not EMA model weights
non-trainable buffers -> handle according to architecture and frameworkBlindly applying the parameter formula to every object in a checkpoint can be wrong. Some buffers are integer counters; some are running statistics with their own update rules; some model families have no such state.
Follow the architecture’s evaluation semantics and make checkpoint loading tests part of the implementation.
Evaluate the EMA model without corrupting training
There are two common patterns.
The cleanest is to maintain a separate shadow model or parameter collection for EMA evaluation. Training uses theta; validation loads or references theta_ema.
Another pattern temporarily swaps EMA values into the training model:
save current theta
copy theta_ema into model
evaluate
restore saved thetaThis can save memory compared with a complete second model object, but restoration must be reliable. If training resumes with EMA weights while the optimizer state still corresponds to the original training trajectory, you have silently changed the optimization process.
Checkpointing should therefore make the roles explicit. If you intend to resume training, save the normal training weights and optimizer state. If EMA is used for evaluation or deployment, save the EMA weights too. Do not assume one set can replace the other for every purpose.
Why weight averaging can help
Mini-batch optimization is noisy because each update is based on a sample of the training data. Learning-rate schedules, data order, augmentation, and optimizer dynamics can also make the parameter path fluctuate.
EMA smooths that path in parameter space. When nearby training checkpoints represent similarly useful solutions, averaging them can produce evaluation weights that are less sensitive to the exact final update.
But this explanation has an important boundary: neural-network parameter space is not globally linear in a way that makes arbitrary weight averaging safe. Averaging unrelated models can produce a poor model because hidden units, symmetries, and optimization basins may not align.
EMA avoids part of that problem by averaging successive states from one continuous training trajectory, with exponentially decreasing influence from older states. Even then, improvement is empirical rather than guaranteed.
Measure the trade-offs instead of assuming improvement
EMA adds relatively little arithmetic, but it is not free. Keeping a full shadow copy adds storage proportional to the parameters being averaged. For very large models, that memory can be significant, especially if the EMA copy uses high-precision weights.
EMA can also lag behind rapid improvements. Suppose the learning rate changes and the training model quickly enters a better region. A high-decay EMA still contains substantial weight from earlier states, so validation with EMA may improve more slowly.
The practical comparison is simple: evaluate both the current training weights and EMA weights on the same validation procedure over time. Track the metric that matters for deployment rather than only training loss.
If EMA consistently improves or stabilizes validation performance enough to justify its memory and checkpoint complexity, keep it. If it does not, the simpler training pipeline is preferable.
Avoid common mistakes
Averaging gradients instead of weights. Gradient accumulation and EMA solve different problems. Gradient accumulation combines gradient information before an optimizer update. EMA combines parameter states after updates.
Updating on the wrong clock. The decay only has meaning together with the EMA update frequency. Changing from per-step to occasional updates changes the effective averaging horizon.
Training directly on EMA weights by accident. The optimizer should normally continue from its own current parameters. EMA is a shadow trajectory unless the training algorithm explicitly specifies otherwise.
Forgetting mixed-precision details. Training may use low-precision compute while retaining higher-precision master parameters. Decide which parameter representation feeds EMA and verify the numerical behavior. Framework-specific mixed-precision implementations differ, so do not assume the visible model tensor is always the optimizer’s authoritative copy.
Saving only the deployment weights. EMA weights can be enough for inference, but resuming training generally also requires the non-EMA training state expected by the optimizer and scheduler.
Treating EMA as an ensemble. EMA produces one averaged parameter set. It does not retain multiple independent predictions at inference time, so it does not provide the same diversity or uncertainty information as an ensemble of separately evaluated models.
Know when a simpler checkpoint is enough
EMA is most useful when training is long enough for a meaningful parameter history to accumulate, validation metrics fluctuate across nearby steps, and an additional parameter copy fits the resource budget.
It may add little value when training is already highly stable, the final checkpoint is selected reliably by validation, or model size makes another parameter copy expensive. It is also not a substitute for fixing unstable optimization. If loss diverges, gradients explode, or the learning rate is inappropriate, smoothing the resulting weights does not address the underlying problem.
EMA should be viewed as a checkpoint-quality technique layered on top of a sound training process.
Conclusion
Exponential moving average weights give neural-network training a second, smoother parameter trajectory. The optimizer updates the normal model; after each optimizer step, EMA blends the new parameters into a shadow copy whose memory is controlled by the decay.
The method is easy to write down but should be implemented deliberately. Tie the decay to the update frequency, initialize the shadow state consistently, handle non-parameter state correctly, keep training and evaluation weights separate, and compare both versions on representative validation data. EMA is valuable when that measured stability or quality improvement is worth the extra model state—not because averaging is automatically better than the latest checkpoint.