Backprop needs every layer's activations, so memory grows with depth. Activation checkpointing stores activations only at every k-th layer and recomputes the rest during the backward pass, trading compute for memory.
Model an n-layer network where one layer's activations cost one unit of memory and one unit of forward compute:
ceil(n / k) checkpoints, and while processing one segment in the backward pass you also hold up to k recomputed activations, so peak_memory(n, k) = ceil(n / k) + k. With k = n this is roughly the no-checkpointing cost.recompute_cost(n, k): every activation that isn't a checkpoint is recomputed once: n - ceil(n / k)best_k(n): the k in 1..n that minimises peak memory (ties: the largest k, since it recomputes the least). The optimum is near √n.plan(n, memory_budget): the k with the least recompute whose peak memory fits in the budget, or None