Skip to content

← Flax NNX step 26 of 100

Hard Primitives

NNX Implement BatchNorm

Jump to the task

Parameters and running statistics

BatchNorm uses two kinds of state. Gain gamma and offset beta are trainable parameters. Running mean and variance summarize past training batches and are updated by a separate averaging rule.

This exercise handles input shaped (N, D): N examples with D features each. Each feature gets its own statistics across the N examples, so the reduction axis is 0. All four state vectors have shape (D,). This is a custom implementation for matrices, not a replacement for every option offered by a library BatchNorm layer.

Store gain and offset in nnx.Param, initialized to ones and zeros. Store running mean and variance in nnx.Variable, initialized to zeros and ones. An optimizer restricted to nnx.Param then excludes the running statistics. The built-in NNX BatchNorm uses the more specific nnx.BatchStat wrapper; this exercise uses plain nnx.Variable to practice managing state.

Two forward paths

During training (use_running_average=False), compute each feature's batch mean and population variance. Population variance divides by N, not N - 1. Normalize with those current batch statistics, and also update:

running_mean = 0.9 * running_mean + 0.1 * batch_mean

running_var = 0.9 * running_var + 0.1 * batch_var

Here momentum is the weight on the old value. This running variance is an exponential average of within-batch population variances. It is not an exact variance of all examples seen so far, and this implementation does not apply a sample-variance correction.

During evaluation (use_running_average=True), use the stored running statistics and leave them unchanged. The output still depends on each input value, but not on the other examples in the current evaluation batch. Both branches finish with gamma * (x - mean) / sqrt(variance + eps) + beta.

Positive epsilon keeps the denominator nonzero for a constant feature. Initializing running variance to one gives a fresh model a scale close to identity during evaluation. Initializing it to zero would not cause division by zero with positive epsilon, but would divide by the much smaller sqrt(eps) instead.

A complete train-then-evaluate example

import jax.numpy as jnp
from flax import nnx

class MyBatchNorm(nnx.Module):
    def __init__(self, d, momentum, eps, rngs):
        self.gamma = nnx.Param(jnp.ones((d,), dtype=jnp.float32))
        self.beta = nnx.Param(jnp.zeros((d,), dtype=jnp.float32))
        self.running_mean = nnx.Variable(jnp.zeros((d,), dtype=jnp.float32))
        self.running_var = nnx.Variable(jnp.ones((d,), dtype=jnp.float32))
        self.momentum = momentum
        self.eps = eps

    def __call__(self, x, use_running_average):
        if use_running_average:
            mean = self.running_mean[...]
            variance = self.running_var[...]
        else:
            mean = jnp.mean(x, axis=0)
            variance = jnp.mean((x - mean) ** 2, axis=0)
            self.running_mean[...] = (
                self.momentum * self.running_mean[...]
                + (1 - self.momentum) * mean
            )
            self.running_var[...] = (
                self.momentum * self.running_var[...]
                + (1 - self.momentum) * variance
            )
        normalized = (x - mean) / jnp.sqrt(variance + self.eps)
        return self.gamma[...] * normalized + self.beta[...]

model = MyBatchNorm(1, momentum=0.5, eps=3.0, rngs=nnx.Rngs(0))
x = jnp.array([[1.], [3.]], dtype=jnp.float32)
print(model(x, False))  # [[-0.5], [0.5]]
print(model.running_mean[...])  # [1.0]
print(model.running_var[...])   # [1.0]
print(model(x, True))   # [[0.0], [1.0]]

The example uses a large epsilon and momentum 0.5 for easy arithmetic. The exercise below uses epsilon 1e-5 and momentum 0.9. Retaining the same module between calls is essential to retain its running statistics. All initial values are deterministic, so rngs consumes no key here.

Assigning through variable[...] updates the value held by the wrapper; it does not make JAX arrays mutable. Use NNX-aware transformations for a persistent module whose forward pass changes state. A Python branch on the mode flag needs that flag to be static when JIT-compiling this code. Pass the flag explicitly to this custom module; do not assume that calling model.eval() supplies a required call argument.

The visualization groups rows by feature. For this matrix exercise, each cell is a scalar; there are no additional spatial dimensions to reduce.

Loading visualization…

Problem

Implement batchnorm_forward(seed, x, use_running_average):

  1. Create a fresh MyBatchNorm with d=x.shape[-1], momentum 0.9, epsilon 1e-5, and nnx.Rngs(int(seed)).
  2. Convert the mode input with bool(use_running_average >= 0.5).
  3. Call the module using that flag and return its output flattened in row order, preserving feature order within each example.

x is a nonempty float32 matrix (N, D), with positive N and D. Use finite inputs whose squares and reductions stay within float32 range. seed is a nonnegative integer-valued scalar, possibly supplied as a float. The mode input is 0.0 for training or 1.0 for evaluation. The output is a float32 vector of length N * D. There is no optimizer step.

Each entrypoint call constructs a new module. Thus an evaluation call uses initial running mean zero and variance one, not statistics from an earlier entrypoint call. Training outputs alone do not prove that running statistics were updated: inspect a retained module, as in the example, to verify that.

If JIT-compiling the whole entrypoint, make seed and mode static. Only the input matrix needs to be a dynamic argument.