Throwing work away

Backpropagation needs the forward pass’s intermediates. That is a requirement of the chain rule and not negotiable. What is negotiable is whether those intermediates were stored or are recreated on demand.

The √L result

Chen et al.’s version is elegant enough to derive in a paragraph. Divide $L$ layers into $k$ segments. Keep only the activation at each segment boundary — $k$ tensors. During the backward pass, when a segment is reached, rerun its forward from the stored boundary to regenerate its interior.

$$M \sim k + \frac{L}{k}, \qquad \text{extra compute} \sim \frac{L}{k} \cdot \frac{L}{L} = 1 \text{ forward pass}$$

Minimising the first over $k$ gives $k = \sqrt{L}$ and $M \sim \sqrt{L}$ — memory sublinear in depth, for one extra forward pass. For an 80-layer model that is memory proportional to 9 rather than 80.

The number in practice

One extra forward pass sounds cheap. Put it in the 6ND accounting and it is not.

$$\underbrace{2}_{\text{fwd}} + \underbrace{4}_{\text{bwd}} = 6 \qquad\longrightarrow\qquad \underbrace{2 + 2}_{\text{fwd, twice}} + \underbrace{4}_{\text{bwd}} = 8$$

Thirty-three percent more compute for the entire run. On a month-long job that is ten days. It is a real price and it is routinely paid without being stated — scaling-law plots are drawn in FLOPs and the checkpointing multiplier does not appear anywhere in them.

Why uniform recomputation is the wrong shape

Here is the observation that reframed this for me. The memory in a transformer block is not evenly distributed, and neither is the cost of regenerating it.

Sort the tensors by arithmetic intensity — FLOPs to recompute, divided by bytes stored.

Tensor	Bytes	Recompute cost
attention scores	$5as^2b$	one matmul
softmax output	$2as^2b$	elementwise
dropout mask	$as^2b$	free
QKV projections	$6sbh$	a large matmul
FFN hidden	$8sbh$	a large matmul

The quadratic tensors — the large ones — are almost free to regenerate. The linear tensors are matmul outputs, where recomputing means repeating the expensive part.

So the right move is not “recompute the block”. It is: recompute the attention interior, keep the projections. Korthikanti et al. measured about 70% of full checkpointing’s saving at roughly 4% overhead instead of 36%.

And it composes badly with the free thing

The table above assumes the score matrices are stored. If you took part 1’s advice, they are not — FlashAttention already made that trade, better, and for free.

Which means selective recomputation has much less left to work on than the numbers suggest. Its three big targets are gone, and what remains is the linear term where the exchange rate is poor.

This is the honest state of the technique in 2026: important historically, important in stacks that cannot use fused attention, and considerably less central than the papers make it sound if you are running a modern kernel.

The third option

There is a route neither of these takes: keep the activations and store them in fewer bits. The backward pass reads them once to compute a gradient that is already a noisy minibatch estimate, so a little extra noise is tolerable — ActNN makes this rigorous with stochastic rounding, which keeps the gradient unbiased and preserves convergence guarantees.

It gets about twelve-fold compression for roughly 20% throughput. Sound, and mostly unused — because recomputation is exact, is one function call, and after FlashAttention there is not enough left on the table to justify custom kernels.

I like it as an entry anyway. It is a good example of a technique that is correct, well-motivated, and was overtaken by something that solved the same problem from a different direction.

Part 3 stops trying to fit on one device and splits the batch — where the interesting content is not the splitting but what has to stay identical across the copies.