MALI: A memory efficient and reverse accurate integrator for Neural ODEs
Authors: Juntang Zhuang
Nicha C Dvornek
sekhar tatikonda
James s Duncan
Venue: ICLR 2021 Poster (2021)
Keywords: neural ode
memory efficient
reverse accuracy
gradient estimation

Abstract: Neural ordinary differential equations (Neural ODEs) are a new family of deep-learning models with continuous depth. However, the numerical estimation of the gradient in the continuous case is not well solved: existing implementations of the adjoint method suffer from inaccuracy in reverse-time trajectory, while the naive method and the adaptive checkpoint adjoint method (ACA) have a memory cost that grows with integration time. In this project, based on the asynchronous leapfrog (ALF) solver, we propose the Memory-efficient ALF Integrator (MALI), which has a constant memory cost $w.r.t$ integration time similar to the adjoint method, and guarantees accuracy in reverse-time trajectory (hence accuracy in gradient estimation). We validate MALI in various tasks: on image recognition tasks, to our knowledge, MALI is the first to enable feasible training of a Neural ODE on ImageNet and outperform a well-tuned ResNet, while existing methods fail due to either heavy memory burden or inaccuracy; for time series modeling, MALI significantly outperforms the adjoint method; and for continuous generative models, MALI achieves new state-of-the-art performance. We provide a pypi package: https://jzkay12.github.io/TorchDiffEqPack

Published as a conference paper at ICLR 2021
MALI: A MEMORY EFFICIENT AND REVERSE ACCU- RATE INTEGRATOR FOR NEURAL ODES
Juntang Zhuang; Nicha C. Dvornek; Sekhar Tatikonda; James S. Dunean
(j. zhuang; ni ai.cha.dvorme dvornek; sekhar.tat kkonda; james lrmnesddune duncan) @yale. edu Yale University, New Haven, CT, USA
ABSTRACT
Neural ordinary differential equations (Neural ODEs) are new family of deep- learning models with continuous depth. However, the numerical estimation of the gradient in the continuous case is not well solved: existing implementations of the adjoint method suffer from inaccuracy in reverse-time trajectory, while the naive method and the adaptive checkpoint adjoint method (ACA) have mem- ory cost that grows with integration time. In this project, based on the asyn- chronous leapfrog (ALF) solver, we propose the Memory-efficient ALF Integra- tor (MALI), which has constant memory cost w.r number of solver steps in integration similar to the adjoint method, and guarantees accuracy in reverse-time trajectory (hence accuracy in gradient estimation). We validate MALI in various tasks: on image recognition tasks, to our knowledge, MALI is the first to en- able feasible training of a Neural ODE on ImageNet and outperform well-tuned ResNet, while existing methods fail due to either heavy memory burden or in- accuracy; for time series modeling, MALI significantly outperforms the adjoint method; and for continuous generative models, MALI achieves new state-of-the- art performance. We provide pypi package: https:: http:////kaaa1...gm. github. 1o Lo/oorchDiffE@Pach
INTRODUCTION
Recent research builds the connection between continuous models and neural networks. The theory of dynamical systems has been applied to analyze the properties of neural networks or guide the
design of networks (Weinan 2017 Ruthotto Haber 2019 Lu et al., 2018). In these works, a residual block (He et al.] (016) is typically viewed as a one-step Euler discretization of an ODE; instead of directly analyzing the discretized neural network, it might be easier to analyze the ODE.
Another direction is the neural ordinary differential equation (Neural ODE) (Chen al., 2018 which takes a continuous depth instead of discretized depth. The dynamics a Neural ODE is typically approximated by numerical integration with adaptive ODE solvers. Neural ODEs have been applied in irregularly sampled time-series (Rubanova et af.) 2019), free-form continuous gen- erative models (Grathwohl et al.) 2018: Finlay et al.) (200)), mean-field games (Ruthotto et al. 2020), stochastic differential equations
eevettioon (Li et al. 2020) and physically informed modeling (Sanchez- Gonzalez al. 2019 Zhong et al. 2019).
Though the Neural ODE has been widely applied in practice, how to train it is not extensively stud- ied. The naive method directly backpropagates through an ODE solver, but tracking continuous trajectory requires huge memory. Chen et al. (2018) proposed to use the adjoint method to deter- mine the gradient in continuous cases, which achieves constant memory cost w.r. integration time; however, as pointed out by Zhuang et al.) (2020), the adjoint method suffers from numerical errors due to the inaccuracy in reverse-time trajectory. Zhuang et al. (2020) proposed the adaptive check- point adjoint (ACA) method to achieve accuracy gradient estimation at much smaller memory cost compared to the naive method, yet the memory consumption of ACA still grows linearly with integration time. Due to the non-constant memory cost, neither ACA nor naive method suitable for large scale datasets (e.g. ImageNet) or high- high-dimensional Neural ODEs (c.g. FFJORD (Grath- wohl et al. 2018)).

---

Published as a conference paper at ICLR 2021
In this project, propose the Memory--fficient Asynchronous Leapfrog Integrator (MALI) to achieve advantages of both the adjoint method and ACA: constant memory cost w.Y integration time and accuracy in reverse-time trajectory. MALI is based on the asynchronous leapfrog (ALF) integrator (Mutze) 2013). With the ALF integrator, each numerical step forward in time is reversible. Therefore, with MALI, delete the trajectory and only keep the end-time states, hence achieve constant memory cost w.rl integration time; using the reversibility, we can accurately reconstruct the trajectory from the end-time value, hence achieve accuracy in gradient. Our contributions are:
1. We propose a new method (MALI) to solve Neural ODEs, which achieves constant memory cost w.rff number of solver steps in integration and accuracy in gradient estimation. We provide theoretical analysis.
We validate our method with extensive experiments: (a) for image classification tasks, MALI enables Neural ODE to achieve better accuracy than well-tumed ResNet with the same number of parameters; to our knowledge, MALI is the first method to enable training of Neural ODEs on large-scale dataset such ImageNet, while existing methods fail due to either heavy memory burden or inaccuracy. (b) In time-series modeling, MALI achieves comparable or better results than other methods. (c) For generative modeling, FFJORD model trained with MALI achieves new shate-of-the-eee results on MNIST and Cifar10.
PRELIMINARIES
2.1 NUMERICAL INTEGRATION METHODS
An ordinary differential equation (ODE) typically takes the form
dz(t) = 2()) s.t. 2(to) z, E [to, Loss = L(z(T),
(1)
dt
where =tt) is the hidden state evolving with time, T is the end time, to is the start time (typi- cally 0), x is the initial state. The derivative of x(t) w.r.t t is defined by a function f. and of is defined sequence of layers parameterized by 0. The loss function is I(=(T), y), where y is the target variable. Eq. is called the initial value problem (IVP) because only 2(tp) is specified.
Notations We summarize the notations following Zhuang al. (2020)
Algorithm 1: Numerical Integration
Input initial state start time fo, end time T, error tolerance etol, initial stepsize h.
(144)=((T) hidden state in forward/reverse time trajectory time t//T,.
vh. 2;): the numerical solution time i; + h, starting from 2t), with stepsize h.
Initialize 2(0) = x, ==t==tt to While < T
Ng: Ny is the number of layers in f in Eq. I Ng is the dimension 2.
error est
While error_es >> etol
h + h x Decay Factor (h,
Nr/N;: number discretized points (outer itera-
2, error_est
tions Algo. D in forward / reverse integration. m: average number of inner iterations Algo. to find an acceptable stepsize.
If error_est < etol
+ h Increase actor
t t ++ h
Numerical Integration The algorithm for general
ciapttvve-tepssize numerical ODE solvers is summa-
rized in Algo. [ Wanner & Hairer 1996). The solver repeatedly advances in time by step, which is the outer loop Algo. [ curve in Fig. [ For each step, the solver decreases the stepsize until estimate of error is lower than tolerance, which the inner loop in Algo. (green curve in Fig. D. For fixed-stepsize solvers, the inner loop is replaced with single evaluation of th (t, 2) us- ing predefined stepsize h. Different methods typically use different vb for example different orders of the Runge- Kutta method (Runge 1895).
2.2 ANALYTICAL FORM OF GRADIENT IN CONTINUOUS CASE
We first briefly introduce the analytical form of the gradient in the continuous case, then we compare different numerical implementations in the literature to estimate the gradient. The analytical form
2

---

Published as a conference paper at ICLR 2021
Table 1: Comparison between different methods for gradient estimation continuous case. MALI achieves reverse accuracy, constant memory w.R. number of solver steps in integration, shallow computation graph and low computation cost.
Naive
Adjoint
ACA
MALI
Computation Memory
NEN, X No mxx N=Nς x (Ni N,) X NEN xm
NEN, X N₁ x (m Na + NO) N₁ x N:
NZN X Nr (m NE 1)
NEN, Nj M,
Computation graph depth Reverse accuracy
N₁ X xm
Nr N:
24
8(T3)
3(m2)
23
Search for optimal stepsize
a(T))
Adjoint
22
R
E(Ta)
Step forward with optimal stepsize
MALI/ACA
-. - Naive
Ground-truth
70
trajectory
10
11
12
13
1.4
t
70
71
72
T3
T
Figure 1: Illustration numerical solver forward-pass. For adaptive solvers, for cach step forward-in-time, the stepsize is recursively ad- justed until the estimated error is below prede- fined tolerance; by green curve, and the accepted step (ignore the scarch process) is represented by blue curve.
Figure backward-pass, the adjoint method reconstructs trajectory as separate IVP. Naive, ACA and MALI track the forward-time trajectory, hence accurate. ACA and MALI only back- propagate through the accepted step, while naive method backpropagates through the search pro- cess hence has deeper computation graphs
search process is represented
of the gradient in the continuous case
dL
[ a(t)
Of(z(t),t,))
dt
(2)
de
00
da(t)
0f (f(2(t),,,0) a(t) Vt E (O,T), a(T) 0z(T) dt O=(1)
aL
(3)
where a(t) is the "adjoint state". Detailed proof is given in PPontryggin 1962). In the next section we compare different numerical implementations of this analytical form.
2.3 NUMERICAL IMPLEMENTATION IN THE LITERATURE FOR THE ANALYTICAL FORM
We compare different numerical implementations of the analytical form in this section. The forward- pass and backward-par pass of different methods are demonstrated in Fig. and Fig. 2 respectively. Forward-part forward-pass is similar for different methods. The comparison of backward backward-pass pass among different methods are summarized in Table. We explain methods in the literature below.
Naive method The naive method saves all of the computation graph (including search for optimal stepsize, green curve in Fig.2) in memory, and backpropagates through it. Hence the memory cost is N. N₁ x m and depth of computation graph are N2 x N; X m, and the computation is doubled considering both forward and backward passes. Besides the large memory and computation, the deep computation graph might cause vanishing or exploding gradient (Pascanu et al. 2013).
Adjoint method Note that we use "adjoint state equation" to refer to the analytical form in Eq. 2 and 3 while we use "adjoint method" to refer to the numerical implementation by byeh et al. (2018) As Fig. I and [2] the adjoint method forgets forward-time trajectory (blue curve) to achieve memory cost N Ny which is constant to integration time; takes the end-time state (derived from forward-time integration) the initial state, and solves a separate IVP (red curve) in reverse-time.
Theorem 2.1. iZhuang al. 2020 For ODE solver of order p. the error of the reconstructed initial value by the adjoint method is ( DØT (zx) zx) + h2) 1 DTb'/ (2x) 2x] + O(h+t),, where is the ideal solution, D& is the Jacobian of T. 1(1, 2) and 1(1,, 2) are the local error in forward-time and reverse-time integration respectively.
3

---

Published as a conference paper at ICLR 2021
Theorem 2.1 is stated Theorem 3.2 in Zhuang et al. (2020); please reference paper for de- tailed proof. To summarize, due to inevitable errors with numerical ODE solvers, the reverse-timg trajectory (red curve, 2(7)) cannot match the forward-ting time trajectory (blue curve, z(t)) accurately. The error in propagates to dL de by Eq. hence affects the accuracy in gradient estimation.
Adaptive checkpoint adjoint (ACA) To solve the inaccuracy of adjoint method, Zhuang et al. (020)) proposed ACA: ACA stores forward-time trajectory in memory for backward-pan hence guarantees accuracy; ACA deletes the search process (green curve in Fig. 2). and only back- propagates through the accepted step (blue curve in Fig. 2 hence has shallower computation graph (Ny x N₂ for ACA vs Ny N₂ m for naive method). ACA only stores (= and deletes the computation graph for {ft) t.) } hence the memory cost is (N + N,). Though the memory cost is much smaller than the naive method, it grows linearly with No, and can not handle very high dimensional models. In the following sections, we propose a method to overcome all these
disadvantages of existing methods.
3 METHODS
3.1 ASYNCHRONOUS LEAPFROG INTEGRATOR
In this section we give a brief introduction to the asynchronous leapfrog (ALF) method (Mutze) 2013), and we provide theoretical analysis which is missing in Mutze (2013). For general first- order ODEs in the form of Eq. [ the tuple (z,t) t) is sufficient for most ODE solvers to take step numerically. For ALF, the required tuple is (z, u, t), where v the "approximated derivative". Most numerical ODE solvers such as the Bonne-Kuuta method (Runge, 1895) track state 2 evolving with time, while ALF tracks the "augmented state" (z, v). We explain the details of ALF as below.
Algorithm 2: Forward of 2i in ALF
Algorithm 3:
x- (Inverse of v/) in ALF
Input (Zin, Vin Sin; fi) where 8in is current time, Zin and Usn are correponding values at time Sin, is stepsize.
Input (Zout; Vout, Soat, h)) where Sout is current time, Zont and Vout are corresponding values at 8out, h is stepsize.
Forward
81 Sin h/2
Inverse
S1 Sout h/2
k = Zin + Pin x h/2
kj = Rout Vout X h/2
#1 f(i1, 81)
11 f(ki,, $1)
Cout Vin + 2(u1 - Vin)
-21 Vouut
Zout ki Vout X h/2
Zin = ki - Vin X h
Sout $1 + h/2
81 -1/2
Output
(Zouts Vout: Sout, h))
Output
(Zim: Vin Sin:
Procedure of ALF Different ODE solvers have dif-
ferent da in Algo. T hence we only summarize v for
20 0a
21 11
22
23 Ug
ALF in Algo. Note that for a complete algorithm of integration for ALF, we need to plug Algo. 2 into Algo. The forward-pass is summarized in Algo. Given stepsize h, with input (Zin; Vin, Sin), single step of ALF outputs (Fout, Vout; S.u)).
fo
E1
La
Figure 3: With ALF method, given any tuple (2) By ty) and discretized time points {t. pi we can reconstruct the entire trajectory accu- rately due to the reversibility of ALF.
As in Fig. 3 given (zo, to: 40), the numerical forward- time integration calls Algo. 1ett-2iieaatin iiteraativll:
(zi, vi, li h2) (1 vz-1, 1,,4,441 4-,,, h4)
s.l. h = 44 11-1, i= 1, 2..
(4)
Invertibility of ALF An interesting property of ALF is that a defines bijective mapping: therefore, we can reconstruct (Zin, Vin, Sin, h) from (Zout, Vouts Sout, h), as demonstrated in Algo. 7 As in Fig. reconstruct the entire trajectory given the state (zj, 02) at time ty, and the discretized time points (to....t.. For example, given (ZN, UN.) and ( the trajectory for Eq. reconstructed:
(2i-1, Ui-1, U4-1, 1,-1,44-0 4-
hi) 25 - Ui, li, hi) s.l. h 4, 4i-1, i= No. No 1,.., (5)
4

---

Published as a conference paper at ICLR 2021
In the following sections, we will show the invertibility of ALF is the key maintain accuracy at constant memory cost to train Neural ODEs. Note that "inverse" refers to reconstructing the imput from the output without computing the gradient, hence is different from "back paok-ppoop.attonn.
Initial value For an initial value problem (IVP) such as Eq. [ typically 2g = =(t)) is given while oo is undetermined. We can construct to f(2((t), to), so the initial augmented state is (2o, 00)-
Difference from midpoint integrator The midpoint integrator (Süli & Mayers 2003) is similar to Algo. 2 except that it recomputes U/n f( Sin) for every step. while ALF directly uses the input Cin- Therefore, the midpoint method does not have an explicit form of inverse.
Local truncation error Theorem 3.1 indicates that the local truncation error of ALF is of order O(h"); this implies the global error is O Detailed proof is in Appendix A.3
Theorem 3.1. For a single step in ALF with stepsize h, the local truncation error of is Ohh), and the local truncation error of v is O
A-Stability The ALF solver has a limited stability region, but this can be solved with damping. The damped ALF replaces the update of Vout in Algo. 2 with Vout Vin 2n7uu Vin). where 1) is the "damping coefficient" between O and We have the following theorem on its numerical stability.
Theorem 3.2. For the damped ALF integrator with stepsize h, where o; is the i-th eigenvalue of the Jacobian of, then the solver is A-stable if + :j(ho; 1) + V [2hh, + n/(ho: 1)2] <1, Vi
Proof is in Appendix A..4 and A.5 Theorem 3.2 implies the following: when 17 1, the damped ALF reduces to ALF, and the stability region is empty; when < v) < 1, the stability region non- empty. However, stability describes the behaviour when T goes to infinity; in practice we always use a bounded T and ALF performs well. Inverse of damped ALF is in Appendix A.5.
3.2 MEMORY-EFFICIENT ALF INTEGRATOR (MALI) FOR GRADIENT ESTIMATION
An ideal solver for Neural ODEs should achieve two goals: accuracy in gradient estimation and constant memory cost w.F. integration time. Yet none of the existing methods can achieve both goals. We propose a method based on the ALF solver, which to our knowledge is the first method to achieve the two goals simultaneously.
Algorithm 4: MALI to acheive accuracy at a constant memory cost w.r. integration time Input Initial state 20, start time to. end time T Forward
Apply the numerical integration in Algo. [ with the vi function defined by Algo.
Delete computation graph on the fly, only keep end-time state (EN,, (N,)
Keep accepted discretized time points {ta ) (ignore process to search for optimal stepsize)
Backward
dL Initialize ((T) = 8L by Eq. initialize Jā(T)
= 0 de
For i in (Ne, N₁ - ..., 2, 1}:
Reconstruct (2І-1, 1)-1) from (2,, vz) by Algo.
Local forward (2,, Us, t4, h₂) 20 U/ t--11 h4) Local backward, get 0f(e 1.:--11 and 0f(e:
1.4 1.0)
69
Update a(t) and dL aB by Eq. 2 and Eq. discretized time points 4,-1 and f,
and Delete local computation graph
Output the adjoint state a(to) (gradient wir/ input 20) and parameter gradient
dl de
Procedure of MALI Details of MALI summarized in Algo. For the forward-pass, we only keep the end-time state (ZN., UN,) and the accepted discretized time points (blue curves in Fig. and 2 We ignore the search process for optimal stepsize (green curve in Fig. and 2 and delete other variables to save memory. During the backward pass, we can reconstruct the forward-ting trajectory as in Eq. then calculate the gradient by numerical discretization of Eq. and Eq.
Constant memory cost w.r.f number of solver steps in integration We delete the computation graph and only keep the end-time state memory. The memory cost is N₂ N.(N, 1), where
5

---

Published as a conference paper at ICLR 2021
Amplitude of error dL/dzo
Amplitude of error in dL/da
Memory consumption
1.0
1.0
Eror in gragdieett
0.8
ACA
MALI
Error in gradient
0.8
MALI
ACA
12.5
ACA
Nalve
Nalve
Memory 3
10.0
Naive
0.6
0.6
Adjeint
Adjaint
7.5
Adjaint
0.4
0.4
0.2
0.2
2.5
0.0
0.0
24
26
20
22
24
26
1.0
10
10-3
Integration
Integration
Error Tolerance
Figure 4:
Comparison of error in gradient in Eq. (a) error in
dL in db
(b) error
(c) memory cost.
dzo
Test Acc on CIFAR10
Test ACC Eooch
Test ACC Time (90 Epochs)
90
90
93
Accuracy
Accurscy
92
70
MAUI
70
MALI
ACA
ACA
91
Naive
Maive
50
Adjoint
50
Adjoint
90
MALI
ACA
Naive
Adjoint
ResNet
20
Training spoch
80
Training Time for 90 Epachs (hours)
5
10
Figure 5: Results on Cifar10. From left right: (1) box plot of test accuracy (first columns Neural ODEs, last is ResNet); (2) test accuracy t.ttd v.s. training epoch for Neural ODE; (3) test accuracy te v.s. training time of 90 epochs for Neural ODE.
N is due to evaluating f(z, and is irreducible for all methods. Compared with the adjoint method, MALI only requires extra N₂ memory record UN,, and also has constant memory cost W.E/ time step N,- The memory cost is NN 1).
Accuracy Our method guarantees the accuracy of reverse-time trajectory (e.g. blue curve in Fig. matches the blue curve in Fig. []. because ALF is explicitly invertible for free-form f (see Algo. Therefore, the gradient estimation in MALI more accurate compared to the adjoint method.
Computation cost Recall that on average it takes m steps to find an acceptable stepsize, whose еттоr estimate is below tolerance. Therefore, the forward-pass with search process has computation burden N₂ X 2 x N₂ x m. Note that we only reconstruct and backprop through the accepted step and ignore the search process, hence takes another N₂ N₂ N₂ 2 computation. The overall computation burden is N₂N, x N₂ X (m + 2) as in Table [
Shallow computation graph Similar to ACA. MALI only backpropagates through the accepted step (blue curve in Fig. 2 and ignores the search process (green curve in Fig. 2),, hence the depth of computation graph N x N₁. The computation graph of MALI is much shallower than the naive method, hence is more robust vanishing and exploding gradients (Pascanu et al., 2013).
Summary
The adjoint method suffers from inaccuracy in reverse-time trajectory, the naive method suffers from exploding vanishing gradient caused by deep computation graph, and ACA finds balance but the memory grows linearly with integration time. MALI achieves accuracy in reverse- time trajectory, constant memory integration time, and shallow computation graph.
EXPERIMENTS
4.1 VALIDATION ON A TOY EXAMPLE
We compare the performance of different methods toy example, defined
= =
s.l. x(0) 20; dz(l)/dt = a=a(1)
(6)
The analytical solution is =(t) at L 2aT dL/dzo 2aT dL//do 2T =3 =fe227 20T
We plot the amplitude of error between numerical solution and analytical solution varying with T (integrated under the same error tolerance, rtol 10-5 atol = 10 in Fig 4 ACA and MALI have similar errors, both outperforming other methods. We also plot the memory consumption for

---

Published as a conference paper at ICLR 2021
Table 2: Top-1 test accuracy Neural ODE and ResNet on ImageNet. Neural ODE trained with MALI, and ResNet is trained as the original model; Neural ODE is tested using different solvers without retraining.
Fixed-stepaize solvers of varioss stepsizes
Adaptive-stepsize solver of varioss tolerances
Stepsize
0.15
1,00F-01
1,00F-02
MALI
42.35
66.4
69.59
70.17
69.94
MALI
62.56
69.89
69.87
Neural
Faler
21.94
61.25
68.69
70.02
Huler
69.87
69.88
ODE
RK2
42.33
69
69.72
70.14
69.92
RK23
50.77
69.89
69.93
RK4
12.6
69.91
70.21
69.96
DopriS
52.3
68.58
69.71
ResNet
70.09
Table 3: Top-1 accuracy under FGSM attack. is the perturbation amplitude, For Neural ODE modcls, row names represent the solvers to derive the gradient for attack, and column names represent solvers for inference on the perturbed image.
1/255
2/2nh
Euler
RK23
DopriS
MALI
Euler
RK21
DopriS
MALT
14.09
14.72
14.77
15.71
10.38
10.46
10.62
10.62
Heun Fuler
14.80
ODE
RK23
14.82
14.77
14.79
15.69
10.78
10.53
10:48
10.56
Dupri5
14.79
15.15
10.49
10.48
10.51
Res
13.02
9.57
different methods on a Neural ODE with the same input in Fig. 4 As the error tolerance decreases, the solver evaluates more steps, hence the naive method and ACA increase memory consumption, while MALI and the adjoint method have a constant memory cost. These results validate our analysis in Sec. 3.2 and Table I and shows MALI achieves accuracy at constant memory cost.
4.2 IMAGE RECOGNITION WITH NEURAL ODE
We validate MALI on image recognition tasks using Cifarl0 and ImageNet datasets. Similar to Zhuang et al. (2020). , we modify a ResNet18 into its corresponding Neural ODE: the forward func- tion is y = x + fo(x) and = x + +o fo(2 )dt for the residual block and Neural ODE respectively, where the same fe is shared. We compare MALI with the naive method, adjoint method and ACA.
Results on Cifarl0 Results of 5 independent runs on Cifarl0 are summarized in Fig. 5 MALI achieves comparable accuracy ACA, and both significantly outperform the naive and the adjoint method. Furthermore, the training speed of MALI is similar to ACA, and both are almost two times faster than the adjoint memthod, and three times faster than the naive method. This validates our analysis on accuracy and computation burden in Table 1
Accuracy on ImageNet Due to the heavy memory burden caused by large images, the naive method and ACA are unable to train Neural ODE on ImageNet with 4 GPUs; only MALI and the adjoint method are feasible due to the constant mem- ory. We also compare the Neural ODE to standard ResNet. As shown in Fig. 6) the accuracy of the Neural ODE trained with MALI closely follows ResNet, and significantly outper- forms the adjoint method (top-1 validation: 70% v.s. 63%).
Test ACC on ImageNet
70-
Accuracy
40-
ResNet
Neural ODE (MALI)
Neural ODE (adjoint)
Invariance to discretization scheme A continuous model should be invariant to discretization schemes (e.g. different types of ODE solvers) as long as the discretization is suf- ficiently accurate. We test the Neural ODE using different solvers without re-training; since ResNet is often viewed as one-step Fuler discretization of an ODE (Haber & Ruthotto 2017). we perform similar exper- iments. As shown in Table 2 Neural ODE consistently achieves high accuracy (~70%), while ResNet drops to random guessing (-) because ResNet one-step Euler discretization fails to be a meaningful dynamical system (Queiruga et al. 2020).
20-
20
40
60
80
Epach
Figure 6: Top-1 accuracy on Ima-
geNet validation datasct.
Robustness to adversarial attack Hanshu et al. (2019) demonstrated that Neural ODE is more robust to adversarial attack than ResNet on small-scal scale datasets such as Cifar10. We validate this result on the large-scale ImageNet dataset. The top-1 accuracy of Neural ODE and ResNet under FGSM attack (Goodfellow et al.) 2014) are summarized in Table 3 For Neural ODE, due to its invariance to discretization scheme, derive the gradient for attack using a certain solver (row in
7

---

Published as a conference paper at ICLR 2021
Table 4: Test MSE (x0.01) on Mujoco dataset (lower is better), Results
Table 5: Test ACC on Speech
marked with superscript numbers correspond to literature in the footnote.
Command Dataset
Method
Accuracy (%)
Percentage of training data
RNN
RNN-GRU
Latent-ODE
Adjoint"
92.8 1 0.4
Adjoint
Naive"
ACA*
MALI
SemiNorm
92.9 1 0.4
10%
2.45
1,97
0.47
0.36*
0.31
0.35
Naive
93.2 + 0.2
20%
1.71'
1.42'
0.44
0.30²
0.27²
0.27
ACA
93.2 + 0.2
50%
0.79
0.75
0.40
0.292
0.26
0.26
MALI
93.7 1: 0.3
Table 3). and inference on the perturbed images using various solvers. For different combinations of solvers and perturbation amplitudes, Neural ODE consistently outperforms ResNet.
Summary In image recognition tasks, we demonstrate Neural ODE is accurate, invariant to dis- cretization scheme, and more robust adversarial attack than ResNet. Note that detailed explana- tion on the robustmess of Neural ODE is out of the scope for this paper, but to our knowledge, MALI the first method enable training of Neural ODE on large datasets due constant memory cost.
4.3 TIME-SERIES MODELING
We apply MALI to latent-ODD ODE (Rubanova et al. 2019) and Neural Controlled Differential Equation (Neural CDE) (Kidger et al., 2020a b). Our experiment is based on the official implementation from the literature. We report the mean squared error (MSE) the Mujoco test set in Table 4 which is generated from the "Hopper" model using DeepMind control suite TTssa et al. 2018); for all experiments with different ratios of training data, MALI achieves similar MSE to ACA, and both outperform the adjoint and naive method. We report the test accuracy on the Speech Command dataset for Neural CDE in Table MALI achieves higher accuracy than competing methods.
4.4 CONTINUOUS GENERATIVE MODELS
We apply MALI on FFJORD (Grathwohl et al. 2018), free-from continuous generative model, and compare with several variants in the literature (Finlay al.) 2020 Kidger et al.) 2020a). Our experiment is based on the official implementaion of Finlay al. (2020); for a fair comparison, we train with MALI, and test with the same solver as in the literature (Grathwohl et al.) 2018 Finlay ct al. 2020), the Dopri5 solver with rtol atol 10 from the torchdiffeq package (Chen et al, Bits per dim (BPD, lower is better) on validation set for various datasets are reported in For continuous models, MALI consistently generates the lowest BPD, and outperforms the Vanilla FFJORD (trained with adjoint), RNODE (regularized FFJORD) and the SemiNorm Adjoint
2018
Table 6
Kidgeer at. 2020a). Furthermore, FFJORD trained with MALI achieves comparable BPD state- of-the-art discrete-layer flow models in the literature. Please see Sec. B.3 for generated samples.
RELATED WORKS
Besides ALF, the symplectic integrator (Verlet 1967 Yoshida 1990) is also able to reconstruct tra- jectory accurately, yet it's typically restricted to second order Hamiltonian systems (De Almeida 1990), and are unsuitable for general ODEs. Besides aforementioned methods, there are other meth- ods for gradient estimation such as interpolated adjoint (Daulbaev et al. (020) and spectral method Quaglino al. 2019), yet the implementations are involved and publicly available. Other works tocus on the theoretical properties of Neural ODEs (Dupont et al. 2019: Tabuada & Gharesifard 2020 Massaroli al. 2020). Neural ODE is recently applied to stochastic differential equation (Li et al. 2020). jump differential equation (Jia & Benson 2019) and auto-regressive models (Wehenkel Louppe 2019).
CONCLUSION
Based on the asynchronous leapfrog integrator, we propose MALI to estimate the gradient for Neural ODEs. To our knowledge, our method is the first to achieve accuracy, fast speed and constant memory cost. We provide comprehensive theoretical analysis on its properties. We validate MALI
01. Rubanova et al. (2019) ; 2. Zhuang et al. (020); 3. Kidger et al. (0200a): 4. Chen al. 2018): 5. Finlay
al. 2020); 6. Dinh et al. 2016); Behrmann et al. 2019: 8. Kingma & Dhariwal 2018)
: 9. Ho
2019
10. Chen al. (0019
8

---

Published as a conference paper at ICLR 2021
Table 6: per dim (BPD) of generative models, lower is better. Results marked with superscript numbers correspond to literature in the footnote.
Continuous Flow (FFIORD)
Discrete Flow
Vanilla*
RNODE"
Sem:Nrmm'
MALJ
RealNVP
R-ReNerr
Glow
Flowe*"
Residual Flow
MNIST
0.991
0.97
0.96
0.87
1.06
1.05'
1.05
0.97
CIFARIO
3.40*
3.38
3.35
3.49"
3.457
3.35"
3.28
3.28
ImageNet64
3,83
3,98
3.76
with extensive experiments, and achieved new state-of-the-ar results in various tasks, including image recognition, continuous generative modeling, and time-series modeling.
ACKNOWLEDGEMENT
This research was funded by the National Institutes of Health (NINDS-ROI HANDS-ROINSO33193)
REFERENCES
Jens Behrmann, Will Grathwohl, Ricky TQ Chen, David Duvenaud, and Jörn-Henrik Jacobsen. Invertible residual networks. In International Conference on Machine Learning, pp. 573-582, 2019.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, pp. 6571-6583, 2018.
Ricky TQ Chen, Jens Behrmann, David K Duvenaud, and Jörn-Henris Jacobsen. Residual flows for invertible generative modeling. In Advances in Neural Information Processing Systems, pp- 9916-9926, 2019.
Talgat Daulbaev, Alexandr Katrutsa, Larisa Markeeva, Julia Gusak, Andrzej Cichocki, and Ivan Oseledets. Interpolated adjoint method for neural odes. arXiv preprint arXiv:2003.0527/ 2020.
Alfredo M Ozorio De Almeida. Hamiltonian systems: chaos and quantizaaion. Cambridge Univer- sity Press, 1990.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint 1rXiv:1605.000 08803, 2016.
Emilien Dupont, Arnaud Doucet, and Yee Whye Teh. Augmented neural odes. In Advances in Neural Information Processing Systems, pp. 3140-3150, 2019.
Chris Finlay, neural ode: the world of jacobian and kinetic regularization. International Conference on Machine Learning. 2020.
Jacobsen, Levon Nurbekyan, and Adam M Oberman. How to train your
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. arXiv preprint 4rXiv:1412.6572, 2014.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv preprint arXiv:1810.011367 2018.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse Problems, 34(1):014004, 2017.
YAN Hanshu, DU Jiawei, TAN Vincent, and FENG Jiashi. On robustness of neural ordinary differ- ential equations. In International Conference on Learning Representations, 2019.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.

---

Published as a conference paper at ICLR 2021
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. arXiv preprint arXivv:1022.00275. 2019.
Junteng Jia and Austin Benson. Neural jump stochastic differential equations. In Advances in Neural Information Processing Systems, pp. 9847-9858, 2019.
Patrick Kidger, Ricky T. Q. Chen, and Terry Lyons. "Hey, that's not ODE": Faster ODE Adjoints with 12 Lines of Code. arXiv:2009.0943) 2020a.
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equa- tions for irregular time series. arXiv preprint arXiv:2005.088 08926, 2020b.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In Advances in Neural Information Processing Systems, pp. 10215-10224, 2018.
Xuechen Li, Ting-Kan Kam Leonard Wong, Ricky TQ Chen, and David Duvenaud. Scalable gradients for stochastic differential equations. arXiv preprint arXIv:2001.01328, 2020.
Kuang Liu. Train cifar10 with pytorch. 2017. URL httpp://giithu..com//o pytorch-cifar
Yiping Lu, Aoxiao Zhong, Quanzheng Li, and Bin Dong. Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations. In International Conference on Machine Learming. pp. 3276-3285. PMLR, 2018.
Stefano Massaroli, Michael Poli, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama, Dissecting neural odes. arXiv preprint 200x::0002.0807,,, 2020.
UIrich Mutze. An asynchronous leapfrog method ii. arXiv preprint
2013.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. the difficulty of training recurrent neural networks. In International conference on machine learning, pp. 1310-1318, 2013.
Lev Semenovich Pontryagin. Mathematical theory of optimal processes. Routledge, 1962.
Alessio Quaglino, Marco Gallieri, Jonathan Masci, and Jan Koutník. Snode: Spectral discretization of neural odes for system identification. arXiv preprint arXiv: 1706,0008, 2019.
Alejandro F Queiruga, N Benjamin Erichson, Dane Taylor, and Michael W Mahoney. Continuous- in-depth neural networks. arXiv preprint arXiv:2008.023:0 2020.
Yulia Rubanova, Ricky TQ Chen, and David K Duvenaud. Latent ordinary differential equations for irreuularly-sampledd time series. In Advances in Neural Information Processing Systems, pp- 5320-5330, 2019.
Carl Runge. Über die numerische auflösung von differentialgleichunen. Mathematische Annalen, 46(2):167-178, 1895.
Lars Ruthotto and Eldad Haber.
Deep neural networks motivated by partial differential equations. Journal of Mathematical Imaging and Vision, pp. 1-13, 2019.
Lars Ruthotto, Stanley J Osher, Wuchen Li, Levon Nurbekyan, and Samy Wu Fung, A machine learning framework for solving high-dimensional mean field game and mean field control prob- lems. Proceedings of the National Academy of Sciences, (17(1))9183-9119) 2020.
Alvaro Sanchez-Gonzales Victor Bapst, Kyle Cranmer, and Peter Battaglia. Hamiltonian graph networks with ode integrators. arXiv preprint arXiv: 2019.
John R Silvester. Determinants of block matrices. The Mathematical Gazette, 84(501):400-467, 2000.
Endre Süli and David Mayers, An introduction to numerical analysis. Cambridge university press, 2003.
10

---

Published as a conference paper at ICLR 2021
Paulo Tabuada and Bahman Gharesifard. Universal approximation power of deep neural networks via nonlinear control theory. arXiv preprint wxXiv2000.060007 2020.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite. arXiv preprint arXiv:1801.000 000900 2018.
Loup Verlet. Computer" experiments" on classical fluids. Thermodynamical properties of Lennard-Jones molecules. Physical review, 159(1):98, 1967.
Gerhard Wanner and Ernst Hairer. Solving ordinary differential equations II. Springer Berlin Hei- delberg, 1996.
Antoine Wehenkel and Gilles Louppe. Unconstrained monotonic neural networks. In Advances in Neural Information Processing Systems, pp. 1545-1555, 2019.
E Weinan. A proposal on machine learning via dynamical systems. Communications in Mathematics and Statistics, 5(1):1-11, 2017.
Haruo Yoshida. Construction of higher order symplectic integrators. Physics letters A, 150(5-7): 262-268, 1990.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Symplectic ode-net: Learning hamiltonian dynamics with control. arXiv preprint arXiv: 199.1::909.12077, 2019.
Juntang Zhuang, Nicha Dvorekk Xiaoxiao Li, Sekhar Tatikonda, Xenophon Papademetris, and James Duncan. Adaptive checkpoint adjoint method for gradient estimation neural ode. Inter- national Conference on Machine Learning, 2020.

---

Published as a conference paper at ICLR 2021
CONTENTS (APPENDIX)
Theoretical properties of ALF integrator'
13
A.I
Algorithm of ALF
13
A.2
Preliminaries
13
A.3
Local truncation emror of ALF
13
A.4
Stability analysis
15
A.5
Damped ALF
16
Experimental Details
18
B.1
Image Recognition
18
B.1.1
Experiment Cifar10
18
B.1.2
Experiments on ImageNet
18
B.2
Time series modeling
19
B.3
Continuous generative models
20
B.3.1 Training details
20
B.3.2
Addrional results
20
Error gradient estimation for toy examples when
20
B.5
Results of damped MALI
20
12

---

Published as a conference paper at ICLR 2021
A THEORETICAL PROPERTIES OF ALF INTEGRATOR
A.1 ALGORITHM OF ALE
For the ease of reading, write the algorithm for di in ALF below, which is the same as Algo. in the main paper, but uses slightly different notations for the ease of analysis.
Algorithm 1: Forward of vi ALF
Input (Zeyy Van Sin ft) (žo, 6o, h) where so is current time, 20 and 60 correponding values at time So; stepsize h. Forward
80 + h/2
21 20 + 60 x h/2
(2)
1 8()),81)
(3)
1/2 6 + (0) 06)
2 E 02 x h/2
52 = 51 + h/2
(6)
Output
Fouts Coat; Souts h) (22, U2 82, h)
For simplicity, we can re-write the forward of ALF as
E0 + h hf + l Vo, s0 5
2 + 10 80 + 60
Similarly, the inverse of ALF can be written as
12 11
b 62 4))
(8)
06
12f(22 - h 62,02 82 ()) 62
A.2 PRELIMINARIES
For an ODE of the form
dz(()
f(=((),t)
(9)
dt
We have:
d²=(11 d
[(((1).)) Of(((),() 0f(((1), dz(()
(10)
dt²
di
Ot
Oz
dt
For the ease of notation, we re-write Eq. 10 as
d²z(()
Sc + Szf
(11)
di2
where fi and fa represents the partial derivative of w.n. t and respectively.
A.3 LOCAL TRUNCATION ERROR OF ALF
Theorem A.1 (Theorem 3.1 in the main paper). For a single step in ALF with stepsize h, the local truncation error of is O and the local truncation errof of u O(h2).
Proof. Under the same notation Algo. I denote the ground-truth state of and starting from (oo, 5o) as and respectively. Then the local truncation error is
L, 3(00 h) E2 Ian
7(5o h) 02
(12)
13

---

Published as a conference paper at ICLR 2021
We estimate L2 and Ly in terms of polynomial of h.
Under mild assumptions that f is smooth up to 2nd order almost everywhere (this is typically satis- fied with neural networks with bounded weights), hence Taylor expansion is meaningful for f. By the Taylor expansion
Eq.
around point (2o: Dg. sp) is
[(su + h) = Eu + h d=
h2 d2z
+
+ O
(13)
dt
2 dt
zu + hf so) + h
( su) + fz(Eo, so) sol) + O(h²)
(14)
Next, we analyze accuracy of the numerical approximation. For simplicity, we directly analyze Eq. by performing Taylor Expansion on
h h f + Up, So + 2
hip fz so) + 0
2
= f( s0) + 2
11g5((a0 f10 su) +
(15)
2
h 2 0b 80
h
$ 20 60
2
(16)
Plug Eq.
14 Eq. 15 and E.q. 16 into the definition of L., we get
L [soo + h) E2
(17)
112
20 + hf so) + [f( so) f2 so) ( so)
2
[ + 1h(10 ((((2) su)
h
+2 f2 (20 (n) + 0
hōg
(18)
2
122
2 fz fz( (o) f(20, so) tu) O(h3)
(19)
Therefore, if [( so) 6 is of order O(1), L_ is of order O(h²); f( s0) - 0o is order O(h) or smaller, then L of order (11)). Specifically, at the start time of integration, we have [( 60) - 6g 0/. by induction, Lg at end time (11)).
Next we analyze the local truncation error in v, denoted L. Denote the ground truth ō(to + h), we have
(80 + h) = (3(so + hi), 80 h) = = so) + h-fe Eo, s0) + (3(80 + h) f2).2 80) + 0(h3)
(20)
(21)
Next analyze the error in the numerical approximation. Plug Eq. 15 into Eq.
2 No + h [ 8o) + (J 20) 60) h s0) (ooos((F, O(h3)
6 = 25 +
7/ - 6
(22)
O(()))
(23)
From Eq.
14 Eq. 21 and Eq. 23 we have
I. ū(so h) 02
(24)
[( so) 60)
(((s + h-- - hvm) fa fa s0) + 0(13)
(25)
(f(E), so) 5o) f(E,, so) - 60 fa (2o, so) + O(h²)
(26)
The last equation is derived by plugging in Eq. 14 Note that Eq. 26 holds for every single step forward in time, and at the start time of integration, we have [( Sp) fy = due to our initialization as in Sec. 3.1 of the main paper. Therefore, by induction, L,, of order O(h2) for consecutive steps.
14

---

Published as a conference paper at ICLR 2021
A.4 STABILITY ANALYSIS
A B
Lemma A.1.1. For matrix of the form
if A, C, are square matrices of the same
C D
A B
shape, and CD DC, then we have det
det(AD BC)
C D
Proof. See (Silvester 2000) for a detailed proof.
Theorem A.2. For ALF integrator with stepsize h, if ho; is 0 or is imaginary with norm no larger than I, where o, is the i-th eigenvalue of the Jacobian af of A-stability; otherwise, the solver not A-stable.
then the solver is on the critical boundary dz'
Proof. A solver is A-stable is equivalent to the eigenvalue of the numerical forward has norm below 1. We calculate the eigenvalue & below.
For the function defined by Eq. 7 the Jacobian is
83 0zo
03 Delg
T + h 01
A 0f
0c
3 d2
(27)
06
062 0eu
2
8f 8%
h
320
0s
We determine the eigenvalue of J by solving the equation
h 85 0z
+ (1 A)T
h7t 85 0z
det(J - -1)
(28)
2 82
of 02
(1 + A)
It's trivial to check satisfies conditions for Lemma A.I.I Therefore, we have
det(J A1) det
0f Oz
+ (1 111
(n
Of of 02
(1+ - oz
h² 0f 2
]
0f
(29)
Oz
0f Oz
det
2\h
+ (x=-1)1 1)1
(30)
81 Suppose the eigen-decompostion of 05
can be written
01
aʃ
A
02
A-1
(31)
ON
Note that / = ALX hence we have
Toi
117A-- -(12-117AA-
del(J AT) det A
2Xh
02
+ + -
(32)
NN_
II (xa 2ho 1)
(33)
Hence the eigenvalues are
A4 ho, + v 1207
(34)
A-stability requires |deal < Vi, and has no solution.
The critical boundary is Dia = 1, the solution is: ha; is 0 or on the imaginary line with norm no larger than 1.
15

---

Published as a conference paper at ICLR 2021
A.5 DAMPED ALF
Algorithm 2:
Forward of ab in Damped ALF (n E (),
Input (Z Vin Sinz h) = (Zo, ous 50: h)) where so is current time, Ey and Ey are correponding
values at time Soi stepsize h. Forward
81 Su + h/2
(35)
& Σσ + t0 x h/2
(36)
6i = [( 81)
(37)
62 60 + 2m 61 0%
(38)
52 + 0/ x h/2
(39)
= + h/2
(40)
Output
&outs Coat, Souts h) (22, 82, h2
Algorithm 3: x-1 (Inverse of vb) in Damped ALF (n E (0,
Input (Zout: Vout, Soat, h) where Sout is current time, Zout and Yout are corresponding values at Souts h is stepsize. Inverse
(22, 02 82, 82, h) = (Zout, Voat, Sout, h))
(41)
81 $2 h/2
(42)
& = 22 02 x h/2
(43)
0 81)
(44)
to = (62 2nτe) (1 -2n)
(45)
& 6 h/2
(46)
h/2
(47)
Output
Zins Vnn Sin, h) = (50, Co. So. h))
The main difference between ALF and Damped ALF is marked blue in Algo. 2 In ALF, the update of U2 is 02 (úi 66) 6 2 og) 6o; while in Damped ALF, the update scaled by a factor 1) between 0 and 1, so the update is t2 ( o0) + T.. When rj 1, Damped ALF reduces ALF.
Similar to Sec. A.I
we can write the forward as For simplicity, we can re-write the forward of ALF
2o + nhf 60. so (1) (1 1o
22
(48)
2 + 2 su + ()) + (1 - 27)
Similarly, the inverse of ALF can be written
T3 h
1
f 4/ 82 4
T-2n
(49)
2n 2n
f
b 62, 82 4))
1-
Theorem A.3. For a single step in Damped ALF with stepsize h, the local truncation error of z is (h2),, and the local truncation errof of is (h)).
16

---

Published as a conference paper at ICLR 2021
Proof. The proof is similar to Thm. Eq. 13 and Eq. 14 we have
A.3 By similar calculations using the Taylor Expansion in
h
2 f( 00) + héo fz so)
2 (80 + h) =1 (1-n)hoi + hr/ f(6 (0. su)
[ so) + h 200
h
fe 50) (50, s0)
+ +1h))
(50)
2
= (1 1-- 11) 60 ((50,
7- 112 S1 80)
2
h²
+
noo - f so) fz( su) + O(h²)
(51)
2
Using Eq.
21 Eq.
15 and Eq. 14 we have
0/
(1 2n) + (2n- 1)f(20 80) + (1 - n)) su)
(((5 + It) Eg n)hro) f2(20; 80) O(h²)
(52)
(2n () 80) To} (1 - n)),,((E50 8o)
7 f so) hooo f= so) + O
(53)
Note that when v7 Eq. 51 reduces to Eq. 19 and Eq. 53 reduces Eq. 26 By initialization,
we have [ 50) - 0p. 0 at initial time, hence by induction, the local truncation error for is
Ohh?); the local truncation error for u is (hh when 77 1, and O(h) when 17 =
O
Theorem A.4 (Theorem 3.2 in the main paper). For Dampled ALF integrator with stepsize h,
where og is the i-th eigenvalue of the Jacobian of then the solver is A-stable if + nnhoo
8z
1) + / (2(a, + t/ho: 1)2] 1, Vi.
Proof. The Jacobian of the forward-past ppass of a single step damped ALF is
+ mp oz
(1 n)hl + +1 1
01 d=
J
(54)
2 of
8f 1/h 0=
+ (1- -1
when 17 1, .J reduces to Eq. 27 We can determine the eigenvalue of J using similar techniques.
Assume the eigenvalues for 85 are }or then we have
01 (1 A)I nh
(1 (1- 1)n + 1)2 222 of
O=
det(J AI) = det
(55)
af 2n 05
h/ 0z
+ (1 - 2n - A)T
det
[(1 1) + 1/h
0f
(nh 0f
+ (1 - 2n A)1)
0z
0z
(11 - 1// - 122 a
0
2m Oz)
(56)
2 0z/
II 1/hha 1) + [2hoo n/hh: 1)2]
(57)
when 17 < 1, it's easy to check that 1 1+n - 1) + / m 22oo -nho, - 1)2] < 1 has non-empty solutions for ho.
For quick validation, we plot the region of A-stability on the imaginary plane for a single eigen- value in Fig. [] As 17 increases, the area of stability decreases. When 77 1, the system is no-where A-stable, and the boundary for A-stability is on the imaginary axis [-i, z] where i the imaginary unit.
17

---

Published as a conference paper at ICLR 2021
44
40
Figure 1: Region of A-stability for eigenvalue on the imaginary plane for damped ALF. From left to right, the region of stability for 7) 0.25, 7) 0.7,17 0.8 respectively. As 7 increases to the area of stability region decreases.
B EXPERIMENTAL DETAILS
B.1 IMAGE RECOGNITION
B.1.1 EXPERIMENT ON CIFARIO
We directly modify a ResNet18 into a Neural ODE, where the forward of a residual block (y + fx)) and the forward of an ODE block (y x J f(z, t)dt where T 1) share the same parameterization f. hence they have the same number parameters. Our experiment is based on the official implementation by Zhuang et al. (2020) and an open-source repository (Liu 2017).
All models are trained with SGD optimizer for 90 epochs, with an initial learning rate of 0.01, and decayed by factor of 10 at 30th epoch and 60th epoch respectively. Training scheme is the same for all models (ResNet, Neural ODE trained with adjoint, naive, ACA and MALI). For ACA, we follow the settings in (Zhuang et al. 2020) and use the official implementation torch ACAA and use a Heun-Eule solver with rtol = 10 atol = 10-2 during training. For MALI, we use an adaptive version and set rtol- 10-1 atol 10-2 For the naive and adjoint method, Dopri5 solver from the torchuiiffe package with rtol atol 10-5, We train all models for independent runs, and report the mean and standard deviation across runs.
the default
B.1.2 EXPERIMENTS ON IMAGENET
Training scheme We conduct experiments ImageNet with ResNet18 and Neural-ODE18. All models are trained on 4 GTTX-1800T GPUs with batchsize of 256. All models are trained for 80 epochs, with initial learning rate of O.I. and decayed by factor of 10 at 30th and 60th epoch. Note that due to the large size input 256 x 256, the naive method and ACA requires huge memory, and is infeasible train. MALI and the adjoint method requires constant memory hence is suitable for large-scale experiments. For both MALI and the adjoint menthod, we use a fixed stepsize of 0.25, and integrates from 0 to T 1. As shown in Table. 2 in the main paper, stepsize of 0.25 is sufficiently small to train meaningful continuous model that is robust to discretization scheme.
Invariance to discretization scheme To test the influence of discretization scheme, we test our Neural ODE with different solvers without re-training. For fixed-stepsize solvers, we tested various step sizes including 10.1, 0. 15, 0.25, 0.5, 1.0}; for adaptive solvers, we set rtol=0.1, atol=0.01 for MALI and Heun-Euler method, and set rtol = 10-2 atol 10-3 for RK23 solver, and set rtol = 10-4 atol = 10 for Dopri5 solver. As shown in Table. 2 Neural ODE trained with MALI is robust to discretization scheme, and MALI significantly outperforms the adjoint method in terms of accuracy (70% 63% t accuracy on the validation dataset). An interesting finding is that when trained with MALI which is a second-order solver, and tested with higher-order solver (e.g.
ACA
18

---

Published as a conference paper at ICLR 2021
Training ACC on ImageNet
Test ACC on ImageNet
70
70
60
mn
Accuracy
Accuracy
50
50
40
40
A N
30
ResNet
ResNet
Neural (MALI)
30
Neural ODE (MALI)
20
Neural ODE (adjoint)
Neural ODE (adjoint)
10
20
20
40
60
80
0
20
40 Epoch
60
80
Epoch
(a) Training curve on ImageNet.
(b) Validation curve ImageNet.
Figure 2: Results on ImageNet.
RK4), our Neural ODE achieves 70.21% top-1 accuracy, which is higher than both the same solver during training (MALI, 69.59% accuracy) and the ResNet18 (70.09% accuracy).
Furthermore, many papers claim ResNet to be an approximation for an ODE (Lu al., 2018). However, Queiruga et al. (2020) argues that many numerical discretizations fail to be meaningful dynamical systems, while our experiments demonstrate that our model is continuous hence invariant to discretization schemes.
Adversarial robustness Besides the high accuracy and robustness to discretization scheme, an- other advantage of Neural ODE is the robustness adversarial attack. The adversary robustness of Neural ODE is extensively studied in (Hanshu et al.) 2019), but not only validated on small-scale datasets such as Cifar10. To our knowledge, our method the first enable effectuve training of Neural ODE on large-scale datasets such as ImageNet and achieve a high accuracy, and we are the first to validate the robustness of Neural ODE on ImageNet. We use the advertorch 3 toolbox to perform adversarial attack. We test the performance of ResNet and Neural ODE under FGSM attack. To be more convincing, we conduct experiment on the pretrained ResNet18 provided by the official PyTorch website Since Neural ODE is invariant to discretization scheme, it's possible to derive gradient for attack using one ODE solver, and inference on the perturbed image using another solver. As summarized in Table. 3 Neural ODE consistently achieves higher accuracy than ResNet under the same attack.
B.2 TIME SERIES MODELING
We conduct experiments on Latent-ODE models (Rubanova et al. 2019) and Neural CDE (con- trolled differential equation) (Kidger et al.) 2020a). For all experiments, we use the official imple- mentation, and only replace the solver with MALL. The latent-ODE model is trained on the Majoco dataset processed with code provided by the official implementation, and experiment with dif- ferent ratios (00%%2%%%55055
of training data as described in (Rubanova et al., 2019). All models trained for 300 epochs with Adamax optimizer, with initial learning rate of 0.01 and scaled by 0.999 for each epoch. For the Neural CDE model, for the naive method, ACA and MALI, we perform independent runs and report the mean value and standard deviation; results for the adjoint and seminorm adjoint are from (Kidger et al.] (2200)). For Neural CDE, we use MALI with ALF solver with fixed stepsize of 0.25, and train the model for 100 epochs with an initial learning rate of 0.004.
htm
19

---

Published as a conference paper at ICLR 2021
2
3
4
7
8
7
7
9
4
2
4
7
7
9
0
1
7
3
4
5 4 5
6
0
5
0
8
2
7
5'
A
7
8
0
2
3 9
7
8
3
0
4
0
7
0
8
9
0
8
3
S
8
5
/
3
2
0
5 4 7 9
3
2
6
8
4
7
7
4
3
9
7
2
3
3
S
6
9
2
9
5
8
0
9
8
5 (
/
2
6
7
7
5
8
2
4
9
6
7
8
6
7
3
4
6
8
6
4
7
(a) Real samples from MNIST datasct.
(b) Generated samples from FFJORD.
Figure 3: Results on MNIST dataset.
B.3 CONTINUOUS GENERATIVE MODELS
B.3.1 TRAINING DETAILS
Our experiment based on the official implementation of (Finlay al. 2020), with the only differ- ence in ODE solver. For a fair comparison, we only use MALI for training, and use Dopri5 solver from torchdiffeq package (Chen al. 2018) with rtol atol 10-5, For MALI, ALF solver with rtol 10 atol = 10 and use an initial stepsize of 0.25. Integration time is from 0 to
adaptive
On MNIST and CIFAR dataset, we set the regularization coefficients for kinetic energy and Frobe- nius norm of the derivative function as 0.05. We train the model for 50 epochs with initial learning rate of 0.001.
B.3.2 ADDTIONAL RESULTS
We show generated examples on MNIST dataset in Fig. results for Cifar10 dataset in Fig. and results for ImageNet64 in Fig. 5
B.4 ERROR IN GRADIENT ESTIMATION FOR TOY EXAMPLES WHEN <
We plot the error in gradient estimation for the toy example defined by Eq.6 in the main paper in Fig. 6 Note that the integration time T is set as smaller than 1, while the main paper larger than 20.
We observe the same results, MALI and ACA generate smaller error than the adjoint and the naive method.
B.5 RESULTS OF DAMPED MALI
For all experiments in the main paper, we set 1) and did use damping. For completeness, experimented with damped MALI using different values of 17. As shown in Table. MALI is robust different 7) values.
20

---

Published as a conference paper at ICLR 2021
(a) Real samples from CIFAR10 datasct.
(b) Generated samples from FFJORD.
Figure 4: Results on Cifar 10 dataset.
(a) Real samples from ImageNet64 dataset.
(b) Generated samples from FFJORD.
Figure 5: Results on ImageNet64 dataset.
Table 7: Results of damped MALI with different values. We report the test accuracy of Neural CDE on Speech Command dataset, and the test MSE of latent-ODE on Mujoco data.
1.0
0.95
0.9
0.85
Test Accuracy on Speech Commands
93.7 + 0.3
93.7 = 0.1
93.5 + 0.2
93.7 + 0.3
(Higher is better)
Test MSE of latent ODE on Mujoco (Lower better)
10% training data
0.35
0.36
0.33
0.33
20% training data
0.27
0.25
0.26
0.27
21

---

Published as a conference paper at ICLR 2021
Amplitude of error in dL/dZo
Amplitude of error in dL/da
1.0
MALI
1.0
MALI
Error in gradient
ACA
Error in gradient
0.8
0.8
ACA
Naive
Naive
0.6
Adjoint
0.6
Adjoint
0.4
0.4
0.2
0.2
0.0
0.0
0.2
0.4
0.6
0.8
0.2
0.4
0.6
0.8
Integration Time
Integration Time
(a) Error in the estimation of gradient wint initial condi- tion.
(b) Error in the estimation of gradient w.z/ parameter a.
Figure 6: Comparison of error in gradient estimation for the toy example by Eq.6 of the main paper, when 1.
22


Reviews:
[{"review_id": "MsiMz9fV1d", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 2, "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: \nThis paper proposes a new algorithm for solving neural ODEs. Each numerical solver step of the neural ODE is implemented as an invertible neural network via a variant of the asynchronous leafprog integrator. While still computing an accurate gradient, this allows memory savings by discarding intermediate data from the numerical integration steps since it can be reconstructed using the inverse. A theoretical stability analysis is provided. The experimental results show that the algorithm achieves similar performance to previous methods (e.g. ACA) while using less memory. \n\n\nStrengths:\n+ Identifies a nice connection between invertibility and memory efficiency. Beyond neural ODEs, this could enable use of larger models where invertible networks are useful (e.g. normalizing flows)\n+ The theoretical analysis of stability is useful to build intuition\n\n\nConcerns / weaknesses:\n- Most experiments in the paper use damping_factor $= \\eta = 1$. Since theoretically, this is not stable, it would be nice to see if the empirical improvements hold up for $\\eta < 1$, where stable regions do exist\n- The naive method seems too naive. Why are all results from all $m$ evaluations being saved? It is obvious that $m-1$ of these are unnecessary for gradient computation since they don't affect $z(T)$. The related claim about the computation graph being deeper for the naive method also seems incorrect.\n\n\nOther comments:\n- In Algorithm 1, shouldn’t $error_est = \\inf$ be inside the while loop?\n- In Algorithm 4, shouldn’t $a(T)$ be the partial derivative of $L$ wrt $z(T)$ instead of total derivative?\n- In Theorem 3.2, what is $\\sigma$? Should it be $\\sigma_i$?\n- In various locations, notation like $O(N+1)$ is used. Should this just be $O(N)$ since I assume $N$ is at least $\\Omega(1)$?\n- It seems a bit strange that we have to do a local forward and local backward pass in Algorithm 4. Could this be solved by making each layer of f invertible? In the same vein, it seems that the adjoint method needs to do a separate solve of the reverse ODE because of loss of information. If we were to assume invertibility of the forward map, is there a way to modify the adjoint method to exactly retrace the path backwards?", "extra_scores": {}, "extra_text": {}}, {"review_id": "EtToZWmr-Vd", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary: This paper presents a memory-efficient asynchronous leapfrog integrator for numerically solving neural ODEs, referred to as MALI. The method comes with a constant memory guarantee (like the adjoint method) and also guarantees reverse-time accuracy (like the adaptive checkpoint adjoint (ACA) method). The authors also give a rigorous theoretical analysis of MALI, and also discuss a \"damped\" version with an increased stability region. The method is evaluated on a variety of tasks which includes classification, dynamical modelling and generative modelling.\n\nPros:\n- The theoretical analysis looks correct, noting that I haven't worked out all the details.\n- Experimental evaluation is very exhaustive, and MALI achieves near-the-best performance in all tasks.\n- The method is proven as accurate as the standard numerical ODE solvers. Thanks to its reduced memory cost (compared to ACA), MALI can then be treated as an off-the-shelf replacement.\n\nCons and Questions:\n- Looking at the results, I'm having difficulty seeing any significant improvement upon ACA. Then the main contribution (in addition to the theoretical analysis) is the reduced memory consumption, which makes me rethink whether ICLR is a suitable venue.\n- Although the memory costs of the adjoint method and MALI are $O(N_f)$ and $O(N_f+1)$, this doesn't really reflect in Figure 4c, where the blue bar doubles the red one. I'd be happy if the authors can briefly explain why\n- Looking at Table-2, why does the test performance of a NODE trained with MALI increase when we switch from MALI to RK4? It would be much nicer to see some error variance estimate.\n- I would be happy to see an experimental evaluation of the \"A-stability\". As mentioned by the authors, the stability analysis is asymptotic and T could be arbitrarily small in, e.g., continuous-time flows. However, that's not the case in time-series modelling. So I wonder if the stability claim can be verified on a scenario in which, e.g., 100 observations arrive uniformly in time with 10 secs gaps.\n- To generate Table-2, did you train a ResNet without any differentials/integrals involved and try to evaluate the test performance using an ODE solver (simply using the trained ResNet as the drift function)? If so, I don't think this makes any sense except Euler-1 solver, and the entire ResNet row in Table-2 could go away.\n\nAdditional comments:\n- Figure-4 caption could include some more detail (at least mentioning the experiment name) \n- Why is there a \"local forward\" step within the for loop in the backward routine in Alg.4?\n- It would be nice to see a brief description of the mujoco dataset.\n- Typo in the title of section B.3.2.\n\nNote: After rebuttal, I increase my overall score from 6 to 7.", "extra_scores": {}, "extra_text": {}}, {"review_id": "E6TmZHTdAEz", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "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": "1) Summary\nThe manuscript proposes a reversible integration scheme for approximately estimating the gradient of neural ordinary differential equations. These ODEs represent DNNs with continuous rather than discrete values for the number of layers. The solver is theoretically analysed and empirically compared to other solvers.\n\n2) Strengths\n+ The paper is mostly well written.\n+ The reversibility property of the solver leads to a memory footprint that does not depend on integration time.\n+ The model is applied to standard datasets.\n\n3) Concerns\n- The concept of Neural ODE models, their scope and their expected usefulness should be better motivated. It is not obvious which role these models play and what they offer as potential strengths.\n- The integrations scheme seems already known and well established.\n- It does not seem that the paper makes code and data available to the public.\n\n4) Remarks/Questions\n  a) Algorithm 1: It seems that the stepsize h should be initialized upon every step. Otherwise the steps can only get smaller. \n  b) References: capitalization not correct e.g. \"neural information processing systems\", \"ode-net\", \"lennard-jones\"\n  c) What benefits does the neural ODE model have in the context of image classification? What is the intuition behind the \"continuous depth\" idea in this scenario?", "extra_scores": {}, "extra_text": {}}, {"review_id": "2tYCvPDe7Cm", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Paper summary**\nThere are typically two methods for estimating the gradients with respect to the loss for neural ODEs. The naive method directly backpropagates through the steps of the ODE solver leading to accurate gradients but very large memory cost. The adjoint method in contrast does not store the entire trajectory in memory, but has reverse trajectory errors (i.e. the numerical solution in the reverse direction will not be the inverse of the numerical solution in the forward direction). In this paper, the authors propose a method that is both reverse accurate and has low memory cost.\n\nTo achieve this the authors take advantage of the asynchronous leapfrog solver. This numerical method is reversible: solving an ODE numerically in the reverse direction is the inverse of solving the ODE numerically in the forward direction. This is not generally true for the ODE solvers typically used (RK4 and so on) in the neural ODE literature. As the numerical solver is explicitly invertible, the authors can (from only the final state and not the entire trajectory) locally reconstruct each ODE solver step to get the local gradient of the parameters with respect to the loss. They can then calculate these gradients along the entire reverse trajectory to obtain an accurate estimate of the gradient with respect to the parameters. As each step of the numerical solver is reversible, they do not need to store the entire trajectory. The authors analyse the stability and numerical error of their proposed method and provide a toy example to show how well their method estimates gradients compared the naive, adjoint and adaptive checkpoint methods.\n\nThe authors then perform experiments on a variety of tasks to test their model. They test their model on image classification experiments, both on CIFAR10 and Imagenet and achieve good results compared to the baselines. In addition, they perform adversarial robustness experiments on ImageNet and also show good performance. Finally, the authors test their method both for time series modeling and continuous normalizing flows, again showing good performance compared with naive integration methods.\n\n**Positives**\n- The motivation and core idea of the paper is clear. Numerical solvers are in general not reversible and this can lead to inaccurate gradient estimates when using the adjoint method. The authors explain this clearly and then propose a method that effectively solves this.\n- The experimental results are quite impressive. The model performs on par with the adaptive checkpoint method in terms of accuracy but is much more memory efficient (and notably memory is independent of the number of integration steps). This allows the authors to run their model on large scale datasets like ImageNet which was not previously possible with most neural ODE methods. Further, the authors achieve good performance on quite a wide variety of tasks (image classification, adversarial attacks, time series modeling, generative modeling) which is nice.\n- The authors perform a thorough analysis of the runtime of their integration method compared to others which is very helpful.\n\n**Negatives**\n- The presentation of the method and results is not always very clear. For example, the section about damping for the ALF integrator is not clear. The authors mention that ALF is not stable for eta=1, but (as far as I can tell) never mention what value of eta they use in practice and whether choosing this value is difficult. Further, it is not clear if ALF is still reversible with this eta parameter. Presumably you would have to use 1/eta in the reverse step for it to remain invertible, in which case the reverse is not stable? The authors should be more clear about this.\n- The toy example is confusing. How come the integration time starts from t=20? Is this because the error only grows after t=20? As you use T=1 for all experiments (and the rtol and atol are also roughly the same for all experiments), it would be nice to see if this actually makes a difference also for t<1. In Figure 4, the authors also mention the derivative dL/dy_0 but this derivative is never mentioned in the text. Do you mean dL/dz_0? The plots of memory consumption are nice and clear though.\n- The ALF solver already exists, so the main contribution of the paper is simply to apply the ALF solver to neural ODEs. This means that the novelty of the method is somewhat limited, but I do not think that this is a major issue as the method works well and is clearly motivated.\n- The section about using different numerical solvers for ResNets does not make much sense. ResNets are not designed as flows and do not behave as flows in practice, so we should not expect them to work at all with other numerical solvers than Euler with timestep=1. I don’t really think these experiments show anything interesting and should be removed for clarity.\n\n**Recommendation**\nOverall the paper has a clear motivation, provides a nice and simple solution to an interesting problem and has good experimental results. However, there are some clarity issues which make some aspects of the model and method confusing. I therefore recommend a weak accept but would increase my score if the clarity issues are solved.\n\n**Questions**\nThe model achieves extremely good bits/dim on MNIST (0.87). However, it seems from the appendix that the samples are fairly poor (compared to vanilla FFJORD for example). Log likelihood and sample quality are not always correlated, but the difference seems particularly jarring here. Do you know why this is?\n\n**Typos and small comments**\n- In many places the authors use latex math to write words that should either just be written in italics (w.r.t) or using \\text{} in math mode (e.g. atol, rtol).\n- There are several typos in the script, so I think it would be a good idea for the authors to read through the script again to fix those.\n- In several places, the authors write O(N_f + 1) which instead should be O(N_f)\n- The authors often write “constant memory cost with respect to integration time”. I think it would be more helpful to say “number of solver steps” or something along those lines as integration time typically refers to the upper limit of the integral when solving an ODE.\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper introduced a new ODE integration scheme that allows constant-memory gradient computation.  I was concerned that the low order of convergence of this method would make it impractical, but the authors performed extensive experiments and got impressive results.  Overall the paper addresses one of the main practical difficulties with large neural ODE models.  The authors satisfactorily addressed the reviewers' concerns in the discussion.
Decision: Accept (Poster)