How do you choose activation-checkpoint segment boundaries in a deep network?

## The uniform case, and where the square root comes from Take a network that is a stack of `n` near-identical blocks, each producing an intermediate of about the same size. Cut it into `k` equal segments of `n/k` blocks each. What stays resident? First, the tensor entering each segment: that is `k` tensors, alive for the whole step because each is the starting point for a later recompute. Second, while the backward pass is working on one segment, that segment's interior is reconstructed and live: about `n/k` tensors. Everything else is gone. So retained activation memory scales like `k + n/k` units instead of `n`. Minimising `k + n/k` over `k` gives `k = sqrt(n)`, and the minimum value is `2*sqrt(n)`. That is the classic sublinear-memory result: memory that grew linearly with depth now grows with the square root of depth, for the price of one extra forward pass. For a 48-block text encoder, `sqrt(48)` is about 6.9, so the heuristic says roughly seven segments of seven blocks, retaining about fourteen block-intermediates instead of forty-eight. ## Why the exact number matters less than people expect The function `k + n/k` is flat near its minimum. Checkpointing every fourth block of that same 48-block encoder gives twelve segments of four: twelve boundaries plus a four-block interior, about sixteen units against the optimum's fourteen. A few percent of the activation budget, in exchange for a boundary grid that lines up with the architecture and is easy to reason about. The more important invariant is that `k` does not change the recompute bill at all. Every checkpointed block is recomputed exactly once per step regardless of how the segments are drawn, so the time cost is fixed and `k` only decides how the retained memory is split between many small boundary tensors and one large live interior. Choosing `k` is a memory-shape decision, not a speed decision. ## When the layers are not uniform Real networks are not stacks of identical blocks. A text model has an embedding table lookup, a stack of transformer-style blocks, a pooling step and a head. A vision model has high-resolution early stages holding enormous feature maps and cheap low-resolution late stages. The square-root heuristic silently assumes every block holds the same bytes and costs the same to rerun, and neither is true here. The right ranking is bytes saved per unit of recomputed work. For each candidate segment, ask how many bytes of intermediates you stop retaining, and how many operations you must rerun to get them back. Sort by that ratio and checkpoint from the top until the memory target is met. This ranking produces two clear rules. Checkpoint the parts that hold large intermediates but are cheap to reproduce — elementwise activations, normalisation outputs, high-resolution feature maps early in a vision stack. Do not checkpoint parts whose retained output is already tiny: a pooling layer that collapses a long sequence to a single vector stores almost nothing, so dropping it frees almost nothing while still forcing everything before it in that segment to rerun. ## The failure mode this prevents The common bad configuration is a segment grid drawn where the code was easiest to edit rather than where the bytes are. A team wraps the embedding lookup and the pooling head — the two cheapest, smallest parts of the model — leaves the expensive attention-and-feedforward blocks resident, and reports that checkpointing did not help. It did exactly what it was asked to do; it was pointed at the wrong tensors. The diagnostic is always the same: measure where the activation bytes actually are before drawing any boundary. A second failure mode is drawing boundaries mid-block, so a segment's retained input is the widest tensor in the model rather than the narrowest. Good boundaries sit at natural narrow points in the graph — a block's residual input, a stage transition where resolution drops — because the boundary tensor is exactly what you have to keep. ## The practical procedure Profile the activation footprint per stage. Identify the narrow points in the graph as candidate boundaries. If the stack is uniform, start near the square-root count and round to something architecturally natural. If it is not uniform, rank by bytes freed per recomputed operation and checkpoint down that list only until the batch size you want fits — every extra segment past that is throughput you spent for nothing.