MALI: A memory efficient and reverse accurate integrator for Neural ODEs

Neural ordinary differential equations ( Neural ODEs ) are a new family of deeplearning models with continuous depth . However , the numerical estimation of the gradient in the continuous case is not well solved : existing implementations of the adjoint method suffer from inaccuracy in reverse-time trajectory , while the naive method and the adaptive checkpoint adjoint method ( ACA ) have a memory cost that grows with integration time . In this project , based on the asynchronous leapfrog ( ALF ) solver , we propose the Memory-efficient ALF Integrator ( MALI ) , which has a constant memory cost w.r.t number of solver steps in integration similar to the adjoint method , and guarantees accuracy in reverse-time trajectory ( hence accuracy in gradient estimation ) . We validate MALI in various tasks : on image recognition tasks , to our knowledge , MALI is the first to enable feasible training of a Neural ODE on ImageNet and outperform a well-tuned ResNet , while existing methods fail due to either heavy memory burden or inaccuracy ; for time series modeling , MALI significantly outperforms the adjoint method ; and for continuous generative models , MALI achieves new state-of-theart performance.We provide a pypi package : https : //jzkay12.github . io/TorchDiffEqPack 1 INTRODUCTION . Recent research builds the connection between continuous models and neural networks . The theory of dynamical systems has been applied to analyze the properties of neural networks or guide the design of networks ( Weinan , 2017 ; Ruthotto & Haber , 2019 ; Lu et al. , 2018 ) . In these works , a residual block ( He et al. , 2016 ) is typically viewed as a one-step Euler discretization of an ODE ; instead of directly analyzing the discretized neural network , it might be easier to analyze the ODE . Another direction is the neural ordinary differential equation ( Neural ODE ) ( Chen et al. , 2018 ) , which takes a continuous depth instead of discretized depth . The dynamics of a Neural ODE is typically approximated by numerical integration with adaptive ODE solvers . Neural ODEs have been applied in irregularly sampled time-series ( Rubanova et al. , 2019 ) , free-form continuous generative models ( Grathwohl et al. , 2018 ; Finlay et al. , 2020 ) , mean-field games ( Ruthotto et al. , 2020 ) , stochastic differential equations ( Li et al. , 2020 ) and physically informed modeling ( SanchezGonzalez et al. , 2019 ; Zhong et al. , 2019 ) . Though the Neural ODE has been widely applied in practice , how to train it is not extensively studied . The naive method directly backpropagates through an ODE solver , but tracking a continuous trajectory requires a huge memory . Chen et al . ( 2018 ) proposed to use the adjoint method to determine the gradient in continuous cases , which achieves constant memory cost w.r.t integration time ; however , as pointed out by Zhuang et al . ( 2020 ) , the adjoint method suffers from numerical errors due to the inaccuracy in reverse-time trajectory . Zhuang et al . ( 2020 ) proposed the adaptive checkpoint adjoint ( ACA ) method to achieve accuracy in gradient estimation at a much smaller memory cost compared to the naive method , yet the memory consumption of ACA still grows linearly with integration time . Due to the non-constant memory cost , neither ACA nor naive method are suitable for large scale datasets ( e.g . ImageNet ) or high-dimensional Neural ODEs ( e.g . FFJORD ( Grathwohl et al. , 2018 ) ) . In this project , we propose the Memory-efficient Asynchronous Leapfrog Integrator ( MALI ) to achieve advantages of both the adjoint method and ACA : constant memory cost w.r.t integration time and accuracy in reverse-time trajectory . MALI is based on the asynchronous leapfrog ( ALF ) integrator ( Mutze , 2013 ) . With the ALF integrator , each numerical step forward in time is reversible . Therefore , with MALI , we delete the trajectory and only keep the end-time states , hence achieve constant memory cost w.r.t integration time ; using the reversibility , we can accurately reconstruct the trajectory from the end-time value , hence achieve accuracy in gradient . Our contributions are : 1 . We propose a new method ( MALI ) to solve Neural ODEs , which achieves constant memory cost w.r.t number of solver steps in integration and accuracy in gradient estimation . We provide theoretical analysis . 2 . We validate our method with extensive experiments : ( a ) for image classification tasks , MALI enables a Neural ODE to achieve better accuracy than a well-tuned ResNet with the same number of parameters ; to our knowledge , MALI is the first method to enable training of Neural ODEs on a large-scale dataset such as ImageNet , while existing methods fail due to either heavy memory burden or inaccuracy . ( b ) In time-series modeling , MALI achieves comparable or better results than other methods . ( c ) For generative modeling , a FFJORD model trained with MALI achieves new state-of-the-art results on MNIST and Cifar10 . 2 PRELIMINARIES . 2.1 NUMERICAL INTEGRATION METHODS . An ordinary differential equation ( ODE ) typically takes the form dz ( t ) dt = fθ ( t , z ( t ) ) s.t . z ( t0 ) = x , t ∈ [ t0 , T ] , Loss = L ( z ( T ) , y ) ( 1 ) where z ( t ) is the hidden state evolving with time , T is the end time , t0 is the start time ( typically 0 ) , x is the initial state . The derivative of z ( t ) w.r.t t is defined by a function f , and f is defined as a sequence of layers parameterized by θ . The loss function is L ( z ( T ) , y ) , where y is the target variable . Eq . 1 is called the initial value problem ( IVP ) because only z ( t0 ) is specified . Algorithm 1 : Numerical Integration Input initial state x , start time t0 , end time T , error tolerance etol , initial stepsize h. Initialize z ( 0 ) = x , t = t0 While t < T error est =∞ While error est > etol h← h×DecayFactor ẑ , error est = ψh ( t , z ) If error est < etol h← h× IncreaseFactor t← t+ h , z ← ẑ Notations We summarize the notations following Zhuang et al . ( 2020 ) . • zi ( ti ) /z ( τi ) : hidden state in forward/reverse time trajectory at time ti/τi . • ψh ( ti , zi ) : the numerical solution at time ti + h , starting from ( ti , zi ) with a stepsize h. • Nf , Nz : Nf is the number of layers in f in Eq . 1 , Nz is the dimension of z . • Nt/Nr : number of discretized points ( outer iterations in Algo . 1 ) in forward / reverse integration . • m : average number of inner iterations in Algo . 1 to find an acceptable stepsize . Numerical Integration The algorithm for general adaptive-stepsize numerical ODE solvers is summarized in Algo . 1 ( Wanner & Hairer , 1996 ) . The solver repeatedly advances in time by a step , which is the outer loop in Algo . 1 ( blue curve in Fig . 1 ) . For each step , the solver decreases the stepsize until the estimate of error is lower than the tolerance , which is the inner loop in Algo . 1 ( green curve in Fig . 1 ) . For fixed-stepsize solvers , the inner loop is replaced with a single evaluation of ψh ( t , z ) using predefined stepsize h. Different methods typically use different ψ , for example different orders of the Runge-Kutta method ( Runge , 1895 ) . 2.2 ANALYTICAL FORM OF GRADIENT IN CONTINUOUS CASE . We first briefly introduce the analytical form of the gradient in the continuous case , then we compare different numerical implementations in the literature to estimate the gradient . The analytical form Figure 2 : In backward-pass , the adjoint method reconstructs trajectory as a separate IVP . Naive , ACA and MALI track the forward-time trajectory , hence are accurate . ACA and MALI only backpropagate through the accepted step , while naive method backpropagates through the search process hence has deeper computation graphs . of the gradient in the continuous case is dL dθ = − ∫ 0 T a ( t ) > ∂f ( z ( t ) , t , θ ) ∂θ dt ( 2 ) da ( t ) dt + ( ∂f ( z ( t ) , t , θ ) ∂z ( t ) ) > a ( t ) = 0 ∀t ∈ ( 0 , T ) , a ( T ) = ∂L ∂z ( T ) ( 3 ) where a ( t ) is the “ adjoint state ” . Detailed proof is given in ( Pontryagin , 1962 ) . In the next section we compare different numerical implementations of this analytical form . 2.3 NUMERICAL IMPLEMENTATIONS IN THE LITERATURE FOR THE ANALYTICAL FORM . We compare different numerical implementations of the analytical form in this section . The forwardpass and backward-pass of different methods are demonstrated in Fig . 1 and Fig . 2 respectively . Forward-pass is similar for different methods . The comparison of backward-pass among different methods are summarized in Table . 1 . We explain methods in the literature below . Naive method The naive method saves all of the computation graph ( including search for optimal stepsize , green curve in Fig . 2 ) in memory , and backpropagates through it . Hence the memory cost is NzNf ×Nt×m and depth of computation graph are Nf ×Nt×m , and the computation is doubled considering both forward and backward passes . Besides the large memory and computation , the deep computation graph might cause vanishing or exploding gradient ( Pascanu et al. , 2013 ) . Adjoint method Note that we use “ adjoint state equation ” to refer to the analytical form in Eq . 2 and 3 , while we use “ adjoint method ” to refer to the numerical implementation by Chen et al . ( 2018 ) . As in Fig . 1 and 2 , the adjoint method forgets forward-time trajectory ( blue curve ) to achieve memory cost NzNf which is constant to integration time ; it takes the end-time state ( derived from forward-time integration ) as the initial state , and solves a separate IVP ( red curve ) in reverse-time . Theorem 2.1 . ( Zhuang et al. , 2020 ) For an ODE solver of order p , the error of the reconstructed initial value by the adjoint method is ∑N−1 k=0 [ hp+1k DΦ T tk ( zk ) l ( tk , zk ) + ( −hk ) p+1DΦtkT ( zk ) l ( tk , zk ) ] + O ( hp+1 ) , where Φ is the ideal solution , DΦ is the Jacobian of Φ , l ( t , z ) and l ( t , z ) are the local error in forward-time and reverse-time integration respectively . Theorem 2.1 is stated as Theorem 3.2 in Zhuang et al . ( 2020 ) ; please see reference paper for detailed proof . To summarize , due to inevitable errors with numerical ODE solvers , the reverse-time trajectory ( red curve , z ( τ ) ) can not match the forward-time trajectory ( blue curve , z ( t ) ) accurately . The error in z propagates to dLdθ by Eq . 2 , hence affects the accuracy in gradient estimation . Adaptive checkpoint adjoint ( ACA ) To solve the inaccuracy of adjoint method , Zhuang et al . ( 2020 ) proposed ACA : ACA stores forward-time trajectory in memory for backward-pass , hence guarantees accuracy ; ACA deletes the search process ( green curve in Fig . 2 ) , and only backpropagates through the accepted step ( blue curve in Fig . 2 ) , hence has a shallower computation graph ( Nf × Nt for ACA vs Nf × Nt ×m for naive method ) . ACA only stores { z ( ti ) } Nti=1 , and deletes the computation graph for { f ( z ( ti ) , ti ) } Nti=1 , hence the memory cost is Nz ( Nf + Nt ) . Though the memory cost is much smaller than the naive method , it grows linearly with Nt , and can not handle very high dimensional models . In the following sections , we propose a method to overcome all these disadvantages of existing methods .

Summary: This paper presents a memory-efficient asynchronous leapfrog integrator for numerically solving neural ODEs, referred to as MALI. The method comes with a constant memory guarantee (like the adjoint method) and also guarantees reverse-time accuracy (like the adaptive checkpoint adjoint (ACA) method). The authors also give a rigorous theoretical analysis of MALI, and also discuss a "damped" version with an increased stability region. The method is evaluated on a variety of tasks which includes classification, dynamical modelling and generative modelling.