Skip to content

← Flax step 68 of 100

Hard Primitives

Gradient Accumulation Step

Jump to the task

Gradient accumulation combines gradients from several smaller batches before taking one optimizer step. It can reduce the memory needed for intermediate activations compared with processing all examples in one forward/backward pass. This exercise uses two equally sized micro-batches and a small linear model.

Keep the parameters fixed while accumulating

Compute both gradients from the same parameter values:

grads_a = jax.grad(loss_fn)(params, x_a, y_a)
grads_b = jax.grad(loss_fn)(params, x_b, y_b)
grads = jax.tree_util.tree_map(lambda a, b: (a + b) / 2, grads_a, grads_b)
state = state.apply_gradients(grads=grads)

Updating after the first micro-batch would change the parameters used for the second gradient. That would be two training steps, not one accumulated step. Include both the Dense kernel and its bias in the gradient tree.

Why average the gradients?

Each micro-batch loss is a mean over M examples. With two equally sized micro-batches, the full-batch mean loss is (loss_a + loss_b) / 2, so its gradient is (grads_a + grads_b) / 2. Summing without dividing would double the gradient. For unequal batch sizes, the two means would need weights proportional to their example counts.

This equality holds in exact arithmetic for this Dense-only model and loss. Floating-point reduction order can introduce small differences. Dividing at the end does not make float32 accumulation exact. Batch-dependent operations such as BatchNorm, or different random masks in stochastic layers, also need care before claiming equivalence to one large-batch calculation.

A gradient tree is a nested mapping with one array per parameter. Adding two Python dictionaries does not add their arrays. jax.tree_util.tree_map applies the averaging function to corresponding leaves.

Evaluate after the update

After the single SGD step, combine the two micro-batches and recompute the mean squared error using the updated parameters. This final loss is different from the losses used to compute the gradients. The exercise measures it on the same examples used for the update, not on a separate validation set.

Problem

Implement grad_accumulation_step(seed, x_batches, y_batches, lr):

  1. Define a Linen TinyDense module whose compact __call__ applies one unnamed nn.Dense(1) with default settings and bias. Keep this wrapper so the seeded parameter path is Dense_0.
  2. Initialize it with jax.random.PRNGKey(int(seed)) and x_batches[0], without splitting the key. Create a TrainState with the params collection, model.apply and optax.sgd(float(lr)) using default settings.
  3. Compute the gradient of mean squared error on each micro-batch separately, using the same initial parameters. Reshape predictions from (M, 1) to (M,) before subtracting targets. There is no one-half factor in the loss.
  4. Average corresponding gradient leaves and call apply_gradients once.
  5. Concatenate the inputs to (2*M, D) and targets to (2*M,). Recompute predictions and MSE from the updated state. Return [final_loss] as a JAX array of shape (1,).

x_batches is a finite float32 array of shape (2, M, D) and y_batches is a finite float32 array of shape (2, M), with M >= 1 and D >= 1. seed is a nonnegative integer-valued scalar that may arrive as a float; lr is a finite nonnegative scalar. Leave both input arrays unchanged. With lr=0, parameters stay at their initialized values and the final loss is their full-batch MSE. A positive learning rate does not guarantee that one step lowers the loss.