Skip to content

← Production ML — Training Stack step 6 of 12

Medium Primitives

Exponential Moving Average

An exponential moving average (EMA) keeps a smoothed copy of a model's parameters. Each update blends the previous EMA with the latest parameter values. A larger decay gives the previous EMA more weight, so it changes more slowly.

Your task

Implement update_ema(ema_weights, current_weights, decay).

  • ema_weights is the previous EMA vector, with shape (d,).
  • current_weights is the current parameter vector, also (d,).
  • Both vectors contain finite floating-point values and have matching dtype and device.
  • decay is a scalar in [0, 1].

Compute one elementwise update:

new_ema = decay * ema_weights + (1 - decay) * current_weights

Return a vector of shape (d,) with the same dtype, leaving both inputs unchanged. There is no optimizer step, averaging across vector coordinates, bias correction or parameter initialization to perform.

For ema_weights=[2, -4], current_weights=[10, 4] and decay=0.75, return [4, -2]: keep three quarters of the old vector and add one quarter of the current vector. The endpoint rules are:

  • decay=0 returns the current values.
  • decay=1 returns the old EMA values.
  • decay=0.5 gives equal weight to the two supplied vectors.

What makes it exponential?

When the returned vector is used as the old EMA for the next update, each older contribution is multiplied by decay again. For a fixed decay between 0 and 1, its influence shrinks geometrically with age. This is not an average over a fixed window: old contributions do not disappear at a particular step.

For example, starting from scalar EMA 0 with decay=0.5, a current value of 10 gives EMA 5. A subsequent current value of 6 gives EMA 5.5. The second call uses 5 as its old EMA, not the original 0.

This function only computes that update. Choosing how to initialize the EMA, how often to update it and when to use the smoothed parameters belongs to the surrounding training loop.