Skip to content

← Flax NNX step 55 of 100

Hard Primitives

NNX Train with BatchNorm

Jump to the task

Train parameters and update statistics

A model with BatchNorm has two update rules. The optimizer changes its trainable weights, gains and biases using gradients. Training forward passes update BatchNorm's running mean and variance from the hidden activations. These updates happen on the same retained model, but they serve different purposes.

Build a small network: a linear layer from D_in to 8 hidden features, BatchNorm over those features, ReLU, and a linear layer from 8 to D_out. Train on successive mini-batches, then make one evaluation pass over all the input rows. Return the mean of those final predictions.

Keep both kinds of state

The built-in nnx.BatchNorm stores its running statistics in mean and var, wrapped as nnx.BatchStat. Its gain and bias, and both linear layers' parameters, are nnx.Param objects. Use wrt=nnx.Param on the optimizer. The default nnx.value_and_grad parameter filter also selects nnx.Param. BatchStat values are updated by the forward pass, not by SGD.

For this exercise, BatchNorm uses momentum 0.99 and epsilon 1e-5. Momentum weights the previous running statistic: each update combines 0.99 * old + 0.01 * current_batch_statistic. Compute statistics across examples, independently for each hidden feature. During training, the output uses current batch statistics, not the running averages.

Each loss evaluation must perform exactly one training forward pass. That pass updates the running statistics using activations from the parameters before the SGD step. Evaluating the loss again merely to log it would update the statistics a second time. Keep the loss returned by nnx.value_and_grad if you need to report it.

A complete one-step example

import jax
import jax.numpy as jnp
import optax
from flax import nnx

class BNModel(nnx.Module):
    def __init__(self, d_in, hidden, d_out, rngs):
        self.linear1 = nnx.Linear(d_in, hidden, rngs=rngs)
        self.bn = nnx.BatchNorm(hidden, momentum=0.99, epsilon=1e-5, rngs=rngs)
        self.linear2 = nnx.Linear(hidden, d_out, rngs=rngs)

    def __call__(self, x, use_running_average):
        x = self.linear1(x)
        x = self.bn(x, use_running_average=use_running_average)
        return self.linear2(jax.nn.relu(x))

model = BNModel(2, 8, 1, nnx.Rngs(0))
optimizer = nnx.Optimizer(model, optax.sgd(0.01), wrt=nnx.Param)
x = jnp.array([[1., 0.], [0., 1.]], dtype=jnp.float32)
y = jnp.array([[1.], [-1.]], dtype=jnp.float32)

def loss_fn(model, x, y):
    predictions = model(x, use_running_average=False)
    return jnp.mean((predictions - y) ** 2)

loss, grads = nnx.value_and_grad(loss_fn)(model, x, y)
optimizer.update(model, grads)
mean_before_eval = model.bn.mean[...]
var_before_eval = model.bn.var[...]
predictions = model(x, use_running_average=True)
assert predictions.shape == (2, 1)
assert jnp.array_equal(model.bn.mean[...], mean_before_eval)
assert jnp.array_equal(model.bn.var[...], var_before_eval)

The loss is the mean squared error across both rows and output features. It is measured before the update; the evaluation predictions use the updated parameters and the accumulated running statistics. A zero learning rate still lets training forward passes update BatchStat values.

Batch order and evaluation

Split the input rows into equal contiguous batches. On step s, choose batch s % num_batches. This cycles through the data without shuffling. Changing the order or replacing all mini-batches with one large batch changes both the gradients and the running-statistic history.

After the last update, pass use_running_average=True explicitly. For a fixed trained model, each prediction is independent of the other rows in that evaluation call. The predictions still depend on the input values, and their mean depends on which examples you evaluate.

Here evaluation uses the same inputs supplied for training. The returned prediction mean is a numerical exercise output, not an accuracy score, validation loss, or estimate of performance on held-out data.

Problem

Implement train_with_bn(seed, x_batches_flat, y_batches_flat, batch_size, lr, num_steps):

  1. Construct BNModel with hidden width 8 and a shared nnx.Rngs(int(seed)). Create linear1, BatchNorm, then linear2 in that order, using their default parameter initializers and the BatchNorm settings above.
  2. Create one nnx.Optimizer using plain optax.sgd(lr) and wrt=nnx.Param.
  3. Repeat int(num_steps) times. Select the next batch, compute mean squared error and gradients with one training forward, then update the same model.
  4. Evaluate the final model on all x_batches_flat with running averages. Return jnp.mean(predictions).reshape(1) as a JAX array of shape (1,).

Both input matrices are float32 with shapes (N, D_in) and (N, D_out). All dimensions are positive. batch_size is positive, divides N exactly, and may arrive as an integer-valued float. num_steps is nonnegative and integer-valued; zero steps evaluates the initialized model. seed is also nonnegative and integer-valued, possibly supplied as a float. lr is a finite nonnegative scalar. Use finite data and a learning rate for which intermediate calculations remain finite.

Each entrypoint call starts fresh; state persists only within its loop. When JIT-compiling the whole entrypoint, make seed, batch size and step count static. Keep learning rate and the returned mean in JAX arithmetic instead of converting them to Python floats. A single scalar result cannot prove correct state handling, so also check updates on a retained model.