Stabilize Neural Network Weights with Exponential Moving Averages

A neural network’s parameters rarely move smoothly toward their final values. Mini-batch training produces noisy updates: one batch may push a weight in one direction, while the next pushes it partly back. The final checkpoint therefore represents one point on a noisy training path, not necessarily the most useful point near the end of that path.

An exponential moving average (EMA) of model weights keeps a second set of parameters that changes more slowly than the actively trained model. Recent training states contribute more than old ones, but no single update immediately replaces the averaged weights.

This technique is simple, but several implementation details matter. By the end of this article, you will understand what an EMA is averaging, how its decay controls the time scale, how to update and evaluate EMA weights without corrupting training, and when averaging weights is less useful than fixing the training process itself.

Think of EMA as a shadow model

Suppose training produces a sequence of parameter values:

step 1: 10.0
step 2: 12.0
step 3:  9.0
step 4: 11.0

Using only the last value gives 11.0. A simple arithmetic average gives 10.5, but it treats every historical value equally. During long training runs, very old parameters may no longer describe the region in which the model is currently learning.

EMA instead gives recent values more influence. For a model parameter w, maintain a shadow parameter w_ema:

w_ema = decay * w_ema + (1 - decay) * w

Here, w is the newly updated training weight and decay is a number between 0 and 1. The training optimizer continues to modify w; it does not optimize w_ema directly.

With decay = 0.9, each update keeps 90% of the previous EMA value and takes 10% from the current training value. With decay = 0.999, the EMA changes much more slowly.

That separation is the central mental model: one parameter set learns; the other follows.

Why averaging can help

Mini-batch gradients are estimates computed from subsets of the training data. Different batches contain different examples, so their gradients differ. Learning-rate schedules, augmentation, dropout, and other sources of randomness can add more variation.

Near a useful solution, the trained parameters can therefore move around within a region rather than settling at one exact point. EMA acts as a low-pass filter on that trajectory. Fast changes have less effect on the shadow weights than persistent changes do.

This does not mean EMA magically finds a better optimum. It means the model used for evaluation is less sensitive to the exact final training update. Whether that improves validation quality depends on the model, optimizer, learning-rate schedule, training duration, and task.

Understand the decay as a time scale

The decay controls how quickly old information loses influence. Expanding the recurrence makes this clearer. Ignoring initialization for a moment, a parameter from k updates ago is weighted approximately by:

(1 - decay) * decay^k

So contributions shrink geometrically with age.

A useful rough measure is the effective averaging horizon:

horizon ~= 1 / (1 - decay)

For example:

decay = 0.9    -> about 10 updates
decay = 0.99   -> about 100 updates
decay = 0.999  -> about 1,000 updates

This is a mental model, not a hard cutoff. EMA never suddenly forgets an update after that many steps.

The consequence is important: decay should be interpreted relative to update frequency, not merely epochs. A decay of 0.999 behaves differently when a model receives 200 optimizer updates per epoch than when it receives 20,000.

If gradient accumulation is used, decide what counts as an EMA update. Updating after each optimizer step usually matches the idea of following actual parameter changes. Updating on every micro-batch even though the optimizer weights have not changed only repeats the same value and changes the effective averaging behavior.

Implement the smallest correct version

The framework-independent logic is short:

initialize training model
initialize EMA model as a copy of training model

for each optimizer step:
    compute gradients
    update training model

    for each averaged parameter:
        ema = decay * ema + (1 - decay) * training_parameter

The EMA update belongs after the optimizer update if the intention is to average the sequence of post-update model states.

The shadow parameters should not receive gradients. Treat them as maintained state, similar to a checkpoint copy that is updated by the EMA rule rather than by an optimizer.

For floating-point parameters, the update can also be written as:

ema += (1 - decay) * (training_parameter - ema)

The two forms are mathematically equivalent apart from finite-precision effects.

Keep training and evaluation roles separate

A common pattern is:

training:
    forward/backward with training weights
    optimizer updates training weights
    EMA follows training weights

evaluation:
    run validation with EMA weights

Do not accidentally continue optimization from EMA weights unless that is an explicit training design. Swapping EMA weights into the training model for validation and forgetting to restore the original weights changes the next optimizer step and silently turns the procedure into something else.

Keeping a separate shadow model is easier to reason about, although it consumes additional memory. Another implementation can store only a shadow copy of parameters and temporarily swap them for evaluation, provided the swap and restoration are reliable.

Handle initialization deliberately

The simplest initialization copies the training parameters into the EMA state before training begins:

ema = training_parameter

This avoids starting the average from zero, which would bias early EMA values toward zero.

Even with copy initialization, a very large decay makes the EMA respond slowly during early training. If decay = 0.9999, for example, the shadow model retains a large influence from its initial state for many updates.

Some training systems address this with a warm-up or a step-dependent decay that begins lower and increases later. That can be useful, but it is an additional policy rather than part of the EMA definition. If a framework provides an EMA utility, check its exact initialization and decay schedule instead of assuming it uses the simple recurrence shown here.

Parameters are not the only model state

Neural networks can contain state that is not optimized as a normal parameter. Batch normalization, for example, can maintain running statistics used during inference.

This creates an implementation question: should such buffers be copied, averaged, or recomputed?

There is no universal answer because it depends on the layer and framework. Blindly applying the parameter EMA equation to every tensor in a model can be wrong. Integer counters cannot be meaningfully averaged in the same way as floating-point weights, and inference statistics may need treatment consistent with the model implementation.

A robust EMA implementation therefore distinguishes trainable floating-point parameters from other buffers. If the architecture contains stateful normalization or similar components, verify how the chosen framework or training recipe handles their inference state before relying on the EMA checkpoint.

Account for memory and distributed training

EMA adds state. Keeping a full shadow copy of model parameters requires roughly another model-sized set of weights, though the exact memory cost depends on data types and what is stored. For a small model this may be negligible; for a large model it can be a meaningful constraint.

Distributed training adds another question: where should the EMA live? If every worker already has synchronized model parameters, each worker can in principle maintain the same EMA when updates occur identically. Other systems may keep the shadow weights only on selected devices or offload them to reduce accelerator memory use. Those choices can trade memory for communication or update latency.

The key requirement is consistency. The EMA should follow the parameter sequence that defines the trained model. Averaging stale or differently synchronized copies changes what the shadow model represents.

Evaluate whether EMA actually earns its cost

Treat EMA as a model-selection choice, not an assumed improvement. During development, compare validation results from both parameter sets under the same evaluation procedure:

checkpoint A: current training weights
checkpoint B: EMA weights

Measure the metrics that matter for the application. A small classification gain may be useful when inference cost is unchanged, because the deployed model still uses one set of weights. But the extra training memory and checkpoint complexity may not be justified if validation quality is unchanged.

Also test more than one point near the end of training when practical. A single favorable evaluation can be noise, especially on a small validation set.

EMA generally does not increase inference latency after deployment if you export only the chosen EMA weights. Its main costs occur during training: extra state, extra parameter updates, and more checkpoint bookkeeping.

Avoid common mistakes

Using EMA to hide unstable training

EMA can smooth a noisy parameter path, but it does not repair exploding gradients, an unsuitable learning rate, corrupted data, or a model that fails to learn. If the underlying training loss is unstable or diverging, diagnose that problem first.

Choosing decay without considering update count

Copying a decay value from another project can produce a very different effective horizon. Compare optimizer-step counts and training duration before treating a decay as transferable.

Updating at the wrong moment

If the intended sequence is the model after optimizer steps, update EMA after those steps. Mixing pre-update and post-update conventions makes experiments harder to compare.

Saving only one side unintentionally

If training must resume exactly, the checkpoint may need the live model, optimizer state, scheduler state, EMA state, and any other training state. Saving only EMA weights is sufficient for some inference exports but is not generally a complete training checkpoint.

Assuming averaged weights must outperform the final weights

EMA is a bias toward a smoother recent trajectory. If the model is still improving rapidly, an average can lag behind the newest useful parameters. If training has already converged smoothly, averaging may add little. Validation decides whether the trade-off is worthwhile.

EMA versus other kinds of averaging

EMA is not the only way to combine model states. A plain checkpoint average can combine a fixed set of saved checkpoints with equal weights. Stochastic weight averaging uses a different training and averaging procedure designed to average selected model states, often over a later training phase. Model soups combine weights from multiple compatible fine-tuned models selected after separate training runs.

EMA is distinct because it is an online, exponentially weighted average maintained as training progresses. It needs no list of old checkpoints, but its result depends on the full sequence of updates and the chosen decay.

Use EMA when you want a low-overhead way to smooth recent training states and can afford the shadow parameters. Prefer a simpler final checkpoint when training is already stable and EMA shows no validation benefit. Prefer other averaging methods when you specifically need to combine selected checkpoints or independently fine-tuned models rather than continuously track one training run.

Conclusion

Exponential moving averages turn a noisy sequence of neural network weights into a slowly changing shadow model. The essential rule is simple: train one parameter set normally, then let a second set follow it with an exponential decay.

The practical details determine whether that simplicity survives real training. Interpret decay in optimizer steps, initialize the shadow state deliberately, keep evaluation from mutating training weights, handle non-parameter buffers explicitly, and budget for the extra model state. Most importantly, compare EMA and non-EMA checkpoints on the same validation criteria. EMA is useful when smoothing the training trajectory improves the model you actually care about; otherwise, the simpler checkpoint is the better engineering choice.