Skip to content

← JAX Autodiff step 15 of 25

Medium Primitives

Gradient Checkpointing: Basics

Jump to the task

Reverse-mode differentiation needs values from the forward computation to calculate derivatives. JAX can retain these values, called residuals, or recompute some of them later. It does not necessarily store every activation produced by a model.

jax.checkpoint, also called jax.remat, asks JAX to recompute internal residuals of a wrapped function when needed during differentiation. Inputs can still be saved so the computation can be repeated. This can reduce memory use at the cost of extra work, but the benefit depends on the chosen boundary, operations, compiler and hardware.

A small API exercise

This exercise wraps only the middle matrix multiplication in a three-layer network. The first two products are followed by ReLU, which replaces negative entries with zero. The last product has no activation.

def middle(y, w):
    return y @ w

y1 = jax.nn.relu(x @ w1)
y2 = jax.nn.relu(jax.checkpoint(middle)(y1, w2))
y3 = y2 @ w3
loss = jnp.sum(y3 ** 2)

The second ReLU is outside the checkpointed function. Matrix multiplication derivatives use the operands, not the product itself. Wrapping this one product therefore does not demonstrate a useful reduction in saved activations. Treat it as practice with the API, not a memory benchmark.

For identity weights and x = [[1, 1]], both ReLUs leave the values unchanged, so y3 = [[1, 1]] and the loss is 1 + 1 = 2. Negative values in the final output still contribute their squares.

Check values, then measure memory

Checkpointing preserves the mathematical function and its derivatives. Compare floating-point results with a tolerance; compilation changes do not carry a general bit-for-bit equality guarantee. A forward-only call cannot demonstrate the storage/recomputation tradeoff in a backward pass.

To explore that tradeoff, differentiate a scalar loss, inspect its saved residuals, and profile representative inputs. There is no fixed compute overhead or automatic constant-memory guarantee. Wrapping a larger region or a whole function is valid, but its usefulness depends on which values remain live and what must be recomputed.

Keep transformed functions pure: return results instead of depending on changes to Python lists, counters or global state. Python side effects occur during tracing, not reliably once per numerical execution. The number of traces depends on transformations and caching; it is not a fixed pair of forward and backward Python calls.

Problem

Implement checkpointed_layer_loss(w1, w2, w3, x) using the sequence above. Define the middle multiplication as a helper and wrap that helper with jax.checkpoint. Keep the second ReLU outside the helper.

The three weights are finite float32 JAX matrices of shape (d, d). x is a finite float32 JAX matrix of shape (N, d), with N, d >= 1. There are no biases. Leave all inputs unchanged.

Return a scalar JAX array containing the sum of squared entries of y3. Sum across the entire batch and feature dimensions; do not take a mean or apply ReLU to the final product. Return the loss, not its gradient.

Output tests verify this numerical computation. They cannot establish that the required checkpoint wrapper was used, since an unwrapped implementation can return the same loss.