Heavy Ball Neural Ordinary Differential Equations

Abstract:
We propose heavy ball neural ordinary differential equations (HBNODEs), leveraging the continuous limit of the classical momentum accelerated gradient descent, to improve neural ODEs (NODEs) training and inference. HBNODEs have two properties that imply practical advantages over NODEs: (i) The adjoint state of an HBNODE also satisfies an HBNODE, accelerating both forward and backward ODE solvers, thus significantly reducing the number of function evaluations (NFEs) and improving the utility of the trained models. (ii) The spectrum of HBNODEs is well structured, enabling effective learning of long-term dependencies from complex sequential data. We verify the advantages of HBNODEs over NODEs on benchmark tasks, including image classification, learning complex dynamics, and sequential modeling. Our method requires remarkably fewer forward and backward NFEs, is more accurate, and learns long-term dependencies more effectively than the other ODE-based neural network models. Code is available at https://github.com/hedixia/HeavyBallNODE.

heading: 1 Introduction
text: 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.

heading: 1.1 Contribution
text: 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.

heading: 1.2 Organization
text: 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.

heading: 2 Heavy Ball Neural Ordinary Differential Equations

heading: 2.1 Heavy ball ordinary differential equation
text: 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.

heading: 2.2 Heavy ball neural ordinary differential equations
text: 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.

heading: 3 Generalized Heavy Ball Neural Ordinary Differential Equations
text: In this section, we propose a generalized version of HBNODE (GHBNODE), see (15), to mitigate the potential blow-up issue in training ODE-based models. In our experiments, we observe that h(t) of ANODEs [10], SONODEs [35], and HBNODEs (10) usually grows much faster than that of NODEs. The fast growth of h(t) can lead to finite-time blow up. As an illustration, we compare the performance of NODE, ANODE, SONODE, HBNODE, and GHBNODE on the Silverbox task as in [35]. The goal of the task is to learn the voltage of an electronic circuit that resembles a Duffing oscillator, where the input voltage V1(t) is used to predict the output V2(t). Similar to the setting in [35], we first augment ANODE by 1 dimension with 0-augmentation and augment SONODE, HBNODE, and GHBNODE
with a dense network. We use a simple dense layer to parameterize f for all five models, with an extra input term for V1(t)3. For both HBNODE and GHBNODE, we set the damping parameter γ to be sigmoid(−3). For GHBNODE (15) below, we set σ(·) to be the hardtanh function with bound [−5, 5] and ξ = ln(2). The detailed architecture can be found in Appendix E. As shown in Fig. 2, compared to the vanilla NODE, the `2 norm of h(t) grows much faster when a higher order NODE is used, which leads to blow-up during training. Similar issues arise in the time series experiments (see Sec. 5.4), where SONODE blows up during long term integration in time, and HBNODE suffers from the same issue with same initialization.
To alleviate the problem above, we propose the following generalized HBNODE dh(t)
dt = σ(m(t));
dm(t)
dt = −γm(t) + f(h(t), t, θ)− ξh(t), (15)
where σ(·) is a nonlinear activation, which is set as tanh in our experiments. The positive hyperparameters γ, ξ > 0 are tunable or learnable. In the trainable case, we let γ = · sigmoid(ω) as in HBNODE, and ξ = softplus(χ) to ensure that γ, ξ ≥ 0. Here, we integrate two main ideas into the design of GHBNODE: (i) We incorporate the gating mechanism used in LSTM [19] and GRU [6], which can suppress the aggregation of m(t); (ii) Following the idea of skip connection [17], we add the term ξh(t) into the governing equation of m(t), which benefits training and generalization of GHBNODEs. Fig. 2 shows that GHBNODE can indeed control the growth of h(t) effectively. Proposition 3 (Adjoint equations for GHBNODEs). The adjoint states ah(t) := ∂L/∂h(t), am(t) := ∂L/∂m(t) for the GHBNODE (15) satisfy the following first-order ODE system
∂ah(t)
∂t = −am(t) (∂f ∂h (h(t), t, θ)− ξI ) , ∂am(t) ∂t = −ah(t)σ′(m(t)) + γam(t). (16) 3Here, we exclude an h3 term that appeared in the original Duffing oscillator model because including it would result in finite-time explosion.
Though the adjoint state of the GHBNODE (16) does not satisfy the exact heavy ball ODE, based on our empirical study, it also significantly reduces the backward NFEs.

heading: 4 Learning long-term dependencies – Vanishing gradient
text: It is known that the vanishing and exploding gradients are two bottlenecks for training recurrent neural networks (RNNs) with long-term dependencies [2, 37] (see Appendix C for a brief review on the exploding and vanishing gradient issues in training RNNs). The exploding gradients issue can be effectively resolved via gradient clipping, training loss regularization, etc [37, 11]. Thus in practice the vanishing gradient is the major issue for learning long-term dependencies [37]. As the continuous analogue of RNN, NODEs as well as their hybrid ODE-RNN models, may also suffer from vanishing in the adjoint state a(t) := ∂L/∂h(t) [26]. When the vanishing gradient issue happens, a(t) goes to 0 quickly as T − t increases, then dL/dθ in (2) will be independent of these a(t). We have the following expressions for the adjoint states of the NODE and HBNODE (see Appendix C for details):
• For NODE, we have ∂L ∂ht = ∂L ∂hT ∂hT ∂ht = ∂L ∂hT exp { − ∫ t T ∂f ∂h (h(s), s, θ)ds } . (17) • For GHBNODE4, from (13) we can derive[ ∂L ∂ht ∂L ∂mt ] = [ ∂L ∂hT ∂L ∂mT ] [ ∂hT ∂ht ∂hT ∂mt
∂mT ∂ht ∂mT ∂mt
] = [ ∂L ∂hT ∂L ∂mT ] exp { − ∫ t T [ 0 ∂σ ∂m( ∂f ∂h − ξI ) −γI ] ds︸ ︷︷ ︸
:=M
} .
(18)
Note that the matrix exponential is directly related to its eigenvalues. By Schur decomposition, there exists an orthogonal matrix Q and an upper triangular matrix U , where the diagonal entries of U are eigenvalues of Q ordered by their real parts, such that −M = QUQ> =⇒ exp{−M} = Q exp{U}Q>. (19) Let v> := [ ∂L ∂hT ∂L ∂mT ] Q, then (18) can be rewritten as[
∂L ∂ht ∂L ∂mt ] = [ ∂L ∂hT ∂L ∂mT ] exp{−M} = [ ∂L ∂hT ∂L ∂mT ] Q exp{U}Q> = v> exp{U}Q>. (20)
By taking the `2 norm in (20) and dividing both sides by ∥∥[ ∂L ∂hT ∂L ∂mT ]∥∥ 2 , we arrive at∥∥[ ∂L
∂ht ∂L ∂mt ]∥∥ 2∥∥[ ∂L
∂hT ∂L ∂mT ]∥∥ 2 =
∥∥v> exp{U}Q>∥∥ 2
‖v>Q>‖2 =
∥∥v> exp{U}∥∥ 2
‖v‖2 = ∥∥e> exp{U}∥∥ 2 , (21)
i.e., ∥∥[ ∂L
∂ht ∂L ∂mt ]∥∥ 2 = ∥∥e> exp{U}∥∥ 2 ∥∥[ ∂L ∂hT ∂L ∂mT ]∥∥ 2
where e = v/‖v‖2. Proposition 4. The eigenvalues of −M can be paired so that the sum of each pair equals (t− T )γ.
For a given constant a > 0, we can group the upper triangular matrix exp{U} as follows exp{U} := [ exp{UL} P
0 exp{UV }
] , (22)
where the diagonal of UL (UV ) contains eigenvalues of −M that are no less (greater) than (t− T )a. Then, we have ‖e> exp{U}‖2 ≥ ‖e>L exp{UL}‖2 where the vector eL denotes the first m columns of e with m be the number of columns of UL. By choosing 0 ≤ γ ≤ 2a, for every pair of eigenvalues of −M there is at least one eigenvalue whose real part is no less than (t−T )a. Therefore, exp{UL} decays at a rate at most (t− T )a, and the dimension of UL is at least N ×N . We avoid exploding gradients by clipping the `2 norm of the adjoint states similar to that used for training RNNs.
In contrast, all eigenvalues of the matrix ∫ t T ∂f/∂hds in (17) for NODE can be very positive or negative, resulting in exploding or vanishing gradients. As an illustration, we consider the benchmark Walker2D kinematic simulation task that requires learning long-term dependencies effectively [26, 3]. We train ODE-RNN [42] and (G)HBNODE-RNN on this benchmark dataset, and the detailed experimental settings are provided in Sec. 5.4. Figure 4 plots ‖∂L/∂ht‖2 for ODE-RNN and ‖[∂L/∂ht ∂L/∂mt]‖2 for (G)HBNODE-RNN, showing that the adjoint state of ODE-RNN vanishes quickly, while that of (G)HBNODE-RNN does not vanish even when the gap between T and t is very large.
4HBNODE can be seen as a special GHBNODE with ξ = 0 and σ be the identity map.

heading: 5 Experimental Results
text: In this section, we compare the performance of the proposed HBNODE and GHBNODE with existing ODE-based models, including NODE [4], ANODE [10], and SONODE [35] on the benchmark point cloud separation, image classification, learning dynamical systems, and kinematic simulation. For all the experiments, we use Adam [25] as the benchmark optimization solver (the learning rate and batch size for each experiment are listed in Table 1) and Dormand–Prince-45 as the numerical ODE solver. For HBNODE and GHBNODE, we set γ = sigmoid(θ), where θ is a trainable weight initialized as θ = −3. The network architecture used to parameterize f(h(t), t, θ) for each experiment below are described in Appendix E. All experiments are conducted on a server with 2 NVIDIA Titan Xp GPUs.
Table 1: The batch size and learning rate for different datasets.
Dataset Point Cloud MNIST CIFAR10 Plane Vibration Walker2D
Batch Size 50 64 64 64 256 Learning Rate 0.01 0.001 0.001 0.0001 0.003
NO DE
AN OD E SO NO DE
HB NO DE GH BN OD E
Tol 1e-7: Cloud Separation
0 50 100 150 200 250 3000
100
200
300
Fo rw
ar d
NF Es
0 50 100 150 200 250 3000
100
200
300
Ba ck
wa rd
N FE
s
0 50 100 150 200 250 300 Epochs
0.0
0.5
1.0
1.5
Lo ss
NODE ANODE SONODE HBNODE GHBNODE
Figure 4: Comparison between NODE, ANODE, SONODE, HBNODE, and GHBNODE for two-dimensional point cloud separation. HBNODE and GHBNODE converge better and require less NFEs in both forward and backward propagation than the other benchmark models.

heading: 5.1 Point cloud separation
text: In this subsection, we consider the two-dimensional point cloud separation benchmark. A total of 120 points are sampled, in which 40 points are drawn uniformly from the circle ‖r‖ < 0.5, and 80 points are drawn uniformly from the annulus 0.85 < ‖r‖ < 1.0. This experiment aims to learn effective features to classify these two point clouds. Following [10], we use a three-layer neural network to parameterize the right-hand side of each ODE-based model, integrate the ODE-based model from t0 = 0 to T = 1, and pass the integration results to a dense layer to generate the classification results. We set the size of hidden layers so that the models have similar sizes, and the number of parameters of NODE, ANODE, SONODE, HBNODE, and GHBNODE are 525, 567, 528, 568,
and 568, respectively. To avoid the effects of numerical error of the black-box ODE solver we set tolerance of ODE solver to be 10−7. Figure 4 plots a randomly selected evolution of the point cloud separation for each model; we also compare the forward and backward NFEs and the training loss of these models (100 independent runs). HBNODE and GHBNODE improve training as the training loss consistently goes to zero over different runs, while ANODE and SONODE often get stuck at local minima, and NODE cannot separate the point cloud since it preserves the topology [10].

heading: 5.2 Image classification
text: We compare the performance of HBNODE and GHBNODE with the existing ODE-based models on MNIST and CIFAR10 classification tasks using the same setting as in [10]. We parameterize f(h(t), t, θ) using a 3-layer convolutional network for each ODE-based model, and the total number of parameters for each model is listed in Table 2. For a given input image of the size c× h× w, we first augment the number of channel from c to c+ p with the augmentation dimension p dependent on each method5. Moreover, for SONODE, HBNODE and GHBNODE, we further include velocity or momentum with the same shape as the augmented state.
NFEs. As shown in Figs. 1 and 5, the NFEs grow rapidly with training of the NODE, resulting in an increasingly complex model with reduced performance and the possibility of blow up. Input augmentation has been verified to effectively reduce the NFEs, as both ANODE and SONODE require fewer forward NFEs than NODE for the MNIST and CIFAR10 classification. However, input augmentation is less effective in controlling their backward NFEs. HBNODE and GHBNODE require much fewer NFEs than the existing benchmarks, especially for backward NFEs. In practice, reducing NFEs implies reducing both training and inference time, as shown in Figs. 1 and 5.
Accuracy. We also compare the accuracy of different ODE-based models for MNIST and CIFAR10 classification. As shown in Figs. 1 and 5, HBNODE and GHBNODE have slightly better classification accuracy than the other three models; this resonates with the fact that less NFEs lead to simpler models which generalize better [10, 35].
NFEs vs. tolerance. We further study the NFEs for different ODE-based models under different tolerances of the ODE solver using the same approach as in [4]. Figure 6 depicts the forward and backward NFEs for different models under different tolerances. We see that (i) both forward
5We set p = 0, 5, 4, 4, 5/0, 10, 9, 9, 9 on MNIST/CIFAR10 for NODE, ANODE, SONODE, HBNODE, and GHBNODE, respectively.
and backward NFEs grow quickly when tolerance is decreased, and HBNODE and GHBNODE require much fewer NFEs than other models; (ii) under different tolerances, the backward NFEs of NODE, ANODE, and SONODE are much larger than the forward NFEs, and the difference becomes larger when the tolerance decreases. In contrast, the forward and backward NFEs of HBNODE and GHBNODE scale almost linearly with each other. This reflects that the advantage in NFEs of (G)HBNODE over the benchmarks become more significant when a smaller tolerance is used.

heading: 5.3 Learning dynamical systems from irregularly-sampled time series
text: In this subsection, we learn dynamical systems from experimental measurements. In particular, we use the ODE-RNN framework [4, 42], with the recognition model being set to different ODE-based models, to study the vibration of an airplane dataset [33]. The dataset was acquired, from time 0 to 73627, by attaching a shaker underneath the right wing to provide input signals, and 5 attributes are recorded per time stamp; these attributes include voltage of input signal, force applied to aircraft, and acceleration at 3 different spots of the airplane. We randomly take out 10% of the data to make the time series irregularly-sampled. We use the first 50% of data as our train set, the next 25% as validation set, and the rest as test set. We divide each set into non-overlapping segments of consecutive 65 time stamps of the irregularly-sampled time series, with each input instance consisting of 64 time stamps of the irregularly-sampled time series, and we aim to forecast 8 consecutive time stamps starting from the last time stamp of the segment. The input is fed through the the hybrid methods in a recurrent fashion; by changing the time duration of the last step of the ODE integration, we can forecast the output in the different time stamps. The output of the hybrid method is passed to a single dense layer to generate the output time series. In our experiments, we compare different ODE-based models hybrid with RNNs. The ODE of each model is parametrized by a 3-layer network whereas the RNN is parametrized by a simple dense network; the total number of parameters for ODE-RNN, ANODE-RNN, SONODE-RNN, HBNODE-RNN, and GHBNODE-RNN with 16, 22, 14, 15, 15 augmented dimensions are 15,986, 16,730, 16,649, 16,127, and 16,127, respectively. To avoid potential error due to the ODE solver, we use a tolerance of 10−7.
In training those hybrid models, we regularize the models by penalizing the L2 distance between the RNN output and the values of the next time stamp. Due to the second-order natural of the underlying dynamics [35], ODE-RNN and ANODE-RNN learn the dynamics very poorly with much larger training and test losses than the other models even they take smaller NFEs. HBNODE-RNN and GHBNODE-RNN give better prediction than SONODE-RNN using less backward NFEs.

heading: 5.4 Walker2D kinematic simulation
text: In this subsection, we evaluate the performance of HBNODE-RNN and GHBNODE-RNN on the Walker2D kinematic simulation task, which requires learning long-term dependency effectively [26]. The dataset [3] consists of a dynamical system from kinematic simulation of a person walking from a pre-trained policy, aiming to learn the kinematic simulation of the MuJoCo physics engine [48]. The dataset is irregularly-sampled with 10% of the data removed from the simulation. Each input consists of 64 time stamps fed though the the hybrid methods in a recurrent fashion, and the output is passed to a single dense layer to generate the output time series. The goal is to provide an auto-regressive forecast so that the output time series is as close as the input sequence shifted one time stamp to the right. We compare ODE-RNN (with 7 augmentation), ANODE-RNN (with 7 ANODE style augmentation), HBNODE-RNN (with 7 augmentation), and GHBNODE-RNN (with 7 augmentation) 6 The RNN is parametrized by a 3-layer network whereas the ODE is parametrized by a simple dense
6Here, we do not compare with SONODE-RNN since SONODE has some initialization problem on this dataset; the ODE solver encounters failure due to exponential growth over time.
network. The number of parameters of the above four models are 8,729, 8,815, 8,899, and 8,899, respectively. In Fig. 8, we compare the performance of the above four models on the Walker2D benchmark; HBNODE-RNN and GHBNODE-RNN not only require significantly less NFEs in both training (forward and backward) and in testing than ODE-RNN and ANODE-RNN, but also have much smaller training and test losses.

heading: 6 Related Work
text: Reducing NFEs in training NODEs. Several techniques have been developed to reduce the NFEs for the forward solvers in NODEs, including weight decay [15], input augmentation [10], regularizing solvers and learning dynamics [12, 23, 13, 36], high-order ODE [35], data control [29], and depthvariance [29]. HBNODEs can reduce both forward and backward NFEs at the same time. Second-order ODE accelerated dynamics. It has been noticed in both optimization and sampling communities that second-order ODEs with an appropriate damping term, e.g., the classical momentum and Nesterov’s acceleration in discrete regime, can significantly accelerate the first-order gradient dynamics (gradient descent), e.g., [39, 32, 5, 45, 53]. Also, these second-order ODEs have been discretized via some interesting numerical schemes to design fast optimization schemes, e.g., [44]. Learning long-term dependencies. Learning long-term dependency is one of the most important goals for learning from sequential data. Most of the existing works focus on mitigating exploding or vanishing gradient issues in training RNNs, e.g., [1, 54, 22, 51, 30, 18, 47]. Attention-based models are proposed for learning on sequential data concurrently with the effective accommodation of learning long-term dependency [50, 8]. Recently, NODEs have been integrated with long-short term memory model [19] to learn long-term dependency for irregularly-sampled time series [26]. HBNODEs directly enhance learning long-term dependency from sequential data. Momentum in neural network design. As a line of orthogonal work, the momentum has also been studied in designing neural network architecture, e.g., [31, 47, 27, 43], which can also help accelerate training and learn long-term dependencies. These techniques can be considered as changing the neural network f in (1). We leave the synergistic integration of adding momentum to f with our work on changing the left-hand side of (1) as a future work.

heading: 7 Concluding Remarks
text: We proposed HBNODEs to reduce the NFEs in solving both forward and backward ODEs, which also improve generalization performance over the existing benchmark models. Moreover, HBNODEs alleviate vanishing gradients in training NODEs, making HBNODEs able to learn long-term dependency effectively from sequential data. In the optimization community, Nesterov acceleration [32] is also a famous algorithm for accelerating gradient descent, that achieves an optimal convergence rate for general convex optimization problems. The ODE counterpart of the Nesterov’s acceleration corresponds to (9) with γ being replaced by a time-dependent damping parameter, e.g., t/3 [45] or with restart [52]. The adjoint equation of the Nesterov’s ODE [45] is no longer a Nesterov’s ODE. We notice that directly using Nesterov’s ODE cannot improve the performance of the vanilla neural ODE. How to integrate Nesterov’s ODE with neural ODE is an interesting future direction. Another interesting direction is connecting HBNODE with symplectic ODE-net [57] through an appropriate change of variables.

heading: 8 Acknowledgement
text: This material is based on research sponsored by NSF grants DMS-1924935 and DMS-1952339, DOE grant DE-SC0021142, and ONR grant N00014-18-1-2527 and the MURI grant N00014-20-1-2787.

Review summaries:
This work proposes a reparameterization of a Neural ODE (NODE) as a Heavy-Ball Neural ODE (HBNODE). This is inspired by Heavy-Ball Momentum, which is the continuous limit of classical momentum. The authors make the following claims: The adjoint of an HBNODE is also an HBNODE. HBNODE reduces forward and backward NFE. This is even more apparent at lower tolerances. The spectrum of HBNODE is well-structured. This implies effective learning of long-range dependencies by avoiding vanishing gradients. HBNODEs are more accurate The authors verify these claims on image classification, complex dynamics, and sequential modelling.  

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.

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.