What the backward pass costs, in arithmetic and in memory

STATEMENT

Account for both resources a backward pass consumes. Show that the arithmetic is a small multiple of the forward pass and that the memory is not, find the batch size at which stored activations overtake parameters, and quantify what checkpointing buys.

GIVEN

A stack of $L = 48$ square affine layers of width $d = 1024$, in fp32. Two working regimes: a batch of 32 rows, and a batch of 32 sequences of length 512, which presents $32 \times 512 = 16\,384$ rows to every layer.

FIND

Forward and backward FLOPs and their ratio; activation memory against parameter memory in both regimes; the batch at which the two are equal; and the memory and arithmetic under $\sqrt{L}$ checkpointing.

STRATEGY

Count matrix products, since everything else is lower order. Forward, each layer does one. Backward, by (I.6.3) and (I.6.4), each layer does two — one to send the adjoint onward and one to form the weight gradient. For memory, count what the backward pass will read and cannot recompute for free.

SOLUTION

Step 1 — arithmetic. A product of a $B \times d$ matrix with a $d \times d$ matrix costs $2Bd^2$ floating-point operations, counting a multiply and an add for each of the $Bd^2$ pairs. Forward:

Backward does two such products per layer, so $2F$, and the step costs $3F$:

The ratio is 3 in both regimes and would be 3 for any $B$, $d$ and $L$: it is a property of the graph, not of the size. This is the cheap gradient principle (I.6.T1) at work, and the constant here — three — sits comfortably inside the $c \le 4$ the theorem allows.

Step 2 — memory. The parameters are $Ld^2 = 50\,331\,648$ values, $0.201$ GB. The activations kept for the backward pass are $LBd$ by (I.6.7):

The same model, the same weights, the same arithmetic per row — and the memory picture inverts completely. At batch 32 the parameters dominate by thirty times; at batch 16 384 the activations dominate by sixteen.

Step 3 — the crossover. Setting $LBd = Ld^2$ gives

The layer count cancels, and so does everything else. Activations overtake parameters when the batch reaches the width, which for this model is 1024 rows — a number a single sequence of moderate length already exceeds. That is why the second regime is the normal one and the first is the exception.

Step 4 — checkpointing. Keep every $\sqrt{L}$-th activation and recompute the rest during the backward pass. With $L = 48$, $\sqrt{48} \approx 7$: hold 7 segment boundaries, plus the 6 tensors of whichever segment is currently being recomputed, so 13 tensors instead of 48.

The price is recomputing each segment’s forward pass once, which adds one forward pass to a step that cost three:

A third more arithmetic for nearly four times less memory. On a device where memory is the binding constraint and arithmetic is not — which describes most training — that is not a close decision.

Check — sanity

The FLOP ratio is exactly 3 in both regimes. It has to be — it is a count of matrix products per layer, one forward and two back, and no dimension enters it. A measured ratio far from 3 on real hardware is a statement about memory bandwidth and kernel efficiency, not about the arithmetic.

The crossover has no free parameters. $LBd = Ld^2$ reduces to $B = d$ with $L$ cancelling, so the answer cannot depend on depth. A deeper model has more activations and more parameters, in the same proportion.

Checkpointed memory is near $2\sqrt{L}$. $2\sqrt{48} \approx 13.9$ against the 13 tensors counted, and $48/13 = 3.7$ matches the measured reduction. The $\sqrt{L}$ of (I.6.7) is where boundary storage and segment storage balance; away from it, one term or the other grows.