The Symplectic Adjoint Method: Memory-Efficient Backpropagation of Neural-Network-Based Differential Equations

Abstract

The combination of neural networks and numerical integration can provide highly accurate models of continuous-time dynamical systems and probabilistic distributions. However, if a neural network is used <inline-formula> <tex-math notation="LaTeX">$ \boldsymbol {n}$ </tex-math></inline-formula> times during numerical integration, the whole computation graph can be considered as a network <inline-formula> <tex-math notation="LaTeX">$ \boldsymbol {n}$ </tex-math></inline-formula> times deeper than the original. The backpropagation algorithm consumes memory in proportion to the number of uses <italic>times</italic> of the network size, causing practical difficulties. This is true even if a checkpointing scheme divides the computation graph into subgraphs. Alternatively, the adjoint method obtains a gradient by a numerical integration backward in time; although this method consumes memory only for single-network use, the computational cost of suppressing numerical errors is high. The symplectic adjoint method proposed in this study, an adjoint method solved by a symplectic integrator, obtains the exact gradient (up to rounding error) with memory proportional to the number of uses <italic>plus</italic> the network size. The theoretical analysis shows that it consumes much less memory than the naive backpropagation algorithm and checkpointing schemes. The experiments verify the theory, and they also demonstrate that the symplectic adjoint method is faster than the adjoint method and is more robust to rounding errors.