Skip to content

← Flax step 69 of 100

Medium Primitives

EMA of Parameters

Jump to the task

An exponential moving average (EMA) keeps a smoothed copy of model parameters. Each update blends the previous average with the current parameters:

new_ema = decay * old_ema + (1 - decay) * current

The operation is elementwise. It does not average across different parameters or replace the optimizer's update. Keep the live parameters and their moving average as separate values.

What decay controls

A larger decay retains more of the previous average; a smaller decay follows the current parameters more quickly. For example, with decay=0.9, an old value of 2 and a current value of 6 produce 2.4.

For a fixed decay d, a parameter value from i updates ago has weight (1-d) * d**i. After t updates, the initial average still contributes d**t * initial_ema. Older contributions decrease geometrically; there is no hard cutoff after a fixed number of steps.

EMA can be useful for evaluation, but it does not guarantee better accuracy. Compare the live and averaged parameters on validation data. The decay and initialization affect how quickly the average responds to training changes.

Initialization matters

Starting the average from the initial model parameters avoids introducing an all-zero starting buffer. If you instead start at zero, repeated updates retain a contribution from that zero initialization. Some algorithms apply a bias correction by dividing by 1 - decay**t; this exercise does not. It takes the existing average as an input and performs exactly one update.

For a full parameter tree, the same arithmetic can be applied to corresponding leaves. optax.incremental_update(current, old_ema, step_size=1-decay) expresses that operation. Here the parameters are already flattened, so no tree traversal or Flax module is needed.

Problem

Implement ema_apply(seed, params_a_flat, params_b_flat, decay).

  • params_a_flat is the existing average, with shape (K,).
  • params_b_flat contains the current parameters in the same order and shape.
  • Both arrays contain finite values, and K >= 1. Convert them to float32.
  • decay is a scalar with 0 < decay < 1.
  • seed is retained for signature compatibility and is unused. Do not sample random values or initialize a model.

Return decay * params_a_flat + (1 - decay) * params_b_flat as a float32 JAX array of shape (K,). Leave both inputs unchanged. Do not add bias correction, normalize the result, or average over the parameter axis.