Faster Training of Neural ODEs Using Gauß–Legendre Quadrature
Authors: Alexander Luke Ian Norcliffe
Marc Peter Deisenroth
Venue: Accepted by TMLR ()

Abstract: Neural ODEs demonstrate strong performance in generative and time-series modelling. However, training them via the adjoint method is slow compared to discrete models due to the requirement of numerically solving ODEs. To speed neural ODEs up, a common approach is to regularise the solutions. However, this approach may affect the expressivity of the model; when the trajectory itself matters, this is particularly important. In this paper, we propose an alternative way to speed up the training of neural ODEs. The key idea is to speed up the adjoint method by using Gauß-Legendre quadrature to solve integrals faster than ODE-based methods while remaining memory efficient. We also extend the idea to training SDEs using the Wong-Zakai theorem, by training a corresponding ODE and transferring the parameters. Our approach leads to faster training of neural ODEs, especially for large models. It also presents a new way to train SDE-based models.

Published in Transactions on Machine Learning Research (08/2023)
Faster Training of Neural ODEs Using GauB-Legendre Quadrature
Alexander Norcliffe* University of Cambridge
alin2ccom.ae uk
Mare Peter Deisenroth University College London
m.deisenroothhuccl.ccc uk
Reviewed on OpenReview: htips: // openreview. net/ met/ forum? id- JOFSDAy1bU
Abstract
Neural ODEs demonstrate strong performance generative and time-series modelling. However, training them via the adjoint method is slow compared to discrete models due to the requirement of numerically solving ODFs. To speed neural ODEs up, a common ap- proach is regularise the solutions. However, this approach may affect the expressivity of the model; when the trajectory itself matters, this is particularly important. In this paper, propose alternative way to speed up the training of neural ODEs. The key idea is to speed up the adjoint method by using Gauß Legendre quadrature to solve integrals faster than ODE-based methods while remaining memory efficient. We also extend the idea to training SDEs using the Wong Pakaii theorem, by training a corresponding ODE and trans- ferring the parameters. Our approach leads faster training of neural ODEs, especially for large models. also presents new way to train SDE-based models.
Code is available at:
sqjadii
An associated video presentation is available at: http::/wwwwyouuubb.o pKbLlsqy8aM
1 Introduction
Neural ODEs E 2017 Chen et al. 2018) make an explicit connection between deep feedforward neural networks and dynamical systems. They take inspiration from the residual network architecture (He et al. 2016), where the le + 1-th hidden layer related the k-th by
2k-1 zx + f(zk, 0x),
where 0N parameterises the learnable function f at the k-th layer. Here, zu E IR for l,..., K, iee, the dimensionality of the hidden units does not change as we progress through the K layers.
We interpret equation 1 as the Fuler discretisation of contimuous-time dynamical system with step size of 1. Neural ODEs are the limit of infinitesimally small step sizes, by directly parameterising the instantancous rate of change of state with neural network at timee t via
dz dt = [22,,,
(2)
Equation defines continuous path between the input x(tn) at an initial time to and the state =(tx) at given later time tg, which is evaluated using black-box ODE solvers.
MMaoriity of work done student University College London September 2020 September 2021.

---

Published in Transactions on Machine Learning Research (08/2023)
Due to their continuous nature, neural ODEs are suited to dealing with irregular time series (Chen et al, 2018 Kidger et al. 2020; Rubanova et al. 2019) allowing for analysis to be carried out without binning data. Neural ODEs also include the inductive bias that the data is generated from a dynamical system. This is particularly applicable in the natural sciences, for example, predator prey dynamics described by the Lotka Volterra equations (Norclifle et al. 2021 Dandekar el al. 2020). Additionally, neural ODEs can be applied to generative modelling. Continuous normalising flows (Chen et al. 2018 Grathwohl et al.] 2019) build on normalising flows (Rezende & Mohamed 2015) by allowing the flow be defined by a vector field rather than discrete change. Stochastic differential equations have also been used for score-based generative
modelling (ong et al. 2020 Vahdat et al. 2021) showing strong results.
Neural ODEs are trained by calculating gradients of a scalar loss function L and using first-order optimiser, such as stochastic gradient descent, Adam (Kingmaa & Ba 2015) or RMSProp (Tieleman & Hinton 2012)
Gradients can be calculated by directly backpropagating through the operations the ODE solver has taken; we call this the 'direct method'. An alternative method the adjoint method (Pontryagin) 1987 LeCun al.
1988 Chen et. al.) 2018), which solves a second ODE backwards in time. The direct method is faster but more memory intensive than the adjoint method (we include complexities of time and memory in Appendix C in practice a tradc-off must be made. The focus in this work is building on the adjoint method.
Solving these ODFs in general can be slow due to the computation demand of running an ODE solve [Lehtiäkii et al. 2022). Speed is of particular interest when compute is limited, or more broadly for environmental concerns. We focus on this problem in this paper. Specifically, we focus on speeding up the training procedure of neural ODEs by providing a faster way to carry out the adjoint method while retaining its favourable memory footprint. The key idea is to use numerical integration, Gauß GauB-Legendre quadrature, to solve one-dimensional definite integrals, which appear when computing gradients via the adjoint. These integrals are usually solved as differential equations, which slows down training considerably, compared with the GauB- Legendre quadrature, which we propose using. This is different way to approximate the solution to the same adjoint equation and should therefore produce the same gradients within numerical precision.
Related Work There are various methods for speeding up neural ODFs. We can directly speed up the forward ODE solve by running it in a lower dimensional space using Model Order Reduction (Lehtimäki et al. 2022). We can also use regularisation techniques that the trajectorics are less complex, allowing the ODE solve to be faster (Onken et al. 2021). Kinetic regularisation (Finlay el al., 2020) includes terms in the loss function that penalise large velocities and Jacobians, Another method is to penalise large higher-order derivatives, such as the acceleration of the state, so that the learnt dynamics are easier to solve (Kelly et al., 2020j. We can also include heuristics calculated by the solver in the loss function (Pal et al.) 2021). This reduces the complexity of the dynamics by penalising small step sizes or large st.iffness estimates. Further to this, we can also regularise the solutions without including any new terms in the loss. This can be done by uniformly sampling a terminal time tx (Ghosh et al., 2020). This encourages the model to reach the final state quickly (if the sampled end time is early), and then stay there (if the sampled end time is later), regularising the dynamics. Finally, we can directly restrict the dynamics to be faster to solve, for example (Xia et al., 2021 Nguyen et al.] 2022), introduce damping style terms to the dynamics to regularize the ODE solutions. All these regularisation methods make inference faster, because solving the ODE is easier with simpler dynamics. However, this introduces restrictions on the trajectory, reducing expressivity of the model. This is particularly important in time-series applications, where the trajectory matters.
Instead of regularising the solutions, we can speed up the training process directly. Daulbaev et al. (020)) use barycentric Lagrange interpolation to approximate the trajectory to solve the adjoint equations. This does not require a backwards ODE solve of 2, which therefore speeds up the training process. However, this does require checkpoints of the trajectory z to be stored during the forward solve, making the met less memory efficient than the adjoint method.
Seminorms Kidger et al., 2021a) are the closest method to our work (that remains memory-effiiien) by recognising the adjoint equations contain integrals. During the backward solve only subset of the overall state is used by the ODE solver to estimate the truncation error, which is used to select the stop size. This enables larger step sizes and faster solves. However, the adjoint integrals are still solved with ODE methods, which uses more computation than required. Rackauckas et al. (2020) use GauB -ornroo quadrature in place
2

---

Published in Transactions on Machine Learning Research (08/2023)
of this to solve the integrals. However, this firstly requires more terms than other Gaussian quadratures and secondly requires dense solutions of the state and the adjoint state to carry out the quadrature, i.e., many checkpoints must be taken. This makes the method very memory intensive, which removes the advantage of using an adjoint based method over the direct method.
Contributions (i) We introduce the GQ method, memory efficient and fast way to carry out the adjoint method using GauB Legendre quadrature, speeding up the training of neural ODEs, (ii) We show that this method can also be used to train neural SDEs using the Wong Wong-Zakai theorem.
2 Neural ODEs
A general neural ODE consists of dynamics function with learnable model parameters 0, an encoder with learnable parameters and decoder with learnable parameters Φ. For input æ and given measurement times fto. f1...,,xx txh, the prediction at time & is given by (ttx) by solving the ODE
dz
z =(4) = ho ()),
dt (,,,,,))
a(xx) = ha(=(()),,)))
where the encoder he and decoder ha are general functions. They can be identity operations so that the ODE can be thought of running in observation space, or they can be learnable functions so the ODE is latent [Rubanova et al. 2019
Neural ODEs are trained by obtaining gradients of a scalar loss function L and using a first-order optimiser of choice. ODE solvers, such as the Runge Kutla solvers, consist of algebraic operations, all of which are differentiable. Therefore, one can directly backpropagate through an ODE solve to obtain gradients. However, when the dynamics are complex, this can lead to an arbitrarily large number of function evaluations for adaptive solvers, storing all of the intermediate activations during the solve, and the method becomes prohibitively memory intensive.
A less memory intensive method for training is using the adjoint method (Pontryagin 1987 LeCun et al.
1988
Chen et al. 2018 This introduces the adjoint state az, which obeys the (backward-inn-imme) ODE
da --a= dt
-s Oz
0s
OL
(x(t)) dz(tk)
(4)
dl.) is given by (=(()) and is used to train the encoder. A third state ag is used to caleulate the gradients for the dynamics function. It obeys the ODE
dag dt a de
Of
at(t))= = 0.
The gradients of the loss L with respect to the model parameters O are given by (to) In the standard implementation of the method, these gradients are found by first solving the forward ODE for =(1) up to tx, then solving the following concatenated ODE backwards in time
f(z, t, 0) a V2l
z(tx) V
d
a
(tk) = az [g
V((tx) L
(6)
di
ag
az Vof
0
No intermediate activations are stored and so the method is memory efficient in the integration time. A fourth state ag can be used to calculate gradients associated with measurement times. However, measurement times are typically not learnable functions"] and gradients are not required. Further detail can be found in the original neural ODE paper (Chen et al. 2018). The method developed in this work directly applics to the case where we learn measurementt times, and this is included in our code.
An important observation made by Kidger et al. (2021a) is that the differential equation for ag does not actually contain ag. When the trajectorics of z and a are known, åg a Vof only depends on time.
A noteable exception is adaptive depth neural ODFs Massaroli al. 2020
3

---

Published in Transactions on Machine Learning Research (08/2023)
They introduce seminorms take advantage of this; this is when the ODE solver during the backward solve does not consider dg when calculating the error to choose step size, because the error in ao does not grow significantly compared to 2 or az (Kidger et al. 2021a).. While this allows for larger step sizes and faster training it still requires a= Vof to be calculated at every step. Instead we can rewrite the ODE solve in a way that it calculates the gradients associated with the parameters as a definite integral
dL - -
0f de - 00
f.e rtx
af(t) 00 dt.
T
a²(())
(7)
d0
xx
We can then use more advanced methods for solving definite integrals to compute the desired gradient faster and to the same level of accuracy. We only need to solve a smaller ODE for [z, az). This will become significant when there are many parameters, making ag large. Importantly, despite having many parameters, the integration variable t for computing the gradient in equation 77 is one-dimensional. Hence, we can solve many (casy) 1D integrals in parallel (one for each paramcter), allowing us to use fast and accurate methods for solving ID definite integrals. That is the key idea behind our approach, which we detail in the following,
3 Faster Training of Neural ODEs
In the following, we describe the methodology for faster training of neural ODEs. The key idea is to use more appropriate methods to caleulate the definite intograls in the adjoint method given by equation 7 specifically we use GauB Legendre quadrature.
3.1 GauB-Legendre Quadrature
Gaussian quadrature is numerical integration method for calculating low-dimensional integrals as weighted sum of function values
n _d w (()) Σ (n)(
in the integration domain for given weights wi: locations Ti, and integrand f(t). Weights and locations are determined based on the locations of zeros of given (orthogonal) polynomials (Stoer & Bulirsch 2002).
Gaussian quadrature using terms in equation of polynomials of up to degree 2n - Therefore, if the integrand can be well approximated by such a polynomial, this shall produce a good approximation to the integral. Gaussian quadrature is the fastest method for 1-D integrals outside of analytical solutions, it converges significantly faster than methods such as a Riemann Sum or the trapezoid rule; in particular solving with a differential equation solver requires (t) afit) 80 to be calculated more times than when using Gaussian quadrature (we give error bounds for the different integration methods in Appendix D). Therefore, we use Gaussian quadrature.
( is guaranteed to obtain the exact result for integrals
Why Gauß Legendre quadrature? There are various flavours of Gaussian quadrature such as GauB GauB--Jcobi. We specifically use Gauß-Legendre quadrature over other Gaussian quadratures for multiple reasons. Firstly, our integration intervals are both finite, ruling out GauB-Laguee quadrature which works on the [0, 20) interval and GauB Hermite quadrature which works on the (-00, oo) interval. Secondly and more specifically our integration interval is [to, tx]. it includes the start and end time, rather than (to, tx), this excludes SauB-Jacobi quadrature and GauB Chebyshev quadrature of the first kind. Thirdly, in general, Gaussian quadrature uses weight function so that I Shd (t)dd.
Laguerre,
Why do this? Consider the integral S2 v1 - t² f(t), this can be solved as f2 w(t) - t2 ftt)tt so that w(t) 1 and we use Gaug-Legendre or we can solve it as [_1 f()) where w(t) -1-t² and we use GauB Chebychev of the second kind. If we can write the integrand more simply as the product of a weight function and other function we may use other quadratures. In our case the dynamics function is a neural network and there no clear way to write the dynamics as this product so we let w(t) = 1, which rules out GauB CheeySee quadrature of the second kind. We do not use GauB Kronrod quadrature, an adaptive quadrature, because this requires the integrand to be stored at many locations in order to take advantage of
4

---

Published in Transactions on Machine Learning Research (08/2023)
the adaptive nature of the scheme. This removes the memory efficiency of the method, and if memory is not constraint it is often more useful to use direct backpropagation over any adjoint based method. Finally, we do not use other schemes such as GauB Lobatto quadrature because it does not converge as quickly, rather than n points exactly solving a degree 2n - polynomial it only accurately solves up to degree 2n 3. After ruling out all of these option,s GauB Legendre quadrature is the only choice left.
3.2 Memory Efficiency
Ii is vital when using an adjoint-based method to be memory elficient in integration time. Otherwise directly backpropagating through the solver is the best method. This is the key flaw with the GauB Kronrod implementation (Rackauckas et al.) (202)), in that it requires dense solutions of and azy consuming large amounts of memory. We can achieve memory efficiency by using a running total g and add terms in the quadrature during the solve, rather than all at the end. That is, we initialize g as vector of zeros O, and at each point in the quadrature sum Tr, we update to be g + w;aa Vef(r.) that way only needing two vectors at any point, the running total and the updale. Without dense solutions, it is non-trivial lo arbitrarily add terms to the sum because an ODE must be solved to query the integrand at time t. Therefore, we have to determine how many terms to use before starting the backward solve.
Preferably, the number of terms chosen adapts to the complexity of the problem, so that the smallest number of terms is used to compute accurate gradient. Additionally, any information used to determine this must be gathered during the forward solve, and must be cheap to gather. Otherwise we do not reduce the computation, we just move it to a different process.
We opt for an empirical approach. During the forward solve, for effectively no extra computation, we count how many times the dynamics function is ovaluated (NFE). When using an adaptive solver, the larger the NFE, the more complicated the trajectory. We make the assumption that if the (forward) trajectory of z is complex then the corresponding gradient trajectory Vof is also complex, and more terms are required in the quadrature calculation. Considering the integration interval, larger intervals [ts-1, Lx] likely require more quadrature points to achieve a high accuracy. Knowing how these quantities should affect the number of terms, we propose the heuristic
n C NFE x (1& -ts-1)
(tx- to)
where C' is user chosen constant that tells us how the number of GQ terms scales with the NFE. If C is larger, then require more terms. We apply the ceiling function to ensure is both integer and rounded up rather than down. We also enforce an upper bound of 64 quadrature points, which allows us to prescribe a polynomial of degree 127, to prevent a large number of terms slowing down the training. This does not negatively affect model performance in our experiments, see Appendix E.2 for an ablation study.
3.3 Algorithm
Algorithm [ outlines an implementation of the GQ method. After initialising the state, the adjoint state and the running total, we then calculate the number of terms required using equation ( The weights and locations for Gauß Legendre quadrature are obtained by accessing a lookup table. We then loop through the weights and locations, solving the ODE up to that point, calculating a='Vo VoJ and adding to the running weighted sum. Finally we complete the backwards solve to obtain (to).
If the loss L(t(tu), (t.))...., #(tx)) depends on the state at multiple measurement times, there is a discon- tinuous change in the adjoint stale az at those times:
ØL
a=(tx) az (tt) dz((x)
(10)
This discontinuity is accounted for by splitting the large integral into smaller integrals where the integrand is continuous and stepping through (for pseudocode we refer to Appendix
AK
T 0s dt
K le
as dt.
a 00
=
a
(11)
00
k=
12-1
5

---

Published in Transactions on Machine Learning Research (08/2023)
Algorithm Memory efficient implementation of the GQ method.
al 2 z(1x), az Jatk)) ) g g
n - GotNTerms(to, tx. NFE, C. C)
w, T + GerGGWWiighsssooctttoons//// to, tx)
j +- n
prev f tk
while j 2 1 do
next Tj
f T 0f 8z
ODESolve(
'prevs )
0
az
a
prev
f fnxx
Tr
+ W/ x az 80
1
end while
f
ODESolve(
Aprev, to)
T o5 0= 85
ag
0x
return 9, az
Overall, have devised fast and memory-effieient way to compute gradients in neural ODFs via the adjoint method. The key idea was to formulate the gradicnt of the loss w.r.t. the model paramcters definite integral, and then to use Gaussian quadrature to solve this integral efficiently.
4 Faster Training of Neural SDEs
Neural ODEs have also been extended to learning stochastic diffcrential equations (SDEs) in the form of
neural SDEs (Tzen & Raginsky] 2019: Liu et al. 2019 Xu et al. 2021 Li et al. 2020 Kidger el al. 2021c b). A Stratonovich SDE is given by
d= O))tt + g(z,t,0) dW,
(12)
where f and g are learnable functions of the drift and the diffusion and W, is a Wiener process. Neural SDEs can be applied to situations where there is inherent noise in the system, an underlying random process affecting the nymamiss, for example in mathematical finance [Black & Scholes 1973).
Training neural SDEs suffers from similar speed and memory problems as neural ODEs. The GQ method from Section 3 can be extended to training neural SDEs by training corresponding ODE and transferring
the parameters to SDE solver at test time. The Wong Zakai Wong-Zakai theorem (Wong& & Zakai 1965) tells that
can approximate an SDE using an ODE. For the Stratonovich SDE in equation 12 consider the ODE
dzm
dBm.. dt
dt f(zm. t, 0) + g( Zm t,
(13)
where Bm (I) is a smooth approximation a Wiener process determined by integer m that improves as m gets larger. An example could be polynomial with degree m (Foster et al. 2020). We include subscript on 2m to show this is an approximation to the SDE. The Wong Zakai theorem states that if Bm, (1) + W do, then the solution *t) to the ODE converges to solution =(t) of the SDE This allows
to approximate an SDE using an ODE (Londo & Villegas 2016 Filip et al. 2019) and use the methods for training neural ODEs, such as our GQ method, to train neural SDEs and transfer the parameters at test time (Hodgkinson et al. 2021
Our application requires the approximation of the Wiener process to be smooth and memory efficient integration time; this prevents using interpolations through m points, for example. We therefore use the Karhunen Loève theorem, which says a Wiener process on the domain (No, [x] can be approximated by the Fourier series
sin( )) -to)/c)
W. = 2
&
(14)
(i - 1
1=1

---

Published in Transactions on Machine Learning Research (08/2023)
where c =tx-- to and &. ~ I). We then differentiate this approximation, which yields
dB,, dt
2 1=1
(2 ( to)
E cos
(15)
This is not significantly large memory overhead, and & do not have to be resampled for each evaluation. However, if memory is scarce, can further make this memory efficient by calculating equation (15) using running sum, so that we iteratively sample & and then add the cosine, rather than storing all &:- We use pseudo-random numbers and re-seed a generator to guarantee we recover the same (t) throughout the forwards and backwards solves.
5 Experiments and Results
In this section, we compare the GQ method against the direct, adjoint and seminorm methods. The main aim of this work is to speed up the training of neural ODEs and SDEs, without compromising performance. Therefore, the key metric to compare methods is the training time, provided the final performances are the same, discuss the number of function evaluations and why this is not the most reliable metric in our case in Appendix E.1 The GQ method is implemented in PyTorch building on the torchdiffeq library (Chen) et al.
2018 Code is publicly available at: htttps://githb. com/ wooorliiffettorrc_see..eeee
All experiments were run on NVIDIA GeForce RTX 2080 GPUs. We use the Dormand-Prince 5(4) solver with an absolute and relative tolerance of 1 x 10 and = 0.1 for the GQ method. We found that this value worked well, giving accurate gradicnts quickly, we run an ablation on this value in Appendix E.2 We train using the Adam optimiser (Kingma & Ba 2015). For classification lasks we use constant integration times of [0, 1]. Additional experimental details such as task-dependent learning rates, information datasets and further experiments (such as test performance against wall-clock time) are given in Appendix
5.1 Accuracy of Gradients
For Neural ODEs, it is difficult to determine the true gradient because the forward pass is approximated. To test the GQ method in an analytical setting, we have recreated and extended the experiment in ACA (Zhuang et al. 2020) and MALI (Zhuang et al. 2021). Here we use an analytical system that we can control with exact expressions for the loss and gradients. We consider the exponential growth system == az + x(1 = 2o exp(al), with a loss L = =())² = z exp(2aT)). The gradients with respect to initial condition 20, parameter and integration time T are: 250 exp(2aT)) 2T 2l exp(2aTT) and 20z exp(2aT") rospectively. In Figure [ we plot the relative errors True-Preddeted in the loss and the different gradients for True
various methods as the integration time T is increased. We plot relative errors since use exponentially growing system and want errors for small T to be equally as visible as errors for large T. We use value of 0.2 for 10.0 for zo and a range of 19-29 for T. The GQ method (blue) produces the same gradients as the standard adjoint method (pink), demonstrating that this is as accurate as the adjoint method, which supports the claim it is solving the same equation in a different way. We also see that the direct, standard and GQ adjoints give the same loss as we expect since they have the same forward solve. Importantly we see that the direct method gives different gradients to the adjoint methods which we expect since they are thereiizz-ttee--otiiiieee vs optimize-then-deeee diseerieetie.. MALI and ACA have been
two different approaches included to show how using different methods for the forward solve can change the results.
5.2 Memory Efficiency
We additionally investigate how much memory is used by each method on the analytical task. We test the direct, standard adjoint, GQ adjoint and ACA methods. We solve the same task up to T = 15 and then take the gradient using each method. The maximum memory consumption is then determined in bytes, and plotted against the solver tolerance (lower tolerance means more steps to solve). We plot these in Figure We see that as expected the GQ method constant in the tolerance, and the number of intogration steps. Adjoint is also efficient, but uses slightly less memory. We hypothesize that this is because the adjoint method uses one function evaluation to calculate ((2,1,0). a² a=V=I, a Vol), whereas GQ uses
g
7

---

Published in Transactions on Machine Learning Research (08/2023)
Relative Error in Loss
Relative Error in dL/dz0
1e-5
1e-6
1.2
Direct
1.0
Adjoint
Relative Error
GQ
Relative Error
MALI
ACA
0.4
0.2
0.0
20
22
24
26
28
20
22
24
26
28
Integration Time
Integration Time
Relative Error in dL/da
Relative Error in dL/dT
3.5
3.0
Relative Error
Relative Error
2.5
2.0
1.5
1.0
0.5
0.5
20
22
24
26
20
22
24
26
Integration Time
Integration Time
Figure Absolute relative errors in losses and gradients for the toy gradient problem. We see that the GQ method performs well relative to the other well established methods, showing that it produces reliable gradients.
one evaluation to calculate so there is a slightly higher memory usage for GQ, the main takeaway is that it is constant the tolerance. ACA and the direct method scale poorly as expected.
-a? -(2 and another to calculate a Vol at the quadrature point,
5.3 Nested Spheres
The Nested Spheres experiment (Dupont et al. 2019 Massaroli et al. 2020 is an illustrative classification task. The function d dimensions consists an inner sphere entirely surrounded by an outer sphere
0,
0
<
x S r ac 1a S r3
g(x)
(16)
T2
for 0 <rr < < rs. Vanilla neural ODEs cannot solve this problem, because any mapping preserves topology, preventing the two regions being linearly separable (Dupont et al. 2019). This can be seen in a plot of example 2D data in Figure
We train an augmented neural ODE Dupont et al. 2019) with three augmented dimensions for 100 epochs, with a batch size of 200 using the cross-entropy loss. We use a time-dependent,, multilayer perceptron with softplus activations and two hidden layers as the dynamics function. We use different hidden widths test the GQ method with varying numbers of parameters. We train across 10 seeds with results given in Figure Training times recorded the time taken to complete 100 epochs.
8

---

Published in Transactions on Machine Learning Research (08/2023)
1.00
0.75
log(Memory Usage/bytes)
Direct
12.5
ACA
0.50
12.0
Adjoint
0.25
GQ
0.00
11.5
11.0
-0.75
10.5
-1.00
3.0
3.5
4.0
4.5
5.0
5.5
6.0
6.5
7.0
-0.75
-0.50
-0.25
0.00
0.25
1.00
-tol(tol)
Figure 3: Example 2D data from the nested spheres training set. The two regions can- not be linearly separated by vanilla neural ODE.
Figure Maximum memory usage by each method on the analytical task.
110
Direct
Training Time (/s)
500
Adjoint
Test Accuracy
105
400
Seminorm
100
GQ
300
100
500
1500
2500
3000
500
1000
1300
3000
Hidden Width
Hidden Width
(a) Training times.
(b) Test accuracy.
Training times and test accuracies on the nested spheres task. The GQ method's training time scales well with model size compared to the adjoint and seminorm. The test accuracies are all comparable showing the method produces accurate gradients.
Figure 4:
Directly backpropagating has the fastest training times, as expected. When the number parameters is low, the methods have comparable training times. However, as the number of model parameters increases, the training times for the adjoint and seminorm methods increase more significantly than the GQ or direct methods. We also see that the models have the same accuracies on the test set (within standard deviation) showing that the gradients produced by the GQ method are accurate.
5.4 Time Series
To test neural ODEs trained using the GQ method on time series, consider sine curves -- where the underlying ODE is second order. We train second-order neural ODE (Norclifee al. 2020 Massaroli et al. 2020), where the initial position and velocity are given. Therefore, the model only has to learn the acceleration. We train using MSE as the loss for 250 epochs and a batch size of 15. This is repeated over 5 seeds to obtain means and standard deviations; results are given in Figure
We see similar results as those the nested spheres experiment. The direct method is the fastest. The time taken to train is about the same for small model between the /smminorm methods and the GQ method. However, the models get larger, the time taken scales far more favourably for the GQ method
9

---

Published in Transactions on Machine Learning Research (08/2023)
Direct
Training Time (/s)
Adjoint
300
Seminorm
Test MSE
GQ
100
250
500
1500 1750 2000 Hidden Width
750
1000 1250
250 500
1500 1750 2000 Hidden Width
750
1000
1250
(a) Training times.
(b) Test MSE.
Figure 5: Training times and test MSEs on the sines task. The GQ method's training time scales well with model size compared to the adjoint and seminorm. The test MSEs are comparable showing the method produces accurate gradients.
than the adjoint and seminorm methods. We also see that the final test MSEs are the same, with the same standard deviations, showing the methods produce the same gradients.
5.5 Image Classification
We now test the GQ method on more difficult classification. We consider image classification using the MNIST (LeCun et al. 1998) dataset. Further results are given on the CIFAR-10 (Krizhessky 2009) and SVHN (Ntzerr et al.J 2011) datasets in Appendix E We use convolutional initial downsampling layers and dynamics function and a fully connected final set of layers to obtain logits. We train for 15 epochs with batch size of 16. We train using three random seeds to obtain means and standard deviations. We do not consider the direct method due to memory consumption. Results are given in Figure 0
17500
Adjoint
98.5
Training Time (/s)
15000
Seminorm
Test Accuracy
98.0
12500
97.5
10000
97.0
7500
5000
96.5
2500
96.0
100
200
300
400
500
100
200
300
400
500
Hidden Channels
Hidden Channels
(a) Training times.
(b) Test accuracy.
Figure 6: Training times and test accuracies the MNIST task. The GQ method's training time scales well with model size compared to the adjoint and seminorm. The test accuracies are all comparable showing the method produces accurate gradients.
As in the previous experiments we see that as the model (and the corresponding number of parameters) becomes large, the training time scales more favourably for the GQ method. The accuracy does not change significantly between the methods either, that the method produces accurate gradients.
10

---

Published in Transactions on Machine Learning Research (08/2023)
5.6
Process
Finally, consider how method extends to training neural SDEs, by considering the Orastein-Uhlenbet (OU) process (Uhlenbeck & Ornstein 1930 governed by the SDE
Kidger et al. 2021c A onc-dimensional OU process is
dr frddt + dW..
(17)
We consider a harder variation of the SDE with time-dependent drift and diffusion
dx = (μt 0x)dt + (a + Φt) o dW.
(18)
for scalar parameters μr, 0, and Φ and one-dimensional Wiener process W₁. We train and evaluate using the KL divergence and a batch size of 40. We do not use an encoder or decoder; the model acts in observation space. The drift and diffusion are separate time dependent multilayer perceptrons, with two hidden layers. We use 10 cosines to approximate the Wiener process when training the corresponding ODE, we run an ablation over this value in Appendix Fl We compare our method to directly backpropagating through an SDE solver and with the SDE adjoint method (Li al., 2020), using the reversible Heun method (Kidger et al. 2021b) allowing for fast and reversible SDE solves. We use relatively large step size of 0.01 for the SDE solvers to reduce their training time. We use the same SDE solver to evaluate parameters trained by
the GQ method. Training times and test KL divergences are given in Figure i
2000
Test KL Divergence
0.10
SDE Direct
Training Time (/s)
1800
SDE Adjoint
0.09
GQ
0.08
1400
1200
0.07
1000
0.06
800
0.05
100
100
Hidden Width
Hidden Width
(a) Training times.
(b) Test KL divergence.
Figure 7: Training times and lest. KL divergences for different widths the OU experiment. The SDE direct method is fast, but runs out of memory, the GQ method is faster than the SDE adjoint method. The methods have comparable test KL divergences.
The SDE direct method does not record values for hidden width of 100. This is because the GPU ran out of memory. We see before this, that the direct method is the fastest, due to not having to solve a backwards SDE (or ODE) to obtain gradients. The GQ method is second fastest with the SDE adjoint method being slowest due to the time taken to solve the forward and backward SDEs. We do not see the previous effect of model size on training time because we are still in the 'small parameter regime', which is why the training time fluctuates rather than increasing with model size. The KT. divergences for each method are approximately the same and within a standard deviation of each other, showing that the GQ method produces accurate gradients for neural SDEs well ODEs. The small difference likely comes from training corresponding ODE, rather than the SDE directly.
6 Discussion and Practical Guidelines
We saw that the GQ method a promising alternative to the standard adjoint method for training neural ODEs. One of the significant advantages of the GQ method that it speeds up the training of neural ODEs without affecting the dynamics. Tn this way the GQ method is orthogonal to regularisation schemes, such
11

---

Published in Transactions on Machine Learning Research (08/2023)
STEER Ghosh et al. 2020 kinetic regularisation Finlay al. 2020), regularising higher-order terms Kelly et al. 2020) and training with solver heuristics PPal et al. 2021
When the model has many parameters, the GQ method can be significantly faster than the standard adjoint method depending on the state size. Additionally, the GQ method enforces no further restrictions, so it can be used to replace any existing models trained with the standard adjoint method, and it will speed up models without limiting their expressivity. We can also apply this method to train SDEs and, depending on the step size, the GQ method can be faster than the SDE adjoint method.
However, when the state is large, the improvement in speed is either negligible compared to the adjoint method or it can even be slower as shown in Appendix D There are also no theoretical guarantees on the heuristic used to calculate the number of terms in equation ( We did not experience cases where the method does not produce accurate gradients; however, that does not mean they do not exist. Hence, the GQ method should be used in situations with large models, and small states. memory capacity is large, then the direct method will be the fastest method. In the case where the model is small and the stale is large, the adjoint method will be the best method. The same holds for the SDE versions of these methods.
7 Conclusion and Future Work
We looked at how the training of neural ODEs can be made faster. We showed that the adjoint method, the current method of choice, inelficiently solves a definite integral using ODE methods, contributing to the slow training. We showed that we can use SaauB-Legenrre quadrature to speed up the training and showed how to make this method memory efficient, a key requirement of neural ODE training at scale. Following this, we extended the method to training neural SDEs using the Wong- Wong-Zakai and the Karhmnnn- thcorems.
We tested both the speed of the GQ method and the reliability of its gradients on classification and time-series regression, In all of our experiments, we saw that training using the GQ method produces the same model performance as training with the direct or adjoint methods for ODEs and SDEs. We also saw that when the model size is large and the state is small, the GQ method outperforms the adjoint method significantly in terms of training time.
One of the significant advantages of the GQ method is that it speeds up the training of neural ODEs without affecting the dynamics. In this way the GQ method is orthogonal to regularisation schemes, such as STEER Ghosh et al. 2020), kinctic regularisation (Finlay et al. 2020), regularising higher order terms (Kllyy et al. 2020) and training with solver heuristics (Pal et al.j 2021). Therefore, it would be interesting to see how much the method would speed up if these techniques were used in conjunction with the GQ method.
Broader Impact Statement
In our work we investigate speeding up neural ODEs using Cauß-Legendre quadrature. We also extend the method to training neural SDFs. Our work is incremental and focuses on the speed of training rather than model performance. Therefore, we envision any large societal impacts coming more from applying differential equation models, rather than using our specific method to train.
Applications Neural ODEs and neural SDEs are still relatively new models and are yet to be used in any unethical situations or on any large scale to our knowledge. Differential equation based models have seen success in both time-series modelling and generative modelling. They will likely be applied further in these areas, such predicting the evolution of dynamical system in natural sciences, or generating images using coore-based modelling. The models are still at risk of learning biases within a dataset. However, they are no more at risk of this than discrete models. In fact, there is evidence to show that they may be more robust in many cases than discrete models Yan et al. 2020). due to the property of no crossing trajectorics and preserving topology (Dupont et al. 2019).
Environmental Implications The contributions of this paper are speed up existing methods for train- ing neural ODEs. Therefore we anticipate a positive impact on the environment because less computation is needed to produce the same results for the same model.
12

---

Published in Transactions on Machine Learning Research (08/2023)
Datasets We do not use any datasets that could be considered to be sensitive. All experiments used synthetic data, with the exception of the well-known image datasets: MNIST, CIFAR-10 and SVHN.
Code Release
Our code is publicly available at: httes///ggttiiacccc/rr/ooo/rr//r
Author Contributions
Alexander Norcliffe: Lead author. Developed the idea and full implementation. Ran the experiments. Joint effort writing the manuscript. Wrote the author responses during review.
Mare Deisenroth: Senior author. Initial broad conceptualisation of the project and provided guidance on research directions. Joint effort writing the manuscript. Provided significant input to author responses. Provided access to the hardware to run the experiments.
Acknowledgments
At the time of publication, Alexander Norcliffe is supported by a grant from GlaxoSmithKlinc. We thank So Takao for his help understanding the Wong-Zaks theorem for the SDE adaptation of the GQ method. We would like to thank the anonymous reviewers and action editor Kevin Swensky for their time and efforts to review and constructively critique the paper. Our implementation relied heavily the torchdiffeq library, we thank the authors Ricky Chen, Yulia Rubanova, Jesse Bettencourt and David Duvenaud for their work on this library and the Neural ODE paper.
13

---

Published in Transactions on Machine Learning Research (08/2023)
References
Fischer Black and Myron Scholes. The pricing of options and corporate liabilities. Journal of Political Economy, 81(3)6337-654 1973.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, 2018.
Raj Dandekar, Karen Chung, Vaibhav Dixit, Mohamed Tarek, Aslan Garcia- Valadez, Krishna Vishal Vemula, and Chris Rackauckas. Bayesian neural ordinary differential equations. arXiv preprint orXiv:2012.07244 2020.
Talgat Daulbaev, Alexandr Katrutsa, Larisa Markeeva, Julia Gusak, Andrzej Cichocki, and Ivan Oseledets. Interpolation technique to speed up gradients propagation in neural ODEs. Advances in Neural Information Processing Systems, 2020.
Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural ODEs, Advances in Newral Information Processing Systems, 2019.
Weinan E. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 5(1):1-11, 3 2017.
Silviu Filip, Aurya Javeed, and Lloyd N Trefethen. Smooth random functions, random ODEs, and Gaussian processes. STAM Review, (1(11:185-20 205, 2019.
Chris Finlay, Horn-Henrik Jacobsen, Levon Nurbekyan, and Adam Oberman. the world of Jacobian and kinetic regularization. International Conference on Machine Learning, 2020.
How to train your neural ODE:
James Foster, Terry Lyons, and Harald Oberhauser. An optimal polynomial approximation of Brownian motion. STAM Journal on Numerical Analysis, 58(3):1393-1422 14211 2020.
Amir Gholami, Kurt Keutzer, and George Biros. Anode: Unconditionally accurate memory-efficient gradients for neural odes. arXiv preprint arXiv:1902, 10208, 2019.
Arnab Ghosh, Harkirat Singh Behl, Emilien Dupont, Philip HS Torr, and Vinay Namboodiri. Steer: Simple temporal regularization for neural ODEs. Advances in Neural Information Processing Systems, 2020.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Πya Sutskever, and David Duvenaud. FFJORD: Free-
form continuous dynamics for scalable reversible generative models. International Conference on Learning Representations, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition, 2016.
Liam Hodgkinson, Chris van der Heide, Fred Roosta, and Michael W Mahoney. Stochastic continuous normalizing flows: Training SDEs as ODEs. Conference on Uncertainty in Artificial Intelligence, 2021.
Jacob Kelly, Jesse Bettencourt, Matthew James Johnson, and David Duvenaud. Learning differential equa-
tions that are casy to solve. Advonces in Neural Information Processing Systems, 2020.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equations for irregular time scries. Advances in Neural Information Processing Systems, 2020.
Patrick Kidger, Ricky TQ Chen, and Terry Lyons. "Hey, that's not an ODE": Faster ODE adjoints via seminorms. In International Conference on Machine Learning, 2021a.
Patrick Kidger, James Foster, Xuechen Li, and Terry Lyons. Efficient and accurate gradients for neural SDEs. Aduances in Neural Information Processing Systems, 2021b.
Patrick Kidger, James Foster, Xuechen Li, Harald Oberhauser, and Terry Lyons Neural SDEs infinite- dimensional GANs. In International Conference on Machine Learning, 2021c
14

---

Published in Transactions on Machine Learning Research (08/2023)
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Represenlations, 2015.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Yann LeCun, D Touresky, Hinton, and T Sejnowski. A theoretical framework for back propaaation. Connectionist Models Summer School, 1988.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- based learning applied to docu- ment recognition. Proceedings of the IEEE, 86(11))2278-25 2324, 1998.
Mikko Lehtimäki, Lassi Paumonen, and Marja-Leena Linne. Accelerating neural ODEs using model order reduction. IEEE Transactions on Neural Networks and Learning Systems, 2022.
Xuechen Li, Ting-Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. In International Conference on Artificial Intelligence and Statistics, 2020.
Xuanqing Liu, Si Si, Qin Cao, Sanjiv Kumar, and Cho-Jui Hsich. Neural SDE: Stabilizing neural ODE networks with stochastic noise. arXiv preprint arxiv: 02355, 2019.
Jaime A Londo and Andrés M Villegas. Numerical performance of some Wong Wong-Zakai type approximations for stochastic differential equations. International Journal of Pure and Applied Malhemalics, 107(2):301 331, 2016.
Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Dissecting neural ODEs. In Advances in Newral Information Processing Systems, 2020.
Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Ng. Reading digits in natural images with unsupervised feature learning. NeurIPS Workshop Deep Learning and Unsupercised Feature Learning, 2011.
Ho Huu Nghia Nguyen, Tan Nguyen, Huyen Vo, Stanley Osher, and Thieu Vo. Improving neural ordi- nary differential equations with Nesterov's accelerated gradient method. Advances in Neural Information Processing Systems, 35:7712 7726, 2022.
Alexander Norcliffe, Cristian Bodnar, Ben Day, Nikola Simidjievski, and Pietro Liò. On second order behaviour in augmented neural ODEs. In Advances in Neural Information Processing Systems, 2020.
Alexander Norclifle, Cristian Bodnar, Ben Day, Jacob Moss, and Pietro Liò. Neural ODE processes. In International Conference on Learning Representations, 2021.
Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. OT-flow: Fast and accurate continuous normalizing flows via optimal transport. In AAAI Conference on Artificial Intelligence, 2021.
Avik Pal, Yingbo Ma, Viral Shah, and Christopher Rackauckas. Opening the blackbox: Accelerating neural diffcrential equations by regularizing internal solver houristics. International Conference on Machine Learning, 2021.
Lev Semenovich Pontryagin. Mathematical Theory of Optimal Processes. CRC press, 1987.
Christopher Rackauckas, Yingbo Ma, Julius Martensen, Collin Warner, Kirill Zubov, Rohit Supekar, Do- minic Skinner, Ali Ramadhan, and Alan Edelman. Universal differential equations for scientific machine learning. arXiv preprint arXiv:2001.04530 2020.
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In International Con- ference on Machine Learning, 2015.
Yulia Rubanova, Ricky TQ Chen, and David Duvenaud. Latent ODEs for irregularly-sampled time series. Advances in Neural Information Processing Systems, 2019.
15

---

Published in Transactions on Machine Learning Research (08/2023)
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference Learning Representations, 2020.
Josef Stoer and Roland Bulirsch. Introduction to Numerical Analysis. Vprilger-Verlag, 2002.
T. Tieleman and G. Hinton. Lecture 6.5 RMM-Poppppp Divide the gradient by a running average of its recent magnitude. COURSERA: Neural Networks for Machine Learning, 2012.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent Gaussian models in the diffusion limit. arXiv preprint arXiv:1905.0955 2019.
George E Uhlenbeck and Leonard Ornstein. On the theory of the Brownian motion. Physical Review, 36 (5):823, 1930.
Arash Vahdat, Karsten Kreis, and Jan Kautz. Score-based generative modeling latent space. arXiv preprinl Xxxi:2206..00 05931, 2021.
Eugene Wong and Moshe Zakai. On the convergence of ordinary integrals to stochastic integrals. The Annals of Mathematical Statistics, 36(5):1500 (665):1560-1564 1965.
Hedi Xia, Vai Suliafu, Hangjie Ji, Tan Nguyen, Andrea Bertozzi, Stanley Osher, and Bao Wang. Heavy ball neural ordinary differential equations. Advances in Newral Information Processing Systems, 34:18646- 18659, 2021.
Winnie Xu, Ricky TQ Chen, Xuechen Li, and David Duvenaud. Infinitely deep Bayesian neural networks with stochastic differential equations. arXiv preprint ar arXi::2102.06559 2021.
Yewci Xu, Shi Chen, Qin Li, and Stephen J Wright. Correcting auto-differentiation in ncural-ODE training. arxiv preprinl X30: 02192, 2023.
Hanshu Yan, Jiawei Du, Vincent YF Tan, and Jiashi Feng. On robustness of neural ordinary differential equations. International Conference on Learning Representations, 2020.
Juntang Zhuang, Nicha Dvornek, Xiaoxiao Li, Sekhar Tatikonda, Xenophon Papademetris, and James Dun- can. Adaptive checkpoint adjoint method for gradient estimation in neural ODE. In International Con- ference on Machine Learning, 2020.
Juntang Zhuang, Nicha C Dvornek, Sekhar Tatikonda, and James S Duncan. MALI: A memory efficient and reverse accurate integrator for neural ODEs. International Conference on Learning Representations, 2021.
16

---

Published in Transactions on Machine Learning Research (08/2023)
A Multiple Measurement Times
In Section 3.3 we described the GQ method using an intial and terminal time only, here we consider mul- tiple measurement times, The integrand is given by a Vof which can be integrated using GauB Legendre quadrature. However, when the loss depends on > mcasurement times L(b(to), #(t1), .... .(tx)) there is discontinuily in the adjoint state at the intermediate times given by
ØL dz(tz))
a a= az (tx)
(19)
The integrand is now not continuous, so Gaussian quadrature does not apply well to the integration domain to. tx] However, quadrature still applies well to the sub-domains given by (k-1,t%l, for k = {1,2. K}: because the integrand is continuous there. Therefore, we split the whole integration into many smaller ones and step through
AK Os dt = 00
Σ
tx
T
05
dt,
(20)
00
km
12-1
applying quadrature to the smaller integrals. This more general implementation is given in Algorithm
Algorithm 2 Memory efficient implementation of the GQ method with K measurement times.
az + 0, k + K
g g+0
while do
There is now an outer loop over k, which sums the smaller integrals.
2(1x), a₂ f 0₂
81 (
n - GetNTerms((4
NFE,
C)
w, T + GeGGQWeigtssLootiionnn 1--1, 4b)
f-1 n
fprev + tk while j 2 1 do
Inext
Tj
f- ODESolve
T 0f 0z
(prev (nex.)
az
Qx
tprev
f- tnext
x a= 0f
+- + Wj
50
j j
end while
z
ODESolve
z
f af 0s
Iprevv, kk-1)
ag
a2
k + 1-1
end while return 9, Oz
Algorithm 2 is very similar to Algorithm 1 however there is now an outer loop over the mcasurement times indexed by The adjoint state is initialised at O, and the discontinuity applied the beginning of each smaller integral. The smaller integral is then solved as in Algorithm I but using the integration domain [t&-1,ta] tx] rather than [to. tx].
NOTE: There an important implementation detail. We do not restart a new solve except at the measure- ment times where there is the discontinuous change. That is, in the inner loop of Algorithm [ and the only loop of Algorithm D we continue the ongoing solve and do not start new onc. This is because starting a solve involves initial steps with fixed computational cost that we both wish and are able avoid. We recommend looking at code to exactly how this is done.
17

---

Published in Transactions on Machine Learning Research (08/2023)
B Training Neural SDEs
The final gencralization of the GQ method is to stochastic differential equations. We train the drift f and diffusion by training an ODE approximate solution of the SDE, using a Karhunen Loeve expansion. We can then use the GQ adjoint method to train this. The algorithm is exactly the same as Algorithm 2 for time-series or Algorithm when only the end point matters. The main difference is that those algorithms we have
dz
dt (2
whereas for the SDE we have
dz
(i- --()((( 1o) (
dt = f( ,(2,t,0)
V tk to
& cos
tK to
1=1
where &. ~ NOO, Following this we use an SDE solver with the trained drift and diffusion functions at test time.
C Adjoint vs Direct
In our Introduction, Section ! we stated that the direct method is faster than the adjoint but also more memory intensive. To make the statements more precise, consider a solve with state size N2, dynamics function with N, hidden layers, number of function evaluations evaluations in the backward pass N,. Then the memory usage for the direct method is (NN=-V whereas for the standard adjoint method it is O(N2N,), which is why it is often referred to constant with respect to the integration time, since it does not depend N₁. This is because the direct method has to store every activation every time the dynamics function is evaluated, whereas the adjoint method only has to store the activations from one function evaluation. The computational complexity is then O(2N,N N,),) for the direct method and ONN_NN + N2)) for the adjoint method, but the state in the backward solve is far larger increasing the computation. These values are taken from Table from Zhuang et al. (2021) and they represent the dominant terms in the scaling.
the forward pass N, and number of
It should be noted that the direct and adjoint methods do give different results, certainly for the gradients. Depending on the task and the parameter loss landscape this may not result in a difference in the final performance, the models may settle in different but equally good local optima. Or the loss landscape might be sufficienlly shaped such that even with different gradients they still reach the same minimum, for example if the loss landscape is convex (this very rarely the case). However it is still the case that the gradients are different, this is evident in our analytical experiment Figure [ where the errors in the loss between direct and adjoint are the same but the errors in the gradient are different. Interestingly for this experiment the errors are larger for the direct than the adjoint method, this is in contrast to (Gholami et al. 2019). We hypothesize this is because in ANODE the dynamics are learnt that images converge to be turned into distribution over classes, whereas in our exponent ial experiment the dynamics diverge, therefore for practical purposes where we do not want diverging dynamics it is likely that the direct method will produce preferable gradients. One example where this not the case is Xu et al. (2023) which shows using the leapfrog ODE solver can produce incorrect gradients when directly backpropagating which oscillate around the true ODE gradient. Whilst comparisons with the direct method are not this main paper's aim, we do stress that care should be taken when selecting method, since speed and memory consumption are factors but also in some cases the gradients will be different well, which can lead to different results.
final value
D Numerical Integration
Error Rates of Numerical Integration Method In Section 3.1 we stated that Gaussian Quadrature is the fastest method for numerically calculating 1D integrals. We give precise error bounds here demonstrate this. Assume we wish solve the integral given by Sº f())dt, below we give three common techniques for numerically solving the integral and the error bounds.
18

---

Published in Transactions on Machine Learning Research (08/2023)
Newton-Cotes Newton-Cord integration techniques split the interval [a, into N equidistant intervals. Then a polynomial of a given degree is fit on each interval which can be integrated. Classic examples are:
Trapezoid rule, this fits a linear interpolant to each interval. The error is O (ie)), refers to the largest m-th derivative on the interval [a, bj located at &.
N2 where fz)
Simpson's rule, this fits quadratic interpolants on the intervals. The error O N+
Boole's rule, this fits quartic interpolants on the intervals. The error is O No
In the above, (E) is a constant, it has been included in the error complexity to show that if the true function is exactly a polynomial of a given degree the error can be zero. For example, if the function is exactly linear the second derivative everywhere is zero, so the Trapezoid and other Newton-Cotes rules give the exact results. However the crucial point is that these are still constant with respect. to N so the errors for the Trapezoid, Simpson and Boole's rules are O (*t), Of (+)) and (( Therefore for Newton-Cotes the error is ((N=) where k depends on the rule for interpolation.
Gaussian Quadrature Gaussian Quadrature does not use equidistand poinls, the points are given by principled locations on the interval the roots of orthogonal polynomials which we can look up in a table, This allows us to approximate 2N 1 degroe polynomial with only N points in the quadrature sum. The error using Gaussian Quadrature with weight function 1 is bounded by (6-o)aN+1(NY)) f2N (E). There are two (2N+1)(2N))3
key terms this error, the first is f2N (E), as we use more points we model a higher degree polynomial unlike Newton-Cotes which stays constant. As mentioned we are therefore already able exactly solve a polynomial of degree 2N - 1, so if the function is well approximated by such a polynomial we will have an accurate approximation of the integral. The other key term is ((2N))) in the denominator, which dominates the other terms in the error bound showing the error will shrink very quickly with N, significantly faster than if the error is (++))
Note that the above assumes that f is differentiable, if there are discontinuities in f the integral is split up to make up for that (as described Appendix
Monte Carlo Integration Monte Carlo methods approximate the integral by sampling many points and f()) caleulating the mean of these. fa [(()dt fe p(!) (1t) di & N Ei p(t] 4 ~ p(t). By sampling the points they are not equidistant but also not in principled locations. The distribution of these Monte Carlo estimates follows a normal distribution with the mean being the true integral and the variance is ( so the error using Monte Carlo is ((+x) Monte Carlo has the worst error rates of the described methods which is why it is only used for high dimensional integrals where it is not possible to apply the other methods.
Therefore we use Gaussian Quadrature for these integrals since it has the best error rate. All of this information has been taken from Stoer & Bulirsch (002)
Shifting Integration Domain The weights and locations in GauB-Legendre quadrature are only defined on the interval [-1, 1]. And so these can only be used to approximate integrals of the form L dt) whereas our integral is on the general domain given by the start and end times of the ODE solve [to, [x]. Fortunately with simple change of variables we can scale and shift our domain to be [-1, Consider the integral given by
fx f(I)dt,
I =
to
by using the change of variables
7 ix to
to we can rewrite the integral as ix -to
I = 1' p*
f(t)dt= = [-]
fk
to
tx
tu
tk to
di
10
2
2
2
19

---

Published in Transactions on Machine Learning Research (08/2023)
And this can then be approximated using Gauß Legendre quadrature
Lx f
f(t)dt R3 Σ wif
1K - 1o
1x
+
4o
l to
Ti
to
2
2
2
allowing us to use GauB Legendre quadrature in our setting.
E Experimental Details
In this section we provide further details and results for our experiments. For all experiments we use the Dormand Prince 5(4) solver with an absolute and relative tolerance of 1 x 10-3 For the GQ method we use C= 0.1. We train using the Adam optimiser (Kingma & Ba 2015). For classification tasks we use constant integration times of [0, 1]. For fair comparison, all experiments were run on NVIDIA GeForce RTX 2080 GPUs. Minor hyperparameter tuning of the batch size and learning rate was carried out initially on each task to obtain acceptable model performance. However, the key metric to compare methods is training time, provided the model performances are the same or very similar. We first calculate the approximate number of function evaluations below and justify why it is not the most indicative metric of speed in our scenario.
E.1 Number of Function Evaluations
Unlike in other methods for speeding up the adjoint method, the number of function evaluations is not indicative of wall-clock time for the GQ method. This is because in previous methods during the back- ward solve one function evaluation is used each time to calculate the velocity of [z, az, agl which is [[,,,,, 0), -a Vo.lf. And the gradient is still found as an ODE. However in our method we use one function evaluation to calculate the velocity of [z, az] which [f(,, i, 0), -a"V= v=f and another function evaluation to calculate the term in the quadrature given by Vef. So number of function eval- uations is not indicative of time to obtain the gradient. The reason are faster is because the number of function evaluations does not take into account the time to backpropagate to calculate a V or at Vos; as well as generally solving a differential equation with state size 2|z + [0] (accounting for [z, az agl) versus one with size 2| (accounting for [z, az]). Nevertheless we can make approximate calculations for how many function evaluations are calculated:
The standard adjoint method produces approximately half the number of function evaluations during the backward solve were performed during the forward solve. This is shown in Figure 3 of Chen et al. (2018) (the original Neural ODE paper). Our method then uses the number of evaluation points introduced by heuristic [C
(fx-ta) NFB(tt-ta-1))
In our case we use C' 0.1 giving us NFE 10
NFE ((x-fe-1)) and if we assume we aren't 't 10
(fx-ta)
working with time scries, then this simplifies to function evaluations and the GQ adjoint uses approximately I 5 the number of evaluations as the standard adjoint.
And so the standard adjoint method uses NFE 2 so in relative terms the GQ method uses
NFE
10
Note that this is in fact likely an overestimate. In Kidger al. (2021a) it is demonstrated that if treating 0g integral rather than differential equation far fewer function evaluations during the backward solve are required. That paper only uses the error estimates of [z, az] during the backward solve using an adaptive solver, the error estimate of ag is not included. Similarly since we are only solving the ODE for [z, azl the error estimate in our case only comes from there as well. In Kidger et al. (2021a) Figure and Table 1 it is shown that the number of backward function evaluations significantly fewer in this scenario (40% 62% fewer steps as quoted), and so we can expect even fewer evaluations.
By using range of 40%-60% fewer evaluations given we can put all of this together. If we use NFE function evaluations during the forward solve:
The standard adjoint uses 0.5NFE evaluations of (f(z,t,0) -o Vofl
The seminorm adjoint 1(2,1,0). - V. -a
uses
between
0.2NFE and 0.3NFE
evaluations of
20

---

Published in Transactions on Machine Learning Research (08/2023)
The GQ adjoint uses between 0.2NFE and 0.3NFE evaluations of [f(z, 0)), V=f and 0.1NFE evaluations of Vol
And so we can see how the GQ method uses fewer evaluations compared to both the standard and seminorm adjoints, its just that this is not captured by the standard NFE measure.
E.2
Nested Spheres
Dataset We consider a two-dimensional nested spheres task, where the classification is given by
0, 0 < | S m
g(x) =
(21)
1, r2 | S 13
We use the values 7i = 0.4, 12 = 0.7 and r3 = 0.9, 140 example points from the training set are given in Figure 8 The training set made of 1000 randomly sampled points from each class (making 2000 data points in total). The test set made of randomly sampled points from each class.
Model Our model is an augmented neural ODE (Dupont et al. 2019) with 3 augmented dimensions, augmentation (Massaroli et al.) 2020) and final linear layer to obtain logits. The dynamics function is a time-dependent multilayer perceptron with two hidden layers and softplus activation. The two hidden layers have the same width which we vary to test the methods on different sized models.
Training We train the model 100 epochs, with a batch size of 200, cross entropy loss and learning rates given in Table We repeat the experiment with 10 seeds to obtain means and standard deviations.
Table 1: Learning rates for the nested spheres experiment.
Hidden Width
Learning Rate
5 x 10-3
20
5 x 10-3
100
1 X 10-3
500
1 X 10-3
1000
2 x 10-4
1500
2 x 10-4
2000
1 10
2500
1x 10
3000
1 10-4
Loss Curves We saw in Sention 5.3 that. the GQ method's training times scale well with model size compared to the adjoint and seminorm methods, as expected. We also saw that all methods have similar final accuracies demonstrating the gradients produced by the GQ method are accurate. We look further into this by plotting test loss curves during training, these are given Figure S
The loss curves match each other, showing further that the gradicnts computod with the GQ method are accurate. We also that in wall clock time the GQ method reduces the loss faster than the adjoint and seminorm methods.
Next we carry out an ablation study to see the effect of C on the training time and accuracy. We train the same model, with a hidden width of 1500 using the GQ method. We train with different values of C' and repeat with seeds to obtain means and standard deviations. The results are given Table ?
Effect of C
We see that the time train lowest when C' is lowest, because the fewest terms are used, however we also see that the final accuracy is compromised result. This is what would expect, the number of terms
21

---

Published in Transactions on Machine Learning Research (08/2023)
Direct
0.7
0.6
Adjoint
0.6
Test Loss
Seminorm
Test Loss
GQ
0.4
0.4
0.3
0.2
0.1
0.0
0.0
60
80
100
120
Epoch
Wall Clock Time (/s)
(a) Test loss epoch.
(b) Test loss vs wall clock time.
Figure 8: Test loss curves during training of the largest model the nested spheres task. We see that the loss curves approximately match for all methods and that the GQ method reduces the loss faster than the adjoint and seminorm methods.
Table The effect of the user defined C on training time and final accuracy on the nested spheres experiment. The best values are in bold. When C x 102 and C x 105, both regimes reached the maximum number of quadrature points (64), resulting in the same performance.
Time Train (/s)
Final Accuracy (%)
1 10-3
61.4 + 3.0
79.0 + 26.9
1 X 10
89.8 3.1
95.6 8.8
X 10a
141.0 2.4
95.8 8.4
1 X 10
141.7 = 2.4
95.8 8.4
used is not enough so the gradients are inaccurate. We see that using very high C leads to slightly higher accuracy than when C is 0.1. However, this is insignificant compared to the increase in training time, and the accuracies are casily within standard deviation of cach other. There is no increase in training time between C' being x 10 and 1 x 10 because the number of terms in the quadrature is capped at 64.
E.3
Time Series
Dataset Our dataset consists of sine curves, given by the second order ODE x., For given initial position =0)) xi and initial velocity v(0) vo the position function of time
(t) = To cos(t) + ug sin(t).
(22)
To create trajectory for the dataset, a random position and random velocity are sampled uniformly on [-1, 1], equation (22) is then used to give the position at given time between 0 and 2r. The training set consists of 15 trajectories and the test set contains For the main experiment we use 50 uniformly spaced measurement. times the range [0, 2T)] including the end points.
Model We use a second order neural ODE as our model (Norcliffe et al. 2020 Massaroli et al. 2020 this case the model given the initial position and the initial velocity and predicts the acceleration. The acceleration learnt by a time-independent multilayer perceptron, that takes both position and velocity input. The MLP has two hidden layers of varying width and softplus activations.
Training We train the model for 250 epochs, with batch size of 15, MSE loss and learning rates given in Table We repeat the experiment with 5 seeds to obtain means and standard deviations.
22

---

Published in Transactions on Machine Learning Research (08/2023)
Table 3: Learning rates for the sines experiment.
Hidden Width
Learning Rate
5
2 x 10 2
20
8 x 10-3
200
x 10
1000
3 10-5
2000
x 10-5
Loss Curves To further test if the GQ method produces accurate gradients, we plot the test MSE curves during training in Figure g
2.00
Direct
1.75
1.75
Adjoint
1.50
1.50
Test Loss
Seminorm
Test Loss
1.25
GQ
1.25
1.00
0.75
0.75
0.50
0.50
0.25
0.25
50
100
150
200
250
100
200
300
400
Epoch
Wall Clock Time (/s)
(a) Test loss vs epoch.
(b) Test loss vs wall clock time.
Figure 9: Test loss curves during training of the largest model on the sines task. We that the loss curves match for all methods and that the GQ method reduces the loss faster than the adjoint and seminorm methods.
The loss curves match, showing further that the gradients computed with the GQ method are accurate. We also see that in wall clock time the GQ method reduccs the loss faster than the adjoint and seminorm methods.
Effect of Regularity Here carry out an ablation study to investigate the effect of having regular and irregular measurement times during training. We also test the case where there are only 10 measurement times. In order to use batches, all of the initial states are concatenated to one large state, and the ODE runs on the larger system. For time series, the solver outputs a prediction for the whole concatenated state at every unique measurement. time and a mask is applied if necessary, so that only the relevant. times for each batch are used calculate the loss. Because we use batch sizes greater than one, having irregular times requires the model to output the predictions at many times, so would expect this to be slower than when using regular measurement times. We repeat the above experiment using model with hidden width 1000, using different numbers of measurement times, and changing whether the times are regularly spaced on the domain or not. The results are given in Table 4
We see that having irregular times increases the time to train significantly, as expected. The final MSEs are not significantly different with the standard deviations overlapping. The irregular MSEs are slightly lower because using irregular training times effectively giving the model many different measurement times train However, the decrease in MSF is negligible compared to the increase in training time.
23

---

Published in Transactions on Machine Learning Research (08/2023)
Table 4: The effect of regular/irregulat times on the sines experiment on training time and final test MSE. The best values are in bold.
No. Times
Regularity
Time Train (/s)
Final MSE
Regular
48.5 + 1.3
0.08 + 0.07
10
Irregular
239.2 2.9
0.07 + 0.06
Regular
96.7 + 4.9
0.11 = 0.09
50
Irregular
1880.5 = 607.6
0.09 + 0.08
E.4 Image Classification
Datasets We use standard image datasets for the image recognition experiment. We use the MNIST dataset (LeCun et al. 1998), which consists of handwritten digits 0-9. We use the CTFAR-10 dataset [Krzhevskky] (009) which consists of natural images of 10 classes: Airplane, Automobile, Bird, Cat, Deer, Dog, Frog, Horse, Ship and Truck. And use the SVHN dataset (Netzer al.) 2011), which is harder version of MNIST consisting of natural images of the digits 0-9, obtained from house numbers in Google
Street View.
Model For each dataset, our model consists of initial downsampling layers, a time-independent neural ODE with time domain [0, 1] and fully connected layers, to produce a vector of logits. We vary the number of hidden channels (nhidden) in the neural ODE dynamics function to vary model size. The downsampling layers are given in order below:
Batch Normalisation
ReLU activation
3 x convolution going to 10 channels with no padding
2 x 2 max pool
ReLU activation
3 x convolution going to 10 channels with no padding
The neural ODE dynamics function layers are then given by:
3 x 3 convolution going from 10 channels to nhidden channels with zero padding
Softplus activation
3 x 3 convolution going from nhidden channels to nhidden channels with zero padding
Softplus activation
3 x 3 convolution going from nhidden channels to 10 channels with zero padding (must have padding and go to 10 channels so the ODE layers preserve the size of the state)
An important implementation detail is that our implementation only works with vectors. So included in the downsampling layers is flattening operation the end, additionally the ODE function has unflattening operation at the start and flattening operation at the end.
The fully connected layers at the end of the classifier are given by a one hidden layer multilayer perceptron, with hidden width 50 and a ReLU activation.
Training We train the models for 15 epochs, with a batch size of 16 and cross entropy loss. The learning rate used for MNIST was 10-5 and the learning rate used for CIFAR-10 and SVHN was 10-5 We run each experiment using 3 seeds obtain means and standard deviations. We do consider the direct method due to memory limitations.
24

---

Published in Transactions on Machine Learning Research (08/2023)
Additional Results In addition to the results on MNIST in Section 5.5 we also present further results CIFAR-10 and SVHN here; this includes loss curves and results for different batch sizes. We first plot the loss curves from the MNIST experiment in Figure 10
1.0
Adjoint
Seminorm
Test Loss
GQ
Test Loss
0.6
0.4
0.4
0.2
0.2
2
14
2500
5000
7500
10000 12500 15000 17500
Epoch
Wall Clock Time (/s)
(a) Test loss vs epoch.
(b) Test loss vs wall clock time.
Figure 10: Test loss curves during training of the largest model on the MNIST task using batch size of 16. We see that the loss curves match for all methods and that the GQ method reduces the loss slightly faster than the adjoint and seminorm methods.
We see that the loss curves match, providing further evidence that the GQ method produces accurate gradients. The GQ method also reduces the loss faster than the adjoint and seminorm methods.
We also carry out the same experiments on the CIFAR-10 and SVHN datasets, results for the CIFAR-10 dataset are given in Figure 11
We see that the GQ method scales more favourably than the adjoint and seminorm methods for large models, however the effect far less noticeable than for the nested spheres and sines tasks. The final accuracies are comparable showing the gradients are accurate. In further support this, the test loss curves match for the methods, with the GQ method reducing the loss faster than the adjoint and seminorm methods.
25

---

Published in Transactions on Machine Learning Research (08/2023)
12000
61
Adjoint
Training Time (fs)
10000
Seminorm
Test Accuracy
60
GQ
59
8000
58
6000
57
4000
56
200
100
200
300
400
500
Hidden Channels
Hidden Channels
(a) Training time vs model size.
(b) Test accuracy hidden channels.
Adjoint
1.8
S Seminorm
1.8
Test Loss
GQ
Test Loss
1.4
1.4
1.2
2
2000
4000
6000
8000
10000
Epoch
Wall Clock Time (/s)
(c) Test loss epoch.
(d) Test loss vs wall clock time.
Figure 11: Training times and test accuracies for models of varying size the CIFAR-10 dataset with batch size of 16. Loss curves for the largest model are also plotted in the second row. The GQ method scales well to large models, the loss curves in the second row are approximately the same for each method and the GQ method reduces the loss faster than the adjoint and seminorm methods.
26

---

Published in Transactions on Machine Learning Research (08/2023)
We carry out the same experiment the SVHN dataset with results given in Figure 12
92
Adjoint
Training Time (/s)
20000
Seminorm
Test Accuracy
GQ
89
88
10000
87
5000
100
200
300
400
500
Hidden Channels
Hidden Channels
(a) Training time vs model size.
(b) Test accuracy vs hidden channels.
Adjoint
Seminorm
1.2
1.2
Test Loss
GQ
Test Loss
1.0
1.0
0.6
0.6
0.4
0.4
0
2
4
6
10
12
14
0.00 0.25
0.50 0.75 1.00 1.25 1.50 1.75
2.00
Epoch
Wall Clock Time (/s)
(c) Test loss vs epoch.
(d) Test loss vs wall clock time.
Figure 12: Training times and test accuracies for models of varying size on the SVHN dataset with batch
size of 16. Loss curves for the largest model also plotted in the second row. The GQ method scales well to large models, the loss curves in the second row approximately the same for each method and the GQ method reduces the loss faster than the adjoint and seminorm methods.
We see that the GQ method scales more favourably than the adjoint and seminorm methods for large models. The final accuracies are comparable showing the gradients are accurate. In further support of this, the test loss curves match for the methods, with the GQ method reducing the loss faster than the adjoint and seminorm methods.
Effect of Batch Size In this ablation study we see how the state size affects the training time. The adjoint state is the same size as the state and so as this increases, the effect of using the GQ method will be expected to be less noticeable, as it takes longer to solve the ODE regardless of the model size. We test this by using a larger batch size the image recognition tasks, we use a batch size of 32 doubling the size of the state. We also look at MNIST with a comparably large batch size of 256. We plot the training times for models of different sizes, Figure 13
We that for larger state sizes the advantage of using the GQ method significantly decreases, only scaling slightly more favourably than the adjoint and seminorm methods. This is because the computational cost to solve the differential equation for [z, [z,az] az] is significantly larger than solving the integral of a Vof. regardless of hee it is solved with quadrature or as an ODE. In the extreme case, using a batch size of 256, we actually see that the GQ method is slower than the adjoint and seminorm methods. This is likely because the time taken to calculate terms the quadrature sum becomes significant when the state is large. Whercas they are automatically calculated when solving as an ODE, because the backward state concatenated, allowing Vol to be calculated along with V= in backward pass of
27

---

Published in Transactions on Machine Learning Research (08/2023)
16000
Adjoint
Adjoint
Training Time (/s)
14000
10000
Seminorm
Seminorm
12000
Training Time
GQ
10000
6000
4000
2000
2000
100
200
300
400
500
100
200
300
400
500
Hidden Channels
Hidden Channels
(a) MNIST batch size
(b) CIFAR-10 batch size 32.
20000
12000
Adjoint
Adjoint
Training Time (/s)
Seminorm
Training Time (/s)
10000
Seminorm
15000
GQ
GQ
8000
10000
6000
4000
5000
2000
0
100
200
300
100
500
0.5
1.5
2.0
2.5
Hidden Channels
Parameters
166
(c) SVHN batchsize 32.
(d) MNIST batch size 256.
Figure 13: Training times on the image classification tasks for larger batch sizes. For a batch size of 32 the GQ method still scales more favourably, however the effect is significantly reduced. For a batch size of 256, the GQ method is actually slower than the adjoint and seminorm methods.
Therefore, as discussed in Section the GQ method is best applied to situations where memory is limited (otherwise the direct method best), where the state size is small and the model has many parameters. An example time-series analysis with complex dynamics.
F
Training SDEs
Dataset We consider the one-dimensional time-dependent OU process given by the Stratonovich SDE
dx (ut dx) (a Φt) dW,,
(23)
where we have time dependent drift and diffusion. We the values μ 0.2, A 0.1, 0.6, o 0.15. The data is made by sampling initial condition uniformly [-3, 3]. and then solving the SDE for times between [o, 10] with time intervals of 0.1. This is done using the sdeint Python library. We solve the SDE 45 times for each initial condition to generate enough samples, the training data consists of 200 initial conditions, and the test data consists of 20 initial conditions.
Model Our neural SDE runs in observation space, so do not an encoder or decoder. The drift and diffusion are learnt. by separate, Lime-dependent, multilayer perceptrons. Each MLP has two hidden layers of varying width and uses softplus activations.
Training We train the model for 100 epochs with a batch size of 40 initial conditions and the approximate KL divergence as the loss as well as evaluation metric. The learning rates for each hidden width are given in Table [ We run each experiment using 5 seeds to obtain means and standard deviations. We use the reversible- Heus Heun solver (Kidger et al. 2021b) to produce reversible solves (Zhuang al. (2021). with step size of 0.01. We the same solver to evaluate the parameters trained by the GQ method.
28

---

Published in Transactions on Machine Learning Research (08/2023)
Table Learning rates for the Ornstein Uhlenbeck experiment.
Hidden Width Learning Rate
5
1 x 10-3
20
5 x 10-4
50
1 x 10
100
1 x 10
Loss Curves To further test if the GQ method produces accurate gradients for neural SDEs, we plot the test KL divergences during training in Figure 14
0.5
0.5
Test KL. Divergence
SDE Adjoint
Test KL Divergence
0.4
GQ
0.4
0.2
0.2
0.1
0.1
20
40
60
80
100
250
500
750
1000 1250 1500 1750
Epoch
Wall Clock Time (/s)
(a) Test loss vs epoch.
(b) Test loss vs wall clock time.
Figure 14: Test loss curves during training of the largest model on the OU task. We see that the loss curves match for both methods and that the GQ method reduces the test loss slightly faster than the SDE adjoint method.
We see that the loss curves approximately match, they don't exactly match due to the stochastic element of the task. The approximate match shows that the GQ method and the Wong-Zakai theorem produces accurate gradients and can be used to train neural SDEs. Additionally, the GQ method reduces the loss faster than the SDE adjoint method.
Effect of Number of Cosines Next we consider how the number of cosines affects the results, with fewer cosines the ODE solution worse approximation to the SDE solution, so we would expect a larger KI. divergence. However, with fewer cosines each function evaluation faster, so we would also expect training be faster. The important case is when there are no cosines used, because only the drift is being trained, We test this by carrying out the previous experiment with hidden width 20, using different numbers of cosines to represent the Wiener process. The training times and test KL divergences given in Table 6
We see what is expected, the training time increases with the number of cosines used and the KL. divergence decreases with the number of cosines. We see that for more than 5 cosines, the improvement in performance is small compared to the increase in training time.
Learnt Solutions Finally, to see how the trained model performs, we visualise some of the learnt solutions in Figure 15 when the hidden width is 20. The learnt distribution matches the true distribution well. The learnt samples appear like they could have been sampled from the true distribution, providing qualitative evidence that the GQ method works for training SDEs.
29

---

Published in Transactions on Machine Learning Research (08/2023)
Table 6: Effect of number of cosines the OU experiment. The best values are in bold.
No. Cosines
Training Time (/minutes)
Mean KL Divergence (x10-3
0
9.62 + 0.09
75.55 + 32.30
1
11.34 + 0.04
28.01 = 2.48
5
12.97 + 0.06
7.49 + 0.86
10
19.88 + 6.23
6.50 + 0.74
25
33.13 + 0.96
6.18 + 0.68
50
51.97 0.69
6.12 + 0.66
t 0.3
t 5.0
10.0
Tree Prod P
True
True
Pred
800
x(0) -2.420
x(0) -2.108
x(0) 0.470
18.
125
Predicted
Prodicted
Predicted
10.0
0.8
-25
-2.
-3.0
-5
10 -5.0
10
10
t
Figure 15: SDE predictions using the trained model, with hidden width of 20, the OU experiment. In the top row plot Gaussian distributions, where the mean and standard deviation are calculated from the samples. These are plotted for given times from the same initial condition. We see that the predictions match well. In the bottom row are samples from three different initial condit ions. We see that the predicted samples are similar to the true samples from the SDE.
30

---

Published in Transactions on Machine Learning Research (08/2023)
G
Crossing Trajectories
Finally we consider the additional task of the crossing trajectorics problem Massaroli et al. (2020), originally called 9id in Dupont et al. (0019). It represents a simple problem that vanilla Neural ODEs cannot solve.
The problem is to map [++1,111] to [[1],, using an ODE. In one dimension, this is impossible because
ODE trajectories cannot cross.
The standard solution this problem to augment the state Dupont et al. 2019 this is where the initial condition has its dimensionality increased with zeros
Zo
X0 0
(24)
The new state z is then evolved according to new ODE and can be projected back into the observation space by directly selecting the relevant dimensions, alternatively a projection could be learnt as a single linear layer or a more complex function.
We use state with augmented dimensions, 4 in total. We train the model for 200 epochs and record the time well the MSF during training, using the MSE loss function. For each model this is repeated 10 times to obtain means and standard deviations, using the same seeds between methods for fair comparison.
The ODE function is a multi-layer perceptron, with two hidden layers of varying width, and a softplus activation. After the ODE evolution we then select the first dimension of the augmented state. We use different learning rates depending on the width of the hidden layers which are given in Table
Table 7: Learning rates for the crossing trajectories experiment.
Hidden Width
Learning Rate
5
10
20
10
100
1 x 10-3
500
1 X 10
1000
7 10-4
1500
1 1x 10-4
2000
1 1 10-1
2500
1 10
3000
1 10
The results of the time laken and final MSE for each hidden width, are presented in Figure T6 We can see that the direct method is the fastest, followed by the GQ method. The seminorm method slightly faster the whole than the standard adjoint method as expected. In particular method scales well to the number of parameters, whereas the other two adjoint methods do not. This is because during the backwards solve, the two adjoint methods calculate a a6 at every function evaluation (with the seminorm method taking fewer steps). The GQ method on the other hand only calculates this term for evaluation points in the quadrature calculation, and therefore scales better to larger models. This also explains why the training time for the adjoint methods appears to scale approximately linearly with the number of parameters. For small models the time improvement negligible.
can see that the GQ
T
We that the final MSEs the same, which are all low showing that the problem has been solved. The only exception being for the largest model where the learning rate may have been too high. We see that the losses during training match, so the gradients produced by the GQ method are (approximately) the same as the other methods. We also see that the loss is lowered in shorter wall clock time using the GQ met than the other adjoint methods as expected.
31

---

Published in Transactions on Machine Learning Research (08/2023)
Training Time
MSE
0.0030
Direct
Training Time (/s)
40
Adjoint
0.0025
Seminorm
0.0020
GQ
30
0.0015
MSE
0.0010
20
0.0005
0.0000
10
-0.0005
500
1000
1500
2000
2500
3000
500
1000
1500
2000
2500
3000
Hidden Width
Hidden Width
Loss vs Epoch
Loss vs Wall Clock Time
4.0
4.0
3.5
3.5
3.0
3.0
Test Loss
Test Loss
2.5
2.5
2.0
2.0
1.5
1.5
1.0
1.0
0.5
0.5
0.0
0.0
Epoch
Wall Clock Time (/s)
Figure 16: Training times and MSEs for the crossing trajectorics task. The GQ method's training time scales well with model size compared to the adjoint methods. The test MSEs are the same. In the second row we plot the loss during training for the largest model. We the loss curves all match and that the GQ method reduces the loss faster than the adjoint methods.
32


Reviews:
[{"review_id": "xNoFXGWy1s", "reviewer": "Reviewer_WaJb", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n- As far as I understand, the work is theoretically sound. \n- The findings clearly favor the proposed approach. \n\n### Weaknesses\n- The approach is thoroughly evaluated on standard NODE architectures + benchmarks. I'm curious if more recent NODE variants would also benefit from the approach and how the approach would perform on more challenging setups.\n- There are many claims without citations.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces a new approach to approximate the reverse-time integral that is needed to compute gradients of a neural ordinary differential equation (NODE). The authors propose to replace standard adjoints with a Gauß–Legendre quadrature. This is doable as the rather complicated reverse integral can be expressed in terms of multiple one-dimensional definite integrals, which are then computed in parallel. On a series of benchmarks, the model is shown to increase speed without hurting model performance.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. I think we need to see more comparisons. For instance, would more complicated NODE variants, e.g. latent NODEs, neural controlled differential equations, heavy-ball NODEs, etc, benefit from the proposed construction?  \n2. Comparisons against a more competitive baseline would strengthen the paper. The authors claim \"regularizing solutions may affect the expressivity of the model\". Putting these two together, I'm wondering if, e.g., Jacobian regularization, would help reduce memory footprint or execution time without adversely affecting the model performance.\n3. \"Preferably, the number of terms chosen adapts to the complexity of the problem, so that the smallest number of terms is used to compute an accurate gradient.\" Adjoints also can trade off the ODE accuracy with speed. So, would a less accurate adjoint solver (e.g., rtol=1e-2) lead to competitive performance? Related to this, isn't rtol=atol=1e-3 too high compared to the default values in MATLAB and torchdiffeq implementation?\n4. \"We make the assumption that if the (forward) trajectory of $z$ is complex then the corresponding gradient trajectory $a^T_z \\nabla_\\theta f$ is also complex\". This should be empirically demonstrated.\n5. A comparison of $n$ and model performance would be nice.\n6. Why do we see memory efficiency on the analytical task but not the others?\n7. The approach is tested on neural nets with varying widths but not depths. Why is this?\n8. References are needed for the following claims:\n    - Solving these ODEs in general can be slow, which is one of the major barriers preventing neural ODEs from being used more widely.\n    - \"Because the error in $a_\\theta$ does not grow significantly compared to $z$ or $a_z$.\n    - \"Gaussian quadrature is the fastest method for 1-D integrals outside of analytical solutions\"\n\n### Minor:\n1. \"They introduce seminorms to take advantage of this\" What is \"seminorms\"?\n2. \"Instead we can rewrite the ODE solve in a way that it calculates the gradients associated with the parameters as a definite integral\" In what way (6) and (7) are different?\n3. A discussion on whether the approach suffers from or helps against vanishing/exploding gradients would be nice.\n4. In (9), when $n$ is divided by $t_k-t_0$?\n5. What is $m$ in (13)?\n6. Does the experiment in Sec 5.1 make sense? The approach is proposed for NODEs whereas the experiment considers a setup that is a lot simpler, for which a much smaller $C$ could be sufficient?  ", "broader_impact_concerns": "Everything alright here.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "4IUaG9iGsE", "reviewer": "Reviewer_rgHZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "This trick improves adjoint efficiency quite significantly with clear empirical demonstrations. However, the adjoint is not a practical method in the first place, and naive autodiff is still better in practise.\n\nThe paper is well-written.\n\nThe claims are accurate, and are backed by sufficient evidence. This paper is interesting to the neural ODE community.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper proposes an acceleration technique for computing neural ODE adjoints (backpropagation). The key insight is that the backward pass integral can be approximated using quadrature, and the quadrature points are then used to guide the backward solver. The idea is principled, clever, and novel.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Running comments\n\n- The introduction contrast the ‘direct’ backprop with adjoint. The direct method is said to be “fast” but “memory-intensive”. I find both statements strange: this is just standard neural network backprop over many layers. I don’t really see in what sense do we characterise this as fast or memory intensive. Similarly saying that adjoint is “slow” is also strange: it solves a different problem and one can’t really compare Neural-ODE-adjoint with ResNet-Euler-backprop. Finally, there is also the issue that adjoint method “O(1)” method is biased due to numerical reversibility issues of the neural ODE. The statements need more precision.\n- “arbitrarily large computational graphs”. Not sure I buy this. An ODE takes some amount of time increments, and each takes some linear number of intermediate steps (eg. 4 in RK4). The backpropagation recycles the gradients backwards, so I don’t see why the computational graphs would explode. I also don’t see why we need to store all intermediate states (in practise we probably do). The paper needs to be it more precise what is the complexity of standard backprop wrt number of increments and order of solver.\n- “faster and potentially more accurately”. I don’t see how integral approximations can be more accurate than the true ODE. Surely you are incurring some error.\n- Since w(t)=1, why do we include it at all?\n- “adding each term in the sum..”. I don’t understand this part. This is all vague and I have no idea what terms or sums are we talking about.\n- I don’t understand fig 1. What is integration time? It doesn’t make much sense that we integrate for 20..28 seconds: surely the t_K is constant? I don’t understand y-axis either: relative wrt what? It seems that ACA is the gold standard: how does ACA then incur error? This entire experiment is poorly described, and I can’t follow.\n- In the experiments it’s unclear how many NFEs the backward pass takes between adjoint and GQ. Can you describe this?", "broader_impact_concerns": "no issues", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "xpEHYT2Zpf", "reviewer": "Reviewer_4ho2", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\n- Neural ODEs and similar methods, such as diffusion models, are widely used but can be memory intensive and slow to train. Improvements in training efficiency are very useful, so the paper is well-motivated.\n- The key observation that $a_\\theta$ can be evaluated as a definite integral rather than an ODE is well explained, as is their proposed method.\n\n\n### Concerns\n\n- The authors provide a large set of experiments demonstrating their method, but I think a more focused analysis of a smaller set of experiments would provide more insight.\n- The speedup in this method comes from evaluating $\\frac{da_\\theta}{dt}$ at only a subset of points. As with any quadrature problem, there is a tradeoff between accuracy and the number of function evaluations, but this tradeoff isn't thoroughly explored in the paper. They propose a heuristic for choosing the number of points in equation (9), but given that this is the crux of the method, it would be nice to see the more analysis of speed/accuracy trade-off.\n- Some of the plots are very unclear. Figure 1 in particular is difficult to read---the results from their own method are almost invisible! It's also not clear to me what is being plotted here. What does 'Relative Error in Loss' mean? I think this plot in particular needs a lot of clarification.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose a more efficient method for computing gradients of neural ODEs. They use the observation that the time derivative of one of the terms ($a_\\theta$) in the \"reverse time\" ODE depends only on other terms and not on its own value. As a result, solving for $a_\\theta(t_0)$ can be approached as a quadrature problem rather than a differential equation. This allows them to evaluate $a_\\theta$ only at a subset of well-chosen quadrature points. They apply their method to a range of problems, evaluating its training speed and memory efficiency.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Some of the discussion of quadrature methods is imprecise, e.g. \"Gaussian quadrature is the fastest method for 1-D integrals...\". This should be replaced with a precise statement about the convergence and efficiency of GL quadrature.\n- The speed up in this method is due to fewer evaluations of $\\frac{da_\\theta}{dt}$. How many evaluations does the conventional adjoint method perform? How many does this method save?\n- In figure 2, why does GQ use _more_ memory than the adjoint method? Its state is the same size or smaller during the reverse time solve, isn't it?\n- The authors point out that the true gradient is unknown, so it is difficult to assess accuracy. However, since their method is almost identical to the adjoint method except that it evaluates $\\frac{da_\\theta}{dt}$ at only a subset of points, we can't expect the GQ method to be _more_ accurate than the adjoint in general. It seems reasonable to treat deviation from the adjoint method gradient as the \"error\" caused by the GQ method. Again, it would be nice to see a plot of this error as a function of the number of quadrature points.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "6GRqmo7Ino", "reviewer": "Reviewer_8wu8", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper clearly communicates the basics of Neural ODEs and the existing literature. It clearly explains where existing methods fail, and how their approach tackles some of those issues. \n\nThe experiment on image classification demonstrates their method performs favourably in terms of training time on a larger scale. However, it seems that the GQ method only provides a marginal improvement over other methods. \n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "\nTraining Neural ODEs is slow since they require numerical solvers. A common approach to speed Neural ODEs up is to regularise them, however this can also affect the expressivity of the model (particularly in time-series modelling, where the trajectory itself matters). \n\nThis paper proposes a method to make training Neural ODEs with the adjoint method faster. The method of Gauß–Legendre Quadrature is used for numerical integration, instead of using solvers. This approach is also extended into a new way to train Neural SDEs. \n\nSection 2 of the paper gives a clear explanation of Neural ODEs and training them using the adjoint method. It then goes on to describe the work of Kidger, which points out that the adjoint equation for gradients with respect to the Neural ODE parameters can be seen as an integral. The terms of the integrand are the adjoint with respect to z and the gradient of the dynamics with respect to the parameters. In particular, there is no term in the integrand for the adjoint with respect to the parameters theta. So, we can solve this 1D integral using quadrature methods in-parallel across the dimensions of the parameters. \n\nInstead of requiring a dense solution, the authors propose fixing the number of quadrature points in advance and maintaining a running sum. This avoids high memory consumption that other adjoint quadrature methods suffer from.\n\nThe authors propose a reasonable empirical heuristic for choosing the number of quadrature points. This heuristic requires the choice of a constant, for which an ablation study is done. There could be an experiment done to directly justify the choice of the form of this heuristic (i.e. comparing the NFEs on the forward and backwards passes when using the adjoint method), but I don’t think it is absolutely necessary since similar results have already been shown in other works. \n\nThe paper also extends their approach to training Neural SDEs, where memory requirements can be especially high. They make use of the Wong–Zakai theorem, so that an ODE can be numerically solved instead of an SDE. In particular, this approach requires a smooth and memory efficient approximation of the Wiener process. \n\nThe first experiment in the paper demonstrates that the GQ method is just as accurate as the adjoint method, and has similar (constant) memory requirements. \n\nSecondly, they show these benefits carry over to making the training time faster on a toy spheres experiment.\n\nThere are additional experiments for image classification, time series, and SDEs.\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "The finer details of the approach for Neural SDEs needs further explanation in my opinion. I think having pseudocode in the Appendix (similar to that for the method for ODEs) would be very helpful.\n\nThe authors describe GQ methods in some detail, but I feel there could be further explanation given here. In particular, it would be helpful to describe how the weights need to be shifted to change to a general interval.\n\nThe experiment on time-series models demonstrates the method works well when there are intermediate points of integration. I would be interested to see if there were any ablations done on the number of integration points for GQ for time series modelling tasks, as I imagine the number of integration points might vary more between different examples in the dataset. \n\nIn section 5.6, can you clarify “We do not see the previous effect of model size on training time because we are still in the ‘small parameter regime’, which is why the training time fluctuates rather than increasing with model size.”\n\nAre there any experiments where you compare the NFE on the backwards pass with the number of integration points? I know they are not directly comparable, since the NFE in the adjoint method includes terms for both df/dz and df/dtheta. But it would be nice to see that data directly, since it could disentangle total time spent from the memory requirements needs to solve the adjoint. This seems like it could be related to the ablation experiment done that justifies the choice of constant C in Equation 9, perhaps?\n", "broader_impact_concerns": "I have no concerns here.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper presents a new way of looking at training neural ODEs using the adjoint method with demonstrable improvements in efficiency. There were questions about the practical benefits of the adjoint method itself that were largely resolved during the discussion period. While experiments that test alternative NODE variations and regularization would be very interesting, the approach itself was deemed to be relevant and sound in terms of the TMLR criteria.
Decision: Accept (poster)