Study Plan

Rematerialized MLP Block

JAX

Medium

Compute the loss and the parameter gradients of a residual MLP block, recomputing the block's tanh activations during the backward pass with jax.checkpoint while keeping its matrix multiplication results.

The block maps x of shape (B, D) to

\text{out} = x + \tanh\left( x W_1 + b_1 \right) W_2 + b_2,

where $W_1$ is w1 of shape (D, H), $b_1$ is b1 of shape (H,), $W_2$ is w2 of shape (H, D) and $b_2$ is b2 of shape (D,), all stored in params. The loss is the mean squared error between out and the target y of shape (B, D), averaged over all B times D entries.

Return a dict with keys loss (a float32 JAX array of shape ()) and grads (a dict with keys w1, b1, w2 and b2, each a float32 JAX array with the shape of the matching parameter).

Examples

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

Output: {"loss":0.408815,"grads":{"w1":[[0.332859,-0.197498,-0.033156],[-0.665719,0.394996,0.066313]],"b1":[0.332859,-0.197498,-0.033156],"w2":[[-0.089534,-0.24773],[0.234075,0.647654],[-0.293898,-0.813177]],"b2":[-0.307348,-0.850392]}}

Explanation: The hidden pre-activation x times w1 plus b1 is [0.3, -1.0, 1.9], and its tanh is about [0.291313, -0.761594, 0.956238]. Multiplying that hidden layer by w2 gives about [-0.857348, 0.199608], so the block output, which adds x and b2, is about [0.192652, -1.850392]. The errors against y are about -0.307348 and -0.850392, and the mean of their squares is the loss 0.408815.

Input: params = {"w1":[[1.0,0.0],[0.0,1.0]],"b1":[0.0,0.0],"w2":[[1.0,0.0],[0.0,1.0]],"b2":[0.0,0.0]}, x = [[0.5,-0.5],[1.0,2.0]], y = [[0.0,0.0],[0.0,0.0]]

Output: {"loss":3.435003,"grads":{"w1":[[0.559076,-0.084458],[0.550661,0.398575]],"b1":[0.74824,-0.273622],"w2":[[0.893115,0.906388],[0.626807,1.651008]],"b2":[1.361856,1.000955]}}

Input: params = {"w1":[[4.0,-3.0],[2.5,3.5]],"b1":[1.0,-1.0],"w2":[[0.3,-0.2],[0.1,0.4]],"b2":[0.0,0.2]}, x = [[2.0,1.5],[-1.0,-2.0],[0.5,0.5]], y = [[1.0,1.0],[-1.0,0.0],[0.0,0.5]]

Output: {"loss":1.039321,"grads":{"w1":[[3.7e-05,0.010167],[3.7e-05,0.006994]],"b1":[7.4e-05,0.000838],"w2":[[0.780639,0.623179],[-0.400972,0.681628]],"b2":[0.514078,-0.710164]}}

Hint 1

jax.checkpoint(f, policy=jax.checkpoint_policies.dots_saveable) wraps a function.

Hint 2

jax.value_and_grad(loss)(params, x, y) returns (loss, grads).

Hint 3

print(jax.make_jaxpr(f)(params, x, y)) shows the checkpointed region.

Requirements

Wrap the function that computes the block, including its tanh hidden layer, in jax.checkpoint (also available as jax.remat) with policy=jax.checkpoint_policies.dots_saveable.
Put the checkpoint inside the function you pass to jax.grad or jax.value_and_grad, because a checkpoint outside that function is never differentiated.
The grader traces your function with jax.make_jaxpr, and the traced program must contain a differentiated checkpointed region that holds the tanh and uses the dots_saveable policy.
Use only JAX operations on the inputs, so that the function can be traced.

Constraints

x and y are float32 arrays of shape (B, D), with B between 1 and 4 and D between 2 and 4.
The hidden width H is between 2 and 5.
All parameters are float32 and every value is finite.

Try Similar Problems

Mixed Precision with Loss ScalingMedium Int8 Matmul with Per-Channel ScalesMedium Blockwise Attention with Online SoftmaxHard Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium