Symplectic Adjoint Method for Exact Gradient of Neural ODE with Minimal Memory

1 Introduction . Deep neural networks offer remarkable methods for various tasks , such as image recognition [ 18 ] and natural language processing [ 4 ] . These methods employ a residual architecture [ 21 , 34 ] , in which the output xn+1 of the n-th operation is defined as the sum of a subroutine fn and the input xn as xn+1 = fn ( xn ) + xn . The residual architecture can be regarded as a numerical integration applied to an ordinary differential equation ( ODE ) [ 30 ] . Accordingly , a neural network model of the differential equation dx/dt = f ( x ) , namely , neural ODE , was proposed in [ 2 ] . Given an initial condition x ( 0 ) = x as an input , the neural ODE solves an initial value problem by numerical integration , obtaining the final value as an output y = x ( T ) . The neural ODE can model continuoustime dynamics such as irregularly sampled time series [ 24 ] , stable dynamical systems [ 37 , 42 ] , and physical phenomena associated with geometric structures [ 3 , 13 , 31 ] . Further , because the neural ODE approximates a diffeomorphism [ 43 ] , it can model probabilistic distributions of real-world data by a change of variables [ 12 , 23 , 25 , 45 ] . For an accurate integration , the neural ODE must employ a small step size and a high-order numerical integrator composed of many internal stages . A neural network f is used at each stage of each 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . proposed∗∗ symplectic adjoint method yes { xn } N−1n=0 , { Xn , i } si=1 MN+s L 4MNsL ∗Available only for the asynchronous leapfrog integrator . ∗∗Available for any Runge–Kutta methods . time step . Thus , the backpropagation algorithm consumes exorbitant memory to retain the whole computation graph [ 39 , 2 , 10 , 46 ] . The neural ODE employs the adjoint method to reduce memory consumption—this method obtains a gradient by a backward integration along with the state x , without consuming memory for retaining the computation graph over time [ 8 , 17 , 41 , 44 ] . However , this method incurs high computational costs to suppress numerical errors . Several previous works employed a checkpointing scheme [ 10 , 46 , 47 ] . This scheme only sparsely retains the state x as checkpoints and recalculates a computation graph from each checkpoint to obtain the gradient . However , this scheme still consumes a significant amount of memory to retain the computation graph between checkpoints . To address the above limitations , this study proposes the symplectic adjoint method . The main advantages of the proposed method are presented as follows . Exact Gradient and Fast Computation : In discrete time , the adjoint method suffers from numerical errors or needs a smaller step size . The proposed method uses a specially designed integrator that obtains the exact gradient in discrete time . It works with the same step size as the forward integration and is thus faster than the adjoint method in practice . Minimal Memory Consumption : Excepting the adjoint method , existing methods apply the backpropagation algorithm to the computation graph of the whole or a subset of numerical integration [ 10 , 46 , 47 ] . The memory consumption is proportional to the number of steps/stages in the graph times the neural network size . Conversely , the proposed method applies the algorithm only to each use of the neural network , and thus the memory consumption is only proportional to the number of steps/stages plus the network size . Robust to Rounding Error : The backpropagation algorithm accumulates the gradient from each use of the neural network and tends to suffer from rounding errors . Conversely , the proposed method obtains the gradient from each step as a numerical integration and is thus more robust to rounding errors . 2 Background and Related Work . 2.1 Neural Ordinary Differential Equation and Adjoint Method . We use the following notation . M : the number of stacked neural ODE components , L : the number of layers in a neural network , N , Ñ : the number of time steps in the forward and backward integrations , respectively , and s : the number of uses of a neural network f per step . s is typically equal to the number of internal stages of a numerical integrator [ 17 ] . A numerical integration forward in time requires a computational cost of O ( MNsL ) . It also provides a computation graph over time steps , which is retained with a memory of O ( MNsL ) ; the backpropagation algorithm is then applied to obtain the gradient . The total computational cost is O ( 2MNsL ) , where we suppose the computational cost of the backpropagation algorithm is equal to that of forward propagation . The memory consumption and computational cost are summarized in Table 1 . To reduce the memory consumption , the original study on the neural ODE introduced the adjoint method [ 2 , 8 , 17 , 41 , 44 ] . This method integrates the pair of the system state x and the adjoint variable λ backward in time . The adjoint variable λ represents the gradient ∂L∂x of some function L , and the backward integration of the adjoint variable λ works as the backpropagation ( or more generally the reverse-mode automatic differentiation ) in continuous time . The memory consumption is O ( M ) to retain the final values x ( T ) of M neural ODE components and O ( L ) to obtain the gradient of a neural network f for integrating the adjoint variable λ . The computational cost is at least doubled because of the re-integration of the system state x backward in time . The adjoint method suffers from numerical errors [ 41 , 10 ] . To suppress the numerical errors , the backward integration often requires a smaller step size than the forward integration ( i.e. , Ñ > N ) , leading to an increase in computation time . Conversely , the proposed symplectic adjoint method uses a specially designed integrator , which provides the exact gradient with the same step size as the forward integration . 2.2 Checkpointing Scheme . The checkpointing scheme has been investigated to reduce the memory consumption of neural networks [ 14 , 15 ] , where intermediate states are retained sparsely as checkpoints , and a computation graph is recomputed from each checkpoint . For example , Gruslys et al . applied this scheme to recurrent neural networks [ 15 ] . When the initial value x ( 0 ) of each neural ODE component is retained as a checkpoint , the initial value problem is solved again before applying the backpropagation algorithm to obtain the gradient of the component . Then , the memory consumption is O ( M ) for checkpoints and O ( NsL ) for the backpropagation ; the memory consumption is O ( M + NsL ) in total ( see the baseline scheme ) . ANODE scheme retains each step { xn } N−1n=0 as a checkpoint with a memory of O ( MN ) [ 10 ] . Form each checkpoint xn , this scheme recalculates the next step xn+1 and obtains the gradient using the backpropagation algorithm with a memory of O ( sL ) ; the memory consumption is O ( MN + sL ) in total . ACA scheme improves ANODE scheme for methods with adaptive time-stepping by discarding the computation graph to find an optimal step size . Even with checkpoints , the memory consumption is still proportional to the number of uses s of a neural network f per step , which is not negligible for a high-order integrator , e.g. , s = 6 for the Dormand–Prince method [ 7 ] . In this context , the proposed method is regarded as a checkpointing scheme inside a numerical integrator . Note that previous studies did not use the notation s. Instead of a checkpointing scheme , MALI employs an asynchronous leapfrog ( ALF ) integrator after the state x is paired up with the velocity state v [ 47 ] . The ALF integrator is time-reversible , i.e. , the backward integration obtains the state x equal to that in the forward integration without checkpoints [ 17 ] . However , the ALF integrator is a second-order integrator , implying that it requires a small step size and a high computational cost to suppress numerical errors . Higher-order Runge–Kutta methods can not be used in place of the ALF integrator because they are implicit or non-time-reversible . The ALF integrator is inapplicable to physical systems without velocity such as partial differential equation ( PDE ) systems . Nonetheless , a similar approach named RevNet was proposed before in [ 11 ] . When regarding ResNet as a forward Euler method [ 2 , 18 ] , RevNet has an architecture regarded as the leapfrog integrator , and it recalculates the intermediate activations in the reverse direction . 3 Adjoint Method . Consider a system d dt x = f ( x , t , θ ) , ( 1 ) where x , t , and θ , respectively , denote the system state , an independent variable ( e.g. , time ) , and parameters of the function f . Given an initial condition x ( 0 ) = x0 , the solution x ( t ) is given by x ( t ) = x0 + ∫ t 0 f ( x ( τ ) , τ , θ ) dτ . ( 2 ) The solution x ( t ) is evaluated at the terminal t = T by a function L as L ( x ( T ) ) . Our main interest is in obtaining the gradients of L ( x ( T ) ) with respect to the initial condition x0 and the parameters θ . Now , we introduce the adjoint method [ 2 , 8 , 17 , 41 , 44 ] . We first focus on the initial condition x0 and omit the parameters θ . The adjoint method is based on the variational variable δ ( t ) and the adjoint variable λ ( t ) . The variational and adjoint variables respectively follow the variational system and adjoint system as follows . d dt δ ( t ) = ∂f ∂x ( x ( t ) , t ) δ ( t ) for δ ( 0 ) = I , d dt λ ( t ) = −∂f ∂x ( x , t ) > λ ( t ) for λ ( T ) = λT . ( 3 ) The variational variable δ ( t ) represents the Jacobian ∂x ( t ) ∂x0 of the state x ( t ) with respect to the initial condition x0 ; the detailed derivation is summarized in Appendix A . Remark 1 . The quantity λ > δ is time-invariant , i.e. , λ ( t ) > δ ( t ) = λ ( 0 ) > δ ( 0 ) . The proofs of most Remarks and Theorems in this paper are summarized in Appendix B . Remark 2 . The adjoint variable λ ( t ) represents the gradient ( ∂L ( x ( T ) ) ∂x ( t ) ) > if the final condition λT of the adjoint variable λ is set to ( ∂L ( x ( T ) ) ∂x ( T ) ) > . This is because of the chain rule . Thus , the backward integration of the adjoint variable λ ( t ) works as reverse-mode automatic differentiation . The adjoint method has been used for data assimilation , where the initial condition x0 is optimized by a gradient-based method . For system identification ( i.e. , parameter adjustment ) , one can consider the parameters θ as a part of the augmented state x̃ = [ x θ ] > of the system d dt x̃ = f̃ ( x̃ , t ) , f̃ ( x̃ , t ) = [ f ( x , t , θ ) 0 ] , x̃ ( 0 ) = [ x0 θ ] . ( 4 ) The variational and adjoint variables are augmented in the same way . Hereafter , we let x denote the state or augmented state without loss of generality . See Appendix C for details . According to the original implementation of the neural ODE [ 2 ] , the final value x ( T ) of the system state x is retained after forward integration , and the pair of the system state x and the adjoint variable λ is integrated backward in time to obtain the gradients . The right-hand sides of the main system in Eq . ( 1 ) and the adjoint system in Eq . ( 3 ) are obtained by the forward and backward propagations of the neural network f , respectively . Therefore , the computational cost of the adjoint method is twice that of the ordinary backpropagation algorithm . After a numerical integrator discretizes the time , Remark 1 does not hold , and thus the adjoint variable λ ( t ) is not equal to the exact gradient [ 10 , 41 ] . Moreover , in general , the numerical integration backward in time is not consistent with that forward in time . Although a small step size ( i.e. , a small tolerance ) suppresses numerical errors , it also leads to a longer computation time . These facts provide the motivation to obtain the exact gradient with a small memory , in the present study .

Summary: The paper proposes a new algorithm for propagating gradients through neural ODEs, which minimizes peak memory requirements by checkpointing not only the individual solver steps but also the internal stages. The considered, specific, partitioned Runge--Kutta method has symplectic properties and provides exact gradients. The experiments demonstrate the reduced peak memory cost together with competitive runtimes.