Train with Mutable batch_stats

Train a small Flax Linen model for one step while keeping two kinds of state: learned parameters, updated by SGD, and BatchNorm running statistics, updated by the training forward pass. Then measure the updated model's evaluation loss on the same batch.

The exact model

Define BNDense(nn.Module) with features: int = 4 and this compact method:

@nn.compact
def __call__(self, x, train: bool):
    x = nn.Dense(self.features)(x)
    x = nn.BatchNorm(use_running_average=not train)(x)
    x = nn.relu(x)
    return nn.Dense(1)(x)

Use default layer settings and leave the children unnamed. Their parameter paths are Dense_0, BatchNorm_0 and Dense_1. BatchNorm's default momentum is 0.99 and epsilon is 1e-5. Its scale starts at one and bias at zero; running mean starts at zero and running variance at one.

Split parameters from running state

Initialize BNDense(features=4) with jax.random.PRNGKey(int(seed)), x_batch and train=False. The returned mapping has two collections:

• params contains Dense weights/biases and BatchNorm's learned scale/bias.
• batch_stats contains BatchNorm's running mean and variance.

Initialize optax.sgd(float(lr)) with only params. Collection names do not automatically block gradients or prevent optimizer updates. We keep running statistics outside the optimizer because their update rule is a moving average, not SGD. Optax can accept other array trees; separating these trees expresses which values this training step should optimize.

Return state as auxiliary output

During training, call:

out, updated = model.apply(
    {"params": params, "batch_stats": batch_stats},
    x_batch, train=True, mutable=["batch_stats"]
)
preds = out.reshape(-1)
loss = jnp.mean((preds - y_batch) ** 2)

Reshaping (N, 1) predictions to (N,) avoids accidentally broadcasting against the target vector to form an (N, N) error matrix.

Define loss_fn(params, batch_stats, x, y) to return (loss, updated["batch_stats"]). Then use:

(loss, new_batch_stats), grads = jax.value_and_grad(
    loss_fn, has_aux=True
)(params, batch_stats, x_batch, y_batch)

By default, differentiation targets the first argument, params. The second return value is auxiliary data, not another objective to differentiate. Without has_aux=True, this tuple return is not a valid scalar loss for value_and_grad.

Keep both updates

Use tx.update(grads, opt_state, params) and optax.apply_updates to get the new parameters. Separately retain new_batch_stats from the forward pass. Those statistics describe activations computed with the pre-update parameters; do not run another training pass to recompute them here.

With momentum 0.99, each running value becomes 0.99 * old_value + 0.01 * batch_value. BatchNorm uses the current batch's statistics to normalize training activations. Its stored averages are for evaluation. Even with lr=0, a training forward pass can change these averages while leaving learned parameters unchanged.

Measure the post-step evaluation loss

Apply the model with the updated params and batch_stats, using train=False. Evaluation reads the stored statistics without updating them. It does not need mutable permission, although requesting mutable=["batch_stats"] is valid and returns an unchanged statistics collection alongside the output.

Flatten predictions and compute mean squared error against y_batch again. This is a post-step evaluation loss on the training batch, not a held-out generalization measure. It need not equal the training-mode loss returned by loss_fn, and one SGD step need not decrease it.