Improving Neural Ordinary Differential Equations with Nesterov’s Accelerated Gradient Method
Abstract
We propose the Nesterov neural ordinary differential equations (NesterovNODEs), whose layers solve the second-order ordinary differential equations (ODEs) limit of 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 the number of function evaluations (NFEs) needed to solve the ODEs.
We also prove that the adjoint state of a GNesterovNODEs also satisfies a GNesterovNODEs, thus accelerating both forward and backward ODE solvers and allowing the model to be scaled up for large-scale tasks.
We empirically corroborate the advantage of
GNesterovNODEs on a wide range of practical applications, including point cloud separation, image classification, and sequence modeling.
Compared to NODEs,
GNesterovNODEs require a significantly smaller number of NFEs while achieving better 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 modeling techniques from dynamical systems can be employed to improve neural networks.
Along this research direction, the correspondence between residual networks, a popular class of neural networks with skip connections and the numerical solution of ordinary differential equations (ODEs) have been vastly studied in [15, 57, 47, 4].
The resulting Neural ODEs (NODEs) model when taking the the discretization step to zero have shown great promises in a wide range of applications including scientific discovery [24, 62], irregular time series modeling [46, 5], mean-field games [48], and generative modeling [13, 61].
NODEs model the dynamics of hidden state h ( t )
∈ R N in a neural
[Equation]
where function f captures the dynamics and is chosen to be a neural network with parameters θ that are learned from the data.
Starting from the input h ( t 0 ) at the initial time t 0 , NODEs compute the output h ( T ) at time T by solving the Initial Value Problem in Eq.
(1) for some time T ≥ t 0 .
NODEs are trained by optimizing the loss L ( h ( T )) between the prediction h ( T ) and the ground truth where the parameters θ are updated using the following gradient [42]
[Equation]
where a ( t ) :
=
∂ L /∂ h ( t ) is the continuous adjoint state, which satisfies the continuous adjoint equation
[Equation]
NODEs solve both the ODE Eq.
(1) in its forward pass and the ODEs (2) and (3) in its backward pass using black-box numerical ODE solvers.
The literature refers to this approach as either the continuous adjoint method or optimise-then-discretise [21].
The number of function evaluations (NFEs) that these solvers need in a single forward and backward pass is among the main factors that decide the computational efficiency of the model, i.e. how fast the model is.
Unfortunately, in many applications, NODEs require high NFEs in both training and inference, especially when the error tolerances of the solvers are set to small values for obtaining high accuracy.
Furthermore, the NFEs increase rapidly when training progresses.
High NFEs deteriorate the efficiency of NODEs, reduce the accuracy of the trained model, and results in instability during training, making it difficult to scale up the models to large-scale tasks [14, 8, 30, 37, 11].
1.1 Contribution
We propose the Nesterov Neural ODEs (NesterovNODEs) that leverage the continuous limit of the
Nesterov’s accelerated gradient (NAG) descent [52] and the convergence rate O (1 /k 2 ) of the NAG scheme to enhance NODE training and inference.
Our contributions are four-fold:
We formulate the NesterovNODE that solves Nesterov ODEs, i
e. second-order ODEs with a time-dependent damping term, instead of first-order ODEs (1).
To improve computational efficiency of the model, we convert these second-order ODEs into equivalent systems of first-order differential-algebraic equations that are solved in both forward and backward propagations 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 [16] and gating mechanisms [18] into NesterovNODEs.
In general, GNesterovNODEs form a wide class of neural differential equations which are represented by differential-algebraic systems and contain NesterovNODEs as a subclass.
We prove that the continuous adjoint equation used to compute the gradients for updating the parameters θ in a GNesterovNODE also follows a generalized Nesterov ODE.
Thus, the NFEs in both forward and backward passes of GNesterovNODEs are significantly reduced, especially when the solvers are used with small error tolerances.
We prove that the spectrum of the GNesterovNODE is well-structured.
This property of GNesterovNODEs helps alleviate the vanishing gradient issue during training and allows the model to capture long-term dependencies in the data.
We empirically demonstrate the advantages of the NesterovNODEs/GNesterovNODEs over the baseline 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 gradient descent [60] on a wide range of applications including point cloud separation, image classification, and kinetic simulation.
In all experiments, our proposed models achieve better accuracy and smaller
NFEs than the baselines.
1.2 Organization
We structure this paper as follows:
In Section 2, we review HBNODEs and NesterovODEs.
In Section 3, we present the algorithm and analysis of the NesterovNODEs and GNesterovNODEs.
We study the 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 and analyze our models with ablation studies.
We discuss related works in Section 7.
The paper ends with concluding remarks.
Proofs and additional experimental details are provided in the Appendix.
An Integration of Nesterov ODEs into NODEs
We first 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 among the methods of choice in optimization and machine learning for training complex systems.
Starting from initial point x 0 ∈ R d , GD iterates as x k = x k − 1 − s ∇ F ( x k ) with s > 0 being the step size in order to find a minimum of the function F ( x ) .
Let s → 0 , we obtain the following ODE limit of the GD
GD
HB
NAG
[Figure] [Figure]
[Equation]
Comparing Eq.
(1) and (4), we observe that a NODE solves the ODE limit of the GD where the gradient −∇ F ( x t ) is parameterized by a neural network f
( x ( t ) , t, θ ) .
Heavy ball neural ordinary differential equations
HBNODEs are proposed in [60].
This model takes advantage of the acceleration of heavy ball (HB) momentum [41] to reduce the NFEs needed for solving the ODEs and speed up NODEs.
In particular, HBNODEs replace the first-order ODE limit of GD by the following secondorder ODE limit of heavy ball momentum method:
[Equation]
Similar to NODEs, [60] parameterize −∇ F ( x ( t )) by a neural network f ( x ( t ) , t, θ ) and formulate
HBNODEs as follows
[Equation]
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 convergence and accelerates GD, both GD and HB share the same convergence rate of O (1 /k ) for convex smooth optimization.
A breakthrough due to Nesterov [33] replaces the constant momentum γ with ( k − 1) /
( k + 2) , a.k.a.
NAG momentum, improves the convergence rate to O (1 /k 2 ) , which is proved to be optimal for convex and smooth objective functions [33, 52].
We demonstrate the faster convergence of NAG in comparison with GD and HB on a quadratic optimization 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 [33] takes the following form: given initial points x 0 ∈ R N and y
= x 0 , the sequence { ( x k , y k ) } k is defined inductively as:
[Equation] [Equation]
The continuous limit of the Nesterov scheme is obtained by setting x k = h ( k √ s )
=
h ( t ) with t = k √ s and some smooth function h from R to R N .
According to [52], the function h satisfies the
Nesterov ODE
[Equation]
with the initial conditions h (0) =
h 0 , h ′ (0)
= 0 .
Remark 1 (Nesterov factor) .
The constant 3 in the coefficient of h ′ ( t ) in Eq.
(8) is originally from the approximation ( k − 1) /
( k + 2) = 1 − 3 /k
+ O (1 /k 2 ) .
This constant will be replaced by a constant r if the factor ( k − 1) / ( k + 2) in Eq.
(7) is replaced by ( k − 1) /
( k + r − 1) .
It is proved in [52] that the Nesterov ODE still holds the quadratic convergence rate when 3 is replaced by any number r > 3 .
Remark 2 (Numerical stability) .
If the Euler method is used, then to keep the numerical solution close 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 is the maximum value for which the step size can be chosen so that the numerical solution remains close to the exact solution, reflects the numerical stability of the ODEs.
It is proved in [52] that the maximum stable step size of the Nesterov ODE is much larger than that of the ODE (4) , thus showing the numerical stability advantage of the Nesterov ODE.
Generalize Nesterov ODEs to Differential-Algebraic Systems
One can parameterize ∇ 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
[Figure] [Figure]
[Equation]
This second-order NesterovNODE can be written in term of the first-order NesterovNODE as
[Equation]
However, because of the singularity created by the coefficient 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 ′ ( t ) + 3 t k ( t ) =
k ( t ) .
By dividing both sides of Eq.
(11) by k ( t ) , we obtain:
[Equation]
where
[Equation]
Let m ( t ) = x ′ ( 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 algebraic equation of the above differential-algebraic system.
It is often the case when training ODE-based models that some functions diverge or explode at a finite time.
This phenomenon is called blow-up, which we demonstrate in Fig.
In order to alleviate the blow-up problem, we introduce a generalized version of NesterovNODE, termed Generalized
NesterovNODE (GNesterovNODE).
For the NesterovNODE, the potential blow-up is due to the oscillation inherited from NAG scheme as can be seen in Fig.
The blow-up can also occur because of the singularity caused by the factor t − 32 e t 2 in the algebraic equation.
Following the techniques presented in [60], 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.
The activation function can be any activation function commonly used.
In our experiments, we use tanh and hardtanh .
In addition, the residual term ξ h ( t ) , which stands for a skip connection [16], is also added into the governing equation of m ( t ) , which benefits training and generalization of
GNesterovNODEs.
The GNesterovNODE differential-algebraic system is then given by
[Equation]
with the initial conditions h (0) =
h 0 , x (0) = x 0 , m (0) = m 0 .
It is noted that, from h (0) = hand h ′ (0) = 0 , we must have x 0 = h 0 lim t → 0 σ ( t − 32 e t 2 ) − 1 and m 0 = h 0 lim t → 0 ddt σ ( t − 32 e t 2 ) − 1 . Fig.shows that GNesterovNODE can indeed control the growth of h ( t ) effectively.
In general, GNesterovNODEs form a wide class of neural differential equations which are represented by differential-algebraic systems and contain NesterovNODEs as a subclass.
Eqs.
(9) and (15) define the forward ODE for the (G)NesterovNODE.
To efficiently update the parameters during the training 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 continuous adjoint state a ( t ) =
∂
L ∂ h ( t ) of the NesterovNODE given in Eq.
(9) satisfies the following NesterovNODE
[Equation]
Proposition 2 (Continuous adjoint equation for GNesterovNODE) .
The continuous adjoint state functions a h ( t ) :
= ∂ L ∂ h ( t ) , a x ( t ) :
=
∂ L ∂ x ( t ) and a m ( t ) :
=
∂ L ∂ m ( t ) of the GNesterovNODE system given in Eq.
(15) satisfy the following differential-algebraic adjoint system
[Equation]
) with the final value conditions a h ( T ) =
∂L∂ h T , a x ( T ) =
∂L∂ x T and a m ( T ) =
∂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 [40].
In the cases of NODEs and their hybrid ODE-RNN variants, the vanishing gradient issue may occur when the adjoint state a ( t ) :=
∂ L ∂ h ( t ) goes to 0 quickly as T − t increases.
In this section, we will 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 the algebraic relation:
[Equation]
While the gradients ∂ L ∂ x t and ∂ L ∂ m t satisfy the following proposition.
[Figure]
Proposition 3.
For every t ∈ (0 , T ) , there exist a unit length row vector v ∈ C N and an upper triangular matrix U ∈ 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 =
(cid:20) U large
P
U small (cid:21) 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, exp( U ) =
(cid:20) exp
( U large ) ˜ P exp( U small ) (cid:21) and ∥ v exp( U ) ∥ 2 ≥ ∥ v large exp( U large ) ∥ 2 .
Here, the vector v large is the first m columns of v , and m is the size of U large .
Since the real parts of elements in the diagonal of U large is no less than 12 ( t − T ) , exp( U large ) decays at a rate at most 12 ( t − T ) .
This results in the nonvanishing gradient of (cid:104) ∂ L ∂ x t ∂ L ∂ m t (cid:105) , hence so is (cid:104) ∂
L ∂ h t
[Equation]
To illustrate, we take the Walker2D kinematic simulation task [2] in consideration, which requires learning long-term dependency effectively [27].
We train ODE-RNN [46], GHBNODERNN [60], NesterovNODE-RNN and GNesterovNODE-RNN on this benchmark dataset (More details in Appendix D.4).
Fig.
3 plots (cid:13)(cid:13)(cid:13) ∂
L ∂ h t (cid:13)(cid:13)(cid:13) 2 for ODE-RNN, (cid:13)(cid:13)(cid:13)(cid:104) ∂
L ∂ h t ∂ L ∂ m t (cid:105)(cid:13)(cid:13)(cid:13) 2 for GHBNODERNN and (cid:13)(cid:13)(cid:13)(cid:104) ∂
L ∂ h t ∂ L ∂ x t ∂ L ∂ m t (cid:105)(cid:13)(cid:13)(cid:13) 2 for NesterovNODE-RNN and GNesterovNODE-RNN, showing that when the gap between the final 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 the vanishing gradient issue.
Remark 3.
The gradient exploding problem can be effectively resolved via gradient clipping, training loss regularization, etc.
Therefore, in practice the vanishing gradient problem is the major issue for training deep neural networks [40].
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)
[9], the Second Order NODE (SONODE) [37], HBNODE/GHBNODE [60] on a variety of benchmarks including point cloud separation, image classification, and kinetic simulation which involve different data modalities ranging from point cloud to images and time series.
ANODEs augments the space on which the ODE is solved while SONODEs and (G)HBNODEs solve a secondorder ODE.
We aim to show that: (i)
NesterovNODEs/
GNesterovNODEs require significantly fewer
NFEs while attaining similar or even better accuracy as the baselines; (ii) GNesterovNODEs avoid the blow-up of h ( t ) and thus improve over NesterovNODEs;
(iii) NesterovNODEs/GNesterovNODEs
[Table]
[Figure]
n o r w a r d N F E s
[Figure]
capture better long-term dependencies than the baselines and achieve better results in long-sequence modeling tasks.
For all the experiments, we use Adam [23] as the optimizer and Dormand-Prince 5(4)
[7] as the numerical ODE solver.
We choose the network architecture used to parameterize f ( h ( t ) , t, θ ) so that our proposed models and the baselines have similar numbers of parameters in our experiments as shown in Table 1.
Other training/model/dataset details are provided in Appendix D. All results are averaged over 5 runs with different seeds.
We conduct the experiments on a server with 6 NVIDIA 2080Ti GPUs with 11GB of GPU memory.
Our PyTorch code with documentation can be found at https://github.com/minhtannguyen/NesterovNODE.
5.1 Image classification
We validate the accuracy and efficiency advantage of NesterovNODE/GNesterovNODE for image classification on MNIST [6] and CIFAR10 [25] in comparison with other ODE-based baselines.
We follow the same training and model settings as in [60].
NFEs.
As shown in Fig.
4, our NesterovNODE and GNesterovNODE reduce the NFEs in both the forward 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 methods reduce the NFEs more significantly.
Compared to the other second-order methods i.e. SONODE,
HBNODE, and GHBNODE, GNesterovNODE achieve much better NFE reductions on both MNIST (see Fig 12 in the Appendix) and CIFAR10, indicating the improvement in efficiency and stability of our methods over the other second-order baseline models.
Such advancement is an essential step to scale 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 result from SONODE while being much more efficient than SONODE.
This advantage in terms of accuracy can be associated with the small number of NFEs needed by NesterovNODE and GNesterovNODE above, which reduces the model complexity and leads to better generalization.
Note that GNesterovNODE improves over NesterovNODE in efficiency and accuracy (on larger and more challenging benchmark like CIFAR10).
This justifies the effectiveness of our solution that introduces an additional bounded activation function σ and the residual term ξ h ( t ) to prevent the blow-up of h ( t ) as explained in Section 3. [Figure] [Figure]
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 first set consists of 40 points drawn from a circle with the radius
|| r || < 0 .
5 , while the second set comprises 80 points drawn from an annulus with the inner and outer radius of 0 .
85 and 1 , respectively, i.e.
85 < || r || < 1 . 0 .
Fig.
5 shows that our NesterovNODE and GNesterovNODE models are able to converge to 0 loss consistently while other methods have difficulty to reach 0 loss.
In addition,
NesterovNODE and especially the GNesterovNODE require significantly fewer NFEs in forward and backward passes compared to the baselines.
Thus, NesterovNODE and GNesterovNODE help improve both the training and the efficiency of the model.
We plot the evolution of the point cloud separation 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 good separation between the two classes of point clouds in these experiments while NODE and ANODE fail for this task.
5.3 Walker2D kinematic simulation
In this section, we investigate NesterovNODE/GNesterovNODE when applied on time-series data.
In particular, we use the ODE-RNN framework [46], with the recognition model being set to different
ODE-based models, to study Walker2D kinematic simulation task, which requires learning long-term dependency effectively [26].
As shown in Fig.
6, our NesterovNODE-RNN and GNesterovNODERNN not only reduce the NFEs both in the forward and the backward stages, but also achieve smaller test loss compared to the baseline models that 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-RNN for this task is much more time-consuming compared to our Nesterov-based methods.
5.4 Continuous Normalizing Flows for MNIST
We compare the GNesterovNODE with the baseline NODE and GHBNODE for use in variational inference with the continuous normalizing flow model trained on the MNIST dataset.
The continuous normalizing flow model we use is the FFJORD in [14].
We summarize our results in Figure 7.
Compared to the FFJORD-NODE and FFJORD-GHBNODE, the FFJORD-GNesterovNODE significantly reduces the NFEs in both forward and backward passes while improving the negative ELBO on the test set.
This experiment demonstrates that our method GNesterovNODE accelerates NODE-based models on both discriminative tasks and generative tasks.
Observed Properties of NesterovNODE in GNesterovNODE
In this section, we verify empirically that with the addition of the activation function σ and the skip connection, GNesterovNODEs still preserve important properties of NesterovNODEs as stated in Re-
[Figure]
[Figure]
mark 1 and Remark 2.
This hints that GNesterovNODEs have the same behaviors as NesterovNODEs while enjoying more stable training.
Effect of the Nesterov factor on NesterovNODE (Remark 1)
As stated in [52], 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 test loss eventually despite a smaller r outperforming in the beginning
[Figure]
but their experiments are based on
Nesterov ODE.
This section investigates whether this behavior extends to our generalized method
GNesterovNODE.
As shown in Fig.
8 and Table 2, a larger r also leads to a higher test accuracy, and for r = 3 and r = 5 , the model converges to its best test accuracy sooner than r = 10 .
Although the training loss of r = 10 is larger than r = 3 and r = 5 , we observe that all three values of r reach their best testing accuracy when the training loss is approximately in the range of [0 . 65 , 0 . 75] , which hints that decreasing the loss further leads to overfitting.
One more interesting observation is that a higher value of r increases the forward and backward NFEs.
Intuitively speaking, the term − rt h ′ ( t ) in the NesterovNODE, which is opposite to the product of the damping parameter r/t and the velocity h ′ ( t ) , represents the friction force of the model (see [52, Section 4]).
When r is larger, the friction force resists the movement of h ( t ) along the trajectory stronger, which slows down the convergence of training loss.
Stability of NesterovNODE (Remark 2)
[Table]
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 for the GNesterovNODE in comparison with both NODE and GHBNODE.
To illustrate this fact, we perform experiments on CIFAR10 using Euler solvers with large step sizes (0 . 1 , 0 . 2 , 0 . 5) .
Due to the instability of large step size, GHBNODE moves fast to the maximum accuracy points and then goes down, as shown in Fig.
9, while NODE achieves high train loss and low test accuracy.
Even when the step size is big, GNesterovNODE’s training is stable as the curve evolves gradually.
More interestingly, as shown in Fig. 10, GNesterovNODE solved with rk4 (Fourth-order Runge-Kutta with 3/8 rule) solver using a large step size outperforms the same model solved with the adaptive step size solvers like dopri5 (Dormand-Prince of order 5) solver.
This show the promise of using large step sizes in
GNesterovNODE to obtain models with better accuracy and efficiency.
Related Work
Increasing efficiency of training NODEs.
Several methods have been proposed to reduce the NFEs in NODEs and increase the model efficiency.
Among them are works that use weight decay [14] and [Figure] [Figure] other regularizers applied on the solver and the learned dynamics [11, 20, 12, 39].
Other works employ input augmentation [8], data-control [30] and depth-variance [30, 36] to reduce the NFEs needed to compute the ODE solution.
Another method replaces norms with seminorms in the backward continous adjoint stage to reduce backward NFEs [22].
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 [37] replaces the first-order ODE in Eq.
(1) by a second ODE which can be solved as a system of first-order ODEs.
HBNODE [60] also solves a second-order ODE but with an additional constant damping parameter, which corresponds to the 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 NAG momentum method.
These momentum-based systems have also been employed in designing deep neural networks as in [32, 28, 49, 35].
Learning long-term dependencies.
The ability of a model to learn long-term dependencies is highly needed to scale up the model to large-scale tasks that involve very long sequences.
Existing works try to alleviate exploding or vanishing gradient issues happened during the training of recurrent neural networks, including [1, 59, 19, 55, 31, 17, 34].
Recently, learning long-term dependencies with NODEs has been explored.
For example, [26] integrate a long-short term memory cell into
NODEs.
Among the hallmarks of NesterovNODEs is that our proposed models can directly capture long-term dependencies in long sequences.
Concluding Remarks
In this paper, we propose the NesterovNODE and its generalized version GNesterovNODE that solve 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-based models 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 least are 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 issue and can capture long-term dependencies in long sequences effectively.
It is worth mentioning that
NesterovNODEs
GNesterovNODEs do not introduce any inherently negative societal impact.
The high-resolution ODE in [51, 50] is an alternative way to take a continuous-time limit of the
NAG scheme and heavy-ball momentum method.
This high-resolution ODE introduces a gradient correction that allows the NAG scheme to achieve an inverse cubic rate for minimizing the squared gradient norm, which is better than the inverse square rate in the low-resolution ODE that our method uses.
A limitation of our paper is that we have not incorporated restart schemes [38, 45, 56] into the
NesterovNODE, and we leave this as future work.
Acknowledgements
This material is based on research sponsored by the AFOSR MURI FA9550-18-1-0502, the ONR grant N00014-20-1-2093, the MURI N00014-20-1-2787, and the NSF under Grant# 2030859 to the
Computing Research Association for the CIFellows Project (CIF2020-UCLA-38).
NH acknowledges support from the NSF IFML 2019844 and the NSF AI Institute for Foundations of Machine Learning.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect 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 Appendix 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental 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 they were 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., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you 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 using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
Supplementary materials for Improving Neural Ordinary Differential Equations with Nesterov’s Accelerated Gradient Method A Review of the adjoint equation and the gradient for the first-order NODEs
A NODE for a hidden feature z : R →
R N takes of the form
[Equation]
(18) where g ( z ( t ) , t, θ ) ∈ R N is a neural network with learnable parameters θ .
For a scalar loss function
L , the adjoint state a ( t ) :
=
∂ L ∂ z ( t ) satisfies the following differential equation
[Equation]
B The adjoint equations for the NesterovNODEs and GNesterovNODEs
The adjoint equation for the NesterovNODEs (Proposition 1) is an implication of [37, Proposition 3.1] for Nesterov differential equations.
In this section, we give the proofs for Proposition 2.
The proofs will be intrinsically based on Eq.
Proof of Proposition 2 - the adjoint equation for GNesterovNODE
The GNesterovNODE is formulated as the following differential-algebraic system:
[Equation]
(20) where k ( t ) =
t − 32 e t 2 .
The adjoints of this system are given by a h ( t ) :
= ∂ L ∂ h ( t ) , a x ( t ) :
=
∂ L ∂ x ( t ) and a m ( t ) :
=
∂ L ∂ m ( t ) .
From the first equation in the system, we have
[Equation]
To determine the dynamics of a x ( t ) and a m ( t ) , we rewrite the last two equations in system (20) as the first-order system
[Equation]
Set z ( t ) =
[ x ( t ) m ( t )]
T and
[Equation]
then the first-order NesterovNODE can be rewritten as
[Equation]
In this case, we have
[Equation]
Thus, by using Eq.
(19), we obtain the first-order differential system for the adjoints a x and a m as
[Equation]
Together with Eq.
(21), we obtain the differential-algebraic system for the adjoints of the GNesterovNODE in Proposition 2.
C Proof of Proposition 3 - the nonvanishing gradient for GNesterovNODEs
Following the lines in [60], for a NODE given by z ′ ( t ) =
g ( z ( t ) , t, θ ) , we have
[Equation]
For the GNesterovNODE given in Eq.
(15), the gradient ∂ L ∂ h t can be determined from ∂ L ∂ x t via the algebraic relation:
[Equation]
While the gradients ∂ L ∂ x t and ∂ L ∂ m t can be determined by using Eq.
(23) as
[Equation]
∂
[Equation]
(24) where
[Equation]
Let M =
QUQ ⊤ be a Schur decomposition of M where Q is an orthogonal matrix and U is an upper triangular matrix with eigenvalues of M in the diagonal.
Then from Eq.
(24), we have
[Equation]
∂ L
[Equation]
m
[Equation]
(cid:21)
Q exp
[Equation]
U
[Equation]
Set v =
[Equation]
(cid:104)
[Equation]
m
[Equation]
(cid:105) Q , then v is a unit length vector and
[Equation]
To finish the proof of Proposition 3, we need to prove that at least half of complex numbers in the diagonal of U have the real parts greater than or equal to t − T 2 .
We first claim that the eigenvalues of M can be paired in such a way that each pair has the sum t − T .
To prove the claim, we write M = (cid:20)
A
B ( t − T ) I (cid:21) , where
A = − t (cid:90)
T σ ′ ( m ( s ))
ds, and B = − t (cid:90)
T (cid:20) ∂σ ( f ) ∂ x − ξt − 3 2 e t 2 I (cid:21) ds.
Since the matrices ( t − T ) I and A commute, the characteristic polynomial of M can be determined as
[Equation]
which is the value of the characteristic polynomial of the matrix BA at λ ( λ − t + T ) .
Over the field of complex numbers, the characteristic polynomial of BA is splitted completely and it has the form
[Equation]
where λ BA,i are the eigenvalues of BA .
Then the characteristic polynomial of M has the form
[Equation]
BA
[Equation]
Thus the eigenvalues of M can be paired in such a way that each pair is the roots of the quadratic equation
[Equation]
Such pairs always have the sum t − T .
The claim is proved.
Since the diagonal of U are exactly the eigenvalues of M , the claim implies that at least half of complex numbers in the diagonal of U are greater than or equal to t − T 2 .
The proposition is then proved.
D Implementation details
[Table] [Table]
[Table]
To solve the NesterovNODE in Eq.
(9), we rewrite it in the differential-algebraic form given by Eq.
We solve Eq.
(15) as follows.
First, given h ( t 0 ) , we compute x ( t 0 ) .
Using the computed x ( t 0 ) , we solve the ODE given by the last two equations in Eq.
(15), x ′ ( t ) = m ( t ) and m ′ ( t ) =
− m ( t ) − f
( h ( t ) , t, θ ) by using an ODE solver.
To get h ( t ) for the second differential equation, we compute h ( t ) from x ( t ) .
Alternatively, the second differential equation can be expressed as m ′ ( t ) =
− m ( t ) − f ( x ( t ) , t, θ ) , which we can think of applying a composite function f ( g ( h ( t ) , t, θ ) to h ( t ) .
Finally, we compute h ( t n ) from x ( t n ) .
We list some experimental details that are common to all experiments before presenting more details for each experiment.
For ODE solvers in the experiments ( dopri5 , euler , rk4 , explicit adams ), we use the implementation provided by torchdiffeq 1 .
• fc n is a fully-connected layer with the output dimension of size n .
HTanh: HardTanh( − 5 , 5 ).
• LReLU: LeakyRELU( 0 . 3 ).
• tpad: Padding a set of features with the value of the time t .
Specifically, this is equivalent to augmenting the feature tensor with shape c
× x × y to with a time tensor with shape 1 × x × y filled with the value of t , creating a time-augmented feature tensor with shape ( c + 1) × x × y .
• We use a learnable γ with for HBNODE/GHBNODE for all tasks.
• The values of ξ used for GHBNODE and GNesterovNODE used in all tasks are in Table 5.
• Except for the experiments in Section 6 where we investigate the effect of the variation of the Nesterov factor r , we choose r = 3 for the remaining experiments.
• n ∗ =
n for all models except for SONODE where n ∗ = 2 n .
The bounded activation function σ is chosen to be either tanh or hardtanh in PyTorch.
• The tolerance (for adaptive solvers) and step sizes (for fixed step-size solvers) used in the experiments are in Table 6.
• Other hyperparameters are in Table 3 and Table 4.
The difference in the architecture between first-order NeuralODE methods (NODE, ANODE) and second-order NeuralODE methods (SONODE, HBNODE/GHBNODE, NesterovNODE/GNesterovNODE) happens because of the structural difference in their dynamics function.
In particular, second-order NeuralODE methods use extra states to model the momentum of the original hidden states.
To get roughly similar numbers of parameters for first-order NeuralODE methods and secondorder NeuralODE methods, there must be differences in their architecture.
In choosing the architecture for our method, we have made sure that there is no meaningful architecture difference in our methods compared to the baseline second-order Neural ODE methods (SONODE and HBNODE/GHBNODE).
D.1 Experimental details used in Section 3 Silverbox Initialization Test
The Silverbox dataset [58] is as follows:
Given the input voltage V 1( t ) , the models must predict the output voltage V 2( t ) and the experiment is evaluated over 64 time steps.
In this experiment, we use the dopri5 solver (implemented in torchdiffeq ) with a tolerance of 10 − 7 for adaptive step sizes in both forward and backward passes.
Similar to [60], we parameterize the dynamic f of all methods with a dense layer.
The network architecture is:
• ODE layer: input n ∗ +1 → fc n
D.2 Experimental details used in the Point Cloud benchmark in Section 5.
Following the setting in [60, 9], we use a 3-layer neural network to parameterize the function f
( h ( t ) , t, θ ) on the right hand side of the ODE-based models in our study.
We integrate the ODE from t 0 = 1 to T = 2 , and pass the output h ( T ) at time T to a dense classifier.
We also set the tolerance of the ODE solvers to be 10 − 7 so that the effect of numerical error is minimized.
Visualization of the point cloud evolution for a random run of each model is in Fig.
• Initial Velocity layer: input 2 → fc h → HTanh → fc h → HTanh → fc n • ODE layer: input n ∗ → fc h → ELU → fc h → ELU → fc n • Output layer: input n → fc 1 → Tanh
D.3 Experimental details for MNIST/CIFAR10 in Section 5.
In our experiment, we parameterize f ( h ( t ) , t, θ ) in the NODE-based layer using a 3-layer neural network.
The output of this NODE-based layer is then passed through a dense layer to perform the classification task.
Following the augmentation approach in ANODE [8], we add p additional channels to input image, augmenting the number of image channels from c to c + p where p is differently chosen for each method.
The values of p chosen for each model are specified in Table 3.
For SONODE, HBNODE, GHBNODE, NesterovNODE and GNesterovNODE, we also incorporate velocity or momentum of the same shape as the augment state.
We also present extra experiment results for MNIST in Fig.
The architecture for MNIST:
[Figure]
[Figure]
• Initial Velocity layer: input 1 × 28 × 28 → conv h, 1 → LReLU → conv h, 3 → LReLU → conv 2 n − 1 ,• ODE layer: input n ∗ × 28 × 28 → tpad → conv h, 1 → ReLU → tpad → conv h, 3 → ReLU → tpad → conv n,• Output layer: input n × 28 × 28 →
fc
The architecture for CIFAR10:
• Initial Velocity layer: input 3 × 28 × 28 → conv h, 1 → LReLU → conv h, 3 → LReLU → conv 2 n − 3 ,• ODE layer: input n ∗ × 32 × 32 → tpad → conv h, 1 → ReLU → tpad → conv h, 3 → ReLU → tpad → conv n,• Output layer:
input n × 32 × 32 → fc
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 [53].
Following the procedure in HBNODE [60], we randomly take out 10% of the data to make the time series irregularly-sampled.
Each input sequence consists of 64 timestamps, which are recurrently fed through a hybrid technique, with the output of the hybrid method being transferred to a single dense layer to form the output time series.
The goal is to generate an auto-regressive forecast with an output time series that is as close as the input sequence when shifted one time stamp to the right.
The
RNN and ODE are parameterized by a 3-layer network.
The network architecture is:
• ODE layer: input n ∗ → fc ∗ n → Tanh → fc n → Tanh → fc n • RNN layer: input dn + k → fc h 1 → Tanh → fc h 2 → Tanh → fc dn • Output layer: input n → fc
Experiments
[Table]
Silverbox Initialization
[Table] [Table]
Point Cloud separation (Section 5.2)
[Table]
MNIST/CIFAR10 (
[Table] [Table]
Walker2D (
[Table] [Table]
Stability of NesterovNODE (Section 6)
Euler, RK4,
Explicit Adams, dopri 10 − 5 for dopri 0.1, 0.2, 0.
(for details, refer to Fig.and Fig. 10)
D.5 Experimental details for Continous Normalizing Flow for VAE on the MNIST dataset in Section 5.
Our training and the baseline FFJORD-NODE model follow the setting in Section 4.3 in [14].
We use the dopri5 solver with a tolerance of 10 − 5 .
We use the identity function as σ , and ξ = 2 for
GHBNODE and GNesterovNODE.
Aggregated results about the experiment can be found in Table 7.
[Table]
[Table]
E Additional experiments E.1 Integration time
In the the differential-algebraic version of NesterovNODE
[Table]
we reparametrize from h to x in order to eliminate the time-dependent damping coefficient 3 t in the Nesterov ODE given by Eq. 8.
In
[Figure] [Figure]
particular, we set h ( t ) =
k ( t ) x ( t ) and then find k ( t ) such that Eq. 8 written in terms of x ( t ) matches the Heavy-Ball ODE with a constant damping coefficient given by Eq.
We find that k ( t ) =
t − 32 e tsatisfies this requirement.
Solving this Heavy-Ball ODE with respect to x ( t ) is easier and more stable than solving the Nesterov ODE with respect to h ( t ) .
To strengthen our proposed trick, we conduct two experiments (CIFAR 10 and Walker2d) using the original version of Nesterov in Eq.
10 with the starting integration time t 0 = 1 .
We have to change the starting integration to t 0 =
due to the singularity in t 0 = 0 .
The empirical results on CIFAR10 (Fig. 13) and Walker2d (Fig.
14) benchmarks show that the approach of changing starting integration time shows much worse accuracy and efficiency than using our reparametrization trick.
E.2 The Effect of Nesterov factor regarding the solver’s tolerance
We run another CIFAR10 expreriment using smaller tolerance ( 10 − 7 ) in order to study the effect of
Nesterov factor regarding the solver’s tolerance.
As shown in Fig.
15, when the tolerance are shrunk small enough, the forward NFEs of the smaller factor (3, 5 compared to 10) are still smaller, although the opposite is true for backward NFEs.
Moreover, the larger factor converges to a better solution (higher accuracy) despite converging to its best solution later than the smaller factor.
[Table]
E.3 GPU memory consumption
We present the maximum CUDA memory consumption of the models used in the CIFAR10 experiments in Table 8.
We extract the numbers using the function max memory allocated 2 in
PyTorch.
E.4 Wall-clock time advantage of NesterovNODEs/GNesterovNODEs
[Figure]
In this Figs. 16 and
[Figure]
we show the advantage of NesterovNODEs in wall-clock time.
NesterovNODEs/
GNesterovNODEs uses less time to achieve a comparable accuracy compared to other methods.
E.5 The Effect of using Seminorm for reducing backward NFEs on GNesterovNODE
We tested the effect of the seminorm method [22] on GNesterovNODE for the CIFAR10 classification task and demonstrated the results in Fig.
We observe that using the adjoint seminorm method slightly reduces the NFEs compared to using only GNesterovNODE while also considerably reducing the test accuracy.
This effect is expected because while the adjoint seminorm method helps reduce the NFEs, its semi-strong constraints might affect the final ODE solutions
[Figure]
E.6 Human Activity Dataset
We verify the advantage of the GNesterovNODE over the baseline NODE and GHBNODE for time series data on the Human Activity dataset [29].
This dataset consists of time series from five individuals doing various activities: walking, sitting, lying, etc.
The data contains 3D positions of tags attached to those individuals’ belt, chest and ankles (12 features in total).
After preprocessing, the dataset has 6554 sequences of 211 time points.
In this task, the model classifies each time point
[Figure]
[Figure] [Figure]
into one of seven types of activities (walking, sitting, etc.)
We use the ODE-RNN architecture described in Section 4.5 of [46] as the baseline model, with dopri5 adaptive solver and tolerance 10 − 7 .
For the GNesterovNODE, we use σ = tanh, and ξ is a learnable scalar.
Figure 19 and Tableshow that GNesteroveNODE-RNN achieves the best accuracy and the smallest NFEs in both forward and backward pass.
F Solving Neural Differential-Algebraic Equations
Our implementation uses an ODE to calculate the DAE.
An alternative approach is using DAE solvers to solve the DAE, which are implemented by the DifferentialEquations.jl library [44, 43].
[Table]
Although Dormand-Prince 5(4) [7] is a common choice for adaptive solver, is a more efficient