Use an Exponential Moving Average of Model Weights

Neural network parameters move after every optimizer step. Near the end of training, those updates can leave a model bouncing around a useful region instead of sitting at one especially representative point.

An exponential moving average, usually shortened to EMA, keeps a second set of weights that changes more smoothly. The optimizer still updates the normal model. EMA simply tracks a weighted history of those parameters and can provide a separate set of weights for evaluation or deployment.

The mechanism is small enough to implement in a few lines, but details such as update timing, decay, checkpoint state, and non-parameter buffers determine whether the result is trustworthy.

Keep two versions of the parameters

Let the trainable model parameters after optimizer step (t) be (\theta_t). Maintain another set, (\bar{\theta}_t), for the moving average.

After each optimizer update, compute:

[ \bar{\theta}t = \beta \bar{\theta}{t-1} + (1-\beta)\theta_t ]

where (\beta) is a decay value between 0 and 1.

A larger (\beta) gives older parameter values more influence. A smaller value makes the average follow the current model more closely.

The key separation is simple:

• (\theta_t) is the model the optimizer changes.
• (\bar{\theta}_t) is a tracked copy used as an alternative set of weights.

Backpropagation does not need to run through the EMA update. The moving average is bookkeeping over parameter values, not another optimization objective.

A tiny numerical example

Suppose one scalar parameter has these values after three optimizer steps:

[ \theta_1 = 10,\quad \theta_2 = 14,\quad \theta_3 = 12. ]

Initialize the EMA from the first observed value:

[ \bar{\theta}_1 = 10. ]

Using (\beta = 0.8), the next update is:

[ \bar{\theta}_2 = 0.8(10) + 0.2(14) = 10.8. ]

After the third optimizer step:

[ \bar{\theta}_3 = 0.8(10.8) + 0.2(12) = 11.04. ]

The raw parameter moved from 10 to 14 and then back to 12. The EMA moved from 10 to 10.8 to 11.04. It responds to the same updates but changes more gradually.

A real model applies this operation element by element across every tracked parameter tensor.

EMA is not a replacement for the optimizer

A common conceptual mistake is to feed averaged weights back into the optimizer after every step. Standard EMA usage keeps the two paths separate.

A training loop has this shape:

for batch in data:
    loss = model(batch)
    loss.backward()

    optimizer.step()
    optimizer.zero_grad()

    update_ema(ema_weights, model_weights)

The optimizer owns the live model weights. The EMA state observes the result after the optimizer step.

This ordering matters. If the average is updated before optimizer.step(), it records the previous parameter state instead of the newly updated one. That shifts the sequence being averaged and makes the implementation disagree with the intended update rule.

Some training systems update EMA less often than every optimizer step to reduce overhead. That can be valid, but the decay then applies per EMA update rather than per training step. Decay values are not directly comparable across different update frequencies.

The decay controls the averaging horizon

Repeated substitution of the EMA equation shows that recent parameter states receive geometrically decreasing weights. Ignoring initialization terms, a parameter value from (k) updates ago receives a factor proportional to:

[ (1-\beta)\beta^k. ]

This gives decay a useful interpretation. With (\beta) close to 1, the average spans a longer history. With a lower value, recent updates dominate sooner.

A rough effective horizon is often described as being on the order of

[ \frac{1}{1-\beta} ]

EMA updates. This is a mental model, not a hard cutoff. Exponential weighting never reaches zero at a fixed age.

For example, changing the update frequency changes the amount of training time represented by the same numeric decay. A decay chosen for an EMA updated every optimizer step should not be copied blindly to an EMA updated once every ten steps.

Initialization changes early behavior

One simple strategy is to initialize EMA weights as a copy of the current model weights before tracking begins:

[ \bar{\theta}_0 = \theta_0. ]

This avoids starting the average from an unrelated all-zero parameter set.

Another implementation may initialize accumulators at zero and apply a bias correction during early updates. That approach can also be mathematically coherent, but the correction must match the exact averaging rule.

For most custom training code, copying the model parameters at the point where EMA tracking starts is easier to inspect and harder to misuse.

Tracking can also begin after a warm-up period. If so, initialize from the model state at that point rather than pretending earlier updates were included.

Evaluate raw and EMA weights separately

EMA weights should be treated as a distinct model state during evaluation.

A safe evaluation pattern is:

1. preserve the live training parameters;
2. load or swap in the EMA parameters;
3. run evaluation without optimizer updates;
4. restore the live parameters before training resumes.

An alternative is to keep a separate evaluation model populated from EMA state. That uses more memory but avoids temporary swapping.

Do not overwrite the live training weights permanently unless the training procedure explicitly calls for it. If training continues from averaged weights by accident, the optimizer state may no longer correspond to the parameters it had been updating.

It is also useful to measure both versions. EMA can improve evaluation behavior in some training setups, but it is not guaranteed to outperform the current parameters at every checkpoint. Validation results should decide which state is appropriate for deployment.

Check buffers as well as parameters

Neural network state can contain values that are not optimizer-managed parameters.

Batch normalization, for example, commonly maintains running statistics as buffers. A parameter-only EMA does not automatically average or refresh those values. If EMA parameters are evaluated with stale or mismatched running statistics, the result may not represent the intended model.

Possible strategies depend on the architecture and framework:

• copy suitable buffers from the live model;
• maintain buffer state with an explicit policy;
• recompute running statistics using data before final evaluation when the method supports it.

There is no universal rule that every buffer should receive the same EMA equation as trainable weights. Treat buffers according to their semantics.

Architectures without stateful normalization buffers have fewer moving parts, but checkpoint code should still distinguish parameters from other persistent state.

Account for optimizer-step frequency

Gradient accumulation introduces another easy source of confusion.

Suppose gradients are accumulated over four microbatches before one optimizer step. The model parameters do not change after each microbatch. If EMA is intended to track parameter updates, update it once after the optimizer step, not four times during accumulation.

The same principle applies when an optimizer step is skipped, such as after invalid numerical values are detected by a mixed-precision training system. If the live parameters did not receive the expected update, EMA logic should follow the actual parameter-update semantics of the training loop.

Thinking in optimizer steps rather than batches keeps the averaging rule precise.

Distributed training needs one coherent source

In data-parallel training, model replicas are typically synchronized through the training framework. EMA should track a coherent parameter state after synchronization and optimizer updates.

A simple design is to maintain EMA on every worker if each worker has identical model parameters and executes identical EMA updates. Another design keeps one EMA copy on a designated process.

The right choice depends on the distributed system. The invariant is more useful than a framework-specific recipe: the EMA used for evaluation and checkpointing must correspond to a well-defined sequence of model parameter states.

If only one process owns EMA, checkpoint and evaluation code must know that. If every process owns it, occasional equality checks can expose synchronization mistakes.

EMA has a real memory and compute cost

Maintaining EMA requires another copy of each tracked parameter. For a large model, that extra storage can be substantial.

The update also reads the current parameter, reads the old average, performs arithmetic, and writes the new average. Relative to a full training step this may be modest, but it is not free, especially when parameters live across devices or storage tiers.

Precision is another design choice. Keeping EMA state in a stable floating-point format can avoid accumulating extra rounding error from repeatedly updating a low-precision copy, but it increases memory when the live model uses a smaller format. The appropriate choice depends on the training stack and available memory.

If memory is tight, compare EMA against simpler alternatives before adding another full parameter copy.

Checkpoints must preserve both states

A training checkpoint that supports exact continuation should record enough state to restore:

• live model parameters;
• optimizer state;
• EMA parameters;
• the number of EMA updates or other schedule state if decay behavior depends on it;
• any scheduler and training state already required by the run.

Saving only EMA weights can be sufficient for an inference-only artifact, but it is not equivalent to a resumable training checkpoint.

Use explicit names such as model_state and ema_state rather than storing two anonymous parameter dictionaries. Clear naming reduces the chance that a deployment pipeline exports the wrong version.

After restoring a checkpoint, compare evaluation output against a pre-save reference on a fixed input when practical. That catches missing EMA state and accidental swaps quickly.

Common implementation mistakes

Several bugs recur because the EMA equation itself looks too simple to fail.

Updating before the optimizer step. This tracks the wrong point in the update sequence.

Averaging gradients instead of parameters. Gradient averaging is a different operation with different effects.

Applying EMA once per microbatch during gradient accumulation. The decay then follows microbatch count even though parameters change only on optimizer steps.

Forgetting non-parameter state. Stateful buffers can make EMA evaluation inconsistent even when every trainable tensor is correct.

Resuming without EMA state. Reinitializing the average after a restart changes its history and can create a visible discontinuity.

Assuming EMA must be better. It is a candidate model state, not a guarantee. Compare it with the live model using the same evaluation procedure.

Decide based on the training behavior you need

EMA is most attractive when you want a smoothed parameter trajectory and can afford an additional model-sized state. It is straightforward to add to many neural network training loops because it does not change the loss function or require gradients through the averaging operation.

It is less attractive when memory is already the main constraint, when model state contains difficult-to-handle buffers, or when validation shows no useful benefit. Short training runs can also leave little time for a high-decay average to become representative.

Start with the invariant: the optimizer updates one model, and EMA tracks those completed parameter updates in a second state. Once that separation is correct, decay, update frequency, checkpointing, and evaluation become explicit engineering choices rather than hidden assumptions.