Optax EMA on Params

An exponential moving average (EMA) blends an existing estimate with a new observation. Here the observations are model parameters, and params holds the previous average, not necessarily the model's current weights.

Blend old and new values

Apply this rule independently at every position:

new_average = decay * params + (1 - decay) * new_params

For params = [1, 2], new_params = [3, 6] and decay = 0.9, return [1.2, 2.4]. The old estimate contributes 90 percent and the new values contribute 10 percent. Repeating the update with the same new values gives [1.38, 2.76].

A larger decay retains more of the old estimate. With decay zero, take the new values directly; with decay one, keep the old values. For decay between zero and one, each result lies between its old and new values, up to rounding.

What this exercise includes

Compute one blend directly with JAX array operations. There is no gradient, learning rate, step counter or bias correction. To continue an EMA sequence, the caller would pass the returned average as params on the next call.

Optax also provides optax.ema, a stateful transformation of incoming updates with optional bias correction, enabled by default. It has a different interface from this task's explicitly supplied old estimate. An exponential average also differs from an equal-weight average of saved model snapshots.