Memory-Efficient Adjoints via Graph Partitioning

Abstract

Derivative information plays a crucial role in the correctness and performance of scientific computing in a wide variety of scientific domains such as computational fluid dynamics (CFD), finance engineering and so on. The reverse mode of Algorithmic Differentiation is particularly efficient for the computation of derivatives of multivariate vector functions $F: R^{n}\mapsto R^{m}$, where the number of inputs $n$ far exceeds the number of outputs $m$, frequently appearing as cost functions in numerical optimization kernels. In particular, reverse-mode AD is at the heart of the back propagation algorithm widely used in machine learning. The reverse mode of AD requires that the control flow of the derivative program be reversed, meaning that the results of intermediate computations (in our case, the computational graph of $F$) must be stored either in memory or secondary storage. As a result, this requirement leads to the memory wall problem, especially for large-scale numerical problems, where the results of intermediate computations cannot fit entirely in memory. In this paper, we present an algorithm called Memory-Efficient Adjoints (ME-Adjoints) for solving the memory wall problem by dynamically applying a simple partitioning scheme to the computational graph of the function $F$ at runtime. Our approach employs operator overloading in C++ to achieve a fully automatic adjoining process, whereby derivative programs require only trivial changes to the code as opposed to the use of checkpointing techniques, which require substantial changes to the code.