Skip to content

← Flax NNX step 56 of 100

Hard Primitives

NNX Gradient Accumulation

Jump to the task

Why this matters

Gradient accumulation combines gradients from several smaller batches before applying one optimizer update. This can reduce the activation memory needed for each forward and backward pass, while optimizing a loss over more examples.

Here the model is a deterministic Linear and each loss is a mean squared error. Compute both gradients at the same parameter values, average them, and update once. Updating between batches would be a different algorithm.

Why equal batch sizes matter

Each batch contains M examples with D_out target values per example. Its MSE averages M * D_out squared errors. Because both batches have the same size, the full-batch loss is (loss_a + loss_b) / 2. Differentiation is linear, so its gradient is (grads_a + grads_b) / 2 in exact arithmetic.

For unequal batch sizes, weight batch means by their number of error terms. With matching output widths and batch sizes M_a and M_b, use (M_a * grads_a + M_b * grads_b) / (M_a + M_b). An unweighted average would give the smaller batch too much influence. Unequal sizes are outside this exercise's input contract.

Combine matching gradient trees

NNX gradients contain nested state and Variable wrappers. Use a tree map to combine corresponding array leaves while preserving their structure.

import jax
import jax.numpy as jnp

grads_a = {"kernel": jnp.array([[2.0], [4.0]]), "bias": jnp.array([1.0])}
grads_b = {"kernel": jnp.array([[-2.0], [2.0]]), "bias": jnp.array([3.0])}
average = jax.tree.map(lambda a, b: (a + b) / 2, grads_a, grads_b)
assert jnp.allclose(average["kernel"], jnp.array([[0.0], [3.0]]))
assert jnp.allclose(average["bias"], jnp.array([2.0]))

Both trees must have compatible structures and leaf shapes. The same operation applies to the gradient states returned by the two NNX gradient calls.

Limits of the equivalence

Floating-point accumulation can round or overflow. Neither dividing at the end nor scaling each contribution first is universally exact or more accurate. Compare the two-batch result with a full-batch calculation using a tolerance.

The equivalence also depends on the model and loss. BatchNorm using separate batch statistics does not reproduce a single combined-batch forward pass. Synchronized BatchNorm across devices does not automatically combine statistics across sequential micro-batches. Stochastic layers require attention to random draws and mutable state. This exercise avoids those complications.

The final metric below deliberately concatenates the data and runs a full-batch forward pass. That is convenient for this small exercise, but it means this implementation is not a complete recipe for keeping every operation within micro-batch memory limits.

Problem

Implement grad_accumulation_step(seed, x_a, y_a, x_b, y_b, lr):

  1. Construct nnx.Linear(x_a.shape[-1], y_a.shape[-1], rngs=nnx.Rngs(int(seed))) with default initializers and bias.
  2. Construct one nnx.Optimizer(model, optax.sgd(lr), wrt=nnx.Param). Use plain SGD, without momentum or weight decay.
  3. Define loss_fn(model, x, y) as jnp.mean((model(x) - y) ** 2).
  4. Compute gradients separately for A and B with nnx.value_and_grad, using the unchanged model for both calls.
  5. Average corresponding leaves with a tree map and apply exactly one optimizer.update(model, avg_grads).
  6. Concatenate the inputs and targets along axis 0. Compute MSE over every entry using the updated model and return final_loss.reshape(1).

Keep the learning rate and loss as JAX values, without Python float conversion. The seed must be static if JIT-compiling the entire entrypoint. Numerical tests check the final loss; they cannot prove that you computed two separate gradients instead of one full-batch gradient. Follow the accumulation workflow.

Inputs: both input arrays have shape (M, D_in) and both target arrays have shape (M, D_out), with positive dimensions. All are finite float32 arrays with matching widths and equal batch sizes. seed is a nonnegative integer-valued scalar supplied as a float; cast it to int. lr is finite and nonnegative. Values are small enough for finite intermediate calculations.

Output: JAX array of shape (1,), containing the full-batch MSE after the single update.