Activation Memory, Gradient Recomputation, and the Checkpoint Trade

Training a convolutional network is a memory problem before it is an arithmetic one. The parameters are the part everyone counts, and for image models they are usually the small part. A fifty-layer residual backbone holds roughly $2.5 \times 10^7$ parameters, about 100 MB in 32-bit precision. The activations of a single forward pass at a batch of 32 images of moderate resolution exceed that by a wide margin, and backpropagation requires all of them to be alive simultaneously.

The reason is structural. To compute the gradient of the loss with respect to a convolution’s weights, the backward pass needs both the incoming gradient and the layer’s input. The input was consumed during the forward pass and is not otherwise recoverable, so it must be stored. Every layer in the stack makes the same demand, and their storage intervals overlap: the earliest layer’s input is needed last, and the latest layer’s input is needed first. Peak memory is therefore proportional to depth times the activation size per layer, not to the largest single layer.

Counting the bytes

For a layer producing $C$ channels at extent $H \times W$ for a batch of $B$, the stored activation is $B \cdot C \cdot H \cdot W$ elements. At 4 bytes each, a stride-1 stage with 64 channels at 224×224 for a batch of 32 is $32 \cdot 64 \cdot 224 \cdot 224 \cdot 4 \approx 411$ MB. The layer’s input is stored for its own backward pass, and the next layer stores that output as its input, so two such tensors are live at once and the totals grow quickly.

Two factors moderate the count. Downsampling reduces the spatial term by a factor of 4 per halving, so most of the memory sits in the early, high-resolution, low-channel stages even though those stages carry a small fraction of the parameters. Inverted-residual blocks, which expand channels by a factor of several in the middle of the block and contract again, store the expanded intermediate and therefore cost more than their input and output widths suggest. And the training framework stores more than the mathematical minimum: a convolution’s backward pass typically needs both operands, so each layer’s input is stored once and referenced by the backward node, which is close to the minimum; but a fused block whose internal pointwise stages are not separately stored still has to store the boundary tensors between fusions.

The practical accounting is a per-layer line item of $B \cdot C \cdot H \cdot W \cdot \text{bytes}$ summed over layers, plus the parameter and optimiser state, plus the fragmentation of the allocator. The last is not negligible: an allocator that cannot reuse a free block across differently shaped tensors inflates the peak, and a caching allocator that holds released blocks for reuse trades total footprint for speed.

Quantity	Scales as	Where it lands
weights	$C_{in}C_{out}K^2$ per layer	deep stages, high channel count
activations	$B C H W$ per layer	early stages, full resolution
optimiser state	2–3× weights	independent of batch
gradients	1× weights	independent of batch

Why larger batches do not fix the per-sample cost

Activations scale linearly in batch, so a batch of 64 needs twice the memory of a batch of 32 for the same model. The only reductions available on the activation side are resolution (fewer pixels), width (fewer channels), depth (fewer stored layers), or recomputation. The first three change the model; the fourth changes the arithmetic. Gradient checkpointing is the only one that trades compute for memory without altering the function, and its analysis occupies the rest of what follows.

Checkpointing: storing boundaries instead of everything

The idea is to divide the network into segments, store only the input to each segment, and recompute the interior activations during the backward pass. Backpropagation proceeds from the last segment to the first; when it reaches a segment, it runs a forward pass over that segment from its stored input to reconstruct the interior tensors, then backpropagates through them. After the segment is done, its reconstructed tensors are freed, so the peak is the maximum over segments of the interior footprint plus the stored inputs of all segments.

For $L$ layers divided into $\sqrt{L}$ segments of $\sqrt{L}$ layers each, the stored inputs are $\sqrt{L}$ tensors and each segment’s recomputation holds $\sqrt{L}$ interior tensors, so the peak is $O(\sqrt{L})$ rather than $O(L)$. Chen et al.’s formulation is exactly this: the memory recovers to a square root of the naive cost, and the trade is well-defined.

The arithmetic cost depends on where the checkpoint boundaries are. If a segment’s forward pass costs $F$ and its backward pass costs $2F$ (each layer’s backward is roughly twice its forward: two matrix products, one for the input gradient and one for the weight gradient), then a recomputation adds one extra $F$ per segment, so the total becomes $3F$ plus $F$ recomputed, or $4F$ against the naive $3F$: a 33% increase in flops for a checkpointer with one segment per layer. Coarser segments reduce the number of recomputations but each recomputes more: with $\sqrt{L}$ segments the recomputation cost is the full forward pass, so the total is $3F + F = 4F$ as well, since every layer is recomputed exactly once regardless of segment size. The segment size therefore does not change the flop overhead much; it changes the memory peak, which is minimised at the square root.

import numpy as np

def checkpointed_forward(layers, x, segment):
    """Store only segment inputs; interior activations are dropped."""
    stored, outs = [], []
    for start in range(0, len(layers), segment):
        stored.append(x)                       # kept for the later replay
        for layer in layers[start:start + segment]:
            x = layer(x)                       # interior tensor freed immediately
        outs.append(x)
    return stored, outs

def replay(layers, seg_input):
    """Reconstruct a segment's interior activations from its stored input."""
    trace = [seg_input]
    for layer in layers:
        trace.append(layer(trace[-1]))
    return trace                               # released after the segment's backward

layers = [lambda a: np.tanh(a) for _ in range(6)]
stored, outs = checkpointed_forward(layers, np.ones((4, 4)), segment=2)
print(len(stored), len(outs), len(replay(layers[:2], stored[0])))   # 3 3 3

The replay is the cost: the segment’s forward pass runs a second time, and its interior tensors exist again only for the duration of that segment’s backward. The release is immediate, so the allocator can reuse the space for the next segment.

Where checkpointing fails to help

Three situations defeat the trade. A segment that contains a non-deterministic operation recomputes a different answer, so the reconstructed gradient is inconsistent with the one the forward pass produced; any stochastic operator inside a checkpointed region must have its random state captured and restored. A segment whose internal state is large, such as a block that materialises a very wide intermediate in the expanded channel stage of an inverted-residual block, has a recomputation peak close to its storage peak, so the saving is small.

The third is the interaction with the memory system’s behaviour. Recompute passes read the segment’s weights again, so the weight traffic doubles for the checkpointed portion. On a device whose training step is limited by weight bandwidth rather than by flops, which is more common at small batch sizes where the arithmetic available per weight load falls, recomputation can cost more wall-clock time than the flops suggest. The saving is real but the exchange rate is hardware-dependent.

The remaining alternative to all of this is to reduce what is stored by fusing layers. A convolution followed by a normalisation and an activation stores its output for the backward pass, but the normalisation’s backward needs the pre-normalisation value, and the activation’s needs the pre-activation value. Fusing the three into one kernel stores one tensor instead of three and saves two round trips to memory, without changing the arithmetic at all. That is the most favourable exchange available in the memory system, and it is why training throughput at a fixed batch size is as much a property of the kernel library as of the model.

Precision as memory

The other lever is the width of the stored values. Activations stored in 16-bit floating point halve the memory of the activation stack; mixing 16-bit compute with a 32-bit master copy of the weights keeps the optimiser state accurate while the forward and backward passes run at reduced width. The consequence for a convolutional model is not only memory: the arithmetic throughput of 16-bit multiply-add is higher on hardware that has dedicated paths, so the same step runs faster as well as smaller.

The limit is dynamic range. A sum over $C_{in}K^2$ products accumulates in a narrower exponent range, so the accumulation of a convolution with many channels benefits from a 32-bit accumulator even when the operands are 16-bit. That is why the precision of the accumulator is a separate decision from the precision of the operands, and why the error behaves differently in layers with large reduction depth, which is the subject that the numeric analysis of im2col, FFT, and Winograd runs into from the transform side.

The interaction with checkpointing is adverse. A checkpointed segment must reproduce its forward pass exactly enough that the reconstructed gradient matches, and a lower-precision recomputation introduces a difference between the two passes that the original training step does not have. The difference is generally small and acts like noise, but it removes the guarantee that the gradient is the one computed by the forward pass, which is a real cost in debugging.

There is a fourth lever that changes memory only in appearance: storing the activations of a narrower network and widening later. Factoring a wide convolutional stage into a narrow spatial stage followed by a channel expansion, as in depthwise and grouped convolutions, reduces the arithmetic but can raise the activation footprint, because the expanded intermediate that sits inside the block is stored even though it is never needed outside it.

Choosing among the levers

The order of preference follows from what each costs. Fusion is free: same arithmetic, less traffic, no change to the function, limited only by the kernel library’s coverage of the block. Reduced precision is nearly free in flops and costs numeric range, with a 32-bit accumulation path required for wide reductions. Checkpointing costs a third more arithmetic in the clean accounting and a variable amount in practice, and its saving is exactly the interior activations of the checkpointed region. Reducing the batch changes the model’s optimisation dynamics and is the last resort rather than the first.

The decision boundary sits where the stored activation stack is larger than the parameter and optimiser state put together, which for image models at any reasonable resolution is almost always. At that point the memory budget is an activation budget, and the depth of the network is limited by it rather than by the parameter count. That is the opposite of the picture a parameter-count table suggests, and it is why architectural choices that reduce activation memory, such as stem downsampling at the earliest stage, an aggressive early stride, and narrow high-resolution channels, are worth more at training time than the same flop reduction placed deeper in the network.