Second-Order Neural ODE Optimizer
Year: 2021

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 continuoustime 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 timeseries 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.

text: 1 Introduction
Neural ODEs (Chen et al., 2018) have received tremendous attention over recent years. Inspired by taking the continuous limit of the “discrete” residual transformation, xk+1 = xk + F (xk, θ), they propose to directly parameterize the vector field of an ODE as a deep neural network (DNN), i.e.
dx(t) dt = F (t,x(t), θ), x(t0) = xt0 , (1)
where x(t) ∈ Rm and F (·, ·, θ) is a DNN parameterized by θ ∈ Rn. This provides a 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 et al., 2020; Mathieu & Nickel, 2020).
Due to the continuous-time representation, Neural ODEs feature a distinct optimization process (see Fig. 2) compared to their discrete-time counterparts, which also poses new challenges. First, the
35th Conference on Neural Information Processing Systems (NeurIPS 2021)
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-propagating 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 O(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). 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 in 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 et al., 2018; Liu & Theodorou, 2019), which treats the parameter θ as a control variable, so that the training process, i.e. optimizing θ 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, a 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 sameO(1) memory and a comparable runtime to first-order methods by adopting Kronecker factorization (Martens & Grosse, 2015). The resulting method – called SNOpt – admits superior convergence in 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 “depth” 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 is a well-studied problem for practical applications with a 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 feedback policy that adapts the integration time throughout training. On training CIFAR10, this further leads to a 20% runtime reduction, yet without hindering test-time accuracy.
In summary, we present the following contributions.
• We propose a novel computational framework for computing higher-order derivatives of deep continuous-time models, with a rigorous analysis using continuous-time 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 secondorder feedback policy for adapting the integration horizon and show it further reduces the runtime.

heading: 2 Preliminaries
text: Notation. We use roman and italic type to represent a variable x(t) and its realization x(t) 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(t0), t0, t1, F) where dx(t) dt = F .
Forward and backward computations of Neural ODEs. Given an initial condition x(t0) and integration interval [t0, t1], Neural ODEs concern the following optimization over an objective L,
min θ L(x(t1)), where x(t1) = x(t0) + ∫ t1 t0 F (t,x(t), θ) dt (2)
is the solution of the ODE (1) and can be solved by calling a black-box ODE solver, i.e. x(t1) = ODESolve(x(t0), t0, t1, F). The use of ODESolve allows us to adopt higher-order numerical methods, e.g. adaptive Runge-Kutta (Press et al., 2007), which give more accurate integration compared with e.g. vanilla Euler discretization in residual-based discrete models. To obtain the gradient ∂L∂θ of Neural ODE, one may naively Back-propagate through ODESolve. This, even if it could be made possible, leads to unsatisfactory memory complexity since the computation graph can grow arbitrarily large for adaptive ODE solvers. Instead, Chen et al. (2018) proposed to apply the Adjoint Sensitivity Method (ASM), which states that the gradient can be obtained through the following integration.
∂L ∂θ = − ∫ t0 t1 a(t)T ∂F (t,x(t), θ) ∂θ dt , (3)
where a(t) ∈ Rm is referred to the adjoint state whose dynamics obey a backward adjoint ODE,
−da(t) dt = a(t)T ∂F (t,x(t), θ) ∂x(t) , a(t1) = ∂L ∂x(t1) . (4)
Equations (3, 4) present two coupled ODEs that can be viewed as the continuous-time expression of the Back-propagation (LeCun et al., 1988). Algorithmically, they can be solved through another call of ODESolve (see Fig. 2) with an augmented dynamics G, i.e.x(t0)a(t0)
∂L/∂θ
 = ODESolve( x(t1)a(t1)
0
 , t1, t0, G), where G t, x(t)a(t) ·  , θ  := F (t,x(t), θ)−a(t)T ∂F∂x −a(t)T ∂F∂θ  (5) augments the original dynamics F in (1) with the adjoint ODEs (3, 4). Notice that this computation (5) depends only on (x(t1),a(t1)). 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,1 the computation in (5) only consumes constant O(1) memory cost.
Chen et al. (2018) noted that if we further encapsulate (5) by ∂∂θL = grad(L, θ), one may compute higherorder derivatives by recursively calling ∂nL ∂θn = grad( ∂n−1L ∂θn−1 , θ), starting from n=1. This can scale unfavorably due to its recursive dependence and accumulated integration errors. Indeed, Table 1
Table 1: Numerical errors between ground-truth and adjoint derivatives using different ODESolve on CIFAR10.
rk4 implicit adams dopri5 ∂L ∂θ 7.63×10 −5 2.11×10−3 3.44×10−4 ∂2L ∂θ2 6.84×10 −3 2.50×10−1 41.10 suggests that the errors of second-order derivatives, ∂ 2L ∂θ2 , obtained from the recursive adjoint procedure can be 2-6 orders of magnitude larger than the ones from the first-order adjoint, ∂L∂θ . In the next section, we will present a novel optimization framework that computes these higher-order derivatives without any recursion (Section 3.1) and discuss how it can be implemented efficiently (Section 3.2).

heading: 3 Approach

heading: 3.1 Dynamics of Higher-order Derivatives using Continuous-time Optimal Control Theory
text: OCP perspective is a recently emerging methodology for analyzing optimization of discrete DNNs. Central to its interpretation is to treat the layer propagation of a DNN as discrete-time dynamics, so
1 T̃ is the number of the adaptive steps used to solve (1), as an analogy of the “depth” of Neural ODEs.
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.1.
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 principle (Pontryagin et al., 1962), which is an optimality condition to OCP. Hence, OCP analysis is 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-time OCP analysis.
min θ
[ Φ(xt1) + ∫ t1 t0 `(t,xt,ut)dt ] subjected to { dxt dt = F (t,xt,ut), xt0 = xt0 dut dt = 0, ut0 = θ , (6)
where x(t) ≡ xt, and etc. It should be clear that (6) describes (2) without loss of generality by having (Φ, `) := (L, 0). These functions are known as the terminal and intermediate costs in standard OCP. In training Neural ODEs, ` can be used to describe either the weight decay, i.e. ` ∝ ‖ut‖, or more complex regularization (Finlay et al., 2020). The time-invariant ODE imposed for ut makes the ODE of xt equivalent to (1). Problem (6) shall be understood as a particular type of OCP that searches for an optimal initial condition θ of a time-invariant control ut. Despite seemly superfluous, this is a necessary transformation that enables rigorous OCP analysis for the original training process (2), and it has also appeared in other control-related analyses (Zhong et al., 2020; Chalvidal et al., 2021).
Next, define the accumulated loss from any time t ∈ [t0, t1] to the integration end time t1 as
Q(t,xt,ut) := Φ(xt1) + ∫ t1 t `(τ,xτ ,uτ ) dτ , (7)
which is also known in OCP as the cost-to-go function. Recall that our goal is to compute higherorder derivatives w.r.t. the parameter θ of Neural ODEs. Under the new OCP representation (6), the first-order derivative ∂L∂θ is identical to ∂Q(t0,xt0 ,ut0 )
∂ut0 . This is because Q(t0,xt0 ,ut0) accumulates
all sources of losses between [t0, t1] (hence it sufficiently describes L) and ut0 = θ by construction. Likewise, the second-order derivatives can be captured by the Hessian ∂ 2Q(t0,xt0 ,ut0 )
∂ut0∂ut0 = ∂ 2L ∂θ∂θ ≡ Lθθ.
In other words, we are only interested in obtaining the derivatives of Q at the integration start time t0.
To obtain these derivatives, notice that we can rewrite (7) as
0 = `(t,xt,ut) + dQ(t,xt,ut)
dt , Q(t1,xt1) = Φ(xt1), (8)
since the definition of Q implies that Q(t,xt,ut) = `(t,xt,ut)dt+Q(t+ dt,xt+dt,ut+dt). We now state our main result, which provides a local characterization of (8) with a set of coupled ODEs expanded along a solution path. These ODEs can be used to obtain all second-order derivatives at t0.
Theorem 1 (Second-order Differential Programming). Consider a solution path (x̄t, ūt) that solves the ODEs in (6). Then the first and second-order derivatives of Q(t,xt,ut), expanded locally around this solution path, obey the following backward ODEs:
−dQx̄ dt = `x̄ + F T x̄Qx̄, − dQū dt = `ū + F T ūQx̄, (9a)
−dQx̄x̄ dt = `x̄x̄ + F T x̄Qx̄x̄ +Qx̄x̄Fx̄, − dQx̄ū dt = `x̄ū +Qx̄x̄Fū + F T x̄Qx̄ū, (9b) −dQūū dt = `ūū + F T ūQx̄ū +Qūx̄Fū, − dQūx̄ dt = `ūx̄ + F T ūQx̄x̄ +Qūx̄Fx̄, (9c)
where Fx̄(t)≡ ∂F∂xt |(x̄t,ūt), Qx̄x̄(t)≡ ∂2Q ∂xt∂xt |(x̄t,ūt), and etc. All terms in (9) are time-varying vector-
valued or matrix-valued functions expanded at (x̄t, ūt). The terminal condition is given by
Qx̄(t1) = Φx̄, Qx̄x̄(t1) = Φx̄x̄, and Qū(t1) = Qūū(t1) = Qūx̄(t1) = Qx̄ū(t1) = 0.
The proof (see Appendix A.2) relies on rewriting (8) with differential states, δxt := xt − x̄t, which view the deviation from x̄t as an optimizing variable (hence the name “Differential Programming”). It can be shown that δxt 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 Qx̄(t) ≡ a(t) follows the same backward ODE in (4) and the solution of the second ODE in (9a), Qū(t0) = − ∫ t0 t1 Fū
TQx̄dt, gives the exact gradient in (3). Meanwhile, solving the coupled matrix ODEs presented in (9b, 9c) yields the desired second-order matrix, Qūū(t0) ≡ Lθθ, 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.

heading: 3.2 Efficient Second-order Preconditioned Update
text: Theorem 1 provides an attractive computational framework that does not require recursive computation (as mentioned in Section 2) to obtain higher-order derivatives. It suggests that we can obtain first and second-order derivatives all at once with a single function call of ODESolve:
[xt0 , Qx̄(t0), Qū(t0), Qx̄x̄(t0), Qūx̄(t0), Qx̄ū(t0), Qūū(t0)]
= ODESolve([xt1 ,Φx̄,0,Φx̄x̄,0,0,0], t1, t0, G̃), (10)
where G̃ augments the original dynamics F in (1) with all 6 ODEs presented in (9). Despite that this OCP-theoretic backward pass (10) retains the same O(1) memory complexity as in (5), the dimension of the new augmented state, which now carries second-order matrices, can grow to an 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.
Proposition 2 (Low-rank representation of (9)). Suppose `:=0 in (6) and letQx̄x̄(t1)= ∑R i=1 yi⊗yi be a symmetric matrix of rank R ≤ n, where yi ∈ Rm and ⊗ is the Kronecker product. Then, for all t ∈ [t0, t1], the second-order matrices appeared in (9b, 9c) can be decomposed into
Qx̄x̄(t) = R∑ i=1 qi(t)⊗ qi(t), Qx̄ū(t) = R∑ i=1 qi(t)⊗ pi(t), Qūū(t) = R∑ i=1 pi(t)⊗ pi(t),
where the vectors qi(t) ∈ Rm and pi(t) ∈ Rn obey the following backward ODEs:
−dqi(t) dt = Fx̄(t) Tqi(t), − dpi(t) dt = Fū(t) Tqi(t), (11)
with the terminal condition given by (qi(t1),pi(t1)) := (yi,0).
The proof is left in Appendix A.2. Proposition 2 gives a nontrivial conversion. It 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 R 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 xt in a latent space of higher dimension (Chen et al., 2018; Grathwohl et al., 2018; Kidger et al., 2020b).
Based on Proposition 2, the second-order precondition matrix Lθθ is given by2
Lθθ ≡ Qūū(t0) = R∑ i=1 (∫ t0 t1 Fū Tqi dt ) ⊗ (∫ t0 t1 Fū Tqi dt ) , (12)
where qi ≡ qi(t) follows (11). Our final step is to facilitate efficient computation of (12) with Kronecker-based factorization, which underlines many popular second-order methods for discrete DNNs (Grosse & Martens, 2016; Martens et al., 2018). Recall that the vector field F is represented
2 We drop the dependence on t for brevity, yet all terms inside the integrations of (12, 13) are time-varying.
Algorithm 1 SNOpt: Second-order Neural ODE Optimizer 1: Input: dataset D, parametrized vector field F (·, ·, θ), integration time [t0, t1], black-box ODE
solver ODESolve, learning rate η, rank R, interval of the time grid ∆t 2: repeat 3: Solve x(t1) = ODESolve(x(t0), t0, t1, F), where x(t0) ∼ D. B Forward pass 4: Initialize (Ān, B̄n) := (0,0) for each layer n and set qi(t1) := yi. 5: for t′ in {t1, t1 −∆t, · · · , t0 + ∆t, t0} do 6: Set t := t′ −∆t as the small integration step, then call
[x(t), Qx̄(t), Qū(t), {qi(t)}Ri=1] = ODESolve([x(t′), Qx̄(t′), Qū(t′), {qi(t′)}Ri=1], t′, t, Ĝ), B Backward pass
where Ĝ augments the ODEs of state (1), first and second-order derivatives (9a, 11). 7: Evaluate zn(t), hn(t), F (t,xt, θ), then compute An(t),Bn(t) in (13). 8: Update Ān ← Ān + An(t) ·∆t and B̄n ← B̄n + Bn(t) ·∆t. 9: end for
10: ∀n, apply θn ← θn − η · vec(B̄−1n Qūn(t0)Ā−Tn ). B Second-order parameter update 11: until converges
by a DNN. Let zn(t), hn(t), and un(t) denote the activation vector, pre-activation vector, and the parameter of layer n when evaluating dxdt at time t (see Fig. 3), then the integration in (12) can be broken down into each layer n,∫ t0 t1 ( Fū Tqi ) dt =[· · · , ∫ t0 t1 ( FTūnqi ) dt, · · · ]
=[· · · , ∫ t0 t1 ( zn ⊗ ( ∂F∂hn T qi) ) dt, · · · ],
where the second equality holds by FTūnqi = ( ∂F ∂hn
∂hn ∂un ) Tqi = z n ⊗ ( ∂F∂hn T qi). This is an essential
step towards the Kronecker approximation of the layer-wise precondition matrix:
Lθnθn ≡ Qūnūn(t0) = R∑ i=1 (∫ t0 t1 ( zn ⊗ ( ∂F∂hn T qi) ) dt ) ⊗ (∫ t0 t1 ( zn ⊗ ( ∂F∂hn T qi) ) dt )
≈ ∫ t0 t1
(zn ⊗ zn)︸ ︷︷ ︸ An(t)
dt⊗ ∫ t0 t1 R∑ i=1 ( ( ∂F∂hn T qi)⊗ ( ∂F∂hn T qi) )
︸ ︷︷ ︸ Bn(t)
dt. (13)
We discuss the approximation behind (13), and also the one for (14), in Appendix A.2. Note that An(t) and Bn(t) are much smaller matrices in Rm×m compared to the ones in (9), and they can be efficiently computed with automatic differentiation packages (Paszke et al., 2017). Now, let {tj} be a time grid uniformly distributed over [t0, t1] so that Ān= ∑ jAn(tj)∆t and B̄n= ∑ jBn(tj)∆t approximate the integrations in (13), then our final preconditioned update law is given by
∀n, L−1θnθnLθn ≈ vec ( B̄−1n Qūn(t0)Ā −T n ) , (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 et al. (2018); see Alg. 2 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, ` := 1t1−t0 ‖θ‖2, which is time-independent, it can be separated from the backward ODEs (9) and added after solving the backward integration, i.e.
Qū(t0)← γθ +Qū(t0), Qūū(t0)← γI +Qūū(t0), where γ is the regularization factor. Finally, we find that using the scaled Gaussian-Newton matrix, i.e. Qx̄x̄(t1) ≈ 1t1−t0 Φx̄ ⊗ Φx̄, generally provides a good trade-off between the performance and runtime complexity. As such, we adopt this approximation to Proposition 2 for all experiments.

heading: 3.3 Memory Complexity Analysis
text: 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 O(1) memory as with the first-order adjoint method, their complexity differs in terms of the state and parameter dimension. Starting from our encouraging result in Theorem 1, 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 toO(Rm+Rn) and paves a 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-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 t1. In practice, when problems contain no prior information on the integration, [t0, t1] 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 CIFAR10 for instance (see Fig. 5), the required training time decreases linearly as we drop t1 from 1, yet the accuracy retains mostly the same unless t1 becomes too small. Similar results also appear on MNIST (see Fig. 12 in Appendix A.5). In other words, we may interpret the integration bound t1 as an architectural parameter that needs to be jointly optimized during training.
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 a new variable:
Q̃(t,xt,ut,T) := Φ̃(T,x(T)) + ∫ T t `(τ,xτ ,uτ ) dτ , (15)
where Φ̃(T,x(T)) explicitly imposes the penalty for longer integration time, e.g. Φ̃ := Φ(x(T)) + c 2T
2. 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 (x̄t, ūt, T̄ ). After some tedious derivations, which are left in Appendix A.3, we will arrive at the update rule below,
T← T̄ − η · δT(δθ), where δT(δθ) = [Q̃T̄ T̄ (t0)]−1 ( Q̃T̄ (t0) + Q̃T̄ ū(t0)δθ ) . (16)
Similar to what we have discussed in Section 3.1, one shall view Q̃T̄ (t0) ≡ ∂L∂T as the first-order derivative w.r.t. the terminal time T. Likewise, Q̃T̄ T̄ (t0) ≡ ∂ 2L ∂T∂T , and etc. Equation (16) is a secondorder feedback policy that adjusts its updates based on the change of the parameter θ. Intuitively, it moves in the descending direction of the preconditioned gradient (i.e. Q̃−1
T̄ T̄ Q̃T̄ ), while accounting for
the fact that θ is also progressing during training (via the feedback Q̃T̄ ūδθ). The latter is a 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).
Dataset. We select 9 datasets from 3 distinct applications where N-ODEs have been applied, including image classification (•), time-series prediction (•), and continuous normalizing flow (•; CNF): • MNIST, SVHN, CIFAR10: MNIST consists of 28×28 gray-scale images, while SVHN and
CIFAR10 consist of 3×32×32 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 a speech dataset, whereas ArticularyWordRecognition (ArtWR) and CharacterTrajectories (CharT) are motion-related datasets. Table 3 details their sample sizes.
• Circle, Gas, Miniboone: Circle is a 2-dim synthetic dataset adopted from Chen et 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 a 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-based feature extraction, followed by a Neural ODE and 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 al. (2019), which consist of a 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 6 illustrates this process. We detail other configurations in Appendix A.4.
ODE solver. We use standard Runge-Kutta 4(5) adaptive solver (dopri5; Dormand & Prince (1980)) implemented by the torchdiffeq package. The numerical tolerance is set to 1e-6 for CNF and 1e-3 for the rest. We fix the integration time to [0, 1] whenever it appears as a hyper-parameter (e.g. for image and CNF datasets3); otherwise we adopt the problem-specific setup (e.g. for time series).
Training setup. We consider Adam and SGD (with momentum) as the first-order baselines since they are default training methods for most Neural-ODE applications. As for our second-order SNOpt, we set up the time grid {tj} 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 et 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.

heading: 4.1 Results
text: Convergence and computation efficiency. Figures 1 and 7 report the training curves of each method measured by wall-clock time. It is obvious that our SNOpt admits a superior convergence
3 except for Circle where we set [t0, t1]:=[0, 10] in order to match the original setup in Chen et al. (2018).
Figure 8: Relative runtime and memory of our SNOpt compared to Adam (denoted by the dashed black lines) on all 9 datasets, where ‘Mn’ is the shorthand for MNIST, and etc.
rate compared to the first-order baselines, and in many cases exceeds their performances by a large margin. In Fig. 8, we report the computation efficiency of our SNOpt compared to 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-order 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 the vanilla state dynamics, 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 dataset. However, the actual increase in memory (less than 1GB for all datasets; see Table 10) remains affordable on 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 performance, including the accuracies (%) for image and time-series classification, and the negative log-likelihood (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 our method not only converges to higher
Table 5: Performance of jointly optimizing the integration bound t1 on CIFAR10
Method Train time (%)w.r.t. t1=1.0 Accuracy (%)
ASM baseline 96 76.61 SNOpt (ours) 81 77.82
0 5k 0.5
1.0
t 1
t1 Optimization
ASM baseline SNOpt (ours)
t1=1.0
Train Iteration
Figure 10: Dynamics of t1 over CIFAR10 training using different methods.
Table 6: Measure of implicit regularization on SVHN
# of function Regularization evaluation (NFE) ( ∫ ‖∇xF‖2 + ∫ ‖F‖2)
Adam 42.1 323.88 SNOpt 32.6 199.1
Mn Sv Cf SA AW CT 1 3 5
Relative Runtime ( )
Mn Sv Cf SA AW CT 0 5
10 15
Accuracy Improvement (%) ( )SNOpt RecursiveAdjointRecursiveAdjoint / SNOpt
Figure 11: Comparison between SNOpt and secondorder recursive adjoint. SNOpt is at least 2 times faster and improves the accuracies of baselines by 5-15%.
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-parameters.
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-order feedback policy (16) derived in Section 3.4 to the first-order ASM baseline proposed in Massaroli et al. (2020). It is clear that our OCP-theoretic method leads to substantially faster convergence, and the optimized t1 stably hovers around 0.5 without deviation (as appeared for the baseline). This drops the training time by nearly 20% compared to the vanilla training, where we fix t1 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-time 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), our method may run slightly faster than first-order methods. This is a distinct phenomenon arising exclusively from training these continuous-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, ∫ ‖∇xF‖2, and Kinetic, ∫ ‖F‖2, 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-conditioned ODE dynamics) compared to the one found by Adam. This provides a plausible explanation of the reduction in the NFE when using our method, yet without hindering the test-time performance (see Table 4).

heading: 5 Conclusion
text: 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 in 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 O(T ). 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 et al., 2021b) to deep continuous-time models, yet using a 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/PDEs.
4 In Appendix A.4, we provide some theoretical discussions (see Corollary 9) in this regard.
Acknowledgments and Disclosure of Funding
The authors would like to 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 #1932068, and Tianrong Chen was supported by ARO Award #W911NF2010151.

Review summaries:
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.

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.

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.