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: This is a very nice paper presenting a novel neural ODE with momentum. The idea behind the paper is solid, and the contribution is motivated by some known limitations of Neural ODEs. The authors not only provide us with compelling experimental results but also with a theoretical discussion and a lot of insights. I think the paper would benefit a lot the community.