Running a computation backwards with little memory

A slope without a subtraction ended on the fact that makes modern optimisation possible: the gradient of a function of a million variables costs only a few times as much as evaluating the function once, provided the derivatives are carried backwards through the computation rather than forwards. The chain rule says the derivative of a composition is a product of the derivatives of the steps — each one the flat map that fits a step closest — and a product can be multiplied from either end. Multiplying from the output end, the reverse sweep, needs one pass whatever the number of inputs.

It also needs something the forward sweep does not. To multiply backwards through step $k$, the sweep needs the derivative of step $k$, and that derivative depends on the state the computation was in when it reached step $k$. So the reverse sweep consumes the states in the opposite order to the one in which they were produced: the last state first, the first state last. The obvious way to provide them is to store every one, and for a long computation — a simulation of ten thousand time steps, a network with hundreds of layers applied to large inputs — that can be far more memory than the machine has.

The alternative is to store a few states and recompute the others from the nearest stored one when they are needed. That trades memory for time, and the question is what the best trade is. It has an exact answer, and the answer is Pascal’s triangle.

The problem, stated exactly #

Take a computation of $l$ steps, $x_0 \to x_1 \to \cdots \to x_l$, where each step computes the next state from the current one. The reverse sweep needs $x_{l-1}$, then $x_{l-2}$, and so on down to $x_0$ — the input to every step, in reverse order. The computer has $c$ slots, each able to hold one state, and $x_0$ is in one of them at the start. The allowed actions are to run a step forward from a state in hand, to copy the state in hand into a free slot, to free a slot, and to hand a state to the reverse sweep. The cost is the number of forward steps run, counting the first pass through the chain.

With one slot, which holds $x_0$, there is no choice. To produce $x_{l-1}$ the computation runs $l - 1$ steps from $x_0$; to produce $x_{l-2}$ it starts again from $x_0$ and runs $l - 2$; and so on. The total is $1 + 2 + \cdots + (l - 1) = l(l-1)/2$, the curve that leaves the hero’s frame at once: 19,900 forward steps to reverse a chain of 200. With $l$ slots there is no recomputation at all beyond the first pass. Everything interesting happens in between.

The recursion that finds the best schedule #

With $c$ slots and $x_0$ stored, any schedule must at some point store a second state; say the first one it stores is $x_j$. From then on the work splits cleanly into two independent problems. Everything beyond $x_j$ has to be produced first, with $x_j$ occupying a slot, so it is a problem of reversing $l - j$ steps with $c - 1$ slots. Once that is done, $x_j$’s slot is free again, and the states before $x_j$ have to be produced with all $c$ slots and $x_0$ still stored: a problem of reversing $j$ steps with $c$ slots. Running from $x_0$ to $x_j$ costs $j$ steps. So the least cost $T(l, c)$ satisfies

$$T(l, c) = \min_{1 \le j < l} \big( j + T(l - j, c - 1) + T(j, c) \big),$$

with $T(1, c) = 0$ and $T(l, 1) = l(l-1)/2$. This is a dynamic programme of the most ordinary kind, and the figures solve it exactly for every length up to 200 and every number of slots up to ten — and for the trade figure further down, up to forty slots and ten thousand steps, through the formula the recursion is checked against.

The two-slot schedule runs six steps, stores $x_6$, reverses the last four steps from $x_6$ with what is now a single free slot — three, two and one steps of recomputation — and then reverses the first six from $x_0$ with both slots again, storing $x_3$ on the way. Twenty forward steps in all. With three slots the same chain takes fifteen, and the schedule nests one level deeper.

The shape of an optimal schedule is the shape of the recursion. A stored state divides the chain; everything after it is finished before anything before it is touched; and the slot is released only when the part it governs is complete. The pictures look like the recursion trees of a divide-and-conquer algorithm, and they are, with the twist that the two halves are not equal and do not get equal resources: the part after the stored state gets one slot fewer, and the part before it gets the steps already run.

Pascal’s triangle #

The recursion has an exact solution, and the cleanest way to state it is to turn the question round. Instead of asking for the cost of reversing $l$ steps with $c$ slots, ask: if no step may be run more than $r$ times, the first pass included, how long a chain can $c$ slots reverse? Call the answer $\beta(c, r)$.

Look at the first stored state $x_j$ in a schedule that respects the limit. The steps before it have already been run once, on the way to $x_j$, so in reversing them later each may be run at most $r - 1$ more times, with all $c$ slots available: there are at most $\beta(c, r - 1)$ of them. The steps after $x_j$ have not been run yet, and they are reversed with $c - 1$ slots and the full $r$ runs: at most $\beta(c - 1, r)$ of them. So

$$\beta(c, r) = \beta(c, r - 1) + \beta(c - 1, r),$$

with $\beta(c, 0) = 1$ and $\beta(0, r) = 1$ in the natural conventions, and the solution of that recursion is the binomial coefficient

$$\beta(c, r) = \binom{c + r}{c}.$$

The table is Pascal’s triangle laid on its side, and the addition rule that builds the triangle is the schedule itself: the entry above counts the steps after the first stored state, which get one slot fewer, and the entry to the left counts the steps before it, which get one run fewer. Three slots and three runs per step reverse $\binom{6}{3} = 20$ steps; three slots and five runs reverse 56; six slots and six runs reverse 924. The total cost follows from the same count. If $r$ is the least number with $\binom{c+r}{c} \ge l$, the fewest forward steps are

$$T(l, c) = r\,l - \binom{c + r}{c + 1},$$

which is Griewank’s formula of 1992. The figures compute $T$ by the recursion and check it against this formula at every length and slot count drawn, and the two agree everywhere.

Two slots and the triangular numbers #

The smallest non-trivial case shows the whole mechanism with numbers small enough to check by hand. With two slots, $\beta(2, r) = \binom{r+2}{2}$, the triangular numbers 3, 6, 10, 15, 21: two slots and no step run more than three times reverse a chain of ten steps, and no longer.

The ten-step drawing above is that case exactly. The part before the first stored state has already been run once and has two runs left, with both slots: at most $\beta(2, 2) = 6$ steps. The part after it has three runs and only one slot, which means recomputing from the stored state every time: at most $\beta(1, 3) = 4$ steps. Six plus four is ten, and the schedule stores $x_6$ — exactly where the arithmetic says it must. Its cost is $3 \cdot 10 - \binom{5}{3} = 30 - 10 = 20$, the number the drawing reports.

A chain of eleven steps would need one step run four times, and the formula then gives $4 \cdot 11 - \binom{6}{3} = 44 - 20 = 24$: crossing a triangular number costs one extra run of every step from there on. That is why the hero’s curves are not smooth parabolas but are made of straight pieces whose slopes step up by one at each binomial coefficient — $r$ is the slope, and it changes exactly when the length passes $\binom{c+r}{c}$. On a long chain the pieces are too short to see.

Every step is run about equally often #

The binomial bound is tight in a strong sense. At the length $l = \binom{c+r}{c}$, an optimal schedule does not merely keep its total within the bound: it runs every individual step at most $r$ times.

With three slots and 56 steps — exactly $\binom{8}{3}$ — the optimal schedule runs 21 steps five times, 15 four times, 10 three times, 6 twice, 3 once, and one step not at all. Those counts are themselves binomial coefficients: 21, 15, 10, 6, 3 and 1 are the entries $\binom{k+2}{2}$, the triangular numbers, read downwards. The total is $5 \cdot 56 - \binom{8}{4} = 280 - 70 = 210$ forward steps, an average of 3.75 runs per step, for a reversal that would cost 1,540 with one slot and nothing with 56.

The five-times-run steps are spread through the chain rather than gathered at its start, which is the visible difference between this schedule and the naive one. With one slot the first step is run $l - 1$ times and the last never: all the burden falls on the beginning. With several slots the stored states act as fresh starting points, and the recursion distributes the burden so evenly that no step is run more than the bound requires.

How much a slot is worth #

For a long chain the trade between memory and time becomes a trade between a few numbers. Take ten thousand steps.

One slot costs 5,000 forward steps per step of the chain — the quadratic disaster. Two slots bring it to 93, three to 28, and by ten slots the cost is 5.76 forward steps per step: the whole chain run about six times over, in exchange for storing ten states rather than ten thousand. Forty slots bring it to 2.91. Because $\binom{c + r}{c}$ grows like $c^r$ for fixed $r$ and like $r^c$ for fixed $c$, a chain of length $l$ can be reversed with $c$ and $r$ both around $\log l$, and in practice a memory of a few dozen states makes the reverse sweep cost only a handful of forward sweeps.

The natural alternative to Griewank’s schedule is repeated halving: store the midpoint, reverse the second half, then the first, and do the same inside each half. That needs about $\log_2 l$ slots and runs each step about $\tfrac12\log_2 l$ times on average. For ten thousand steps it uses fifteen slots and costs 6.46 forward steps per step. The binomial schedule with the same fifteen slots costs 4.52. Halving puts the stored state in the middle, and the recursion says the middle is wrong, in a direction that depends on the memory. The part after the stored state is about to lose a slot; the part before it has already used up a run. With two slots the lost slot is decisive and the first state is stored 9,859 steps in, leaving a short tail. With fifteen slots the used run is decisive, and every optimal schedule stores its first state between step 816 and step 3,876, well before the middle.

The schedule most often used in practice is simpler still. Store every hundredth state on a first pass through a chain of ten thousand steps; then, working backwards one segment at a time, recompute each segment of a hundred steps from its stored start and keep all of its states while the reverse sweep consumes them. That uses about two hundred slots — a hundred for the segment starts and a hundred for the segment in hand — and runs the chain about twice. In general it needs about $2\sqrt l$ slots for about two runs, which is why it is known as square-root checkpointing in the training of neural networks. The binomial formula says how close to optimal it is: with two hundred slots the least possible cost for ten thousand steps is $2 \cdot 10{,}000 - \binom{202}{201} = 19{,}798$ forward steps, 1.98 runs per step. The simple scheme is within a few per cent of the best, at that memory. Its weakness is the other end of the table: it has no answer for a machine that can hold fifteen states, and the binomial schedule does.

Where this sits in differentiation #

The argument nowhere used the fact that the reverse sweep is computing derivatives. It is a statement about any process that consumes the states of a forward computation in reverse order, and the same schedules are used to reverse simulations for other reasons — to run a time-reversible system backwards, or to trace the history of an error. In differentiation the reverse sweep is the reason the order matters, and the cost of the derivative is the cost of the reversal: a gradient that is cheap in arithmetic, as the comparison of forward and reverse sweeps showed, is cheap in memory too, at a price that grows only logarithmically.

The counting argument behind $\binom{c+r}{c}$ is the same one that counts paths in a grid: a schedule that respects the limits is a sequence of decisions to spend a slot or to spend a run, and the longest chain is the number of ways to interleave $c$ of one with $r$ of the other. The number of times a value appears in Pascal’s triangle has nothing to do with differentiation, but the entries being counted are the same: the lengths a memory of a given size can handle. The bracketings of a chain of products — the other half of the cost, taken up at the end — are counted by the Catalan numbers, which is another way of saying there are far too many to try one by one. And the arithmetic of the chain rule, the product rule and the rest is untouched by all of it. Only the bookkeeping of when each piece of the product is available has changed.

A model that counts only forward steps #

The model counts forward steps and nothing else. Every step is taken to cost the same, storing and restoring a state to cost nothing, and the backward steps to be a fixed cost outside the count. Real computations violate all three: steps differ in cost, writing a state to slow memory can cost more than recomputing it, and the backward step is typically two or three times the cost of a forward one. Griewank and Andrea Walther extended the analysis to some of these refinements; the clean binomial answer holds only for the uniform model.

The length is known in advance. The recursion places its stored states using $l$, and a computation whose number of steps is decided as it runs — an adaptive solver that stops when it has converged — cannot use it. Online schedules that place states without knowing the length exist and cost somewhat more; how much more, in the worst case, is a separate question with its own bounds.

One processor, one memory. Everything here is sequential. With several processors, recomputation of different segments can run in parallel, and the best trade among memory, recomputation and the number of processors is a three-way problem that the two-way table does not address.

Optimal schedules are not unique. The recursion’s minimum is often reached at several split points, and the drawings show one choice. The run counts in the 56-step figure are those of the schedule drawn; the bound that no step is run more than five times holds for every optimal schedule at that length, and the multiset of counts may differ between them.

Still open: the cheapest order for the whole product #

Checkpointing settles how to supply the states. The other half of the cost is the product itself. Forward and reverse sweeps are the two extreme orders for multiplying the chain of small Jacobian matrices — one from the input end, one from the output end — and any bracketing in between is allowed. For a chain whose matrices have known shapes, the cheapest bracketing is the classical matrix-chain problem, solved by another dynamic programme of the same kind as the one above. But a real computation is not a chain; it is a graph of operations in which intermediate results feed several later ones, and there the cheapest order of elimination is NP-complete, as Uwe Naumann proved in 2008. Nobody has a fast method for finding it, and nobody knows how much the best order can save over the reverse sweep on the computations that matter most. Combining that question with this one — the cheapest order of multiplication when the states needed for it must also be recomputed — is the form in which the problem actually arises, and it has no exact solution.