Memory and checkpointing

A thousand-layer model returns one number, and the backward traversal needs a value from every layer. Keeping all of them costs a thousand stored values, and a training run that cannot hold a thousand must trade something else. Checkpointing keeps a few values and recomputes the rest during the backward traversal.

In the four-block picture below, the blocks run in order and a kept input is a checkpoint. Three kept inputs occupy three slots, each holding one intermediate array. The backward walk reaches the fourth block and needs the value block 3 produced, which was not kept, so it rebuilds that value with one recomputation. The pass therefore makes five forward block calls where four would do. The rule that fixes which inputs are kept before the walk starts is the schedule. !Four blocks run forward, and three of the four inputs are kept; when the backward pass reaches the fourth block it rebuilds that input from the nearest kept value. One extra forward call replaces one stored value: five calls where four would do, 1.25 times the arithmetic of storing everything. The gradient is unchanged because the rebuild repeats the original run exactly. Status: schematic.

Worked example. Scardapane's picture is the four-block case. It keeps the outputs of blocks 1, 2 and 4, so block 3's activations are missing and only the third block is recomputed. The caption's 1.25 is the ratio of forward computations, five against four [S:65:145]. Counted over the whole gradient, forward plus backward, the ratio is 9/8, because the four adjoint steps are unchanged and only the forward part is repeated. !At 1,024 blocks, storing every value holds 1,024 slots and costs 1,023 calls. Recursive halving holds ten slots and costs 5,120 calls, five times the arithmetic; recomputing everything holds one slot and costs 523,776 calls, 512 times. The gradient is the same in all three, so the choice is the point on this curve the machine can afford. Status: schematic.

The counts use one unit throughout. A call is one evaluation of one block charged by the schedule. The convention here is that the initial forward pass is not charged and the block that produces the loss is never evaluated again, so an eight-block chain has seven counted blocks.

Store all keeps every value, which for eight blocks means seven calls and eight slots: one call for each of the seven blocks before the loss, and one stored value for each of those seven plus the input [S:72:222]. Recompute all keeps one value and rebuilds each input from the start. Obtaining the $j$-th input costs $j$ calls, and the backward walk needs $j=1,\dots ,7$, so the total is $1+2+\cdots +7=28$ calls in a single slot [S:72:223]. Goodfellow's three-node case is the smallest instance of the same choice, storing the shared value once against recomputing it at every use [S:66:186]. Neither endpoint is attractive: one stores a value per block, and the other multiplies the work by the depth.

Recursive halving sits between them. Split the chain in half, keep the midpoint, and solve each half the same way. For eight blocks it makes 12 calls and holds three slots, and the slot the input occupied is reused by the first half. For a chain of L blocks that is

for L a power of two.

That line says memory grows like the logarithm of the depth while recomputation grows like the depth times its logarithm. At L = 1024 the formula gives 512 × 10 = 5120 calls in ten slots, against 1023 calls in 1024 slots for storing everything.

The three schedules form a frontier. Store all sits at linear memory and the fewest calls, and recompute all sits at constant memory with a quadratic number of calls. Halving sits between them at logarithmic memory [S:62:302]. A practical system selects a point on that curve and states its assumptions, because the curve is a menu rather than a rule.

Halving is not the best schedule at its own slot count. A schedule computed by dynamic programming can save calls at the same memory, and the saving grows with depth. The frontier is therefore a bound on what a simple rule achieves, not a theorem about the best possible rule.

Block cost varies across a network, and where it is largest depends on the architecture rather than on depth, so no general statement about which end is expensive can be made from the count. The schedule must use measured per-block costs: the analytic optimum assumes a uniform cost and does not take into account varying computational costs, which is why an optimum in calls need not be an optimum in seconds [S:72:227]. A count of calls assumes equal block costs, and a real network rarely supplies that equality.

The ratio is the price of the trade. Recursive halving costs $(K/2){\log }_{2}K$ calls against $K-1$ for store-all, so the ratio is 1.71 at eight blocks, 3.05 at 64, and 5.00 at 1024, in 3, 6 and 10 slots [S:72:222-224]. The memory falls from a value per block to ten values at 1024, and the computation rises by a factor of five. Murphy states the trade-off without counts: one can trade off memory for computation by discarding some prefixes and re-computing them as needed [S:62:302]. The gradient is the same in all three schedules, which is what makes the trade worth making.

Checkpointing does not change the gradient, provided the second run of a block repeats the first. A block that increments a global counter, or draws a fresh random mask, returns a different value when it is recomputed [S:67:267]. The gradient is then wrong with nothing to announce it. The schedule buys memory with time and with two conditions on the program, and both conditions are properties of the code rather than of the mathematics.

A dropout mask is the standard failure. The mask is drawn from a random stream, and the recomputation draws the next mask rather than the saved one, so the second forward pass describes a different network. A library that restores the random state before recomputing avoids the failure, and that restoration belongs to the schedule rather than to the optimisation of the code.

The count also assumes that every block costs the same. A chain whose expensive blocks happen to sit in the recomputed region pays more than the arithmetic suggests, and a count of calls is not a measurement of seconds. What the count does settle is the shape of the frontier: memory from linear to logarithmic, and time from linear to linear times logarithmic, with recompute-all at the quadratic end.

Limitation: Checkpointing moves the memory bill rather than removing it, and the move costs time. At a thousand layers the recomputation is five times the work of storing everything, and the trade assumes equal block costs and a deterministic program. A random draw or a global counter breaks the second assumption silently.

The formal statement. Fix a chain of $L$ blocks, $s_k=f_k(s_{k-1})$ for $k=1,\dots ,L$, and let the backward pass need the inputs $s_1,\dots ,s_{L-1}$. A schedule is a rule fixed before the backward pass, which decides for each needed input whether to store it and, if not, from which value and by which calls to recompute it [S:72:221].

Write $\mathrm{calls}(L)$ for the number of block calls and $\mathrm{slots}(L)$ for the number of values held at the peak. Store-all gives $L-1$ calls in $L$ slots. Recompute-all gives $L(L-1)/2$ calls in one slot. Recursive halving gives $(L/2){\log }_{2}L$ calls in ${\log }_{2}L$ slots, for $L$ a power of two.

The derivation. Recompute-all is a sum. To obtain $s_j$ the schedule calls $f_1,\dots ,f_j$ from $s_0$, costing $j$ calls. The backward pass needs $s_j$ for every $j$ from 1 to $L-1$, so the total is the sum, $L(L-1)/2$ calls with one slot held.

Recursive halving is a recurrence. Split at $L/2$, call the first half, and keep the midpoint in one new slot. Recurse on the second half with one slot fewer, then on the first half in the slot the input holds. The counts satisfy

with both counts zero at $L=1$. That line says each level pays half a chain in calls, doubles the work of the level below, and adds one slot per split.

Solve it by guessing. For $L=2^m$, guess $\mathrm{calls}=m{2}^{m-1}$. Substituting gives $2^{m-1}+(m-1)2^{m-1}$, so the guess satisfies the recurrence. The slot count unrolls to ${\log }_{2}L$. The omitted step is the induction; here substitution is the induction, since both sides reduce to the same expression.

The edge case. Take $L=100$. The closed forms assume an exact split at every level, and 100 is not a power of two, so the last split is uneven and the formulas are wrong. The procedure still runs, because it never tests for a power of two: it halves until a segment holds one block. The formula is a theorem under a hypothesis; the algorithm is not the formula.

Equal block cost is a second hypothesis. A schedule optimal in calls can be wrong in seconds, and the count never sees it [S:72:227].

What the one-paragraph version hides. Logarithmic memory is not free: the call count rises by a factor that grows with depth, so the saving and the price are one object seen from two sides. And the three schedules are not the whole frontier, because a dynamic programme finds fewer calls at the same slot count [S:72:226].

The correction condition. Recomputation returns the same gradient only if the second run of a block is the same run as the first. Each block must be a function of its arguments, and any random draw must repeat. Break either condition and the backward pass multiplies by a different value, with no error raised. This is a condition on the program, not on the chain.

Where the analogy breaks. Memory bought with time suggests a fixed exchange rate. The rate is not fixed: going from $L$ slots to one costs a factor that grows with $L$, so the last unit of memory is the most expensive.

Original sources. The scheduling problem is Griewank's, and Blondel and Roulet review it in §8.5, giving recursive halving as an algorithm and the split recurrence that defines an optimal scheme [S:72:222]. Scardapane states the same trade in a paragraph and draws it: store a few checkpoints, recompute the rest, and pay more computation for less memory [S:65:144]. Goodfellow, Bengio and Courville give the three-node case that motivates the whole question, with one implementation that stores a shared value and one that recomputes it [S:66:186].

An open question. Optimal schedules when blocks do not cost the same. The analytic optimum assumes a cost per block that is the same everywhere, and Blondel and Roulet note that it does not take into account varying computational costs for the functions $f_k$ [S:72:227]. The dynamic programme is the exception: it is generic and could a priori incorporate varying computational costs [S:72:226]. What is missing at scale is therefore measured per-block costs fed into an existing formulation, not a formulation that can price unequal blocks. A schedule optimal in call count can be slower in seconds, which matters because the two quantities are the ones an engineer trades.

What would settle it: per-block timings measured on the target hardware and fed into the schedule, with the resulting calls and seconds compared against the uniform-cost optimum at the same memory budget. A comparison that reports only calls has measured the proxy and not the thing.

A frontier result. Online checkpointing. The optimal fixed-length scheme needs the number of nodes before the run starts, and a while loop does not supply it. Online schemes keep a fixed set of checkpoints and rewrite them as the computation continues, and they are proved nearly optimal [S:72:227]. Reversible layers take the other route: when every block is invertible and its inverse is cheap, the backward pass rebuilds each input from the next one and the activation storage disappears [S:72:227].

What to read next. Blondel and Roulet §8.5 for halving, the dynamic programme and online schemes [S:72:222], then §8.6 for reversible layers [S:72:227]. Scardapane p.144 for the framework view in one paragraph and one figure [S:65:144]. Goodfellow p.186 for the smallest case, where the choice is visible without any counting.