#compute-memory-tradeoff

A network is a chain of 100 layers. Its input and every layer output have the same size; call one such tensor a unit, and count only these activations (ignore weights and gradients). Plain backpropagation keeps the input and all 100 layer outputs until the backward pass has used them: 101 units.

With activation checkpointing, the chain is cut into equal segments. The forward pass keeps only the input of each segment and frees every other activation as soon as the next layer has used it. The backward pass then handles the segments from last to first: it recomputes all the layer outputs of a segment from the segment's saved input, keeps them until that segment's backward pass is done, and then frees them together with the segment's saved input.

With 4 segments of 25 layers, the peak memory is units. The number of equal segments that minimizes the peak is .