"Hey, that's not an ODE'": Faster ODE Adjoints with 12 Lines of Code
Authors: Patrick Kidger
Ricky T. Q. Chen
Terry Lyons
Venue:  (2021)
Keywords: neural differential equations
neural ordinary differential equations
adjoint

Abstract: Neural differential equations may be trained by backpropagating gradients via the adjoint method, which is another differential equation typically solved using an adaptive-step-size numerical differential equation solver. A proposed step is accepted if its error, \emph{relative to some norm}, is sufficiently small; else it is rejected, the step is shrunk, and the process is repeated. Here, we demonstrate that the particular structure of the adjoint equations makes the usual choices of norm (such as $L^2$) unnecessarily stringent. By replacing it with a more appropriate (semi)norm, fewer steps are unnecessarily rejected and the backpropagation is made faster. This requires only minor code modifications. Experiments on a wide range of tasks---including time series, generative modeling, and physical control---demonstrate a median improvement of 40% fewer function evaluations. On some problems we see as much as 62% fewer function evaluations, so that the overall training time is roughly halved.

Under review as a conference paper at ICLR 2021
"HEY, THAT'S NOT AN ODE": FASTER ODE ADJOINTS WITH 12 LINES OF CODE
Anonymous authors Paper under double-blind review
ABSTRACT
Neural differential equations may be trained by backpropagating gradients via the adjoint method, which is another differential equation typically solved using adaptive-step-size numerical differential equation solver. A proposed step is accepted if its error, relative to some norm, is sufficiently small; else it is rejected, the step is shrunk, and the process is repeated. Here, we demonstrate that the particular structure of the adjoint equations makes the usual choices of norm (such L²) unnecessarily stringent. By replacing with more appropriate (semi)norm, fewer steps are unnecessarily rejected and the backpropagation is made faster. This requires only minor code modifications. Experiments on a wide range of dasks-including time series, generative modeling, and physical control-dmmostrrall
median improvement of 40% fewer function evaluations. On some problems we see as much as 62% fewer function evaluations, so that the overall training time is roughly halved.
INTRODUCTION
We begin by recalling the usual set-up for neural differential equations.
1.1 NEURAL ORDINARY DIFFERENTIAL EQUATIONS
The general approach neural ordinary differential equations E 2017 Chen et al.. 2018 is to use ODEs as learnable component of a differentiable framework. Typically the goal is to approximate a map x +> by learning functions ()) (z(. 1p) and f(.,,,)) 0)). which are composed such that
z(t) =(t)= z())+ + f. f(s, 2(s), () ds and y == f2(((T)) v)).
2(T) (zx Φ),
(1)
The variables d. 0, denote learnable parameters and the ODE solved over the interval [7, 7]. We include the (often linear) maps ()) (2(.,
..(....) 1b) for generality. as in many contexts they are important for the expressiveness the model (Dupont et al.) 2019: Zhang al. (020)), though our contributions will be focused around the ODE component and will not depend on these maps.
Here we will consider neural differential equations that may be interpreted neural ODE.
1.2 APPLICATIONS
Neural differential equations have to the best our knowledge three main applications:
1. Time series modeling. Rubanova et al. produce ODEs with jumps. Kidger et al.
interleave Neural ODEs with RNNs to 2020 take f(t, 2, 9(2, A) dx (t), dependent on some time-varying input X, to produce a neural controlled differential equation.
2019)
2. Continuous Normalising Flows in Chen et al. (2018); Grathwohl et al. (2019), in which the overall model acts as coupling or transformation between probability distributions,
3. Modeling or controlling physical environments, for which a differential equation based model may be explicitly desired, see for Hample/Zhong al. (2020).

---

Under review as a conference paper at ICLR 2021
1.3 ADJOINT EQUATIONS
The integral in equation ① may be backpropagated through either by backpropagating through the internal operations of numerical solver, or by solving the backwards-in-tinn adjoint equations with respect to some (scalar) loss I.
aL
0s
dL = (2(T),
(=( = dz(T))
(1) = (T)
(
(8, =(x), 0) ds
and
0z
dz(r)
0f
dL um(
aT(T) 0,
(1) = (T(
T
(6)
00
88 ((), 0) ds
and
d0
dL
05
dL dr (e(T),
(T(T) dT
(t) (T)
a=(s)
(s, 2()), 8) 0) ds
and
T Ös
(2)
These equations typically solved together joint system a(l) (1), ag (g()), (()).. (They are already coupled; the latter two equations depend on az.) As additionally their integrands require (s). and as the results of the forward computation of equation [T] are usually not stored, then the adjoint equations are typically additionally augmented by recovering by solving backwards-in- time
=(1) = = + [ [(s,2()),))d..
1.4 CONTRIBUTIONS
We demonstrate that the particular structure of the adjoint equations implies that numerical equation solvers will typically take many steps, that are too small, wasting time during backpropagation. Specifically, the accept/reject step of adaptive-step-size solvers is too stringent.
By applying correction account for this, demonstrate that the number of steps needed to solve the adjoint equations may be reduced by typically about 40%. We observe improvements on some problems by as much 62%. Factoring the forward pass (which is unchanged), the overall training time is roughly halved. Our method is hyperparameter-ffee and requires no tuning.
We do not observe any change in model performance, and least with the torchdi fee package (our chosen differential equation package), this correction may be applied with only 12 lines of code.
METHOD
2.1 NUMERICAL SOLVERS
Both the forward pass given by equation (I) and the backward pass given by equations and solved by invoking numerical differential equation solver. Our interest here is in alaptive-step- size solvers. Indeed the default choice for solving many equations is the adaptive-step-size Runge- Kutta 5(4) scheme of Dormand- Princee (Dormand & Prince 1980), for example implemented by depri5 in the torchdi ffeq package or ode45 in MATLAB.
A full discussion of the internal operations of these solvers is beyond scope here; the part of interest to us is the accept/reject scheme. Consider the case of solving the general ODE
y(t)= = u()) + S: f( v(s))
with v(t ddd.
Suppose for some fixed the solver has computed some estimate (t) zo y(t). and it now seeks to take step A > 0 to compute 7(t + A) s y(t + A). A step is made, and some candidate Teandiitate (t + A) is generated. The solver additionally produces yen E Rd representing an estimate of the numerical error made in each channel during that step.
2

---

Under review as a conference paper at ICLR 2021
Given some prespecified absolute tolerance ATOL (for example 10 10-9), relative tolerance RTOL
(for example 10-9), and (semi)norm || I|: R + [0, 0o) (for example ||vil / V v the
RMS norm), then an estimate of the size of the equation given by
SCALE ATOL + RTOL. max Deandidate (t A)) E ER""
(4)
where the maximum is taken channel-wise, and the error ratio
Mee
R
(5)
SCALE
then computed. If 1 then the error deemed acceptable, the step is accepted and take (+++A) Veandidane (t- I >> 1 then the error is deemed too large, the candidate Dcandidate A) rejected, and the procedure is repeated with a smaller A.
Note the dependence on the choice of norm || ||: in particular this determines the relative importance of each channel towards the accept/reject criterion.
2.2 ADJOINT SEMINORMS
The key observation is that ag (and in fact also a,) does appear anywhere in the vector fields of equation 2
Not an ODE
This means that (conditioned on knowing & and a.), the integral corresponding as is just an integral-not an ODE. As such, it is arguably inappropriate to solve it with an ODE solver, which makes the implicit assumption that small errors now may propagate create large errors later.
Accept/reject This is made manifest in the accept/reject step of equation norm ||-Ii, |l, such as L will usually weight each channel equally. But we have just established that to solve the adjoint equations accurately, it is far more important that and a₂ be accurate than it is that 0g accurate.
Typical choices of
Seminorms Thus, when solving the adjoint equations equation 2 we propose to use || that scales down the effect in those channels corresponding to ag.
In practice, in experiments, offending channels, so that ||- in fact a seminorm. This means that the integration steps are chosen solely for the computation of ag and 2, and the values of as are computed just by integrating with respect to those steps.
scale || all the way down by applying zero weight to the
Example As an explicit example, note that (T) 0. When solving the adjoint equation numerically, this means for / close to T that the second term in equation 4 is small. As ATOL is typically also small, then SCALE is additionally small, and the error ratio r in equation (5) is large.
This implies that it becomes easy for the eπror ratio to violate < 1, and it is casy for the step to be rejected. Now there is nothing intrinsically bad about a step being rejected-we would like to solve the ODE accurately, after all-the problem is that this is spurious rejection, the rejection occurred to ensure the accuracy of ag, which is as already established unnecessary.
In practice, observe that spurious rejections may occur for any t, not just those near T.
Other channels In fact, essentially the same argument applies to a, as well: this does not affect the value of the vector field either. In continuous normalising flow, the log- oog-rrbbbiylity channel is also only an integral, rather than an ODE, and again the same argument may be applied.
Does this reduce the accuracy of parameter gradients? One obvious concern is that we are typically ultimately interested in the parameter gradients ag, in order train model; with respect to this our approach seems counter-intuitive
However, we verify empirically that models still train without reduction in performance. We explain this by noting that the 2, a channels truly are ODEs, so that small errors now do propagate to create larger emrors later. Thus these are likely the dominant source of error overall.
3

---

Under review as a conference paper at ICLR 2021
2.3 CODE
Depending on the software package, the code for making this change can be trivial. For example, using PyTorch (Paszke al. 2019) and torchdiffed (Chen al., 2018), the standard set-up requires only a few additional lines of code. The additional 12 lines are marked.
import torchdiffeq
def rms-norm (tensor): return tensor porsorrppw(22 (2). mean () sqrt (
#
#
#
def make norm state
#
state size state numel ()
def norm ( aug-state):
#
== [1: state-size]
#
10
adj-y aug- sta_statell state_size:! + state-size]
#
11
return max nox (y), rms_norm (adj-y)) return norm
12
#
13
#
14
torchdiffeq odeint adjoint (func=.... y0=....
adjoint aptions=dics (norm=make norm (y0)) #
This amounts to the extra 12 lines code stated in the title-a number that even includes the additional whitespace and visual indents.
To keep the remainder of this discussion software-agnostic,, we defer further explanation of this specific code to Appendix A
EXPERIMENTS
We compare our proposed technique against conventionally-trained neural differential equations, across multiple tasks-time series, generative, and physics-informed. These are each drawn from the main applications of neural differential equations, discussed in Section 1.2 every case, differential equation solver used is the Dormand-Prinee 5(4) solver "dopri5" The default norm a mixed L /L2 norm used in torchdi ffeq; see Appendix A For the sake an interesting presentation, investigate different aspects of each problem.
In each case see Appendix B for details on hyperparameters, optimisers and so on.
The code for these experiments can be found at [redacted; see supplementary material].
3.1 NEURAL CONTROLLED DIFFERENTIAL EQUATIONS
Consider the Neural Controlled Differential Equation (Neural CDE) model of Kidger et al. (2020).
To recap, given some (potentially irregularly sampled) time series x ((tu, To),.... ..., (In, xx)), with each I; E R the timestamp of the observation E R" let X: [4,, (n.] > R1+ be interpolation such that ((.) = (44, For example X could be a natural cubic spline.
Then take ft,, 2, = g(, 0) d.X (t) in a Neural ODE model, so that changes in x provoke changes in the vector field, and the model incorporates the incoming information x. This may be thought of as continuous-time RNN; indeed Kidger al. (2020) use this learn functions of (irregular) time series.
de
We apply a Neural CDE the Speech Commands dataset (Warden 2020). This is a dataset of one-sccond audio recordings of spoken words such 'left', 'right' and so We take 34975 time series corresponding to 10 words, to produce a balanced classification problem. We preprocess the dataset by computing mel-frequency cepstrum coefficients that each time series is then regularly spaced with length 161 and 20 channels. The data was then normalised to have zero mean and unit variance. We used the torchcde package (Kidger 2020), which wraps torchdiffeee ffeq.
The initial map L (of equation (T) taken to be linear on (t), The terminal map &g is taken to be linear on 2(In).

---

Under review as a conference paper at ICLR 2021
Results for Neural CDEs. The number of function evaluations (NFE) used in the adjoint equations accumulated throughout training- are reported. Test set accuracy is additionally reported. Mean + standard deviation over five repeats. Every model has 88,940 parameters.
Table
Default norm
Seminorm
RTOL, ATOL
Accuracy (%)
Bwd. NFE (10%)
Accuracy (%)
Bwd. NFE (10%)
10-3, 30-3,10 10-6
92.6=00
14.36=1.09
92.5+00
8.67
10-4, 10-7
92.8104
30.67 2,48
92.5100
12.75 2.00
10 9,10
92.4
77.95
92.9#0.4
29.39
1.00
Accuracy
0.75
0.50
Seminorm Default norm
0.25
0.00
25
50
75
100
125
150
175
200
Epoch
Backward NFE
60000
Seminorm Default norm
40000
20000
0
25
50
75
100
125
150
175
200
Epoch
Figure 1: Mcan + standard deviation, for accuracy and the number of function evaluations (NFE) for the backward pass, during training, for Neural CDEs.
We investigate how the effect changes for varying tolerances by varying the pair (RTOL, ATOL) over (10 (10-3.10-; 10-0), (10-4, 4 10-"), and (10- 10 10-8). For each such pair we run five repeated experiments.
ATOLL
10-00-9
See Table T for results on accuracy and number of function evaluations. We see that the accuracy of the model is unaffected by our proposed change, whilst the backward pass uses 40%-6244 62% fewer steps, depending on tolerance.
Next, we investigate how accuracy and function evaluations change during training, See Figure I We see that accuracy quickly gets close to its maximum value during training, with only incremental improvements for most of the training procedure. In particular, this statement is true of both approaches, and we do not see any discrepancy between them. Additionally, we see that the number of function evaluations is much lower for the seminorm throughout training.
3.2 CONTINUOUS NORMALISING FLOWS
Continuous Normalising Flows (CNF) (Chen et al. 2018) are class of generative models that define a probability distribution as the transformation of a simpler distribution by following the vector field parameterised by Neural ODE. Let p(2o) be arbitrary base distribution that we efficiently sample from, and compute its density. Then let =t) be the solution of the initial value problem
=(0) ~ p(=0),
dz(t)
zt = (1.=
log p(g((()))) -t
(f 02
(1, (()) =())
dt

---

Under review as a conference paper at ICLR 2021
Table 2:
Results for Continuous Normalising Flows. The number of function evaluations (NFE)
used in the adjoint equations- accuuuuaccccd throughout training- are reported. Test set bits/dim and number of parameters are additionally reported. Mean = standard deviation over three repeats.
Default norm
Seminorm
Model
Parameters
Test bits/dim
Bwd. NFE (10%)
Test bits/dim
Bwd. NFE (10%)
CIFAR-10
1,200,600
3.3492
11.65
3..0281100021
38.93:1132
(dy. = 64)
MNIST
281,928
0.906
22.50
0.90421000013
37.03
(dy = 32)
MNIST
1,005,384
0.9623=0.00p
44 1.77
0..0031000021
37.85.10.94
(dn 64)
MNIST
3,779,400
0.9439
48.60
0.946611000488
41.84
(dh = 128)
for which the initial point is randomly sampled from p(z0), and for which the change in log probability density also tracked as the sample is transformed through the vector field.
The distribution at an arbitrary time value T can be trained with maximum likelihood to match generative model of data with the distribution p(z(T"). Furthermore,
data samples, resulting
Grathwohl et al. (2019) combines this with a stochastic trace estimator (Hutchinson, 1989) to create an efficient and unbiased estimator the log probability for data samples x.
log((=(TT) = x) = log p(=())) + Eo-A((11) [ o" (of (2 v dt
(6)
In Table 2 we show the final test performance and the total number of function evaluations (NFEs) used in the adjoint method over 100 epochs. We see substantially fewer NFEs in experiments both MNIST and CIFAR-10. 10. Next, we investigate changing model size, by varying the complexity of the vector field f. which is CNN with d, hidden channels. We find that using the seminorm, the backward NFE does not increase as much as when using the default norm. In particular, we can roughly same NFEs the smaller (d,, = 32) model to train larger (dg 128) model and in doing so achieve substantial gain in performance (1.00 + 0.95).
3.3 HAMILTONTAN DYNAMICS
Finally we consider the problem of learning Hamiltonian dynamics, using the Symplectic ODE-Net model of Zhong et al. (2020. Under Hamiltonian mechanics, the vector field parameterised
(g. 0) dp dH
f(t, 2) g(g.
0)u-
dg q. P.
where z (q. p. u) is the state decomposed into (generalised) positions momenta and control w. H a Hamiltonian and y is an input term, both parameterised as neural networks. The input offers way to control the system, which can be understood via energy shaping.
This parameterisation lends itself to being interpretable. H learns encode the physics of the system in the form of Hamiltonian mechanics, whilst g learns to encode the way in which inputs affect the system. This can then be used to construct controls u driving the system to a desired state. In this context (unlike the other problems output of the model. The evolution of z is matched against the observed state at multiple times 2(t)), 2(t2). ., and trained with 2 loss.
consider here), is not hidden state but instead the
We consider the fully fully-actuated double pendulum ("acrobot")) problem. Training data involves small oscillations under constant forcing. We investigate several quantities of interest.

---

Under review as a conference paper at ICLR 2021
Table 3: Results for Symplectic ODE-Net on fully-actuated double pendulum. The number of function evaluations (NFE) used in the adjoint equations--acuueeee throughout training- are reported. Test set loss additionally reported. Mean standard deviation over five repeats. The model has 509,108 parameters.
Default norm
Seminorm
Test T loss (10-²) Bwd. NFE (10")
Test 1,2 loss (10-1) Bwd. NFE (10')
1 1.247=0.520
46.45=0.01
2.995
26.55+e.01
Figure 2: Frames from controlling the fully-actuated double pendulum to the full-upright position.
41
U1
10
15
20
10
15
20
10
42
42
0.5
Sinto
0.0
Figure 3: Evolution of the state (4, Q. u) whilst controlling the fully-actuated double pendulum from the full-down to the full-upright position. gi denotes the angle of the first joint, and d2 denotes angle second joint.
6000
Default norm
Seminorm
NFE
4000
0
-0.1
0.0
0.1
0.3
t
Figure 4: Location of function evaluations for the adjoint equations, during training for the fully- actuated double pendulum problem. All evaluations over five runs are shown.
First we investigate the number of function evaluations. See Table 3 We see that the model successfully learns the dynamics, with very small loss (order O(10~)) ")) both cases. However, under our proposed change the model is trained using 43% fewer function evaluations on the backward pass.
Next, verify that the end goal of controlling the system is achievable. See Figures 2 and in which the double pendulum successfully controlled from the full-down to the full-upright position, using the ciainorm-traineed model.
After that, we investigate the locations the function evaluations. See Figure We see that the default norm makes consistently hypothesis was that there would be more spurious rejections for near the terminal time T, where (t( is small, (T) In fact observe consistent improvements for all
evaluations for every time t. This is interesting our initial

---

Under review as a conference paper at ICLR 2021
400
Steps
Default norm Seminorm
0.00
0.05
0.10
0.15
0.20
Figure 5: Location of rejected steps during training for the fully-actuated double pendulum problem. All steps over five runs are shown. Rejected steps for the seminorm can just barely be seen.
To conclude, we investigate the locations the rejected steps. See Figure 5 Note that the Dormand- Prince solver makes 6 function evaluations per step, which is the reason for the difference in scale between Figures [ We see that the seminorm produces almost no rejected steps. This accounts for the tight grouping seen in Figure 4 for each batch the evaluations times are typically the same locations. Moreover, we observe an interesting structure to the the location of the rejected steps with the default norm. We suspect this is due the time-varying physies of the problem.
3.4 WHERE DO IMPROVEMENTS COME FROM?
Parameter-state ratio We hypothesised that the ratio between size of state and number of parameters may explain the differences between the smaller improvements for continuous normalising flows against the substantial improvements for Hamiltonian neural networks and neural CDEs (corresponding to the proportion of channels that can no longer cause rejections). Given model with p parameters and state 2(t) Rd, then seminorms reduce the number of rejection-cause channels from 1 2d + to just + 2d.
HNN
JCDE
0.35
0:30
0.20
0.15
JVV di du-66
0.10
#MNIST
CIBAR
0.05
10
10
10"
We plot the (log) ratio 1+24+0/1+24 against the percentage improvement backward steps. See Figure 6 We that this hypothesis partly explains things: broadly speaking continuous normalising flows (which have O(1*)) state variables) small improvements. whilst the neural CDEs and Hamiltonian neural networks (with O(10') state variables) see larger improvements.
Figure 6: Parameter-state ratio against improvement, as given by the relative decrease in backward NFE. Only one neural CDE experiment (RTOL==1033 ATOL=10-99 shown as the others use different tolerances. CNFs are averaged over their different blocks.
There are still many differences between problems (tolerances, complexity of vector field, and so on): see for example the variation amongst the continuous normalising flows. Thus for any given problem this serves more as a rule-offthumm.
Number of accepted/rejected steps Next investigate the improvements for each problem, broken down by accepted and rejected steps (over a single backward pass). See Table 4 We begin by noting that the absolute number of accepted and rejected steps nearly always decreased. (In the case of modelling Hamiltonian dynamics, the seminorm actually produced zero rejected steps.) However beyond this, the proportion of rejected steps nearly always decreased.
This highlights that the benefit of seminorms is two-fold. First, that the number of accepted steps reduced indicates that use of seminorms solve, allowing larger step sizes. Second, the reduced proportion of rejected steps means that the differential equation solve is more efficient, with fewer wasted steps.
systems treated smaller and easier-to-
4 RELATED WORK
Several authors have sought techniques for speeding up training of neural differential equations.
8

---

Under review as a conference paper at ICLR 2021
Table Number of rejected and accepted backward steps on one pass over the test set.
Rejected steps (10%)
Accepted steps (10%)
Proportion rejected
Experiments
Default
Seminorm
Default
Seminorm
Default
Seminorm
NCDE
0.13=0.02
0. 0.152
3.45 + 0.26
1.92 + 0.18
3.63% %
7.25%
(small)
NCDE
(medium)
1.01
0.17.10.0
8.38 + 0.86
3.25
0.43
10.76%
4.97%
NCDE
(large)
10.91 1.03
1.16:10.11
28.86
+ 2.15
8.61
0.45
27.43%
11.87%
CIFAR-10
11.04:10.9
9.30:#1.22
23.50
+ 1.19
22.82
1.34
31.96%
28.95%
MNIST
(dA 32)
5.54
4.51
20.91
+ 0.90
19.93
0.42
20.95%
18.45%
MNIST
7.79
4.29
22.95 1 1.02
19.73 0.99
38.53%
17.86%
(dn 64)
MNIST
8.88=0.21
7.40
26.54 + 0.91
25.01 + 1.17
25.07%
22.83%
(d, 128) HNN
0.1 I 10.00
0.00:10.00
0.18 +: 0.00
0.16
0.00
37.93%
0.00%
Ghosh et al. (020)) (much like this work) make a conceptually simple change. They regularise the neural ODE by randomly selecting the terminal integration time. We note that this does not seem applicable to Neural CDEs as well as Neural ODEs.
Finlay et al. 2020) and Kelly et al. (2020) investigate regularising the higher-order derivatives of the model; the idea is that this encourages simpler trajectories that are easier to integrate. However, this improvement must be weighed against the extra cost of computing the regularisation term. Thus whilst Finlay et al. (2000) describe speed improvements for CNFs, Kelly et al. (2020) describe
slower training as result of this extra cost.
Quaglino et al. 2020)) describe speeding up training via spectral approximation, Unfortunately. this method is rather involved, and our knowledge does not have reference implementation.
When backpropagating through the internal operations of the solver (so not the adjoint method used here), Zhuang al. (2020) note that backpropagating through rejected steps is unnecessary.
Massaroli et al. (2020b) discuss hypersolvers, which are hybrids of neural networks and numerical differential equation solvers, trained to efficiently solve a desired Neural ODE. However the Neural
ODE changes during training, making these most useful after training, to speed up inference.
Dupont et al. (2019) and Massaroli al. (2020a) note that adding extra dimensions the Neural ODE improve expressivity and training. Indeed this has arguably now become a standard part of the model; we include this via the linear & in equation (T).
In the context of Continuous Normalizing Flows, Grathwohl et al. (2019), Chen & Duvenaud (0119)
and Onken al. (2020) have proposed specific architectural choices to reduce training cost.
Chen et al. throughout our experiments.
2018 mention regularizing the differential equation using weight decay; we include this
Independent of modifications to the model or training procedure, Rackauckas & Nie (017) and
Rackauckas et al. (2019) claim speed-ups simply through an improved implementation.
5 CONCLUSION
We have introduced a method for reducing the number of function evaluations required to train a neural differential equation, by reducing the number of rejected steps during the adjoint (backward) pass. The method simple to implement, straightforward to integrate into existing codebases, and offers substantial speed-ups across variety of applications with observed downsides.

---

Under review as a conference paper at ICLR 2021
REFERENCES
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural Ordinary Differential Equations. In Advances Neural Information Processing Systems 31, pp. 6571 6583. Curran Associates, Inc., 2018.
Ricky TQ Chen and David Duvenaud. Neural networks with cheap differential operators. In Advances in Neural Information Processing Systems, pp. 9961-9971, 2019.
R. Dormand and P. J. Prince. A family of embedded Runge- Kutta formulae. J. Comp. Appl. Math, 6:19-26, 1980.
Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural odes. In Advances in Neural Information Processing Systems 32, pp. 3140-3150. Curran Associates, Inc., 2019.
Weinan E. A Proposal on Machine Learning via Dynamical Systems. Commun. Math. Stat., 5(1): 1-11, 2017.
Chris Finlay, Jörn-Herrik Jacobsen, Levon Nurbekyan, and Adam M. Oberman. How to train your neural ODE: the world of Jacobian and kinetic regularization. ICML 2020, 2020.
Arnab Ghosh, Harkirat Singh Behl, Emilien Dupont, Philip H. S. Torr, and Vinay Namboodiri. STEER: Simple Temporal Regularization For Neural ODEs. Neural Information Processing Systems, appear, 2020.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. International Conference on Learning Representations, 2019.
Michael F Hutchinson. stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Sottistis-Simmlattion and Computation, 1859)1609- 1989.
Jacob Kelly, Jesse Bettencourt, Matthew James Johnson, and David Duvenaud. Learning Differential Equations that are Easy to Solve. arXiv:2007.04504, 2020.
Patrick Kidger. torchade, 2020.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons Neural Controlled Differential Equations for Irregular Time Series. arXiv:2005.088 08926, 2020.
D Kingma and J Ba. Adam: A method for stochastic optimization. ICLR, 2015.
Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Dissecting Neural ODEs. arXiv:2002.08077, 2020a.
Stefano Massaroli, Michael Poli, Atsushi Yamashita, Hajime Asama, and Jinkyoo Park. Hypersolvers: Toward Fast Continuous-Depth Models. arXiv:2007.09601 2020b.
Derek Onken, Samy Wu Fung, Xingjian Li, and Lars Ruthotto. Ot-flow: Fast and accurate continuous normalizing flows via optimal transport. arXiv preprint arXiv:2006.00104 2020.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vetito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy. Benoit Steiner, Fang, Junjie Bai, and Soumith Chintala. PyTorch: An Imperative Style, High-PPeformance Deep Learning Library. In Advances in Neural Information Processing Systems 32, pp. 8024- 8035. Curran Associates, Inc., 2019.
Alessio Quaglino, Marco Gallieri, Jonathan Masci, and Jan Koutník. Snode: Spectral discretization of neural odes for system identification. International Conference on Learning Representations, 2020.
10

---

Under review as a conference paper at ICLR 2021
Chris Rackauckas, Mike Innes, Yingbo Ma, Jesse Bettencourt, Lyndon White, and Vaibhav Dixit. Diffeqflux.jl-a julia library for neural differential equations. arXiv preprint arXiv:1902.0233 2019.
Christopher Rackauckas and Qing Nie. Differentilleuations---- performant and feature-rich ecosystem for solving differential equations in julia. Journal of Open Research Software, 5(1), 2017.
Yulia Rubanova, Tian Qi Chen, and David K Duvenaud. Latent Ordinary Differential Equations for Irregularly-Sampled Time Series. In Advances in Neural Information Processing Systems 32, pp- 5320-5330. Curran Associates, Inc., 2019.
Pete Warden, Speech commands: A dataset for limited-vocabulary speech recognition. arXiv: 10X4::88040030000 032000 2020.
Han Zhang, Xi Gao, Jacob Unterman, and Tom Arodz. Approximation capabilities of neural odes and invertible residual networks. International Conference on Machine Learning, 2020.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Symplectic ode-net: Learning hamiltonian dynamics with control. International Conference on Learning Representations, 2020.
Juntang Zhuang, Nicha Dvonek, Xiaoxiao Li, Sekhar Tatikonda, Xenophon Papademetris, and James Duncan. Adaptive checkpoint adjoint method for gradient estimation in neural ode. International Conference Machine Learning, 2020.

---

Under review as a conference paper at ICLR 2021
A
CODE FOR TORCHDIFFEQ
Here we discuss in more detail the code used for torchdi f ffee that was introduced Section
23
By default, torchdi ffeq uses a (slightly odd) mixed -M-S norm. Given the adjoint state [0t, 3, uz az). then the norm used is
| 2, a2s ad]ll max [ RMS ||2|| RMS
|| IRMS: || RMS)
where for (I1....,.T.)
|
Respecting this convention, the code provided reduces this to the seminorm || 2, 0g ag)
max ( RMS [[lll
Other similar approaches such pure-RMS over [z, 02) should be admissible well.
B EXPERIMENTAL DETAILS
B.1 NEURAL CONTROLLEE DIFFERENTIAL. EQUATIONS
We use the same setup and hyperparameters as in Kidger et al. 2020). The loss function is cross entropy. The optimiser used was Adam (Kingma & Ba 2015), with learning rate 1.6 x 10 batch size of 1024, and 0.01-weighted I weight regularisation, trained for 200 epochs. The number of hidden channels (the size of 2) is 90, and is parameterised feedforward network, of width 40 with hidden layers, ReLU activation, and tanh final activation.
B.2 CONTINUOUS NORMALISING FLOWS
We follow Grathwohl et al. (2019) and Finlay et al. (2020) The loss function the negative log likelihood - Iog(p(=(())
z)) of equation (6) The optimiser used was Adam, with learning rate 10-3 and batch size 256, trained for 100 epochs. Relative and absolute tolerance the solver are both taken to be 10-3, We used multi-scale architecture as in Grathwohl et al. (2019) and Finlay et al. (2020), with 4 blocks of CNFs at different scales.
B.3 SYMPLECTIC ODE-NET
The optimiser used was Adam, with learning rate 10-3, batch size of 256, and 0.01 0.01-weighted L weight regularisation. Relative and absolute tolerance of the solver are both taken to be 10-4. We use the same architecture in /hhong et al.] (2020): this involves parameterising H as sum of kinetic and potential energy terms. These details of Sympletic ODE-Net are little involved, and we refer the reader to Zhong et al. (2020) for full details.
12


Reviews:
[{"review_id": "Xyu_zl7F8G", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summarizing the paper claims\n------------------------------------------\nThe paper addresses the problem of reducing the number of function evaluations (NFEs) during neural ODE training with adaptive solver and the adjoint method. Namely, the authors claim that for the variety of applications,  NFEs at backward pass can be reduced if the automated step selection will be based only on the part of the information of the adjoint state.\n\nStrong points\n-------------------\nThe paper is clearly written. The performance of the proposed method is evaluated for a wide variety of tasks (time series, generative modeling, physical control). Hyperparameter setups are provided to reproduce the experiments.\n\nWeak points\n-----------------\nThe paper doesn't provide a comparison with the existing techniques for speeding up NFEs (e.g., Finlay et al. (2020)). \nThe paper claims that the proposed approach is complementary to the techniques mentioned in the related paper. However, I haven't found supportive experiments for this assumption.  If I understood correctly,  the authors use weight decay everywhere in their experiments. Therefore, I wonder if the proposed method shows similar behavior when weight decay is omitted.\n\nRecommendation (accept or reject)\n-------------------------------------------------\nFor this moment, I'd recommend rejecting the paper due to the mentioned weak points; however, I consider updating the decision if weak points and questions will be closed during the rebuttal period.\n\nUpdate: Thanks to the authors for the work done during the rebuttal period. The proposed method seems working but it's still not very clear when it is beneficial to use it. The authors provided some intuition about the cases when the method is effective, however, I think it is necessary to more thoroughly explain the applicability of the method before sharing it with the community at the conference, that is why I don't change my score.\n\nQuestions\n--------------\nI'd like to understand why for different applications relative influence of this approach is quite different? For example, for Continuous Normalising Flows (Table 2), the gain is not so significant as for Neural CDEs (Table 1). What are the crucial properties of the task to benefit a lot from the proposed technique? Are there any limitations?\n\nDescribing Symplectic ODE-Net results in Table 3, the authors claim that the loss is almost the same for default norm/seminorm training. However, the reported mean loss for seminorm is twice bigger, and the reported standard deviation is four times bigger than corresponding values for the default norm. Hence, that would be nice to see NFEs comparison when loss statistics are closer, because for now, that is not  clear whether there is a real improvement due to the proposed method or this is a  trade-off between loss statistics and the number of NFEs", "extra_scores": {}, "extra_text": {}}, {"review_id": "Jg5mWp0xDV", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: The paper proposes a modification for the adjoint method, such that to improve the training efficiency of neural ODEs. The proposed idea is that the solution of some terms in the adjoint method can be less accurate, because these are not ODEs but simple integrals, and hence, the error does not propagate. Thus, the solver can utilize bigger steps, and in total to perform less steps. In the experiments the efficiency is demonstrated under different scenarios where neural ODEs are used.\n\nComments: At first, I think that the proposed idea is quite interesting. Even if I am not an expert in the field, the theoretical argumentation seems to be rather reasonable. Also, the proposed trick seems to improve in practice the efficiency of the method.\n\nHowever, I think that the writing of the paper can be improved. In my opinion, at the current stage the paper looks like a collection of (technical) notes, but probably not as a proper conference paper. Since the proposed idea is a rather simple trick, I believe that is necessary a better description of the technical background, for example the adjoint method and the whole neural ODEs procedure. In my opinion, the paper is written in way which assumes that the reader is very familiar with the topic and its corresponding details. Also, I find the coherence of the story a bit lacking. Therefore, I believe, that improving the paper along these line is essential.\n\n===== After rebuttal =====\n\nAs I have mentioned in my review, I am not an expert in the field of neural ODEs. Regarding the technical contribution I think is quite interesting, but probably rather small as far as I understand from the rest of the reviews. In my opinion, the current paper is very much focused to audience which is very familiar with the topic, and its impact on a non-expert probably limited. Therefore, I tend to keep my score and vote for rejection, because I feel that the submission has to become more accessible to general audience.", "extra_scores": {}, "extra_text": {}}, {"review_id": "3JTEobPOfw6", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Update\n\nI have revised my rating based on the updates. The paper has good theoretical insights, however I agree with the other reviewers that more completeness is required to make the proposal stronger. \n\n--------------------------\n\nSummary\nThe paper recognizes that the error calculations in Neural ODEs involve some terms that don’t need to be accounted for. It then shows that by ignoring those terms, the number of function evaluations reduces while approximately maintaining accuracy. The paper demonstrates this in 3 tasks - speech classification using Neural CDEs, generative modeling using FFJORD, and Hamiltonian dynamics using Symplectic ODE-Net.\n\nStrengths\nThe paper is well written, the motivation is clearly explained. There is a diverse set of tasks that are tackled, which is great to check for consistency of the conclusions drawn. A code snippet is provided for one implementation to show how simple the proposed change is to incorporate into existing code. Where possible, helpful graphs have been provided to help visualize the conclusions of the respective experiments. The paper also mentions other works and how this work is not closely related to them.\n\nComments\nDifferent aspects of introducing the paper’s contribution are explored only in one task, and not all. For example, the graph of the Backward NFE vs Epoch is showcased for the 1st task of speech classification, but not the rest of the tasks. It would be important to understand if the conclusion that the semi-norm has better NFE throughout training is consistent in all the tasks explored. For example, the proposed change does not seem to be very effective in the task of generative modelling in terms of BPD, hence it would be important to understand how the proposed change affects training in that task.\n\nThis also holds for locations of rejection steps, as it was mentioned that although intuitively the spurious steps closer to T=0 would be rejected but empirically it holds true for all t, but this aspect is only explored in the 3rd task. It would be important to understand how consistent this is across tasks.\n\nOverall, the paper seems to introduce a minor change that only works on simpler tasks but does not affect higher-dimensional tasks such as generative modeling very effectively. This is probably because the underlying dynamics for the generative modeling tasks are necessarily complex. Works such as Finlay et al (2020) and Ghosh et al (2020) (uncited in the paper!) propose to reduce the complexity of the underlying dynamics. It is pertinent to see how this paper’s proposed change affects the training of those models. Hopefully, the simpler dynamics should implicitly lend to smaller NFEs, which the proposed change can take advantage of and show faster results on.\n\nIt is important to note that terms such as “standard” may not apply to Neural ODEs, since they are hardly 2 years old and so are quite nascent and yet to be widely adopted. The fact that some papers that work on Neural ODEs are not as cited or implemented in software packages as others does not mean they are not “standard” (or any such related term), and vice-versa.\n\nThe “12 lines of code” contribution also only works for the “standard” implementation. This is not a weakness per se, but software-related changes are highly subject to the respective software being used. Something that’s 12 lines in one particular code repository implemented in PyTorch may be a lot more lines in a different implementation or in a different language such as Julia or Jax.\n\nFinlay et al. (2020) := How to Train Your Neural ODE; ICML 2020\nGhosh et al. (2020) := STEER: Simple Temporal Regularization For Neural ODEs ; NeurIPS 2020; https://arxiv.org/pdf/2006.10711.pdf\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "3If4ranM5Cb", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\nThe paper suggests improving the training run-time of neural ODEs when using the adjoint sensitivity method by relaxing the computational steps of $a_{\\theta}$ and $a_{t}$. \n\n**Significance**\nThere are a couple of major issues about the content and the significance of the results presented in this paper:\n- While the paper is beautifully organized and easy to follow, I consider the scientific contributions of the work to be limited.\n- In the CNF setting, Table 2, I do not see a major improvement in NFE. Why is this the case? \n- All neural ODEs have been tested with a single type of ODE solver. I strongly suggest the authors investigate other solvers to see how a model's performance changes when applying the proposed method? \n- Deriving a theoretical ground for the reduction in time complexity and its effect on high- and low-frequency time series modeling performance. I suspect the proposed method to speed up the training of neural ODEs would have a significant effect on the performance of time series modeling with high-frequency fluctuations. Could the authors comment on this please?\n\nThis paper is a 10 out of 10 as a blog post for the users of the adjoint method for training neural ODEs. I believe that the contribution of this paper is limited to a software hack that would benefit the community best if presented as a blog post. \n\n**Clarity**\nAn extremely clear submission. I greatly appreciate the authors' effort in preparing such a clear presentation. \n\n**Originality**\nThe paper is an original contribution. \n\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper focuses on NeuralODE and shows that for the implementation popular among ML community, one of the equations is not an ODE and can be replaced by an integral. This is implemented using "seminorm" (just assigning zero weight to the last equation).

Pros: 
- Well written
- Useful to replace the "standard" implementation
-  Consistent benchmarking

Cons: 
- Contribution is too limited
- Used in several "prior" codes without explicit ICLR submission. 
- (My personal) The title is not good: more on the "hype side" of the story, rather than progressing the field. I don't think we need to put every single minor fact into a ICLR submission.  For example, one can just compute the integral as an alternative by any suitable quadrature rule. That would add 10-15 function evaluations at most, since most of the functions in NeuralODEs are quite smooth.
Decision: Reject