Study Plan

Gradient Accumulation with Scan

JAX

Hard

Compute the full-batch loss and gradient of a small two-layer network by splitting the batch into num_microbatches equal microbatches and accumulating their results. For a batch of $N$ examples the model and its loss are

h = \tanh(x W_1 + b_1), \qquad o = h W_2 + b_2, \qquad L = \frac{1}{N} \sum_{n=1}^{N} \sum_{k=1}^{K} \left( o_{n,k} - y_{n,k} \right)^{2}

Here x has shape (B, D) and y has shape (B, K), and params is a dict holding w1 ( $W_1$ ) of shape (D, H), b1 ( $b_1$ ) of shape (H,), w2 ( $W_2$ ) of shape (H, K) and b2 ( $b_2$ ) of shape (K,). The index $n$ runs over examples and $k$ over the K outputs. With $S$ = B / num_microbatches, microbatch $j$ holds rows $jS$ to $(j + 1)S - 1$ of x and y, counting from 0.

Your function is traced under jax.jit with num_microbatches fixed as a Python int.

Return a dict with keys loss and grads. The value under loss is the full-batch $L$ as a float32 JAX array of shape (), and the value under grads is a dict with keys w1, b1, w2 and b2 holding the gradient of $L$ as float32 JAX arrays shaped like the matching entries of params.

Examples

Input: params = {"w1": [[0.5, -0.5], [0.25, 1.0]], "b1": [0.0, 0.1], "w2": [[1.0], [-1.0]], "b2": [0.2]}, x = [[1.0, 0.0], [0.0, 1.0], [1.0, 1.0], [-1.0, 0.5]], y = [[0.5], [-0.5], [0.0], [1.0]], num_microbatches = 2

Output: {"loss": 1.060169, "grads": {"w1": [[1.155724, -0.689779], [-0.270026, 0.043908]], "b1": [-0.483697, -0.012093], "w2": [[0.588589], [-0.749159]], "b2": [-0.487136]}}

Explanation: The first microbatch holds rows 0 and 1 and the second holds rows 2 and 3. Each microbatch gives its own mean loss and gradient, and because both have two rows, their average equals the loss and gradient of all four rows at once.

Input: params = {"w1": [[0.3, 0.0, -0.2]], "b1": [0.1, -0.1, 0.0], "w2": [[0.5, 0.5], [1.0, -1.0], [0.0, 0.2]], "b2": [0.0, 0.0]}, x = [[0.5], [-1.5], [2.0]], y = [[1.0, 0.0], [0.0, 1.0], [0.5, 0.5]], num_microbatches = 1

Output: {"loss": 0.735669, "grads": {"w1": [[0.245756, -1.287083, 0.158517]], "b1": [-0.72046, -0.369854, -0.116445], "w2": [[-0.219351, 0.189402], [0.102495, 0.065263], [0.088264, -0.165526]], "b2": [-1.028366, -0.654801]}}

Input: params = {"w1": [[0.2, -0.4], [0.6, 0.1]], "b1": [0.0, 0.0], "w2": [[-0.7], [0.9]], "b2": [0.1]}, x = [[0.1, 0.2], [0.3, -0.4], [-0.5, 0.6], [0.7, 0.8], [-0.9, -1.0], [1.1, 0.0]], y = [[0.0], [1.0], [0.0], [1.0], [0.0], [1.0]], num_microbatches = 6

Output: {"loss": 0.942053, "grads": {"w1": [[0.6885, -0.985434], [0.193408, -0.426145]], "b1": [0.627196, -0.786288], "w2": [[-0.475989], [0.420651]], "b2": [-0.96138]}}

Hint 1

x.reshape(num_microbatches, -1, x.shape[1]) stacks the microbatches.

Hint 2

jax.lax.scan(step, (0.0, zero_grads), (xs, ys)) walks them one at a time.

Hint 3

jax.tree.map(jnp.zeros_like, params) builds the zero gradient carry.

Requirements

Accumulate the microbatches with one jax.lax.scan of length num_microbatches, with no Python loop over them.
Call jax.value_and_grad or jax.grad inside the scan body, so that each step runs the forward and backward pass of one microbatch, instead of differentiating a function that contains the scan.
Each scan step must use only the S rows of its own microbatch, and no matrix product may run outside the scan.
The grader traces your function with jax.make_jaxpr and checks these rules in the traced program.

Constraints

B is between 1 and 8, and num_microbatches is between 1 and B and divides B exactly.
D, H and K are each between 1 and 4.
params has exactly the keys w1, b1, w2 and b2, and every value in params, x and y is a finite float32 number.

Try Similar Problems

Global Norm Gradient ClippingMedium AdamW from ScratchMedium SGD with Momentum over a PytreeEasy Blockwise Attention with Online SoftmaxHard Beam Search StepHard Speculative Decoding AcceptanceHard Ring All-Reduce with ppermuteHard Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Traced Learning Rate ScheduleMedium