Skip to content

← Flax NNX step 16 of 100

Medium Primitives

NNX Running Average / EMA Step

Jump to the task

Why this matters

An exponential moving average (EMA) combines the previous estimate with the newest value. It can smooth a changing signal or maintain a separate averaged set of model weights. Whether averaged weights improve a model depends on the training setup; it is not a guaranteed accuracy gain.

For decay d, one update is:

new_ema = d * old_ema + (1 - d) * current_value

Larger decay retains more of the old estimate. With d = 0.9, the new value receives weight 0.1. This differs from an arithmetic average that gives every past sample equal weight.

This exercise implements the array calculation directly. It does not train a model, update an optimizer or require an NNX module.

Worked example

import jax.numpy as jnp

old = jnp.array([0.0, 2.0], dtype=jnp.float32)
current = jnp.array([4.0, -2.0], dtype=jnp.float32)
decay = 0.75
updated = decay * old + (1.0 - decay) * current
assert jnp.array_equal(updated, jnp.array([1.0, 1.0]))

Each coordinate moves one quarter of the way from its old value to the current value. With finite inputs and decay between zero and one, the exact-arithmetic result lies between those two values coordinate by coordinate, apart from floating-point rounding.

If current stays constant and 0 <= d < 1, repeated updates approach it in exact arithmetic. The direction depends on the starting value: a negative target approached from zero is approached from above, not below. In floating point, very small updates can round away.

Storing the average in NNX

A module can store the estimate as a Variable and replace it on each update:

from flax import nnx

class RunningAverage(nnx.Module):
    def __init__(self, initial):
        self.average = nnx.Variable(initial)

    def update(self, current, decay):
        self.average[...] = decay * self.average[...] + (1 - decay) * current
        return self.average[...]

tracker = RunningAverage(jnp.zeros(2))
assert jnp.array_equal(tracker.update(jnp.array([4.0, -4.0]), 0.5),
                       jnp.array([2.0, -2.0]))

Variable identifies state for filtering; it does not automatically block gradients. If an algorithm needs a detached target, its gradient handling must express that choice separately.

Initialization and bias correction

Starting from zero gives total observation weight 1 - d**t after t updates with constant decay. Dividing by that factor can correct this initialization effect under the appropriate assumptions, as in Adam's moment estimates. It is not a universal correction for arbitrary initial values or changing decay. Do not apply it in this exercise.

Problem

Implement ema_step(seed, params_arr, ema_arr, decay) and return:

decay * ema_arr + (1.0 - decay) * params_arr

Params_arr and ema_arr are finite float32 vectors with the same shape. Decay is a scalar in [0, 1]. Return a float32 vector of that same shape. The seed argument is unused and retained for compatibility with the exercise interface. Do not mutate either input or add bias correction.

Common pitfalls

  • Decay multiplies the old estimate. Reversing the weights changes the update.
  • At decay 0, return the current values. At decay 1, preserve the old estimate. These are valid boundary cases, not errors.
  • Keep decay as a JAX value. Converting a dynamic decay to a Python float prevents JIT tracing and differentiation with respect to decay.
  • The arrays must have matching shapes. Accidental broadcasting may return a different shape while still producing a mathematically valid array.
  • Seed does not affect the result; no random sampling is needed.