Improving Neural Ordinary Differential Equations with Nesterov’s Accelerated Gradient Method
Anonymous Author(s)
AfﬁliationAddress email
Abstract
We propose the Nesterov neural ordinary differential equations (NesterovNODEs),whose layers solve the second-order ordinary differential equations (ODEs) limitof Nesterov’s accelerated gradient (NAG) method, and a generalization called
GNesterovNODEs.
Taking the advantage of the convergence rate O (1 /k 2 ) of the
NAG scheme, GNesterovNODEs speed up training and inference by reducing thenumber of function evaluations (NFEs) needed to solve the ODEs.
We also provethat the adjoint state of a GNesterovNODEs also satisﬁes a GNesterovNODEs,thus accelerating both forward and backward ODE solvers and allowing the modelto be scaled up for large-scale tasks.
We empirically corroborate the advantage of
GNesterovNODEs on a wide range of practical applications, including point cloudseparation, image classiﬁcation, and sequence modeling.
Compared to NODEs,
GNesterovNODEs require a signiﬁcantly smaller number of NFEs while achievingbetter accuracy across our experiments.
Introduction
Dynamical systems have been recently integrated into deep neural networks for modeling highdimensional data.
The advantage of this approach is that well-developed mathematical modelingtechniques from dynamical systems can be employed to improve neural networks.
Along this researchdirection, the correspondence between residual networks, a popular class of neural networks with skipconnections, and the numerical solution of ordinary differential equations (ODEs) have been vastlystudied in [Haber & Ruthotto, 2017; Weinan, 2017; Ruthotto & Haber, 2019; Chen et al., 2018].
Theresulting Neural ODEs (NODEs) model when taking the the discretization step to zero have showngreat promises in a wide range of applications including scientiﬁc discovery [K¨ohler et al., 2020;
Zhong et al., 2020], irregular time series modeling [Rubanova et al., 2019;
De Brouwer et al., 2019],mean-ﬁeld games [Ruthotto et al., 2020], and generative modeling [Grathwohl et al., 2019b; Yildizet al., 2019].
NODEs model the dynamics of hidden state h ( t ) 2 R N in a neural network by an ODE
[Equation]
where function f captures the dynamics and is chosen to be a neural network with parameters ✓ thatare learned from the data.
Starting from the input h ( t 0 ) at the initial time t 0 , NODEs compute theoutput h ( T ) at time T by solving the Initial Value Problem in Eq. (1) for some time T  t 0 .
NODEsare trained by optimizing the loss L ( h ( T )) between the prediction h ( T ) and the ground truth wherethe parameters ✓ are updated using the following gradient [Pontryagin, 1987]
[Equation]
where a ( t ) :
= @
L / @
h ( t ) is the continuous adjoint state, which satisﬁes the continuous adjointequation
[Equation]
NODEs solve both the ODE Eq. (1) in its forward pass and the ODEs (2) and (3) in its backward passusing black-box numerical ODE solvers.
The literature refer to this approach as either the continuousadjoint method or optimise-then-discretise
[Kidger, 2022]. The number of function evaluations(NFEs) that these solvers need in a single forward and backward pass is among the main factors thatdecide the computational efﬁciency of the model, i.e. how fast the model is.
Unfortunately, in manyapplications, NODEs require high NFEs in both training and inference, especially when the errortolerances of the solvers are set to small values for obtaining high accuracy.
Furthermore, the NFEsincrease rapidly when training progresses.
High NFEs deteriorate the efﬁciency of NODEs, reducethe accuracy of the trained model, and results in instability during training, making it difﬁcult to scaleup the models to large-scale tasks [Grathwohl et al., 2019a; Dupont et al., 2019b; Massaroli et al. ,2020; Norcliffe et al., 2020; Finlay et al., 2020].
1.1 Contribution
We propose the Nesterov Neural ODEs (NesterovNODEs) that leverage the continuous limit ofthe Nesterov’s accelerated gradient (NAG) descent and the convergence rate O (1 /k 2 ) of the NAGscheme to enhance NODE training and inference.
Our contributions are four-fold:
We formulate the NesterovNODE that solves Nesterov ODEs, i e. second-order ODEs witha time-dependent damping term, instead of ﬁrst-order ODEs (1).
To improve computationalefﬁciency of the model, we convert these second-order ODEs into equivalent systems ofﬁrst-order differential-algebraic equations that are solved in both forward and backwardpropagations of the NesterovNODE.2.
To eliminate the potential blow-up problem in training NesterovNODEs, we further develop the Generalized NesterovNODEs (GNesterovNODEs) by introducing skip connections [ Heet al., 2016] and gating mechanisms [Hochreiter & Schmidhuber, 1997] into NesterovNODEs.
In general, GNesterovNODEs form a wide class of neural differential equationswhich are represented by differential-algebraic systems and contain NesterovNODEs as asubclass.3.
We prove that the continuous adjoint equation used to compute the gradients for updatingthe parameters ✓ in a GNesterovNODE also follows a generalized Nesterov ODE.
Thus, the NFEs in both forward and backward passes of GNesterovNODEs are signiﬁcantly reduced,especially when the solvers are used with small error tolerances.4.
We prove that the spectrum of the GNesterovNODE is well-structured.
This property of
GNesterovNODEs helps alleviate the vanishing gradient issue during training and allowsthe model to capture long-term dependencies in the data.
We empirically demonstrate the advantages of the NesterovNODEs/GNesterovNODEs over thebaseline NODE and the state-of-the-art neural ODE models including the Heavy Ball NODEs(HBNODEs), which solve the continuous limit of the heavy ball momentum accelerated gradientdescent
[Xia et al., 2021] on a wide range of applications including point cloud separation, imageclassiﬁcation, and kinetic simulation.
In all experiments, our proposed models achieve better accuracyand smaller NFEs than the baselines.
1.2 Organization
We structure this paper as follows:
In Section 2, we review HBNODEs and NesterovODEs.
In Section3, we present the algorithm and analysis of the NesterovNODEs and GNesterovNODEs.
We studythe spectrum structure of the adjoint equations of NesterovNODEs/GNesterovNODEs to show that
NesterovNODEs/GNesterovNODEs can learn long-term dependencies effectively in Section 4.
In
Section 5 and 6, we empirically validate the advantages of NesterovNODEs/GNesterovNODEs andanalyze our models with ablation studies.
We discuss related works in Section 7.
The paper endswith concluding remarks.
Proofs and additional experimental details are provided in the Appendix.
An Integration of Nesterov ODEs into NODEs
We ﬁrst establish a connection between NODEs and gradient descent (GD), then review the Heavy
Ball Neural ODEs (HBNODEs), and motivate the integration of Nesterov ODEs into NODEs.
ODE limit of gradient descent and connections to NODEs Gradient descent (GD) has been amongthe methods of choice in optimization and machine learning for training complex systems.
Startingfrom initial point x 0 2 R d , GD iterates as x k = x k  1  s r F ( x k ) with s > 0 being the step size in order to ﬁnd a minimum of the function F ( x ) .
Let s ! 0 , we obtain the following ODE limit of the GD
GD
HB
NAG
[Figure]
[Equation]
Comparing Eq.
(1) and (4), we observe that a NODEsolves the ODE limit of the GD where the gradient r F ( x t ) is parameterized by a neural networkf ( x ( t ) , t, ✓ ) .
Heavy ball neural ordinary differential equations
HBNODEs are proposed in [Xia et al., 2021].
This modeltakes advantage of the acceleration of heavy ball (HB)momentum [Polyak, 1964] to reduce the NFEs neededfor solving the ODEs and speed up NODEs.
In particular, HBNODEs replace the ﬁrst-order ODE limit of GDby the following second-order ODE limit of heavy ballmomentum method:
[Equation]
Similar to NODEs, [Xia et al., 2021] parameterize r F ( x ( t )) by a neural network f ( x ( t ) , t, ✓ ) andformulate HBNODEs as follows
[Equation]
(6) where  > 0 is the damping parameter, which can be a hyperparameter or a learnable parameter.
Nesterov’s accelerated gradient (NAG) momentum Even though HB improves the convergenceand accelerates GD, both GD and HB share the same convergence rate of O (1 /k ) for convex smoothoptimization.
A breakthrough due to Nesterov [Nesterov, 1983] replaces the constant momentum with ( k  1) /
( k + 2) , a.k.a.
NAG momentum, improves the convergence rate to O (1 /k 2 ) , which isproved to be optimal for convex and smooth objective functions [Nesterov, 1983; Su et al., 2014].
We demonstrate the faster convergence of NAG in comparison with GD and HB on a quadraticoptimization problem in Figure 1.
The much faster convergence rate of NAG than that of GD and
HB motivates us to incorporate the second-order ODE limit of NAG into a NODE and propose the
NesterovNODE.
Nesterov acceleration gradient method [Nesterov, 1983] takes the following form:given initial points x 0 2 R N and y
= x 0 , the sequence { ( x k , y k ) } k is deﬁned inductively as:
[Equation]
The continuous limit of the Nesterov scheme is obtained by setting x k = h ( k p s )
= h ( t ) witht =
k p s and some smooth function h from R to R N .
According to [Su et al., 2014], the function hsatisﬁes the Nesterov ODE
[Equation]
with the initial conditions h (0) =
h 0 , h 0 (0) = 0 .
Remark 1 (Nesterov factor) .
The constant 3 in the coefﬁcient of h 0 ( t ) in Eq. (8) is originally from theapproximation ( k  1) /
( k + 2) =
1  3 /k
+ O (1 /k 2 ) .
This constant will be replaced by a constantr if the factor ( k  1) / ( k + 2) in Eq.
(7) is replaced by ( k  1) /
( k + r  1) .
It is proved in [Su et al.
,2014] that the Nesterov ODE still holds the quadratic convergence rate when 3 is replaced by anynumber r > 3 .
Remark 2 (Numerical stability) .
If the Euler method is used, then to keep the numerical solutionclose to the exact solution, the step size chosen in the Euler method must be small enough.
However,the smaller the step size, the more expensive the computation.
The maximum stable step size, which isthe maximum value for which the step size can be chosen so that the numerical solution remains closeto the exact solution, reﬂects the numerical stability of the ODEs.
It is proved in [Su et al., 2014]that the maximum stable step size of the Nesterov ODE is much larger than that of the ODE (4) , thusshowing the numerical stability advantage of the Nesterov ODE.
Generalize Nesterov ODEs to Differential-Algebraic Systems
One can parameterize r F ( h ( t )) in Eq.
(8) by a neural network f ( h ( t ) , t, ✓ ) with learnable parameters✓ in a similar way as NODE.
This results in the following Nesterov Neural ODE (NesterovNODE)
[Equation]
This second-order NesterovNODE can be written in termof the ﬁrst-order NesterovNODE as
[Equation]
However, because of the singularity created by the coefﬁcient 3 t , the training process based directly on Eq.
(9)and (10) will be unstable.
To avoid the instability issue,we set h ( t ) =
k ( t ) x ( t ) with k ( t ) =
t  32 e t 2 .
Then Eq.
(8)becomes
[Equation]
We observe that 2 k 0 ( t ) + 3 t k ( t ) =
k ( t ) .
By dividing both sides of Eq.
(11) by k ( t ) , we obtain:
[Equation] [Equation]
where
[Equation]
Let m ( t ) = x 0 ( t ) .
Then Eq. (8) is equivalent to the following differential-algebraic system
[Equation]
We parameterize f ( h ( t ) , t ) as a neural network, recalled as f ( h ( t ) , t, ✓ ) with learnable parameter ✓ ,and obtain the differential-algebraic version of NesterovNODE
[Equation]
The singularity at t = 0 of the NesterovNODE in Eqs.
(9) and (10) is now moved to the algebraicequation of the above differential-algebraic system.
It is often the case when training ODE-based models that some functions diverge or explode at aﬁnite time.
This phenomenon is called blow-up, which we demonstrate in Fig.
In order to alleviatethe blow-up problem, we introduce a generalized version of NesterovNODE, termed Generalized
NesterovNODE (GNesterovNODE).
For the NesterovNODE, the potential blow-up is due to theoscillation inherited from NAG scheme as can be seen in Fig.
The blow-up can also occur becauseof the singularity caused by the factor t  32 e t 2 in the algebraic equation.
Following the techniquespresented in [Xia et al., 2021], we address these potential blow-up by applying an activation function to the function f
( h ( t ) , t, ✓ ) , the factor t  32 e t 2 , and the momentum state m ( t ) of the NesterovNODE.
In addition, the residual term ⇠ h ( t ) , which stands for a skip connection [He et al., 2016], is also addedinto the governing equation of m ( t ) , which beneﬁts training and generalization of GNesterovNODEs.
The GNesterovNODE differential-algebraic system is then given by
[Equation]
[Equation]
with the initial conditions h (0) =
h 0 , x (0) = x 0 , m (0) = m 0 .
It is noted that, from h (0) =
h 0and h 0 (0)
{+= 0 , we must have x 0 = h 
+}
lim t
[Equation]
and
[Equation]
lim
[Equation]
Fig.
2shows that GNesterovNODE can indeed control the growth of h ( t ) effectively.
In general, GNesterovNODEs form a wide class of neural differential equations which are representedby differential-algebraic systems and contain NesterovNODEs as a subclass.
Eqs.
(9) and (15)deﬁne the forward ODE for the (G)NesterovNODE.
To efﬁciently update the parameters during thetraining process based on (G)NesterovNODEs, we use the continuous adjoint sensitivity given in Propositions 1 and 2 below.
Proposition 1 (Continuous adjoint equation for the second-order NesterovNODE) .
The continuousadjoint state a ( t ) = @
L @ h ( t ) of the NesterovNODE given in Eq.
(9) satisﬁes the following NesterovNODE
[Equation]
Proposition 2 (Continuous adjoint equation for GNesterovNODE) .
The continuous adjoint statefunctions a h ( t ) :
= @
L @ h ( t ) , a x ( t ) :
= @
L @ x ( t ) and a m ( t ) :
= @
L @ m ( t ) of the GNesterovNODE system givenin Eq.
(15) satisfy the following differential-algebraic adjoint system
[Equation]
(17) with the ﬁnal value conditions a h (0) = @
L @ h T , a x (0)
= @
L @ x T and a m (0) = @
L @ m T .  The Effectiveness of GNesterovNODE in Alleviating Vanishing Gradients
In neural networks with many layers, the vanishing gradient is one of the major issues [Pascanu et al.
In the cases of NODEs and their hybrid ODE-RNN variants, the vanishing gradient issue mayoccur when the adjoint state a ( t ) :
= @
L @ h ( t ) goes to 0 quickly as T  t increases.
In this section, wewill prove that this vanishing gradient issue can be avoided in GNesterovNODEs.
For the GNesterovNODE given in Eq.
(15), the gradient @
L @ h t can be determined from @ L @ x t via thealgebraic relation:
[Equation]
While the gradients @
L @ x t and @ L @ m t satisfy the following proposition.
Proposition 3.
For every t 2 (0 , T ) , there exist a unit length row vector v 2 C N and an uppertriangular matrix U 2 C N ⇥ N such that
[Equation]
and at least N 2 complex values in the diagonal of U have the real parts greater than or equal to t  T 2 .
The term v exp( U ) in the above proposition plays the essential role in the nonvanishing gradient issue.
Without loss of generality, we can assume that U = 
U large
P
U small  where all complex numbers in the diagonal of U large (resp.
U small ) have the real parts greater than or equal to (resp.
smaller than) 12 ( t  T ) .
According to Proposition 3, the size of U large is at least N/ 2 .
Then we have,
[Equation]
and
[Equation]
Here, the vector v large is the ﬁrst m columns of v , and m is the size of U large .
Since the real partsof elements in the diagonal of U large is no less than 12 ( t  T ) , exp( U large ) decays at a rate at most [Table] 12 ( t  T ) .
This results in the nonvanishing gradient of h @
L @ x t @
L @ m t i , hence so is h @
L @ h t @
L @ x t @
L @ m t i .
To illustrate, we take the Walker2D kinematic simulation task
[Brockman et al., 2016a] inconsideration, which requires learning long-term dependency effectively [Lechner & Hasani,2020b].
We train ODE-RNN [Rubanova et al., 2019], GHBNODE-RNN
[Xia et al., 2021],
NesterovNODE-RNN and GNesterovNODE-RNN on this benchmark dataset (More details in
Appendix D.4).
Fig.
3 plots  @ L @ h t  2 for ODE-RNN, h @
L @ h t @
L @ m t i 2 for GHBNODE-RNN andh @
L @ h t @
L @ x t @
L @ m t i 2 for NesterovNODE-RNN and GNesterovNODE-RNN, showing that whenthe gap between the ﬁnal time T and intermediate time t becomes larger, the adjoint states of
NesterovNODE-RNN, GNesterovNODE-RNN and GHBNODE-RNN decay much more slowly than
NODE-RNN.
Thus, NesterovNODE-RNN and GNesterovNODE-RNN have the ability to tackle thevanishing gradient issue. Remark 3. The gradient exploding problem can be effectively resolved via gradient clipping, trainingloss regularization, etc.
[Pascanu et al., 2013; Erichson et al., 2020].
Therefore, in practice thevanishing gradient problem is the major issue for training deep neural networks [Pascanu et al.
Experimental Results
In this section, we empirically study the advantages of our proposed NesterovNODE/GNesterovNODE over the baseline NODEs and other popular NODE-based architectures, including the augmented
NODE (ANODE)
[Dupont et al., 2019a], the Second Order NODE (SONODE)
[Norcliffe et al.
,2020], HBNODE/GHBNODE
[Xia et al., 2021] on a variety of benchmarks including point cloudseparation, image classiﬁcation, and kinetic simulation which involve different data modalities ranging from point cloud to images and time series.
ANODEs augments the space on which the ODEis solved while SONODEs and (G)HBNODEs solve a second-order ODE.
We aim to show that:
(i) NesterovNODEs/GNesterovNODEs
require signiﬁcantly fewer NFEs while attaining similar oreven better accuracy as the baselines; (ii) GNesterovNODEs avoid the blow-up of h ( t ) and thusimprove over NesterovNODEs;
(iii) NesterovNODEs/GNesterovNODEs capture better long-termdependencies than the baselines and achieve better results in long-sequence modeling tasks.
For all the experiments, we use Adam [Kingma & Ba, 2015] as the optimizer and Dormand-Prince 5(4)[Dormand & Prince, 1980] as the numerical ODE solver.
We choose the network architecture used toparameterize f ( h ( t ) , t, ✓ ) so that our proposed models and the baselines have similar numbers ofparameters in our experiments as shown in Table 1.
Other training/model/dataset details are providedin Appendix D. All results are averaged over 5 runs with different seeds.
We conduct the experimentson a server with 6 NVIDIA 2080Ti GPUs with 11GB of GPU memory.
5.1 Image classiﬁcation
We validate the accuracy and efﬁciency advantage of NesterovNODE/GNesterovNODE for imageclassiﬁcation on MNIST [Deng, 2012] and CIFAR
[Krizhevsky et al., 2009] in comparison withother ODE-based baselines.
We follow the same training and model settings as in [Xia et al., 2021].
NFEs.
As shown in Fig.
4, our NesterovNODE and GNesterovNODE reduce the NFEs in both theforward and the backward propagations compared to the baseline models.
Although the augmented input dimensions in ANODE help reduce the NFEs compared to NODE, second-order methodsreduce the NFEs more signiﬁcantly.
Compared to the other second-order methods i.e. SONODE,
HBNODE, and GHBNODE, GNesterovNODE achieve much better NFE reductions on both MNIST(see Fig 11 in the Appendix) and CIFAR10, indicating the improvement in efﬁciency and stability ofour methods over the other second-order baseline models.
Such advancement is an essential step toscale GNesterovNODE to larger and more complex practical tasks.
Accuracy.
Table 1 shows that our GNesterovNODE achieves the highest test accuracy on CIFAR10.
On MNIST, NesterovNODE attains the second-highest test accuracy and very close to the best resultfrom SONODE while being much more efﬁcient than SONODE.
This advantage in terms of accuracycan be associated with the small number of NFEs needed by NesterovNODE and GNesterovNODEabove, which reduces the model complexity and leads to better generalization.
Note that GNesterovNODE improves over NesterovNODE in efﬁciency and accuracy (on larger andmore challenging benchmark like CIFAR10).
This justiﬁes the effectiveness of our solution thatintroduces an additional bounded activation function  and the residual term ⇠ h ( t ) to prevent theblow-up of h ( t ) as explained in Section 3.
5.2 Point cloud separation
We perform experiments on a point cloud separation task in order to verify that NesterovNODEs/GNesterovNODEs can learn effective features to separate two sets of point clouds.
The ﬁrst set consistsof 40 points drawn from a circle with the radius
|| r || < 0 .
5 , while the second set comprises 244 80 points drawn from an annulus with the inner and outer radius of 0 .
85 and 1 , respectively, i.e.0 .
85 < || r || < 1 . 0 .
Fig.
5 shows that our NesterovNODE and GNesterovNODE models are able 246 to converge to 0 loss consistently while other methods have difﬁculty to reach 0 loss.
In addition,
NesterovNODE and especially the GNesterovNODE require signiﬁcantly fewer NFEs in forwardand backward passes compared to the baselines.
Thus, NesterovNODE and GNesterovNODE helpimprove both the training and the efﬁciency of the model.
We plot the evolution of the point cloudseparation through 100 epochs for a random run of each model in Appendix D.2 .
Like SONODE,
HBNODE/GHBNODE, NesterovNODE/GNesterovNODE learn effective features that allow goodseparation between the two classes of point clouds in these experiments while NODE and ANODEfail for this task.
5.3 Walker2D kinematic simulation
In this section, we investigate NesterovNODE/GNesterovNODE when applied on time-series data.
Inparticular, we use the ODE-RNN framework [Rubanova et al., 2019], with the recognition modelbeing set to different ODE-based models, to study Walker2D kinematic simulation task, whichrequires learning long-term dependency effectively [Lechner & Hasani, 2020a].
As shown in Fig.6, our NesterovNODE-RNN and GNesterovNODE-RNN not only reduce the NFEs both in theforward and the backward stages, but also achieve smaller test loss compared to the baseline modelsthat we compare with, including (G)HBNODE-RNN, ANODE-RNN, and NODE-RNN.
Although SONODE-RNN achieves much smaller losses, its NFEs are too high, thus training SONODE-RNNfor this task is much more time-consuming compared to our Nesterov-based methods.
Observed Properties of NesterovNODE in GNesterovNODE
In this section, we verify empirically that with the addition of the activation function  and the skipconnection, GNesterovNODEs still preserve important properties of NesterovNODEs as stated in Remark 1 and Remark 2.
This hints that GNesterovNODEs have the same behaviors as NesterovNODEswhile enjoying more stable training.
Effect of the Nesterov factor on NesterovNODE (Remark 1)
[Table]
As stated in [Su et al., 2014], the “magic constant” 3 can be replaced by a constant r > 3 while maintaining a convergence rate of
O (1 /k 2 ) .
In this work’s experiments, a larger r leads to a better testloss eventually despite a smaller r outperforming in the beginning,but their experiments are based on Nesterov ODE.
This section investigates whether this behavior extends to our generalized method
GNesterovNODE.
As shown in Fig.
7 and Table 2, a larger r alsoleads to a higher test accuracy, and for r = 3 and r = 5 , the modelconverges to its best test accuracy sooner than r = 10 .
Although the training loss of r = 10 is largerthan r = 3 and r = 5 , we observe that all three values of r reach their best testing accuracy whenthe training loss is approximately in the range of [0 . 65 , 0 .
75] , which hints that decreasing the lossfurther leads to overﬁtting.
One more interesting observation is that a higher value of r increases theforward and backward NFEs.
Intuitively speaking, the term  r t h 0 ( t ) in the NesterovNODE, which isopposite to the product of the damping parameter r/t and the velocity h 0 ( t ) , represents the frictionforce of the model (see [Su et al., 2014, Section 4]).
When r is larger, the friction force resists themovement of h ( t ) along the trajectory stronger, which slows down the convergence of training loss.
Stability of NesterovNODE (Remark 2)
The NesterovNODE is more numerically stable than NODE in the sense that the step size in the
Euler method for solving the NesterovNODE can be chosen larger while the stability of the numerical solution is still guaranteed (see Remark 2).
We hypothesize that this fact still holds forthe GNesterovNODE in comparison with both NODE and GHBNODE.
To illustrate this fact, weperform experiments on CIFAR10 using Euler solvers with large step sizes (0 . 1 , 0 . 2 , 0 . 5) .
Due tothe instability of large step size, GHBNODE moves fast to the maximum accuracy points and thengoes down, as shown in Fig.
8, while NODE achieves high train loss and low test accuracy.
Evenwhen the step size is big, GNesterovNODE’s training is stable as the curve evolves gradually.
Moreinterestingly, as shown in Fig.
9, GNesterovNODE solved with rk4 (Fourth-order Runge-Kutta with3/8 rule) solver using a large step size outperforms the same model solved with the adaptive step sizesolvers like dopri5 (Dormand-Prince of order 5) solver. This show the promise of using large stepsizes
[Figure]
Related Work
[Figure]
Increasing efﬁciency of training NODEs.
Several methods have been proposed to reduce the NFEsin NODEs and increase the model efﬁciency.
Among them are works that use weight decay [Grathwohl et al., 2019a] and other regularizers applied on the solver and the learned dynamics [Finlay et al.
,2020; Kelly et al., 2020; Ghosh et al., 2020;
Pal et al., 2021].
Other works employ input augmentation [Dupont et al., 2019b], data-control [Massaroli et al., 2020] and depth-variance [Massaroli et al.
,2020] to reduce the NFEs needed to compute the ODE solution.
Another method replaces norms withseminorms in the backward continous adjoint stage to reduce backward NFEs [Kidger et al., 2021].
NesterovNODE solves second-order Nesterov ODE to reduce both forward and backward NFEs.
Second-order dynamical systems.
Second-order ODEs have also been employed to speed up
NODEs.
SONODE [Norcliffe et al., 2020] replaces the ﬁrst-order ODE in Eq.
(1) by a second ODE which can be solved as a system of ﬁrst-order ODEs.
HBNODE [Xia et al., 2021] also solvesa second-order ODE
but with an additional constant damping parameter, which corresponds tothe ODE limit of the HB momentum method.
NesterovNODE solves the second-order Nesterov
ODE with a time-dependent damping parameter, which corresponds to the ODE limit of the NAGmomentum method.
These momentum-based systems have also been employed in designing deepneural networks as in [Moreau & Bruna, 2017;
Li et al., 2018;
Sander et al., 2021]
Learning long-term dependencies.
The ability of a model to learn long-term dependencies ishighly needed to scale up the model to large-scale tasks that involve very long sequences.
Existingworks try to alleviate exploding or vanishing gradient issues happened during the training of recurrentneural networks, including [Arjovsky et al., 2016;
Wisdom et al., 2016;
Jing et al., 2017; Vorontsovet al., 2017; Mhammedi et al., 2017; Helfrich et al., 2018].
Recently, learning long-term dependencieswith NODEs has been explored.
For example, [Lechner & Hasani, 2020a] integrate a long-short termmemory cell into NODEs.
Among the hallmarks of NesterovNODEs is that our proposed models candirectly capture long-term dependencies in long sequences.
Concluding Remarks
In this paper, we propose the NesterovNODE and its generalized version GNesterovNODE thatsolve the second-order ODE limit of NAG.
These models take advantage of the convergence rate O (1 /k 2 ) of the NAG scheme to gain acceleration over NODEs and the existing NODE-basedmodels such as HBNODE by reducing the NFEs in solving both forward and backward ODEs.
Our Nesterov-based NODEs also achieve better accuracy than NODEs and outperform or at leastare on par with other NODE-based models in our experiments while requiring much fewer NFEs.
We also prove that NesterovNODEs and GNesterovNODEs can avoid the vanishing gradient issueand can capture long-term dependencies in long sequences effectively.
It is worth mentioning that
NesterovNODEs/
GNesterovNODEs do not introduce any inherently negative societal impact.
Also,the high-resolution ODE in [Shi et al., 2019, 2021] is an alternative way to take a continuous-timelimit of the NAG scheme and heavy-ball momentum method.
This gradient correction allows the
NAG scheme to achieve an inverse cubic rate for minimizing the squared gradient norm, whichis much better than the inverse square rate in the low-resolution ODE our method use.
This is alimitation of our method and it is interesting to explore how to incorporate this construction into
NesterovNODEs.
Checklist
For
all authors
(a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See
Section
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes
If
are using existing assets (e
g., code, data, models) or curating/releasing new assets.
(a) If your work uses existing assets, did you cite the creators?
[N/A]
(b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you’re
[Table]
using/curating?
[N/A]
[Table]
(e) Did you discuss whether the data you are using/curating contains personally identiﬁable
[Table]
information or offensive content?
[N/A]
[Table]
If you used crowdsourcing or conducted research with human subjects...
[Table]
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
[Table]
(b) Did you describe any potential participant risks, with links to Institutional Review
[Table]
Board (IRB) approvals, if applicable?
[N/A]
[Table]
(c) Did you include the estimated hourly wage paid to participants and the total amount
[Table]
spent on participant compensation?
[N/A]
[Table]