Heavy Ball Neural Ordinary Differential Equations

1 Introduction . Neural ordinary differential equations ( NODEs ) are a family of continuous-depth machine learning ( ML ) models whose forward and backward propagations rely on solving an ODE and its adjoint equation [ 4 ] . NODEs model the dynamics of hidden features h ( t ) ∈ RN using an ODE , which is parametrized by a neural network f ( h ( t ) , t , θ ) ∈ RN with learnable parameters θ , i.e. , dh ( t ) dt = f ( h ( t ) , t , θ ) . ( 1 ) Starting from the input h ( t0 ) , NODEs obtain the output h ( T ) by solving ( 1 ) for t0 ≤ t ≤ T with the initial value h ( t0 ) , using a black-box numerical ODE solver . The number of function evaluations ( NFEs ) that the black-box ODE solver requires in a single forward pass is an analogue for the continuous-depth models [ 4 ] to the depth of networks in ResNets [ 16 ] . The loss between NODE prediction h ( T ) and the ground truth is denoted by L ( h ( T ) ) ; we update parameters θ using the following gradient ∗Co-first author †Please correspond to : wangbaonj @ gmail.com 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . dL ( h ( T ) ) dθ = ∫ T t0 a ( t ) ∂f ( h ( t ) , t , θ ) ∂θ dt , ( 2 ) where a ( t ) : = ∂L/∂h ( t ) is the adjoint state , which satisfies the following adjoint equation da ( t ) dt = −a ( t ) ∂f ( h ( t ) , t , θ ) ∂h . ( 3 ) NODEs are flexible in learning from irregularlysampled sequential data and particularly suitable for learning complex dynamical systems [ 4 , 42 , 56 , 35 , 9 , 24 ] , which can be trained by efficient algorithms [ 40 , 7 , 58 ] . NODE-based continuous generative models have computational advantages over the classical normalizing flows [ 4 , 15 , 55 , 12 ] . NODEs have also been generalized to neural stochastic differential equations , stochastic processes , and graph NODEs [ 21 , 28 , 38 , 49 , 20 , 34 ] . The drawback of NODEs is also prominent . In many ML tasks , NODEs require very high NFEs in both training and inference , especially in high accuracy settings where a lower tolerance is needed . The NFEs increase rapidly with training ; high NFEs reduce computational speed and accuracy of NODEs and can lead to blow-ups in the worst-case scenario [ 15 , 10 , 29 , 35 ] . As an illustration , we train NODEs for CIFAR10 classification using the same model and experimental settings as in [ 10 ] , except using a tolerance of 10−5 ; Fig . 1 shows both forward and backward NFEs and the training time of different ODE-based models ; we see that NFEs and computational times increase very rapidly for NODE , ANODE [ 10 ] , and SONODE [ 35 ] . More results on the large NFE and degrading utility issues for different benchmark experiments are available in Sec . 5 . Another issue is that NODEs often fail to effectively learn long-term dependencies in sequential data [ 26 ] , as discussed in Sec . 4 . 1.1 Contribution . We propose heavy ball neural ODEs ( HBNODEs ) , leveraging the continuous limit of the classical momentum accelerated gradient descent , to improve NODE training and inference . At the core of HBNODE is replacing the first-order ODE ( 1 ) with a heavy ball ODE ( HBODE ) , i.e. , a second-order ODE with an appropriate damping term . HBNODEs have two theoretical properties that imply practical advantages over NODEs : • The adjoint equation used for training a HBNODE is also a HBNODE ( see Prop . 1 and Prop . 2 ) , accelerating both forward and backward propagation , thus significantly reducing both forward and backward NFEs . The reduction in NFE using HBNODE over existing benchmark ODE-based models becomes more aggressive as the error tolerance of the ODE solvers decreases . • The spectrum of the HBODE is well-structured ( see Prop . 4 ) , alleviating the vanishing gradient issue in back-propagation and enabling the model to effectively learn long-term dependencies from sequential data . To mitigate the potential blow-up problem in training HBNODEs , we further propose generalized HBNODEs ( GHBNODEs ) by integrating skip connections [ 17 ] and gating mechanisms [ 19 ] into the HBNODE . See Sec . 3 for details . 1.2 Organization . We organize the paper as follows : In Secs . 2 and 3 , we present our motivation , algorithm , and analysis of HBNODEs and GHBNODEs , respectively . We analyze the spectrum structure of the adjoint equation of HBNODEs/GHBNODEs in Sec . 4 , which indicates that HBNODEs/GHBNODEs can learn long-term dependency effectively . We test the performance of HBNODEs and GHBNODEs on benchmark point cloud separation , image classification , learning dynamics , and sequential modeling in Sec . 5 . We discuss more related work in Sec . 6 , followed by concluding remarks . Technical proofs and more experimental details are provided in the appendix . 2 Heavy Ball Neural Ordinary Differential Equations . 2.1 Heavy ball ordinary differential equation . Classical momentum method , a.k.a. , the heavy ball method , has achieved remarkable success in accelerating gradient descent [ 39 ] and has significantly improved the training of deep neural networks [ 46 ] . As the continuous limit of the classical momentum method , heavy ball ODE ( HBODE ) has been studied in various settings and has been used to analyze the acceleration phenomenon of the momentum methods . For the ease of reading and completeness , we derive the HBODE from the classical momentum method . Starting from initial points x0 and x1 , gradient descent with classical momentum searches a minimum of the function F ( x ) through the following iteration xk+1 = xk − s∇F ( xk ) + β ( xk − xk−1 ) , ( 4 ) where s > 0 is the step size and 0 ≤ β < 1 is the momentum hyperparameter . For any fixed step size s , let mk : = ( xk+1 − xk ) / √ s , and let β : = 1− γ √ s , where γ ≥ 0 is another hyperparameter . Then we can rewrite ( 4 ) as mk+1 = ( 1− γ √ s ) mk − √ s∇F ( xk ) ; xk+1 = xk + √ smk+1 . ( 5 ) Let s→ 0 in ( 5 ) ; we obtain the following system of first-order ODEs , dx ( t ) dt = m ( t ) ; dm ( t ) dt = −γm ( t ) −∇F ( x ( t ) ) . ( 6 ) This can be further rewritten as a second-order heavy ball ODE ( HBODE ) , which also models a damped oscillator , d2x ( t ) dt2 + γ dx ( t ) dt = −∇F ( x ( t ) ) . ( 7 ) In Appendix E.6 , we compare the dynamics of HBODE ( 7 ) and the following ODE limit of the gradient descent ( GD ) dx dt = −∇F ( x ) . ( 8 ) In particular , we solve the ODEs ( 7 ) and ( 8 ) with F ( x ) defined as a Rosenbrock [ 41 ] or Beale [ 14 ] function . The comparisons show that HBODE can accelerate the dynamics of the ODE for a gradient system , which motivates us to propose HBNODE to accelerate forward propagation of NODE . 2.2 Heavy ball neural ordinary differential equations . Similar to NODE , we parameterize −∇F in ( 7 ) using a neural network f ( h ( t ) , t , θ ) , resulting in the following HBNODE with initial position h ( t0 ) and momentum m ( t0 ) : = dh/dt ( t0 ) , d2h ( t ) dt2 + γ dh ( t ) dt = f ( h ( t ) , t , θ ) , ( 9 ) where γ ≥ 0 is the damping parameter , which can be set as a tunable or a learnable hyperparmater with positivity constraint . In the trainable case , we use γ = · sigmoid ( ω ) for a trainable ω ∈ R and a fixed tunable upper bound ( we set = 1 below ) . According to ( 6 ) , HBNODE ( 9 ) is equivalent to dh ( t ) dt = m ( t ) ; dm ( t ) dt = −γm ( t ) + f ( h ( t ) , t , θ ) . ( 10 ) Equation ( 9 ) ( or equivalently , the system ( 10 ) ) defines the forward ODE for the HBNODE , and we can use either the first-order ( Prop . 2 ) or the second-order ( Prop . 1 ) adjoint sensitivity method to update the parameter θ [ 35 ] . Proposition 1 ( Adjoint equation for HBNODE ) . The adjoint state a ( t ) : = ∂L/∂h ( t ) for the HBNODE ( 9 ) satisfies the following HBODE with the same damping parameter γ as that in ( 9 ) , d2a ( t ) dt2 − γ da ( t ) dt = a ( t ) ∂f ∂h ( h ( t ) , t , θ ) . ( 11 ) Remark 1 . Note that we solve the adjoint equation ( 11 ) from time t = T to t = t0 in the backward propagation . By letting τ = T − t and b ( τ ) = a ( T − τ ) , we can rewrite ( 11 ) as follows , d2b ( τ ) dτ2 + γ db ( τ ) dτ = b ( τ ) ∂f ∂h ( h ( T − τ ) , T − τ , θ ) . ( 12 ) Therefore , the adjoint of the HBNODE is also a HBNODE and they have the same damping parameter . We can also employ ( 10 ) and its adjoint for the forward and backward propagations , respectively . Proposition 2 ( Adjoint equations for the first-order HBNODE system ) . The adjoint states ah ( t ) : = ∂L/∂h ( t ) and am ( t ) : = ∂L/∂m ( t ) for the first-order HBNODE system ( 10 ) satisfy dah ( t ) dt = −am ( t ) ∂f ∂h ( h ( t ) , t , θ ) ; dam ( t ) dt = −ah ( t ) + γam ( t ) . ( 13 ) Remark 2 . Let ãm ( t ) = dam ( t ) /dt , then am ( t ) and ãm ( t ) satisfies the following first-order heavy ball ODE system dam ( t ) dt = ãm ( t ) ; dãm ( t ) dt = am ( t ) ∂f ∂h ( h ( t ) , t , θ ) + γãm ( t ) . ( 14 ) Note that we solve this system backward in time in back-propagation . Moreover , we have ah ( t ) = γam ( t ) − ãm ( t ) . Similar to [ 35 ] , we use the coupled first-order HBNODE system ( 10 ) and its adjoint first-order HBNODE system ( 13 ) for practical implementation , since the entangled representation permits faster computation [ 35 ] of the gradients of the coupled ODE systems .

Summary: The authors propose a new neural ordinary differential equation (NODE) model that is inspired by a second order ODE $\ddot{x} + a \dot{x} + b \nabla F(x)=0$. The discretization of this ODE leads to a two step procedure known as heavy ball method. Using this idea to model a NODE has the advantage that the corresponding adjoint equation takes a form that reduces the number of function evaluations during the forward and backward pass. Further, the authors show that the spectrum of the model is well-structured which in turn mitigates the vanishing gradient problem. In addition, they formulate a model that integrates skip connections and gating mechanisms. The empirical results demonstrate the favorable performance of the proposed models as compared to other NODEs.