System and method for divide-and-conquer checkpointing
A system and method which allows the basic checkpoint-reverse-mode AD strategy (of recursively decomposing the computation to reduce storage requirements of reverse-mode AD) to be applied to arbitrary programs: not just programs consisting of loops, but programs with arbitrarily complex control flow. The method comprises (a) transforming the program into a formalism that allows convenient manipulation by formal tools, and (b) introducing a set of operators to allow computations to be decomposed by running them for a given period of time then pausing them, while treating the paused program as a value subject to manipulation.
1. A method for computing the gradient of a function implemented as a computer program using a processor and a memory, comprising:
a. representing a computation as a sequence of computational steps involving primitive computations;
b. splitting the sequence of steps at a split point which is either a midpoint, or a point selected based on a trade-off between computation and storage, to generate at least a first portion of the computation and a second portion of the computation;
c. computing an intermediate state of the computation at the split point by applying the first portion of the computation prior to the split point;
d. recursively performing the splitting and the computing with respect to both of the first portion of the computation and the second portion of the computation to determine a plurality of portions of the computation and a plurality of intermediate states of the computation at a respective plurality of split points, the recursive performance of the splitting and the computing continuing until all of plurality of portions of the computation have less than a predetermined number of computational steps; and
e. determining the the gradient of the computation using automatic differentiation based on the plurality of portions of the computation and the plurality of intermediate states of the computation.
2. The method of claim 1 , wherein the functionality is exposed to the user through an application programmer interface comprising an operator or higher-order function that takes a function and an argument of the function as input and returns as output the value of the function and the gradient of the function.
3. The method of claim 2 , wherein the operator or higher-order function is nested.
4. The method of claim 1 , wherein the ability to divide a computation sequence into two portions is implemented via a general-purpose checkpointing mechanism, the general-purpose checkpointing mechanism comprising:
a. evaluating a function at an input of the function and returning both a value at an output of the function and a number of steps required to compute the output;
b. evaluating a first specified number of steps of a computation of the-function applied to an argument and returning an intermediate state of the computation after that specified number of steps as a reusable checkpoint; and
c. resuming a computation at the intermediate state and returning the result of the computation.
5. The method of claim 3 , where the general-purpose checkpointing mechanism is provided by an intepreter written in continuation passing style.
6. The method of claim 3 , where the general-purpose checkpointing mechanism is provided by a compiler that generates code in continuation passing style.
7. The method of claim 3 , where the general-purpose checkpointing mechanism is provided using the POSIX fork( ) primitive.
8. The method of claim 1 , where the checkpointing process proceeds without knowledge of the amount of computation in the primitive computations, and functions in an online fashion, discarding previously acquired checkpoints so as to maintain the logarithmic overhead in both time and space regardless of when the primal computation terminates.
9. A method for applying an arbitrary adjoint computation to the reversal of the computation graph of an arbitrary computer program using a processor and a memory, comprising:
a. representing the computation as a sequence of steps involving primitive computations;
b. splitting the sequence of steps at a split point which is either a midpoint, or a point selected based on a trade-off between computation and storage, to generate at least a first portion of the computation and a second portion of the computation;
c. computing an intermediate state of the computation at the split point, and applying the first portion of the computation prior to the split point;
d. recursively applying the second portion of the computation beginning at the split point to compute both the value of the entire computation and the gradient of an adjoint computation at the split point; and
e. recursively applying the first portion of the computation up to the split point to compute the value of the adjoint computation at the input if the number of computational steps exceeds a predetermined minimum.
10. The method of claim 9 , where the functionality is exposed to the user through an application programmer interface consisting of an operator or higher-order function that takes a function and its argument as input and returns as output the value of the function and the value of the adjoint computation.
11. The method of claim 10 , where the operator or higher-order function is nested.
12. The method of claim 9 , where the ability to divide a computation sequence into two portions is implemented via a general-purpose checkpointing mechanism, the general-purpose checkpointing mechanism comprising:
a. evaluating a function at its input and returning both a value at its output and a number of steps required to compute the output;
b. evaluating a first specified number of steps of a computation of the function applied to an argument and returning an intermediate state of the computation after that specified number of steps as a reusable checkpoint; and
c. resuming a checkpointed computation at the intermediate state and returning the result of its computation.
13. The method of claim 12 , where the general-purpose checkpointing mechanism is provided by an intepreter written in continuation passing style.
14. The method of claim 12 , where the general-purpose checkpointing mechanism is provided by a compiler that generates code in continuation passing style.
15. The method of claim 12 , where the general-purpose checkpointing mechanism is provided using the POSIX fork( ) primitive.
16. The method of claim 9 , where the checkpointing process proceeds without knowledge of the amount of computation in the primal computation, and functions in an online fashion, discarding previously acquired checkpoints so as to maintain the logarithmic overhead in both time and space regardless of when the primal computation terminates.
17. A method of dividing a computation sequence into two portions that is implemented via a general-purpose checkpointing mechanism, comprising:
a. evaluating a function at an input of the function and returning both a value at an output of the function and a number of steps required to compute the output;
b. evaluating a first specified number of steps of a computation of the-function applied to an argument and returning an intermediate state of the computation after that specified number of steps as a reusable checkpoint, the first specified number of steps determined based on the number of steps required to compute the output; and
c. resuming a checkpointed computation at the intermediate state and returning the result of the computation.
18. The method of claim 17 , wherein the general-purpose checkpointing mechanism is provided by an interpreter written in continuation passing style.
19. The method of claim 17 , wherein the general-purpose checkpointing mechanism is provided by a compiler that generates code in continuation passing style.
20. The method of claim 17 , wherein the general-purpose checkpointing mechanism is provided using the POSIX fork( ) primitive.