Second-Order Neural ODE Optimizer

1 Introduction Neural ODEs ( Chen et al. , 2018 ) have received tremendous attention over recent years . Inspired by taking the continuous limit of the “ discrete ” residual transformation , xk+1 = xk + F ( xk , θ ) , they propose to directly parameterize the vector field of an ODE as a deep neural network ( DNN ) , i.e . dx ( t ) dt = F ( t , x ( t ) , θ ) , x ( t0 ) = xt0 , ( 1 ) where x ( t ) ∈ Rm and F ( · , · , θ ) is a DNN parameterized by θ ∈ Rn . This provides a powerful paradigm connecting modern machine learning to classical differential equations ( Weinan , 2017 ) and has since then achieved promising results on time series analysis ( Rubanova et al. , 2019 ; Kidger et al. , 2020b ) , reversible generative flow ( Grathwohl et al. , 2018 ; Nguyen et al. , 2019 ) , image classification ( Zhuang et al. , 2020 , 2021 ) , and manifold learning ( Lou et al. , 2020 ; Mathieu & Nickel , 2020 ) . Due to the continuous-time representation , Neural ODEs feature a distinct optimization process ( see Fig . 2 ) compared to their discrete-time counterparts , which also poses new challenges . First , the 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) forward pass of Neural ODEs involves solving ( 1 ) with a black-box ODE solver . Depending on how its numerical integration is set up , the propagation may be refined to arbitrarily small step sizes and become prohibitively expensive to solve without any regularization ( Ghosh et al. , 2020 ; Finlay et al. , 2020 ) . On the other hand , to prevent Back-propagating through the entire ODE solver , the gradients are typically obtained by solving a backward adjoint ODE using the Adjoint Sensitivity Method ( ASM ; Pontryagin et al . ( 1962 ) ) . While this can be achieved at a favorable O ( 1 ) memory , it further increases the runtime and can suffer from inaccurate integration ( Gholami et al. , 2019 ) . For these reasons , Neural ODEs often take notoriously longer time to train , limiting their applications to relatively small or synthetic datasets ( Massaroli et al. , 2020 ) until very recently ( Zhuang et al. , 2021 ) . To improve the convergence rate of training , it is natural to consider higher-order optimization . While efficient second-order methods have been proposed for discrete models ( Ba et al. , 2016 ; George et al. , 2018 ) , it remains unclear how to extend these successes to Neural ODEs , given their distinct computation processes . Indeed , limited discussions in this regard only note that one may repeat the backward adjoint process recursively to obtain higher-order derivatives ( Chen et al. , 2018 ) . This is , unfortunately , impractical as the recursion will accumulate the aforementioned integration errors and scale the per-iteration runtime linearly . As such , second-order methods for Neural ODEs are seldom considered in practice , nor have they been rigorously explored from an optimization standpoint . In this work , we show that efficient second-order optimization is in fact viable for Neural ODEs . Our method is inspired by the emerging Optimal Control perspective ( Weinan et al. , 2018 ; Liu & Theodorou , 2019 ) , which treats the parameter θ as a control variable , so that the training process , i.e . optimizing θ w.r.t . some objective , can be interpreted as an Optimal Control Programming ( OCP ) . Specifically , we show that a continuous-time OCP methodology , called Differential Programming , provides analytic second-order derivatives by solving a set of coupled matrix ODEs . Interestingly , these matrix ODEs can be augmented to the backward adjoint ODE and solved simultaneously . In other words , a single backward pass is sufficient to compute all derivatives , including the original ASM-based gradient , the newly-derived second-order matrices , or even higher-order tensors . Further , these higher-order computations enjoy the sameO ( 1 ) memory and a comparable runtime to first-order methods by adopting Kronecker factorization ( Martens & Grosse , 2015 ) . The resulting method – called SNOpt – admits superior convergence in wall-clock time ( Fig . 1 ) , and the improvement remains consistent across image classification , continuous normalizing flow , and time-series prediction . Our OCP framework also facilitates progressive training of the network architecture . Specifically , we study an example of jointly optimizing the “ integration time ” of Neural ODEs , in analogy to the “ depth ” of discrete DNNs . While analytic gradients w.r.t . this architectural parameter have been derived under the ASM framework , they were often evaluated on limited synthetic datasets ( Massaroli et al. , 2020 ) . In the context of OCP , however , free-horizon optimization is a well-studied problem for practical applications with a priori unknown terminal time ( Sun et al. , 2015 ; De Marchi & Gerdts , 2019 ) . In this work , we show that these principles can be applied to Neural ODEs , yielding a novel second-order feedback policy that adapts the integration time throughout training . On training CIFAR10 , this further leads to a 20 % runtime reduction , yet without hindering test-time accuracy . In summary , we present the following contributions . • We propose a novel computational framework for computing higher-order derivatives of deep continuous-time models , with a rigorous analysis using continuous-time Optimal Control theory . • We propose an efficient second-order method , SNOpt , that achieves superior convergence ( in wall-clock time ) over first-order methods in training Neural ODEs , while retaining constant memory complexity . These improvements remain consistent across various applications . • To show that our framework also enables direct architecture optimization , we derive a secondorder feedback policy for adapting the integration horizon and show it further reduces the runtime . 2 Preliminaries . Notation . We use roman and italic type to represent a variable x ( t ) and its realization x ( t ) given an ODE . ODESolve denotes a function call that solves an initial value problem given an initial condition , start and end integration time , and vector field , i.e . ODESolve ( x ( t0 ) , t0 , t1 , F ) where dx ( t ) dt = F . Forward and backward computations of Neural ODEs . Given an initial condition x ( t0 ) and integration interval [ t0 , t1 ] , Neural ODEs concern the following optimization over an objective L , min θ L ( x ( t1 ) ) , where x ( t1 ) = x ( t0 ) + ∫ t1 t0 F ( t , x ( t ) , θ ) dt ( 2 ) is the solution of the ODE ( 1 ) and can be solved by calling a black-box ODE solver , i.e . x ( t1 ) = ODESolve ( x ( t0 ) , t0 , t1 , F ) . The use of ODESolve allows us to adopt higher-order numerical methods , e.g . adaptive Runge-Kutta ( Press et al. , 2007 ) , which give more accurate integration compared with e.g . vanilla Euler discretization in residual-based discrete models . To obtain the gradient ∂L∂θ of Neural ODE , one may naively Back-propagate through ODESolve . This , even if it could be made possible , leads to unsatisfactory memory complexity since the computation graph can grow arbitrarily large for adaptive ODE solvers . Instead , Chen et al . ( 2018 ) proposed to apply the Adjoint Sensitivity Method ( ASM ) , which states that the gradient can be obtained through the following integration . ∂L ∂θ = − ∫ t0 t1 a ( t ) T ∂F ( t , x ( t ) , θ ) ∂θ dt , ( 3 ) where a ( t ) ∈ Rm is referred to the adjoint state whose dynamics obey a backward adjoint ODE , −da ( t ) dt = a ( t ) T ∂F ( t , x ( t ) , θ ) ∂x ( t ) , a ( t1 ) = ∂L ∂x ( t1 ) . ( 4 ) Equations ( 3 , 4 ) present two coupled ODEs that can be viewed as the continuous-time expression of the Back-propagation ( LeCun et al. , 1988 ) . Algorithmically , they can be solved through another call of ODESolve ( see Fig . 2 ) with an augmented dynamics G , i.e.x ( t0 ) a ( t0 ) ∂L/∂θ = ODESolve ( x ( t1 ) a ( t1 ) 0 , t1 , t0 , G ) , where G t , x ( t ) a ( t ) · , θ : = F ( t , x ( t ) , θ ) −a ( t ) T ∂F∂x −a ( t ) T ∂F∂θ ( 5 ) augments the original dynamics F in ( 1 ) with the adjoint ODEs ( 3 , 4 ) . Notice that this computation ( 5 ) depends only on ( x ( t1 ) , a ( t1 ) ) . This differs from naive Back-propagation , which requires storing intermediate states along the entire computation graph of forward ODESolve . While the latter requires O ( T̃ ) memory cost,1 the computation in ( 5 ) only consumes constant O ( 1 ) memory cost . Chen et al . ( 2018 ) noted that if we further encapsulate ( 5 ) by ∂∂θL = grad ( L , θ ) , one may compute higherorder derivatives by recursively calling ∂nL ∂θn = grad ( ∂n−1L ∂θn−1 , θ ) , starting from n=1 . This can scale unfavorably due to its recursive dependence and accumulated integration errors . Indeed , Table 1 Table 1 : Numerical errors between ground-truth and adjoint derivatives using different ODESolve on CIFAR10 . rk4 implicit adams dopri5 ∂L ∂θ 7.63×10 −5 2.11×10−3 3.44×10−4 ∂2L ∂θ2 6.84×10 −3 2.50×10−1 41.10 suggests that the errors of second-order derivatives , ∂ 2L ∂θ2 , obtained from the recursive adjoint procedure can be 2-6 orders of magnitude larger than the ones from the first-order adjoint , ∂L∂θ . In the next section , we will present a novel optimization framework that computes these higher-order derivatives without any recursion ( Section 3.1 ) and discuss how it can be implemented efficiently ( Section 3.2 ) . 3 Approach . 3.1 Dynamics of Higher-order Derivatives using Continuous-time Optimal Control Theory . OCP perspective is a recently emerging methodology for analyzing optimization of discrete DNNs . Central to its interpretation is to treat the layer propagation of a DNN as discrete-time dynamics , so 1 T̃ is the number of the adaptive steps used to solve ( 1 ) , as an analogy of the “ depth ” of Neural ODEs . that the training process , i.e . finding an optimal parameter of a DNN , can be understood like an OCP , which searches for an optimal control subjected to a dynamical constraint . This perspective has provided useful insights on characterizing the optimization process ( Hu et al. , 2019 ) and enhancing principled algorithmic design ( Liu et al. , 2021a ) . We leave a complete discussion in Appendix A.1 . Lifting this OCP perspective from discrete DNNs to Neural ODEs requires special treatments from continuous-time OCP theory ( Todorov , 2016 ) . Nevertheless , we highlight that training Neural ODEs and solving continuous-time OCP are fundamentally intertwined since these models , by construction , represent continuous-time dynamical systems . Indeed , the ASM used for deriving ( 3 , 4 ) originates from the celebrated Pontryagin ’ s principle ( Pontryagin et al. , 1962 ) , which is an optimality condition to OCP . Hence , OCP analysis is not only motivated but principled from an optimization standpoint . We begin by first transforming ( 2 ) to a form that is easier to adopt the continuous-time OCP analysis . min θ [ Φ ( xt1 ) + ∫ t1 t0 ` ( t , xt , ut ) dt ] subjected to { dxt dt = F ( t , xt , ut ) , xt0 = xt0 dut dt = 0 , ut0 = θ , ( 6 ) where x ( t ) ≡ xt , and etc . It should be clear that ( 6 ) describes ( 2 ) without loss of generality by having ( Φ , ` ) : = ( L , 0 ) . These functions are known as the terminal and intermediate costs in standard OCP . In training Neural ODEs , ` can be used to describe either the weight decay , i.e . ` ∝ ‖ut‖ , or more complex regularization ( Finlay et al. , 2020 ) . The time-invariant ODE imposed for ut makes the ODE of xt equivalent to ( 1 ) . Problem ( 6 ) shall be understood as a particular type of OCP that searches for an optimal initial condition θ of a time-invariant control ut . Despite seemly superfluous , this is a necessary transformation that enables rigorous OCP analysis for the original training process ( 2 ) , and it has also appeared in other control-related analyses ( Zhong et al. , 2020 ; Chalvidal et al. , 2021 ) . Next , define the accumulated loss from any time t ∈ [ t0 , t1 ] to the integration end time t1 as Q ( t , xt , ut ) : = Φ ( xt1 ) + ∫ t1 t ` ( τ , xτ , uτ ) dτ , ( 7 ) which is also known in OCP as the cost-to-go function . Recall that our goal is to compute higherorder derivatives w.r.t . the parameter θ of Neural ODEs . Under the new OCP representation ( 6 ) , the first-order derivative ∂L∂θ is identical to ∂Q ( t0 , xt0 , ut0 ) ∂ut0 . This is because Q ( t0 , xt0 , ut0 ) accumulates all sources of losses between [ t0 , t1 ] ( hence it sufficiently describes L ) and ut0 = θ by construction . Likewise , the second-order derivatives can be captured by the Hessian ∂ 2Q ( t0 , xt0 , ut0 ) ∂ut0∂ut0 = ∂ 2L ∂θ∂θ ≡ Lθθ . In other words , we are only interested in obtaining the derivatives of Q at the integration start time t0 . To obtain these derivatives , notice that we can rewrite ( 7 ) as 0 = ` ( t , xt , ut ) + dQ ( t , xt , ut ) dt , Q ( t1 , xt1 ) = Φ ( xt1 ) , ( 8 ) since the definition of Q implies that Q ( t , xt , ut ) = ` ( t , xt , ut ) dt+Q ( t+ dt , xt+dt , ut+dt ) . We now state our main result , which provides a local characterization of ( 8 ) with a set of coupled ODEs expanded along a solution path . These ODEs can be used to obtain all second-order derivatives at t0 . Theorem 1 ( Second-order Differential Programming ) . Consider a solution path ( x̄t , ūt ) that solves the ODEs in ( 6 ) . Then the first and second-order derivatives of Q ( t , xt , ut ) , expanded locally around this solution path , obey the following backward ODEs : −dQx̄ dt = ` x̄ + F T x̄Qx̄ , − dQū dt = ` ū + F T ūQx̄ , ( 9a ) −dQx̄x̄ dt = ` x̄x̄ + F T x̄Qx̄x̄ +Qx̄x̄Fx̄ , − dQx̄ū dt = ` x̄ū +Qx̄x̄Fū + F T x̄Qx̄ū , ( 9b ) −dQūū dt = ` ūū + F T ūQx̄ū +Qūx̄Fū , − dQūx̄ dt = ` ūx̄ + F T ūQx̄x̄ +Qūx̄Fx̄ , ( 9c ) where Fx̄ ( t ) ≡ ∂F∂xt | ( x̄t , ūt ) , Qx̄x̄ ( t ) ≡ ∂2Q ∂xt∂xt | ( x̄t , ūt ) , and etc . All terms in ( 9 ) are time-varying vector- valued or matrix-valued functions expanded at ( x̄t , ūt ) . The terminal condition is given by Qx̄ ( t1 ) = Φx̄ , Qx̄x̄ ( t1 ) = Φx̄x̄ , and Qū ( t1 ) = Qūū ( t1 ) = Qūx̄ ( t1 ) = Qx̄ū ( t1 ) = 0 . The proof ( see Appendix A.2 ) relies on rewriting ( 8 ) with differential states , δxt : = xt − x̄t , which view the deviation from x̄t as an optimizing variable ( hence the name “ Differential Programming ” ) . It can be shown that δxt follows a linear ODE expanded along the solution path . Theorem 1 has several important implications . First , the ODEs in ( 9a ) recover the original ASM computation ( 3,4 ) , as one can readily verify that Qx̄ ( t ) ≡ a ( t ) follows the same backward ODE in ( 4 ) and the solution of the second ODE in ( 9a ) , Qū ( t0 ) = − ∫ t0 t1 Fū TQx̄dt , gives the exact gradient in ( 3 ) . Meanwhile , solving the coupled matrix ODEs presented in ( 9b , 9c ) yields the desired second-order matrix , Qūū ( t0 ) ≡ Lθθ , for preconditioning the update . Finally , one can derive the dynamics of other higher-order tensors using the same Differential Programming methodology by simply expanding ( 8 ) beyond the second order . We leave some discussions in this regard in Appendix A.2 .

Summary: This paper proposes a method to derive second-order derivatives for Neural ODEs, which requires O(1) memory in theory. For practical feasibility, the authors propose to use a low-rank approximation of the (dynamic) Hessian matrix, together with Kronecker-based factorization to further reduce the memory cost. The authors validated their methods in image classification, time-series analysis and continuous generative modeling.