Second-Order Neural ODE Optimizer
Authors: Guan-Horng Liu
Tianrong Chen
Evangelos Theodorou
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: Neural ODEs
second-order optimization
continuous-time optimal control

Abstract: We propose a novel second-order optimization framework for training the emerging deep continuous-time models, specifically the Neural Ordinary Differential Equations (Neural ODEs). Since their training already involves expensive gradient computation by solving a backward ODE, deriving efficient second-order methods becomes highly nontrivial. Nevertheless, inspired by the recent Optimal Control (OC) interpretation of training deep networks, we show that a specific continuous-time OC methodology, called Differential Programming, can be adopted to derive backward ODEs for higher-order derivatives at the same O(1) memory cost. We further explore a low-rank representation of the second-order derivatives and show that it leads to efficient preconditioned updates with the aid of Kronecker-based factorization. The resulting method – named SNOpt – converges much faster than first-order baselines in wall-clock time, and the improvement remains consistent across various applications, e.g. image classification, generative flow, and time-series prediction. Our framework also enables direct architecture optimization, such as the integration time of Neural ODEs, with second-order feedback policies, strengthening the OC perspective as a principled tool of analyzing optimization in deep learning. Our code is available at https://github.com/ghliu/snopt.

Second-Order Neural ODE Optimizer
Guan-Horng Liu, Tianrong Chen, Evangelos A. Theodorou
Georgia Institute of Technology, USA
{ghliu, cianrong.chen evangelos..heodoruuuee
theodorou] hetprrouu@@gateh..eee....
Abstract
We propose a novel second-order optimization framework for training the emerg- ing deep continuous-timg models, specifically the Neural Ordinary Differential Equations (Neural ODEs). Since their training already involves expensive gradient computation by solving backward ODE, deriving efficient second-order methods becomes highly nontrivial. Nevertheless, inspired by the recent Optimal Control (OC) interpretation of training deep networks, we show that a specific continuous- time OC methodology, called Differential Programming, can be adopted to derive backward ODEs for higher-orde, derivatives at the same (1)) memory cost. We further explore low-rank representation of the second-ord -oderr derivatives and show that it leads to efficient preconditioned updates with the aid of Kronecker- factorization. The resulting method named SNOpt converges much faster than first-order baselines in wall-clock time, and the improvement remains consistent across various applications, e.g. image classification, generative flow, and time- series prediction. Our framework also enables direct architecture optimization, such as the integration time of Neural ODEs, with second-order feedback policies, strengthening the OC perspective as principled tool of analyzing optimization deep learning. Our code is available at https: /tte::/gitttttet //ggtuuu.b.... com/ com/ghliug igoottt
Introduction
Image Classification (CIFAR10)
Time: Series Prediction (ArtWR)
Accuracy 1%6)
30r
80
Train
Adam
SGD
SNOpt
0.5k
1k
0.5k
1k
1.5k
1k
5.4k
1.kk
2.6k
5.4K
(ours)
Wall-Clock Time (sec)
Wall-Clock Time (sec)
Figure Our second-order method (SNOpt: green solid curves) achieves superior convergence
compared to first-order methods (SGD, Adam) on various Neural-OD ODE applications.
Neural ODEs (Chen al., 2018) have received tremendous attention over recent years. Inspired by taking the continuous limit the "discrete" residual transformation, X;-1 = X + cF(x;, 0), they propose directly parameterize the vector field of an ODE as deep neural network (DNN), Le.
dx(t) F(t,x(1), d/
x(lo) Ftg
(1)
where xtt R and F(:::,0 0) is DNN parameterized by R This provides powerful paradigm connecting modern machine learning to classical differential equations (Weinan, 2017) and has since then achieved promising results on time series analysis (Rubanova et al., 2019; Kidger et al., 2020b), reversible generative flow (Grathwohl et al., 2018: Nguyen et al., 2019), image classification (Zhuang et al., 2020, 2021). and manifold learning (Lou al., 2020; Mathieu & Nickel, 2020).
Due to the continuous time representation, Neural ODEs feature a distinct optimization process (see Fig. 2) compared their discrete-time counterparts, which also poses new challenges. First, the
35th Conference on Neural Information Processing Systems (NcurIPS 2021)

---

Forward ODE
Backward Adjoint ODE
x(fo)
ODESolve
x(11)
x(ta) ac (te( ec 30
ODESolve
(<f2)
Compulation llo
Parameter N-ODE
(l,x(())
(x/r
9C (x((1)
ODE function cal
Quary time dertvative
deffined in ())
G defined (5)
0
Forward vector field
Backwaard vector
Figure 2: Neural ODE features distinct training process: Both forward and backward passes parame- terize vector fields so that any generic ODE solver (which can be non-differentiable) can query time dx(()) to solve the ODEs (1, 5). In this work, we extend to second-order training.
derivatives, egg. dt
forward pass of Neural ODEs involves solving (1) with a black-box ODE solver. Depending on how its numerical integration is set up, the propagation may be refined to arbitrarily small step sizes and become prohibitively expensive to solve without any regularization (Ghosh et al., 2020; Finlay et al., 2020). On the other hand, to prevent Back-propese propaggttng through the entire ODE solver, the gradients are typically obtained by solving a backward adjoint ODE using the Adjoint Sensitivity Method (ASM; Pontryagin et al. (1962)). While this can be achieved at a favorable ((1) memory, it further increases the runtime and can suffer from inaccurate integration (Gholami et al., 2019). For these reasons, Neural ODEs often take notoriously longer time to train, limiting their applications to relatively small or synthetic datasets (Massaroli et al., 2020) until very recently (Zhuang et al., 2021).
To improve the convergence rate of training, it is natural to consider higher-order optimization. While efficient second-order methods have been proposed for discrete models (Ba et al., 2016; George et al., 2018). it remains unclear how to extend these successes to Neural ODEs, given their distinct computation processes. Indeed, limited discussions in this regard only note that one may repeat the backward adjoint process recursively to obtain higher-order derivatives (Chen et al., 2018) 8). This is, unfortunately, impractical as the recursion will accumulate the aforementioned integration errors and scale the per-iteration runtime linearly. As such, second-order methods for Neural ODEs are seldom considered practice, nor have they been rigorously explored from an optimization standpoint.
In this work, we show that efficient second-order optimization is in fact viable for Neural ODEs. Our method is inspired by the emerging Optimal Control perspective (Weinan al., 2018; Liu & Theodorou, 2019), which treats the parameter @ as a control variable, so that the training process. i.e. optimizing 0 w.r.t. some objective, can be interpreted as an Optimal Control Programming (OCP). Specifically, we show that a continuous- time OCP methodology. called Differential Programming, provides analytic second-order derivatives by solving a set of coupled matrix ODEs. Interestingly, these matrix ODEs can be augmented to the backward adjoint ODE and solved simultaneously. In other words, single backward pass is sufficient to compute all derivatives, including the original ASM-based gradient, the newly-derived second-order matrices, or even higher-order tensors. Further, these higher-order computations enjoy the same ((1) memory and comparable runtime to first-order methods by adopting Kronecker factorization (Martens & Grosse, 2015). The resulting method- called SNOpt admits superior convergence wall-clock time (Fig. 1), and the improvement remains consistent across image classification, continuous normalizing flow, and time-series prediction.
Our OCP framework also facilitates progressive training of the network architecture. Specifically, we study an example of jointly optimizing the "integration time" of Neural ODEs, in analogy to the "'epth" of discrete DNNs. While analytic gradients w.r.t. this architectural parameter have been derived under the ASM framework, they were often evaluated on limited synthetic datasets (Massaroli et al., 2020). In the context of OCP, however, free- horizon optimization a well-studied problem for practical applications with priori unknown terminal time (Sun et al., 2015; De Marchi & Gerdts, 2019). In this work, we show that these principles can be applied to Neural ODEs, yielding a novel second-order feedbuck policy that adapts the integration time throughout training. On training CIFAR1O, this further leads to a 20% runtime reduction, yet without hindering test-time accuracy.
In summary, we present the following contributions.
We propose novel computational framework for computing higher-orde derivatives of deep continuous models, with a rigorous analysis using continuous -imeuuuuus--meeeeeee Optimal Control theory.
We propose an efficient second-order method, SNOpt, that achieves superior convergence (in wall-clock time) over first-order methods in training Neural ODEs, while retaining constant memory complexity. These improvements remain consistent across various applications.
To show that our framework also enables direct architecture optimization, we derive a second- order feedback policy for adapting the integration horizon and show it further reduces the runtime.

---

2 Preliminaries
Notation. We roman and italic type represent variable x(t) and its realization ætt) given an ODE. ODESolve denotes a function call that solves an initial value problem given an initial condition, start and end integration time, and vector field, i.e. ODESolve (x(to), to, F1 where dx(()
dt = F.
Forward and backward computations of Neural ODEs. Given an initial condition x(to) and integration interval [to, 4.], Neural ODEs concern the following optimization over an objective L,
mjn ((x(t))), where x(1)) = x(to) + I"" F(,,x((), 0) dt
is the solution of the ODE (1) and can be solved by calling black-box box ODE solver, iee x((1) ODESolve (x(to), t,, F). The use of ODESolve allows to adopt higher-order numerical meth- ods, e.g. adaptive Runge-Kutta (Press et al., 2007), which give more accurate integration compared with e.g. vanilla Euler discretization in saaduall-bsed discrete models. To obtain the gradient 50 Neural ODE, one may naively Back-propagate through ODESolve. This, even if could be made possible, leads to unsatisfactory memory complexity since the computation graph can grow arbitrarily large for adaptive ODE solvers. Instead, Chen et (2018) proposed to apply the Adjoint Sensitivity Method (ASM), which states that the gradient can be obtained through the following integration.
f,,
ăc of
00 =
(t,.FF((((())))) dt, a(1) (()))(F((t,((t))
00
where a() E R. is referred to the adjoint state whose dynamics ohey a backward adjoint ODE, OC
da(t)
a(t) 0F (t)),
a(t2) 0x(t1)
(4)
dt
0x(t)
Equations (3, 4) present two coupled ODEs that can be viewed as the continuous-time expression of the Back-propagation (LeCun al., 1988). Algorithmically, they can be solved through another call of ODESolve (see Fig. 2) with an augmented dynamics G, i.e.
x (44)
x(41)
[x(1)
(t)) 8)
a(fo) = ODESolve (
a(t))
to, G). where G
t,
. = - at) T a(t)
8x
(5)
0C/00]
af)
augments the original dynamics F in (1) with the adjoint ODEs (3, 4). Notice that this computation (5) depends only on (x(4),, a(4)). This differs from naive Back-propagation, which requires storing intermediate states along the entire computation graph of forward ODESolve. While the latter requires O(T) memory cost,' the computation in (5) only consumes constant O(1) memory cost.
Chen et al. (2018) noted that if further encapsulate (5) by C C = grad 0), 0), one may compute higher- order derivatives by recursively calling a" 80" grad ( @" - 4,0), 9), starting from n=1. This can scale unfavorably due to its recursive dependence and accumu- lated integration errors. Indeed, Table 1
Table I: Numerical errors between ground-truth and ad- joint derivatives using different ODESolve on CIFARIO.
rk4
implicit adams
dopri5
0c
7.63 2.11x10-3
3.44x10-4
00 8- ag
6.84x10
2.50x10-1
41.10
suggests that the errors of second-order derivatives,
d'c obtained from the recursive adjoint procedure dc In the next
787 be 2-6 orders of magnitude larger than the ones from the first-order adjoint, 20 section, we will present a novel optimization framework that computes these higher-order derivatives without any recursion (Section 3. and discuss how can be implemented efficiently (Section 3.2).
3 Approach
3.1 Dynamics Higher-order Derivatives using Continuous-tinn Optimal Control Theory
OCP perspective is recently emerging methodology for analyzing optimization of discrete DNNs. Central to interpretation is to treat the layer propagation of a DNN as discrete-time dynamics, so
T' is the number of the adaptive steps used solve (1), as analogy of the "depth" of Neural ODEs.
3

---

that the training process, i.e. finding an optimal parameter of a DNN, can be understood like an OCP, which searches for an optimal control subjected to a dynamical constraint. This perspective has provided useful insights on characterizing the optimization process (Hu et al., 2019) and enhancing principled algorithmic design (Liu et al., 2021a). We leave a complete discussion in Appendix A.I.
Lifting this OCP perspective from discrete DNNs to Neural ODEs requires special treatments from continuous-time OCP theory (Todorov, 2016). Nevertheless, we highlight that training Neural ODEs and solving continuous-time OCP are fundamentally intertwined since these models, by construction, represent continuous- time dynamical systems. Indeed, the ASM used for deriving (3, 4) originates from the celebrated Pontryagin's 's principle (Pontryagin et al., 1962), which is an optimality condition to OCP. Hence, OCP analysis not only motivated but principled from an optimization standpoint.
We begin by first transforming (2) to a form that is easier to adopt the continuous-ting OCP analysis.
f(t,x,, x1, de)dt subjected to min ((xx))+ 1.¹ E(t,
dx; di dul 0,
F F(t, X Wt), Xto to
(6)
Uto
dr
where x(t) = xs. and etc. It should be clear that (6) describes (2) without loss of generality by having (#,4) == (c,0) 0). These functions are known as the terminal and intermediate costs in standard OCP. In training Neural ODEs, l can be used to describe either the weight decay, i.e. e xx ||ur]], or more complex regularization (Finlay et al., 2020). The time- invariant ODE imposed for u, makes the ODE of x, equivalent to (1), Problem (6) shall be understood as particular type of OCP that searches for an optimal initial condition 6 of a time- invariant control u, Despite seemly superfluous, this is a necessary transformation that enables rigorous OCP analysis for the original training process (2), and has also appeared in other reltrel-related analyses (Zhong et al., 2020; Chalvidal al., 2021).
Next, define the accumulated loss from any time i E [to, (1] to the integration end time / as
Qtt, X u.) == ((x,)) + J.² ((T,x;, dr,
which is also known in OCP as the cost-to-go function. Recall that our goal is to compute higher- order derivatives w.r.t. the parameter 0 of Neural ODEs. Under the new OCP representation (6), the
first-order derivative #C is identical to xty This is because Qt Xta+ u.,)) accumulates Buso
all sources of losses between [ta f1] (hence it sufficiently describes L) and Uso = 0 by construction. &fC = Les.
Likewise, the second-order derivatives can be captured by the Hessian 0'Q(to:x...")))
au,, ou,, In other words, are only interested in obtaining the derivatives of Q at the integration start time to- To obtain these derivatives, notice that we can rewrite (7) as
3006
0 = E(t,, Xx, ut) + dQ(I, Xx, Ut) Q(1,,xn) = =
dt
since the definition of Q implies that Q(t, Xx, u.) = E(t, Xx, u,)dt + Q(t + dt, Xe+di: Uctde). We now state our main result, which provides local characterization of (8) with set of coupled ODEs expanded along a solution path. These ODEs can be used to obtain all second-order derivatives at to
Theorem 1 (Second-order Differential Programming). Consider solution path (Ø., ūe) that solves the ODEs in (6). Then the first and second-order derivatives of Q x1, uz). expanded locally around this solution path, obey the following backward ODEs:
dQ= = la + F Qa:
dQu = la + F.
(9a)
dt
dt
dQua las + F Qax + QarFa
dQaa lea + QuxFa + F Q
(9b)
dt
dl
dQun dt.
dQuz lus + F.. + Qua
laa + Qau Qua dt
(9c)
8Q (t)- (a and etc. All terms in (9) are time-varying vector- 8x (a)) (t) ox,0x, where (x OF valued matrix-valued functions expanded at (æx, ūe). The terminal condition is given by
Qx( 2z. (t) Daa, and Qa(ti) Qua (t1) = (t) Qau 0.
4

---

The proof (see Appendix A.2) relies on rewriting (8) with differential states, 8x; := x; Et. which view the deviation from ã, as an optimizing variable (hence the name "Differential Programming"), It can be shown that 8x, follows a linear ODE expanded along the solution path. Theorem 1 has several important implications. First, the ODEs in (9a) recover the original ASM computation (3,4), as one can readily verify that (t) at)) follows the same backward ODE in (4) and the solution of the second ODE in (9a), Qu(4)) - - Fe Qell, gives the exact gradient in (3). Meanwhile, solving the coupled matrix ODEs presented in (9b, 9c) yields the desired second-order matrix, Qua (to) Lee. for preconditioning the update. Finally, one can derive the dynamics of other higher-order tensors using the same Differential Programming methodology by simply expanding (8) beyond the second order. We leave some discussions in this regard in Appendix A.2.
3.2 Efficient Second-order Preconditioned Update
Theorem provides attractive computational framework that does not require recursive computa- tion (as mentioned in Section 2) to obtain higher-order derivatives. It suggests that we can obtain first and second-order derivatives all once with single function call of ODESolve:
[ Qu (4)), (40), (0w)(1)) Qu (4v), Qwa (4)), Qua (40)
(10)
= -DDESolve ( (w Le 0, Daz 0, 0, 0J. f1, G),
where G augments the original dynamics F in (I) with all 6 ODEs presented in (9). Despite that this OCP-theereete backward pass (10) retains the same (1)) memory complexity as in (5), the dimension of the new augmented state, which unfavorable size that dramatically slows down the numerical integration. Hence, we must consider other representations of (9), if any, in order to proceed. In the following proposition, we present one of which that transforms (9) into a set of vector ODEs, so that we can compute them much efficiently.
carries second-order matrices, can grow to
Proposition 2 (Low-rank representation of (9)). Suppose f:=0 (6) and let Qa (11)= Σ be a symmetric matrix of rank R S where y; E Ih and @ is the Kronecker product. Then, for all t E [to ti]. the second-ords matrices appeared in (9b, 9c) can be decomposed into
Y: S y:
R
R Σ (1)
Qua = @ ((t),
Qza(t)=
& pt(), Qua(t)= Σ Pr(t) @ (.()).
where the vectors (t) Rn and pi(t)= R² obey the following backward ODEs:
(1) = (a(t) (t),, dp.(() (t( ((()))
(11)
dt
dt
with the terminal condition given by (q((tt)..(ttt) = O).
The proof is left Appendix A.2. Proposition gives nontrivial conversion. indicates that the coupled matrix ODEs presented in (9b, 9c) can be disentangled into a set of independent vector ODEs where each of them follows its own dynamics (11). As the rank determines the number of these vector ODEs, this conversion will be particularly useful if the second-order matrices exhibit low-rank structures. Fortunately, this is indeed the case for many Neural-ODE applications which often propagate x, in latent space of higher dimension (Chen et al., 2018; Grathwohl et al., 2018; Kidger et al., 2020b).
Based on Proposition 2, the recond-order precondition matrix Lee is given by²
Lag = (ua = Σ n 1m1
Fa La dt & F
q dt
(12)
where q: = (t) follows (11). Our final step is to facilitate efficient computation of (12) with Kronecker-based factorization, which underlines many popular second-orde methods for discrete DNNs (Grosse Martens, 2016; Martens al., 2018). Recall that the vector field F is represented
2 We drop the dependence on for brevity, yet all terms inside the integrations of (12, 13) are time-varying varying.
5

---

Algorithm SNOpt: Second-order Neural ODE Optimizer
1: Input: dataset D, parametrized vector field F(.: " 0), integration time (lo, 41), black-box ODE solver ODESolve, learning rate v,, rank R. interval the time grid A/
2: repeat
3: Solve x(t1) = ODESolve ( to. f1, F), where x(to) ~ D.
D Forward pass
4: Initialize (An, B,) := (0 for each layer and set q: (:(1)) y
fo 5: for t' in {t,, t Åt, ... .to +
At, to) do
6: Set t == At the small integration step, then call (x(t), Qa (t), Qu(t), (t)
= ODESolve (x((t), (t(t), Qa(t), (tt
Backward pass
where G augments the ODEs of state (I), first and second-order derivatives (9a, 11).
7: Evaluate (t), (t), Ft, x,, f), then compute (t). (t) (13).
8: Update A. A, A. (t) At and B, + B, + (t) At.
9:
end for
Vn, apply 0" 0 1 ved 1 A,
D> Second-order parameter update
11: until converges
x(1o) (a)1)
ODESolve
x(t)) (1)) (1)
ODE solution path
Query time derivalives
Qa((o)
sampled matrices
(y) )=u
Sampled grid
A. E, (m, (t))
ODE solver function call
B. Σ, (t))
defined with (1.9a.11)
Backwvard vector derivotives
Figure 4: Our second-order method, SNOpt, solves a new backward ODE, i.e. the G appeared in line of Alg. 1, which augments second-order derivatives, while simultaneously collecting the matrices (t,)) and (t,t)) on sampled time grid {tý} for computing the preconditioned update in (14).
by DNN. Let (t), (t),,, and (f) denote the acti- vation vector, pre-activation vector, and the parameter of layer when evaluating dx at time (see Fig. 3), then the integration in (12) can broken down into each layer n,
FG 0
(x((t) dt
(Lx(())) (()
(*) f(x" (5).
(8)
t0
12--1 (4) ((h
Fa
T 9:
dt =(... S1 (FT q:) .t....]
Figure 3: The layer propagation inside the vector field F, where and denote
or
q:) dt, ..., affine and nonlinear activation functions.
thn
where the second equality holds by F.. q: = BF th" Tq. = z" @ BF T step towards the Kronecker approximation of the layer-wise precondition matrix:
Th q:). This an essential
Th
(S/
R
(
to
Lange = Qunan (lo) = Σ
2
OF Sh
(3)) dt
2" 39
Sh q:) di
oF T
1=1
to (z" 2
di
to
Σ
Th
BF T
q
@
HF Uh
qi)
dl.
(13)
A.(t)
B.(())
We discuss the approximation behind (13), and also the one for (14), in Appendix A.2. Note that
A, (t) and B,, (t) are much smaller matrices in I mm compared to the ones in (9), and they can be efficiently computed with automatic differentiation packages (Paszke et al., 2017). 7). Now, let {t,) be time grid uniformly distributed over [fo, t1] that A, E, A. At, (t,)) At and B,= E, B, (t (t,) approximate the integrations in (13), then our final preconditioned update law is given by
Vn, Con vec (B. ()))
(14)
where vec denotes vectorization. Our second-order method - named SNOpt - is summarized in Alg. 1, with the backward computation (i.e. line 4-9 in Alg. 1) illustrated in Fig. 4. In practice, we also adopt eigen-based amortization with Tikhonov regularization (George al. (2018); see Alg. in Appendix A.4), which stabilizes the updates over stochastic training.

---

Remark. The fact that Proposition 2 holds only for degenerate & can be easily circumvented in practice. As typically represents weight decay, :== I ||0||2- which is iimeinnepenneet, it can be separated from the backward ODEs (9) and added after solving the backward integration, i.e.
(to) c + + Qu())), Qua(to) < -- + Qua (fo),
where y is the regularization factor. Finally, we find that using the scaled Muussian-Newtonn matrix, L.e. Qua(()) (41) ~ ē 2 4m generally provides good trade-off between the performance and
runtime complexity. As such, we adopt this approximation to Proposition 2 for all experiments.
3.3 Memory Complexity Analysis
Table 2:
Memory complexity at different stages of our derivation in terms of x2 E IR",, 0 E RR" and the rank R. Note that all methods have (11) in terms of depth.
Theorem 1 Eqs. (9,10)
Proposition 2
SNOpt (Alg. 1)
first-order adjoint Eqs. (3,4)
Eqs. (11,12)
Eqs. (13,14)
backward storage
((mm + n)
O Rm + Rn)
O(Rm + 2m)
O(m + n) O(n)
parameter update
Onn²)
O(n²)
((2m)
Table 2 summarizes the memory complexity of different computational methods that appeared along our derivation in Section 3.1 and 3.2. Despite that all methods retain (1)) memory as with the first-order adjoint method, their complexity differs in terms of the state and parameter dimension. Starting from encouraging result in Theorem which allows one to compute all derivatives with a single backward pass, we first exploit their low-rank representation in Proposition 2. This reduces the storage to O Rm + Rn) and paves way toward adopting Kronecker factorization, which further facilitates efficient preconditioning. With all these, our SNOpt is capable of performing efficient second-order updates while enjoying similar memory complexity (up to some constant) compared to first-order adjoint methods. Lastly, for image applications where Neural ODEs often consist of convolution layers, we adopt convolution-bial based Kronecker factorization (Grosse & Martens, 2016; Gao et al., 2020), which effectively makes the complexity to scale w.r.t. the number of feature maps (i.e. number of channels) rather than the full size of feature maps.
3.4 Extension to Architecture Optimization
Let us discuss an intriguing extension of our OCP framework to optimizing the architecture of Neural ODEs, specifically the integration bound f,. In practice, when problems contain no prior information on the integration, [fu, (1] is typically set to some trivial values (usually [0, 1)) without further justification. However, these values can greatly affect both the performance and runtime. Take CIFARIO for instance (see Fig. 5), the required training time decreases linearly as we drop & from yet the accuracy retains mostly the same unless & becomes too small. Similar results also appear MNIST (see Fig. 12 in Appendix A.5). In other words, we may interpret the integration bound t as an architectural parameter that needs to be jointly optimized during training.
Relative train time (%)
Accuracy (%)
100
50
40
0.0
1.0
0.0
1.0
6
Figure 5: Training performance of CIFAR 10 with Adam when using different 41, which moti- vates joint optimization of t Experiment setup is left in Ap- pendix A.4.
The aforementioned interpretation fits naturally into our OCP framework. Specifically, we can consider the following extension of Q, which introduces the terminal time T as new variable:
Q(t,x, x4, u, T) = ((I, + J." XT, u,) dr,
(15)
where (T, x(T)) explicitly imposes the penalty for longer integration time, e.g. := (x( - ST² Following a similar procedure presented in Section 3.1, we can transform (15) into its ODE form (as in (8) then characterize its local behavior (as in (9)) along a solution path (æ., ū. T). After some tedious derivations, which are left in Appendix A.3, we will arrive at the update rule below,
T イ- T 7 ST(50), where 8T(50) (²rr(1))) (Q+(1o) + @ra((o)60).
(16)
7

---

as the first-order Similar to what we have discussed in Section 3.1, one shall view @t((a) derivative w.r.t. the terminal time T. Likewise, (to) = o² вTØт. and etc. Equation (16) is a second- order order feedback policy that adjusts its updates based on the change of the parameter 6. Intuitively. it moves in the descending direction of the preconditioned gradient (i.e. Q7) while accounting for the fact that 0 is also progressing during training (via the feedback 00) The latter is distinct feature arising from the OCP principle. As we will show later, this update (16) leads to distinct behavior with superior convergence compared to first-order baselines (Massaroli et al., 2020).
0c ōT
4 Experiments
lime- seri-seres observation
Table 3: Sample size of time-series datasets (input dimension, class label, series length)
Neural ODE
GRU cell
Uiner mapping
SpoAD
AntWR
CharT
prediction
(27, 10, 93) (19, 25, 144) (7, 20, 187)
Figure 6: Hybrid model for time-series prediction.
Dataset. We select 9 datasets from distinct applications where N-ODEs have been applied, including image classification (+), time-series prediction (o), and continuous normalizing flow (o; CNF):
MNIST, SVHN, CIFAR10: MNIST consists of 28x28 gray-scale images, while SVHN and CIFAR1O consist of 3 x2 colour images. All 3 image datasets have 10 label classes
SpoAD, ArtWR, CharT: We consider UEA time series archive (Bagnall et al., 2018). Spoke- nArabicDigits (SpoAD) is speech dataset, whereas Articulary WordRecognition (ArtWR) and Character (CharT) are motion-related datasets. Table details their sample sizes.
Circle, Gas, Miniboone: Circle is 2-dim synthetic dataset adopted from Chen al. (2018). Gas and Miniboone are 8 and 43-dim tabular datasets commonly used in CNF (Grathwohl et al., 2018; Onken et al., 2020). All 3 datasets transform multivariate Gaussian to the target distributions.
Models. The models for image datasets and CNF resemble standard feedforward networks, except now consisting of Neural ODEs as continuous transformation layers. Specifically, the models for image classification consist of convolution linear mapping. Meanwhile, the CNF models are identical to the ones in Grathwohl et al. (2018), which consist of 1-5 Neural ODEs, depending on the size of the dataset. As for the time-series models, we adopt the hybrid models from Rubanova et (2019), which consist of Neural ODE for hidden state propagation, standard recurrent cell (e.g. GRU (Cho et al., 2014)) to incorporate incoming time-series observation, and a linear prediction layer. Figure illustrates this process. We detail other configurations in Appendix A.4.
feature extraction, followed by Neural ODE and
ODE solver. We use standard Runge-Kutee Kutta 4(5) adaptive solver (dopri5; Dormand Prince (1980) implemented by the torchdiffes package. The numerical tolerance is set to le-6 for CNF and 1e-3 for the rest. We fix the integration time [0, 1] whenever it appears as a hyper- papper-prramterreeeee (e.g. for image and CNF datasets'); ); otherwise we adopt the problem-specific setup (e.g. for time series).
Training setup. We consider Adam and SGD (with momentum) the first-order baselines since they are default training methods for most Neural- -oural-DDE applications. As for our second-order SNOpt, we set up the time grid (t,} such that it collects roughly 100 samples along the backward integration to estimate the precondition matrices (see Fig. 4). The hyper-parameters (e.g. learning rate) are tuned for each method on each dataset, and we detail the tuning process in Appendix A.4. We also employ practical acceleration techniques, including the semi-norm (Kidger et al., 2020a) for speeding up ODESolVe, and the Jacobian free estimator (FFJORD; Grathwohl al. (2018)) for accelerating CNF models. The batch size is set to 256, 512, and 1000 respectively for ArtWord, CharTraj. and Gas. The rest of the datasets use 128 as the batch size. All experiments are conducted on a TITAN RTX.
4.1 Results
Convergence and computation efficiency. Figures 1 and 7 report the training curves of each method measured by wall-clock time. is obvious that our SNOpt admits a superior convergence
3 except for Circle where we set [to, [1]:=(0, 10] in order match the original setup in Chen et al. (2018).
8

---

SpoAD and Accuracy
SVHN Lossard Accuracy
Train Loss
1.09
Accuracy (%%
Train Loss
Accuracy
0.%
Chart Training Loss and Accuracy
2.3k
0.7
1
232
13c
Miniboone
Accuracy (56)
Train Loss
MI
NLL 13
SHOpt
1Ik
20k
20k
2k Wall Clock Time (5ec)
12c
2.4c
3tk
Call-Clock Time [sec)
Wall-Clock Time (sec)
Wall-Clock Time (sec)
Figure 7: Training performance in wall-clock runtime, averaged over trials. Our SNOpt achieves
faster convergence against first-order baselines. See Fig. 14 in Appendix A.5 for MNIST and Circle.
Table 4: Test-time performance: accuracies for image and time-series datasets; NLL for CNF datasets
MNIST
SVHN
CIFARIO
SpoAD
ArtWR
CharT
Circle
Gas
Miniboone
Adam
98.83
91.92
77.41
94.64
84.14
93.29
0.90
-6.42
13.10
SGD
98.68
93.34
76.42
97.70
85.82
95.93
0.94
-4.58
13.75
SNOpt
98.99
95.77
79.11
97.41
90.23
96.63
0.86
-7.55
12.50
CIFARIO
ArtWR
Computation Efficiency of SNOpt
Rel, Momory Rel. Runtime
2.5
X
X
Accuracy (%)
1.0
Adam SGD
Mn
Sv
Cf
SA AW
CT
CI
Ga
Mi
SNOpt
1.5
(ours)
60
10-3
10
10
10-2
10
10a
1.0
Train Loss
Train Loss
Learing Rabe (Adan)
Learning Rate 1S0D)
Leamiing Rate (SNOJO
Mn Sv C
SA AW
CT
C
Ga M
Figure 8: Relative runtime and mem- ory of our SNOpt compared to Adam (denoted by the dashed black lines) all 9 datasets, where 'Mn' is the shorthand for MNIST, and etc.
Figure 9: Sensitivity analysis where each sample repre- sents a training result using different optimizer and learn- ing rate (annotated by different symbol and color). Our SNOpt achieves higher accuracies and insensitive to hyper- parameter changes. Note that x-axes are in log scale.
rate compared to the first-order baselines, and in many cases exceeds their performances by large margin. In Fig. 8, report the computation efficiency of our SNOpt compared Adam on each dataset, and leave their numerical values in Appendix A.4 (Table 9 and 10). For image and time-series datasets (i.e. Mn-CT), our SNOpt runs nearly as fast as first-order methods. This is made possible through a rigorous OCP analysis in Section 3, where we showed that second- matrices can be constructed along with the same backward integration when we compute the gradient. Hence, only a minimal overhead is introduced. As for CNF, which propagates the probability density additional to vanilla state dynamies, our SNOpt is roughly 1.5 to 2.5 times slower, yet it still converges faster in the overall wall-clock time (see Fig. 7). On the other hand, the use of second-order matrices increases the memory consumption of SNOpt by 10-40%, depending on the model and datasct. However, the actual increase in memory (less than 1GB for all datasets; see Table 10) remains affordable standard GPU machines. More importantly, our SNOpt retains the O(1) memory throughout training.
Test-time performance and hyper-parameter sensitivity. Table 4 reports the test-time perfor- mance, including the accuracies (%) for image and time-series classification, and the negative log- ligliiillhood (NLL) for CNF. On most datasets, our method achieves competitive results against standard baselines. In practice, we also find that using the preconditioned updates greatly reduce the sensitivity to hyper-parameters (e.g. learning rate). This is demonstrated in Fig. 9, where we sample distinct learning rates from a proper interval for each method (shown with different color bars) and record their training results after convergence. It is clear that method not only converges to higher

---

Table 5: Performance of joindy optimizing the integration bound t on CIFARIO
Table 6: Measure of implicit regularization on SVHN
# of function
Regularization
Train time (%) Accuracy (%)
evaluation (NFE)
(f V + + [|FII
Method
w.L.t f =110
Adam SNOpt
42.1
323.88
ASM baseline
96
76.61
32.6
199.1
SNOpt (ours)
81
77.82
Relative Runtime Recorsivedddiin SNOpt)
Accuracy Improvement (%) (SMDpt - Recorsivinn
Opcimization
x
15
1.0
4-10
10-
3
5
ASM baseline
x
0.5
SNOpt (ours)
Mn
Sv
Cf
SA
AW
CT
Mn Sv SA AW CT
0
Train Iteration
5k
Figure 11: Comparison between SNOpt and second- order recursive adjoint. SNOpt is at least 2 times faster and improves the accuracies of baselines by 5-15%.
Figure 10: Dynamics of 4 over CIFAR1O
training using different methods.
accuracies with lower losses, these values are also more concentrated on the plots. In other words, our method achieves better convergence in a more consistent manner across different hyper-parametenn
Joint optimization of the integration bound t1. Table 5 and Fig. 10 report the performance of optimizing t1 along with its convergence dynamics. Specifically, we compare our second-orde order feedback policy (16) derived in Section 3.4 to the first-order ASM baseline proposed in Massaroli al. (2020). It clear that our OCP-theoretic method leads substantially faster convergence, and the optimized ( stably hovers around 0.5 without deviation (as appeared for the baseline). This drops the training time by nearly 20% compared the vanilla training, where we fix 4 to 1.0, yet without sacrificing the test-time accuracy. A similar experiment for MNIST (see Fig. 13 in Appendix A.5) shows a consistent result. We highlight these improvements as the benefit gained from introducing the well-established OCP principle to these emerging deep continuous models.
Comparison with recursive adjoint. Finally, Fig. 11 reports the comparison between our SNOpt and the recursive adjoint baseline (see Section 2 and Table 1). It is clear that our method outperforms this second-order baseline by a large margin in both runtime efficiency and test-time performance. Note that we omit the comparison on CNF datasets since the recursive adjoint simply fails to converge.
Remark (Implicit regularization). In some cases (e.g. SVHN in Fig. 8), method may run slightly faster than first-order methods. This is a distinct phenomenon arising exclusively from training these continuous-timg time models. Since their forward and backward passes involve solving parameterized ODEs (see Fig. 2), the computation graphs are parameter-dependent hence adaptive throughout training. In this vein, we conjecture that the preconditioned updates in these cases may have guided the parameter to regions that are numerically stabler (hence faster) for integration." 4 With this in mind, we report in Table 6 the value of Jacobian, f FII and Kinetic, I || regularization (Finlay et al., 2020) in SVHN training. Interestingly, the parameter found by our SNOpt indeed has a substantially lower value (hence stronger regularization and better-conditionn compared to the one found by Adam. This provides plausible explanation of the reduction in the NFE when using our method, yet without hindering the test-time performance (see Table 4).
ODE dynamics)
Conclusion
We present an efficient higher-order optimization framework for training Neural ODEs. Our method named SNOpt - differs from existing second-order methods in various aspects. While it leverages similar factorization inherited Kronecker-based methods (Martens & Grosse, 2015), the two methodologies differ fundamentally in that we construct analytic ODE expressions for higher-order derivatives (Theorem 1) and compute them through ODESolve. This retains the favorable O(1) memory as opposed to their (TT). It also enables a flexible rank-based factorization in Proposition 2. Meanwhile, our method extends the recent trend of OCP-inspired methods (Li et al., 2017: Liu al., 202 to deep continuous-time models, yet using rather straightforward framework without imposing additional assumptions, such as Markovian or game transformation. To summarize, our work advances several methodologies to the emerging deep continuous-time models, achieving strong empirical results and opening up new opportunities for analyzing models such as Neural SDEs/PDESs
In Appendix A.4, we provide some theoretical discussions (see Corollary 9) in this regard.
10

---

Acknowledgments and Disclosure of Funding
The authors would like thank Chia-Wen Kuo and Chen-Hsuan Lin for the meticulous proofreading. and Keuntaek Lee for providing additional computational resources. Guan-Horng Liu was supported by CPS NSF Award #1 11320068, and Tianrong Chen was supported by ARO Award #N911N22010151.
References
Almubarak, H., Sadegh, N., and Taylor, Infinite horizon nonlinear quadratic cost regulator, In 2019 American Control Conference (ACC), pp. 5570-5555555 5575. IEEE, 2019.
Amari, S.-i. and Nagaoka, H. Methods of information geometry, volume 191. American Mathematical Soc., 2000.
Ba, J.. Grosse, R., and Martens, J. Distributed second-order optimization using kronecker-factored approximations. 2016.
Bagnall, A., Dau, H. A., Lines, J., Flynn, M., Large, J., Bostrom, A., Southam, P., and Keogh, E. The uea multivariate time series classification archive, 2018. arXiv preprint arXiv:1811 1.1.0005, 2018.
Botev, A., Ritter, H., and Barber, D. Practical gauss-newton optimisation for deep learning. In Proceedings of the 34th International Conference on Machine Learning- Volumee 70, pp. 557-565. JMLR. org, 2017.
Chalvidal, M., Ricci, M., VanRullen, R., and Serre, T. Go with the flow: Adaptive control for neural odes. 2021.
Chen, T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, pp. 6572-6583, 2018.
Cho, K., Van Merriënboer, B., Gulcehre, C., Bahdanau, D., Bougares, F., Schwenk, H., and Bengio, Y. Learning phrase representations using rnn encoder-decoder for statistical machine translation. arXiv preprint arXiv: 1406. 1078, 2014.
De Marchi, A. and Gerdts, M. Free finite horizon lqr: bilevel perspective and its application to model predictive control. Automatica, 100:299-311, 2019.
Desjardins, G., Simonyan, K., Pascanu, R., and Kavukcuoglu, K. Natural neural networks. arXiv preprint 1507.002 10, 2015.
Dormand, J. R. and Prince, P. J. A family of embedded runge-kutta formulae. Journal of computational and applied mathematics, 6(1):19-26 1980.
Finlay, C., Jacobsen, J.-H., Nurbekyan, L., and Oberman, A. How to train your neural ode: the world of jacobian and kinetic regularization. In International Conference on Machine Learning. pp. 3154 3164. PMLR, 2020.
Gao, K.-X., Liu, X.-L., Huang, Z.-H., Wang, M., Wang, Z., Xu, D., and Yu, F. A trace-restricted Kronecker-factoree approximation to natural gradient. arXiv preprint arX1v:2011.10741, 2020.
George, T., Laurent, C., Bouthillier, X., Ballas, N., and Vincent, Fast approximate natural gradient descent in a kronecker factored eigenbasis. In Advances in Neural Information Processing Systems, 2018.
pp.
Gholami, A., Keutzer, K., and Biros, G. Anode: Unconditionally accurate memory-efficient gradients for neural odes. arXiv preprint arXiv: 1902. 10298, 2019.
Ghosh, A., Behl, H. S., Dupont, E., Torr, P. H., and Namboodiri, V. Steer: Simple temporal regularization for neural odes. arXiv preprint arXiv:2006.10711, 2020.
Grathwohl, W. Chen, R. T., Betterncourt, J., Sutskever, I., and Duvenaud, D. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint 1rXivv18100.000 2018.
11

---

Grosse, R. and Martens, J. A froteceer-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning, pp- 573-582, 2016.
Gupta, V., Koren, T., and Singer, Y. Shampoo: Preconditioned stochastic tensor optimization. In International Conference Machine Learning, pp. 1842-1850. PMLR, 2018.
Hu, K., Kazeykina, A., and Ren, Mean-field langevin system, optimal control and deep neural networks. arXiv preprini 1rXiv:1909.0777 2019.
loffe, S. and Szegedy, C. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference on machine learning, pp. 448-456. PMLR, 2015.
Kelly, J., Bettencourt, J., Johnson, M. J., and Duvenaud, D. Learning differential equations that are easy to solve. arXiv preprint wrxiv:2007.04504 2020.
Keskar, N. S., Mudigere, D., Nocedal, J., Smelyanskiy, M., and Tang, P. T. On large-batch training for deep learning: Generalization gap and sharp minima. arXiv preprint arXiv:1609.04480 2016.
Kidger, P., Chen, R. T., and Lyons, T. hey, that's not ode": Faster ode adjoints with 12 lines of code. arXiv preprint arXiv:2009.008 09457, 2020a.
Kidger, P., Morrill, J., Foster, J., and Lyons, T. Neural controlled differential equations for irregular time series. arXiv preprint arXiv:2005.0880 089266 2020b.
Laurent, C., George, T., Bouthillier, X., Ballas, N., and Vincent, P. An evaluation of fisher approxi- mations beyond kronecker factorization. 2018.
LeCun, Y., Touresky, D., Hinton, G., and Sejnowski, T. theoretical framework for back back- toock-roppggtion. In Proceedings of the 1988 connectionist models summer school, volume 1, pp. 21-28. CMU, Pittsburgh, Morgan Kaufmann, 1988.
Li, Q., Chen, L.., Tai, C., and Weinan, E. Maximum principle based algorithms for deep learning. The Journal of Machine Learning Research, 18(1):599866026, 2017.
Liu, G.-H. and Theodorou, E, A. Deep learning theory review: An optimal control and dynamical systems perspective. arXiv preprint arXiv: 10920, 2019.
Liu. G.-H., Chen, T., and Theodorou, E. A. Ddpnopt: Differential dynamic programming neural optimizer. In International Conference on Learning Representations, 2021a.
Liu, G.-H.. Chen, T., and Theodorou, E. A. Dynamic game theoretic neural optimizer. In International Conference on Machine Learning, 2021b.
Lou, A., Lim, D., Katsman, I., Huang, L.. Jiang, Q., Lim, S.-N., and De Sa, C. Neural manifold ordinary differential equations. arXiv preprint arXiv:2006. 10t200000..2244 10254, 2020.
Ma, L., Montague, G., Ye, J., Yao, Z., Gholami, A., Keutzer, K., and Mahoney, M. W. Inefficiency of k-fac for large batch size training. arXiv preprint arXiv:1903.0623; 2019.
Martens, J. New insights and perspectives on the natural gradient method. arXiv preprint 11X1 2014.
Martens, J. and Grosse, R. Optimizing neural networks with cronecker-factordd approximate curvature. In International conference on machine learning, pp. 2408-24024 2417, 2015.
Martens, J., Ba, J., and Johnson, M. Kronecker-factored curvature approximations for recurrent neural networks. International Conference on Learning Representations, 2018
Massaroli, S., Poli, M., Park, J., Yamashita, A., and Asama, H. Dissecting neural odes. arXiv preprint wrXi::2002.08071, 2020.
Mathieu, and Nickel, M. Riemannian continuous normalizing flows. arXiv preprint 706055 2020.
12

---

Nguyen, T. M., Garg, A., Baraniuk, R. G., and Anandkumar, A. Infocnf: An efficient conditional continuous normalizing flow with adaptive solvers. arXiv preprint arXiv:1912.03978, 2019.
Onken, D., Fung, W., Li, X., and Ruthotto, L. Or-flow: Fast and accurate continuous normalizing flows via optimal transport. arXiv preprint arXiv:2006.00104, 2020.
Paszke, A., Gross, S., Chintala, S., Chanan, G., Yang. E., Detito, Z., Lin, Z.,, Desmaison, A., Antiga, L., and Lerer, A. Automatic differentiation in pytorch. 2017.
Pontryagin, I.. S., Mishchenko, E., Boltyanskii, V., and Gamkrelidze, R. The mathematical theory of optimal processes. 1962.
Press, W. H., William, H., Teukolsky, A., Vetterling, W. T., Saul, A., and Flannery, P. Numerical recipes 3rd edition: The art of scientific computing. Cambridge university press, 2007.
Rubanova, Y., Chen, R. T., and Duvenaud, D. Latent odes for iregularly-samplee time series. arXiv preprint arXiv: 1XXiv:1977.03007 2019.
Santurkar, S., Tsipras, D., Ilyas, A., and Madry, A. How does batch normalization help optimization? arXiv preprint arXiv: 1805. 1604, 2018.
Schacke, K. On the kronecker product. Master's thesis, University of Waterloo, 2004.
Sun, W., Theodorou, E., and Tsiotras, P. Model based reinforcement learning with final time horizon optimization. arXiv preprint arXiv:1509 1et.115990011111 086,0,11860000 2015.
Tassa, Y., Mansard, N., and Todorov, E. Control-limited differential dynamic programming. In 2014 IEEE International Conference on Robotics and Automation (ICRA), pp. 1168-1175. IEEE, 2014.
Theodorou, E., Tassa, Y., and Todorov, E. Stochastic differential dynamic programming. In Proceedings of the 2010 American Control Conference, pp. 1125-1 1132. IEEE, 2010.
Todorov, E. Optimal control theory. Bayesian brain: probabilistic approaches neural coding, pp- 269-298, 2016.
Weinan, E. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 5(1):1-11, 2017.
Weinan, E., Han, J., and Li, Q. A mean-field optimal control formulation of deep learning, arXiv preprint arXiv: 10.7.100700088 2018.
Wu, Y., Zhu, X., Wu, C., Wang, A., and Ge, R. Dissecting hessian: Understanding common structure of hessian neural networks. arXiv preprint arXiv:20/0.04261, 2020.
Zhang, G., Martens, J., and Grosse, R. Fast convergence of natural gradient descent for overparame- terized neural networks. arXiv preprint arXiv: 10961, 2019.
Zhong, Y. D., Dey, B., and Chakraborty, A. Symplectic ode-net: Learning hamiltonian dynamics with control. 2020.
Zhuang. J., Dvoreek, N., Li, X., Tatikonda, S., Papademetris, X., and Duncan, J. Adaptive checkpoint adjoint method for gradient estimation neural ode. In International Conference on Machine Learning, pp. 1 1639-116499 PMLR, 2020.
Zhuang, J., Dvomekk N. C., Tatikonda, S., and Dunean, J. Mali: A memory efficient and reverse accurate integrator for neural odes. arXiv preprint arXiv:2102.04666 2021.
13


Reviews:
[{"review_id": "vgOe7t-JX1z", "reviewer": "Reviewer_w3BT", "summary": "This paper presents a novel second-order optimization technique for neural ordinary differential equations (NODEs), derived from the optimal control perspective of NODEs. In a nutshell, the authors consider the second-order Taylor expansion of the loss term, leading to matrix-valued ODEs that give second-order derivatives. A series of approximations are proposed to compute a low-rank representation of the second-order preconditioning matrices. The experiments demonstrate that the proposed optimization method converges faster than a first-order optimization routine in all neural ODE applications at the cost of slightly higher memory usage. The method also allows optimizing the integration time in continuous-time normalizing flows.", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The integration operation in continuous-time models can be overwhelmingly time-consuming, which oftentimes makes it difficult to apply neural ODEs to even simple, low-dimensional problems. Therefore, the proposed approach addresses an overlooked but fundamental restriction of a rather recent model family. The paper is well-written. Despite the heavy technical details, the presentation is very clear and I had no difficulty grasping the motivation and ideas. The experiments and the discussion of the results are exhaustive. Experiments indeed show that the presented optimization scheme improves the NODE training in all three applications of NODEs. \n\nA few concrete suggestions and questions:\n- Since the approximations are very involved, having a broader picture of them would be nice. For instance, a cartoon illustration of the series of approximations and what they aim to address could help the reader gain insights more easily. Also, additional (simple) ablation studies that compare approximations with the corresponding exact expressions might reveal how good/tight the approximations are.\n- A direct comparison against a Quasi-Newton method (I suggest LBFGS) on a simple problem would be very interesting to see as such methods can scale to low/mid-size neural nets.\n- Line 243: Which tolerance values (absolute/relative) does this refer to? Also, any justification for not using the default values?\n- To support the claims in Line 272, it would be nice to test NODE with a regularization scheme, e.g., Finlay et al., 2020. Also, one may further check the number of function evaluations.\n- Looking at Fig13 (train time), I wonder why it first drops and then goes up. Maybe because of a change in the stiffness of the vector field? Can this be experimentally demonstrated?\n\nMinor details:\n- I suggest writing \"back-propagation\" instead of \"Back-propagation\". \n- Line117: ',and etc' seems unnecessary. Also, what is w.l.o.g?\n- An important reference to [1] is missing.\n\n[1] Stapor, Paul, Fabian Fröhlich, and Jan Hasenauer. \"Optimization and profile calculation of ODE models using second order adjoint sensitivity analysis.\" Bioinformatics 34.13 (2018): i151-i159.\n\n**Post-rebuttal comment:** Thanks authors for the rebuttal. As noted in my original review, this is a very nice paper and should be definitely accepted.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations are clearly described. No potentially negative societal impact I can think of.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "dpviFdIhOTy", "reviewer": "Reviewer_HLDZ", "summary": "This paper proposes a method to derive second-order derivatives for Neural ODEs, which requires O(1) memory in theory. For practical feasibility, the authors propose to use a low-rank approximation of the (dynamic) Hessian matrix, together with Kronecker-based factorization to further reduce the memory cost. The authors validated their methods in image classification, time-series analysis and continuous generative modeling.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n\nThis method, to my knowledge, is novel in terms of an efficient method to estimate the second-order gradient information for Neural ODE models. The paper is in general well-written with solid theoretical analysis and experimental validations. It would be better if the authors could release code for such a complicated method.\n\nCons:\n\nI have a few comments and concerns rather than critiques for the authors to clarify.\n\n1. I quickly went through the proof, found the key result rely on two important assumptions (line 526 in appendix), a) $z(t) \\otimes g(t)$ are uncorrelated in time $t$, b) $z$ and $g_i$ are pairwise independent. Assumption a) looks very strong even unpractical to me, since in the update it's always $z(t+\\delta (t))=z(t)+\\Delta(t)$, omit the details of $\\Delta(t)$ here. It's clear that $z(t)$ and $z(t+\\delta (t))$ are highly-correlated.\n\n2. The author did not specify how to choose a proper rank $R$\n\n3. In terms of memory consumption, If I get it correctly, Full rank second-order adjoint > Low-rank second-order adjoint > first-order adjoint, so the memory cost of the proposed method is in-between. Even though it's O(1) in time $t$, I'm concerned the memory cost has a large constant. \\\nSpecifically, suppose the parameter number is $n$ and activation dimension is $m$, first-order adjoint needs $O(n+m)$ memory, while the proposed method needs $O(n+m+Rm+Rn)$, where $R$ is the rank. This memory cost is required to store the first and second-order gradients. \\\nIn terms of update, the method requires a much higher peak memory of $O(n^2+m^2)$, though it's reusable, because Eq.13 requires the inversion of two matrices of size $n^2$ and $m^2$, this inversion could also be very expensive.  \\\nTable 7 in appendix shows that the proposed method requires 40% more memory than first-order methods. Note that this is for small-size image ($m=32\\times 32$) on CIfar, on large images such as ImageNet ($m=224\\times 224$), I expect the memory would be even much larger. \\\nSo the scalability to large scale networks especially with high-resolution images is one of my concerns.\n\nSuggestions\n\nIt's shown that numerical errors in gradient for Neural ODE is the cause of bad empirical performances, so I wonder how does the proposed method compare with first-order methods that tackle the numerical issues carefully [1,2,3].\n\n[1] Gholami, Amir, Kurt Keutzer, and George Biros. \"Anode: Unconditionally accurate memory-efficient gradients for neural odes.\" arXiv preprint arXiv:1902.10298 (2019).  \\\n[2] Zhuang, Juntang, et al. \"Adaptive checkpoint adjoint method for gradient estimation in neural ode.\" International Conference on Machine Learning. PMLR, 2020.  \\\n[3] Zhuang, Juntang, et al. \"MALI: A memory efficient and reverse accurate integrator for Neural ODEs.\" arXiv preprint arXiv:2102.04668 (2021).", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "4Tb-jTay-g6", "reviewer": "Reviewer_ikHp", "summary": "This work proposes an alternative to the recursive adjoint method for computing higher order gradients of neural ODEs based on optimal control theory that promises to be more efficient and more accurate. ", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "My main concern is that the authors do not take sufficient care in statements about the memory and time costs. The adjoint method is O(1) memory *in depth* i.e. in time, but it is O(n) in the number of parameters, i.e. $dim_R(\\theta) = n$. Naively computing higher-order derivatives by recursively applying the adjoint method has O(n) memory footprint and O(n) time, but these calculations are independent and can be parallelised to give O(n^2) memory and O(1) time (as described by the authors L98). The proposed method maintains O(1) in depth, but has a footprint of O(n^2) ($Q_{uu}$ is matrix valued $\\in R^{n\\times n}$), and is O(1) in time.\n\nSo there isn’t actually any change in the memory or time costs as compared to recursively applying the adjoint method. There may be a difference in the error, as stated in L41-43, but there is no evidence presented in favour of this claim nor evidence that this is of practical consequence in the training of neural ODEs, which is the evaluation criteria used in this work.\n\nL97-101 make some arguments against the recursive method, stating that the method scales unfavourably due to recursion and errors. That k=2 is apparently the useful limit of this recursion (‘*second*-order neural ODE optimizer') diminishes the significance of the first complaint and it is diminished further when we note that the time cost is linear in the order of the optimization, which is a fixed hyperparameter. Evidence is not presented to support the second complaint, and this is a missed opportunity as whether the accumulation of errors is of any significance in training is an interesting scientific question. Recursively applying the adjoint method is the essential baseline to compare the method against, particularly given the proposals from section 3.2 onwards to approximate the optimisation and L194 which introduces a further approximation explicitly to trade off complexity for performance, such that both methods run with some error.\n\nThis is a shame as the work is otherwise interesting and highly creative, and seeing this comparison seems essential.\n\nI’d also be interested in some discussion about how the number of terms in equation 9 grows in the order of the optimisation. I followed the discussion in the appendix to Almubarak 2019, though the appendix indicates that there is a further reference within Almubarak 2019 that is to be followed next, though doesn’t specify which. Combinatorially, we expect to get 2^(n+1) - 2 (i.e. 2 for first order, 6 for second order, 14 for third order…) This seems to suggest that there is a further term missing from the memory footprint. Including the optimisation order, k, the costs seem to be\nRecursive adjoint: memory O(n^k) time O(k)\nProposal: memory O(2^k n^k) time O(1)\nGiven that k is generally a fixed hyperparameter and ≤ 2 I doubt this is a pressing problem, but some comment would be appreciated.\n\nSection 3.3 would benefit from including the choice of initialisation in the discussion. We can freely vary the integration time without modifying the output by making corresponding changes to the magnitude of the weights in the derivative function i.e. if you want to double the integration time, half the weights in the final layer of the derivative network. So discussions about the integration interval are missing important context. Furthermore, given that the model uses a CNN feature extractor (Appendix A.4 gives this as a 3 layer conv net) prior to the NODE, and that output of the integration is passed through a linear layer to get the prediction, the results in Figure 5 are explained equally as the ODE being superfluous. The correct baselines for comparison here would be 0 integration, which shows 50% reduction in train time for what appears to be a 2-3 percentage point reduction in accuracy, and a model that has the same capacity as the NODE added. In this case the ODE function is a 2 layer conv net, so a 5 layer conv net in total would be appropriate. Figure 13 in the appendix, showing the same experiment for MNIST, presents similarly inconclusive results, in that case the change in accuracy is even less significant across the range of values tested.\n\nThe results in Figure 7 are mixed, which isn’t accurately reflected in the discussion in the **Convergence and computational efficiency** paragraph. On SpoAD, SGD trains as quickly and to the same performance. On CharT the three methods reach similar levels of performance. The absence of errors in Table 2 makes it difficult to compare the significance of the bolded results (and appropriate baselines discussed above are absent).\n\nTable 3 does not select the best training time model from Figure 5. There is a higher-performing and faster-training model on that plot, something in the region of t = 0.4 or 0.5 looks to have performance at least as good as that of the t = 1.0 model at a train time percentage of 50-75%. \n\nFigure 9 is a useful plot, though the caption doesn’t seem to match the data: the model is sensitive to the choice of learning rate, with the training loss varying by approximately ~2x between high and low learning rates on CIFAR10. A higher learning rate is not investigated for Adam, which would appear to be on trend to match the performance of the proposal if a higher learning rate were used (perhaps it wouldn’t, the evidence isn’t provided.) Better would be to provide training plots as in Figure 7 for varying learning rates for the different methods.\n\nAlso absent is a discussion of the effect of batch size on the training plots. With the first order methods having a lower memory overhead, there should be more scope for increasing the batch size to accelerate training. Has this been investigated?\n\nOverall, this is an exciting project that is let down by its evaluation. The absence of obvious baselines, most critically the recursive adjoint method, and seemingly cherry picked examples in the comparisons that are present, make it difficult to evaluate the impact of this work.\n\n--- minor points\nThe presentation is clear and of high quality.\nL91 should be ‘the backpropgation algorithm’\nL122 seemly should be seeming\nL160 ‘much more efficiently’\n\n\n### Update\nFollowing the rebuttal and back and forth with the authors, my concerns have been addressed and I'm happy to raise my score (4->7) to reflect this.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Impact is adequately addressed.", "review_text": "", "extra_scores": {"ethical_concerns": "There are no ethical concerns.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "0hvwpKkoT5a", "reviewer": "Reviewer_1V2i", "summary": "This method proposes a second-order optimization network for neural ODEs. Inspired by OC, this paper adopts differential programming to derive backward ODE for higher-order derivatives at constant memory cost. This paper also explores using low-rank representation for the second-order optimization to make more efficient updates. The paper shows their method achieves faster convergence than first-order methods across three different tasks. This framework also enables the optimization of the architecture with the example of integration time of neural ODE.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality\n- This work proposes a novel second-order optimization network specifically for neural ODEs. It adopts ideas from OCP to neural ODE to achieve this goal. It is clear how this work differs from previous work, and the related works are appropriately cited. I like how it introduces the OCP perspective to neural ODE and it can treat integration time as an optimization variable to optimize the architecture directly.\n\nQuality\n- This work is technically sound. It makes sense that using a second-order optimization approach can make neural ODE converge faster and this claim is well-supported by three different tasks, which demonstrate faster convergence and higher accuracy using their method than using the first-order method.\n\n- I would like to see the reasoning behind Figure 9, which claims that this method is less sensitive to hyper-parameter changes. It is not very obvious based on the theoretical ground and is worth exploring.\n\nClarity\n- The overall structure is well-organized. \n\n- Different fonts of x are present in lines 19 and 21 before it is clarified in preliminary.\n\n- On line 118: what are “These functions are known as...” Which are the functions here?\n\nSignificance\n- Other researchers may use their method to run neural ODE, but I doubt the significance of developing even higher-order optimization based on this work. Also, I am not sure if it is worth implementing the method given its complexity and limited improvement on computational time. Please discuss these in Section 5.\n\nAdditional Questions\n-What is the relationship between step sizes and integration time? Does increasing the time step size have the same effect as decreasing the integration time for image and CNF datasets?\n\n-Why do you choose a Runge-Kutta 4 adaptive solver instead of an implicit solver such as the implicit Adams method? Do the benefits depend on the specific solver you choose?\n\n-For Eq. 14, why do you call this specific formula a preconditioned update law? Why do you choose this to update theta?\n\n-Is there any mathematical formulation that supports how much the second-order optimization improves the convergence rate of Neural ODE?\n\n-Why using second-order optimization can also improve the accuracy of the experiments as demonstrated in figure 9?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations were adequately discussed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3 hours"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes a second-order optimization algorithm for neural ODEs. Inspired by optimal control theory, the paper derives backward ODEs for higher-order derivatives at constant memory cost. The authors shows their method achieves faster convergence than first-order methods across three different tasks.

The contribution is novel and relevant to NeurIPS, and is appreciated by the reviewers. The main criticism that surfaced during review is the absence of certain baselines, raised by reviewer ikHp: After back and forth discussion between the authors and reviewer this concern was addressed sufficiently for the reviewer to raise their review score. The authors also replied extensively to comments by the other reviewers, and no further serious concerns were identified. The reviewers formed a consensus recommendation to accept the paper.
Decision: Accept (Spotlight)