Symplectic Adjoint Method for Exact Gradient of Neural ODE with Minimal Memory
Authors: Takashi Matsubara
Yuto Miyatake
Takaharu Yaguchi
Venue: NeurIPS 2021 Poster (2021)
Keywords: ordinary differential equation
neural ODE
adjoint method
backpropagation
checkpointing scheme
symplectic integrator

Abstract: A neural network model of a differential equation, namely neural ODE, has enabled the learning of continuous-time dynamical systems and probabilistic distributions with high accuracy. The neural ODE uses the same network repeatedly during a numerical integration. The memory consumption of the backpropagation algorithm is proportional to the number of uses times the network size. This is true even if a checkpointing scheme divides the computation graph into sub-graphs. Otherwise, the adjoint method obtains a gradient by a numerical integration backward in time. Although this method consumes memory only for a single network use, it requires high computational cost to suppress numerical errors. This study proposes the symplectic adjoint method, which is an adjoint method solved by a symplectic integrator. The symplectic adjoint method obtains the exact gradient (up to rounding error) with memory proportional to the number of uses plus the network size. The experimental results demonstrate that the symplectic adjoint method consumes much less memory than the naive backpropagation algorithm and checkpointing schemes, performs faster than the adjoint method, and is more robust to rounding errors.

Symplectic Adjoint Method for Exact Gradient of Neural ODE with Minimal Memory
Takashi Matsubara
Yuto Miyatake
Osaka University
Osaka University
Osaka, Japan matsubara@sys.c es www.ss...u.ua.....oo
Osaka, Japan
miyatake@cas. cmc. osaka-u. ac jp
Takaharu Yaguchi
Kobe University
Kobe, Japan 657-850
Raguchi@pearl.. kobe-u. ac.
Abstract
A neural network model differential equation, namely neural ODE, has enabled the learning of continuous-time dynamical systems and probabilistic distributions with high accuracy. The neural ODE uses the same network repeatedly during a numerical integration. The memory consumption of the backpropagation algorithm is proportional to the number of uses times the network size. This is true even if a checkpointing scheme divides the computation graph into sub-graphs Otherwise, the adjoint method obtains gradient by a numerical integration backward in time. Although this method consumes memory only for a single network use, it requires high computational cost to suppress numerical errors. This study proposes the symplectic adjoint method, which is an adjoint method solved by symplectic integrator. The symplectic adjoint method obtains the exact gradient (up to rounding eπror) with memory proportional to the number of uses plus the network size. The experimental results demonstrate that the symplectic adjoint method consumes much less memory than the naive backpropagation algorithm and checkpointing schemes, performs faster than the adjoint method, and is more robust to rounding errors.
Introduction
Deep neural networks offer remarkable methods for various tasks, such as image recognition [8 and natural language processing [4]. These methods employ residual architecture [2 34], in which the output 2n+1 of the n-th operation is defined the sum of a subroutine f.. and the input In as In+1 = fxx) + En. The residual architecture can be regarded as a numerical integration ordinary differential equation (ODE) [30. Accordingly, neural network model of the differential equation dæ/dl = f [(x), namely, neural ODE, was proposed in [2]. Given an initial condition .0) x as input, the neural ODE solves initial value problem by numerical integration, obtaining the final value as an output
applied
xx (T)) The neural ODE can model continuous- time dynamics such as irregularly sampled time series (24). stable dynamical systems (37 42), and
physical phenomena associated with geometric structures [3 I3 31. Further, because the neural ODE approximates diffccomorphism (43), it can model probabilistic distributions of real-world data
by change of variables 22 23 25
45
For accurate integration, the neural ODE must employ small step size and high-order numerical integrator composed of many internal stages. neural network f is used at each stage of each
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Table 1: Comparison of the proposed method with existing methods
Methods
Gradient Calculation
Exact
Checkpoints
Memory Consumption Computational Cost
checkpoint
hackprop.
NODE
2
adjoint method
no
JNN
M
L
M(N-2N)6
NODE 2
backpropagation
yes
MN
2M
baseline scheme
backpropagation
yes
M
NsL
3 SL
ACA
46
backpropagation
yes
{I Jn EN
1
MN
SL
3MNsL.
MALI
47
backpropagation
yes
M
SL
4M
proposed**
symplectic adjoint method
yes {Zm) N-1
MN+s
L
4M
"vaailable only for the asynchronous leapfrog integrator. "Available for any Runge-Kutta Kutta methods.
time step. Thus, the backpropagation algorithm consumes exorbitant memory to retain the whole computation graph 39
10 46). The neural ODE employs the adjoint method to reduce memory consumption- method obtains a gradient by backward integration along with the state z, without consuming memory for retaining the computation graph over time [8 41 44]. However, this method incurs high computational costs to suppress numerical errors. Several previous works employed a checkpointing scheme [0C 46 47. This scheme only sparsely retains the state a as checkpoints and recalculates computation graph from each checkpoint to obtain the gradient. However, this scheme still consumes a significant amount of memory to retain the computation graph between checkpoints.
To address the above limitations, this study proposes the symplectic adjoint method. The main advantages of the proposed method are presented as follows.
Exact Gradient and Fast Computation: In discrete time, the adjoint method suffers from numerical errors or needs a smaller step size. The proposed method uses a specially designed integrator that obtains the exact gradient in discrete time. It works with the same step size as the forward integration and is thus faster than the adjoint method in practice.
Minimal Memory Consumption: Excepting the adjoint method, cxisting methods apply the back- propagation algorithm to the computation graph of the whole or a subset of numerical integra- tion 10 46 47). The memory consumption is proportional to the number of steps/stages in the graph times the neural network size. Conversely, the proposed method applies the algorithm only to each of the neural network, and thus the memory consumption is only proportional to the number of steps/stages plus the network size.
The backpropagation algorithm accumulates the gradient from each use of the neural network and tends to suffer from rounding errors. Conversely, the proposed method obtains the gradient from each step numerical integration and is thus more robust rounding errors.
Robust to Rounding Error:
2 Background and Related Work
2.1 Neural Ordinary Differential Equation and Adjoint Method
We use the following notation.
M: the number of stacked neural ODE components,
L: the number of layers in a neural network,
N, N: the number of time steps in the forward and backward integrations, respectively, and the number of uses of neural network f per step.
is typically equal to the number of internal stages of numerical integrator [7]. A numerical integration forward in time requires computational cost of sMMNsL). It also provides compu- tation graph over time steps, which is retained with memory of ( N LL); the backpropagation algorithm is then applied to obtain the gradient. The total computational cost is O(2M sL), where we suppose the computational cost of the backpropagation algorithm equal to that of forward propagation. The memory consumption and computational cost are summarized in Table [

---

To reduce the memory consumption, the original study on the neural ODE introduced the adjoint method 2 17 4144 This method integrates the pair of the system state and the adjoint variable &C backward in time. The adjoint variable represents the gradient de of some function C, and the backward integration of the adjoint variable À works as the backpropagation (or more generally the reverse-mode automatic differentiation) in continuous time. The memory consumption is ((M) to retain the final values x<(") of M neural ODE components and O()) to obtain the gradient of a neural network f for integrating the adjoint variable A. The computational cost is at least doubled because of the re- re-integration of the system state r backward in time. The adjoint method suffers from numerical errors
10. To suppress the numerical errors, the backward integration often requires a smaller step size than the forward integration (i.e., N > N), leading to an increase in computation time. Conversely, the proposed symplectic adjoint method uses specially designed integrator, which provides the exact gradient with the same step size as the forward integration.
2.2 Checkpointing Scheme
The checkpointing scheme has been investigated to reduce the memory consumption of neural networks [3), where intermediate states are retained sparsely as checkpoints, and a computation graph is recomputed from each checkpoint. For example, Gruslys et al. applied this scheme to recurrent neural networks [15]. When the initial value x0) of each neural ODE component is retained checkpoint, the initial value problem is solved again before applying the backpropagation algorithm to obtain the gradient of the component. Then, the memory consumption is O(M) for checkpoints and ONNs [L) for the backpropagation; the memory consumption is O(M + N&L. in total (see the baseline scheme). ANODE scheme retains each step {zn ) as a checkpoint with a memory of O(MN) [10]. Form each checkpoint Zn, this scheme recalculates the next step In+1 and obtains the gradient using the backpropagation algorithm with memory of OsL)); the memory consumption is OMMN N + sL) in total. ACA scheme improves ANODE scheme for methods with adaptive time-stepping by discarding the computation graph to find an optimal step size. Even with checkpoints, the memory consumption is still proportional to the number of uses of a neural network per step, which is not negligible for a high-order integrator, e.g., 8 = 6 for the Dormand-Prince method [7]. In this context, the proposed method is regarded as checkpointing scheme inside numerical integrator. Note that previous studies did not use the notation 8.
Instead of a checkpointing scheme, MALI employs an asynchronous leapfrog (ALF) integrator after the state x is paired up with the velocity state v [47]. The ALF integrator is time-reversible, i.e., the backward integration obtains the state x equal to that in the forward integration without checkpoints [17. However, the ALF integrator a second-order integrator, implying that requires a small step size and high computational cost to suppress numerical errors. Higher-order Runge Kutta methods cannot be used in place of the ALF integrator because they are implicit or con-time-revente The ALF integrator is inapplicable to physical systems without velocity such as partial differential equation (PDE) systems. Nonetheless, similar approach named RevNet was proposed before in [I] When regarding ResNet as forward Fuler method 2 18], RevNet has an architecture regarded as the leapfrog integrator, and it recalculates the intermediate activations in the reverse direction.
3 Adjoint Method
Consider a system
d dt x [xx,,,0)
where x, i, and 0, respectively, denote the system state, an independent variable (e.g., time), and parameters of the function f. Given an initial condition (0) Tin the solution xx() is given by
x(t) = xo + L. (oc(T).T.fddk.
(2)
The solution r=t) is evaluated at the terminal = T by a function L as [(x(T)) Our main interest is in obtaining the gradients of ((x(T)) with respect the initial condition Ti and the parameters 0.
Now, we introduce the adjoint method omit the parameters 0. The adjoint method based on the variational variable 5(t) and the adjoint
44 We first focus on the initial condition ao and
3

---

variable A(t). The variational and adjoint variables respectively follow the variational system and adjoint system as follows.
d Of dt 8(1) Or (z((), ()(() for 8(0) -1 T. A(1) - dx (z, 1) X(1) for X(T) - AT-
d
0f
(3)
The variational variable 5()) represents the Jacobian ax(()) of the state s=() with respect to the initial condition žo: the detailed derivation is summarized in Appendix A Remark 1. The quantity A 6 is time- invarina,r,,, i.e., X(t) 8(t) x(O) 8(0).
The proofs of most Remarks and Theorems in this paper are summarized in Appendix B
Remark 2. The adjoint variable X(t) represents the gradient ( if the final condition XT
of the adjoint variable À is set to (zx(T) (T(
(
This is because of the chain rule. Thus, the backward integration of the adjoint variable X(1) works as reverse-mode automatic differentiation. The adjoint method has been used for data assimilation, where the initial condition To is optimized by gradient-basee based method. For system identification (i.e.. parameter adjustment), one can consider the parameters as part of the augmented state [x 0) of the system
d ग. i(,,0 f(2,1) = fxx t. 0)
=(0) = 20
(4)
The variational and adjoint variables are augmented in the same way. Hereafter, we let a denote the state augmented state without loss of generality. See Appendix for details.
According to the original implementation of the neural ODE (2), the final value xT)) of the system state x is retained after forward integration, and the pair of the system state x and the adjoint variable is integrated backward in time to obtain the gradients. The right-hand sides of the main system in Eq. ① and the adjoint system in Eq. (3 are obtained by the forward and backward propagations of the neural network f. respectively. Therefore, the computational cost the adjoint method twice that of the ordinary backpropagation algorithu.
After a numerical integrator discretizes the time, Remark does not hold, and thus the adjoint variable A(t) is not equal to the exact gradient [10 41). Moreover, in general, the numerical integration backward in time consistent with that forward time. Although a small step size (i.e., small tolerance) suppresses numerical errors, it also leads to longer computation time. These facts provide the motivation obtaín the exact gradient with a small memory, in the present study.
4 Symplectic Adjoint Method
4.1 Runge-Kutta Method
We first discretize the main system in Eq. and state, respectively, where hn fn+1 f.-- Previous studies employed one of the Runge- Kutta methods, generally expressed
[ Let bn, hns and &n denote the n-th time step, step size,
In+1 In + hn Σ bykinnis
i=1
kn,i = f(Xm,is in + crhn),
(5)
X == Jn + hn Σ oiijking
j=1
The coefficients ai.j by, and e; are summarized the Butcher tableau [16 17 41 If a.
for j 2 i, the intermediate state X,,s is calculable from i = 1 to 8 sequentially; then, the Runge- Kutta method considered explicit. Runge-Kun Kutta methods are not time-reversible in general, i..., the numerical integration backward in time is not consistent with that forward in time.
Remark 3 (Bochev and Scovel [I]. Hairer et al. (16). When the system in Eq. ① is discretized by the Runge-Kull Kutta method in Eq. (5), the variational system in Eq. (3) is discretized by the same Runge-Kiun Kutta method.
Therefore, it is not necessary to solve the variational variable 8(t) separately.

---

4.2 Symplectic Runge-Kutta Method for Adjoint System
We assume bj * for s. We suppose the adjoint system be solved by another Runge-Kutee method with the same step size as that used for the system state z, expressed as
Au+1 An + hn Σ
Bihn,i
1=1
L,,i == 0f
(Xniil bn + Cshaa) Angi
(6)
Os
An, = An + hn Aiphnis
j=1
The final condition AN set to (aC(eN))² TT Because the time evolutions of the variational variable and the adjoint variable are expressible by two equations, the combined system is considered as partitioned system. A combination of two Runge-Kutta methods for solving partitioned system called partitioned Kunge-Kutta method, where C, C for i 1, 1,...,8. 8. We introduce the following condition for partitioned Runge- Kutta method.
Condition 1. b,A,, Byag, B, for i,j 1, ,...,,, and B b, + and C C: for i 1= 8. Theorem (Sanz-Serna 41). The partitioned Runge-Kutte method in Eqs. and 6 conserves bilinear quantity S(6, X) if the continuous- time system conserves the quantity ( X) and Condition
7 holds.
Because the bilinear quantity S (including σ)) is conserved, the adjoint system solved by the Runge-Kurte Kutta method Eq. (6) under Condition providees the exact gradient as the adjoint variable [0((z)) The Dormand-Prince method, one of the most popular Runge-Kutta methods, has du.
An
b/
7 For such methods, the Runge-Kutta method under Condition in Eq. is generalized
Ant1 = An + hn Σ beln,s
k=1
0ʃ
In, =
(Xm,is + cihin) An.is
Ox
An + Nn Eÿ=1 by (1 bi
has i E To E To 1o =0= I,..,,,., b, 0}.
azii
An,i =
Ej= bragstnis
where
= bb if iig To
6.
hn if i Io.
5,
(8)
Note that this numerical integrator is longer Runge Kutta method and is an alternative expression for the "fancy" integrator proposed [41]. The combination of the integrators in Eqs. 5 and conserves bilinear quantity
Theorem 2.
((5, A) the continuous-tiie syslem conserves the quantity S(ö, X)-
Remark 4. The Runge-Kiun Kutta method Eq. (6) under Condition I and the numerical integrator in
Eq. 7 are explicit backward in time if the Runge-Kuute method in Eq. 5 is explicit forward in time.
We emphasize that Theorems T aan hold for any ODE systems even if the systems have disconti- nuity [19,, stochasticity 229 or physics constraints [13]. This is because the Thcorems are not the properties of system but of Runge- Kutta methods.
partitioned Hange-Kuuta, method that satisfies Condition is symplectic [7 [6]. is known that, when symplectic integrator is applied to Hamiltonian system using fixed step size, it conserves modified Hamiltonian, which an approximation the system energy of the Hamiltonian system. The bilinear quantity S is associated with the symplectic structure but not with Hamiltonian. Regardless of the step size, a symplectic integrator conserves the symplectic structure and thereby conserves the bilinear quantity S. Hence, we named this method the symplectic adjoint method. For integrators other than Runge-Kutta methods, one can design the integrator for the adjoint system so that the pair of integrators is symplectic (see [32 for example).
5

---

Algorithm 1 Forward Integration
Algorithm 2 Backward Integration
Input: Za
Input: øN )==
Output: 2N (xh 1N-1
Jn=0 I: for to N do
Output: Au
1: for = N to do
Retain n. as a checkpoint
2: Load checkpoint According Eg.
According to Eg. 5
3: for = to do
3: for to do
4: Get X,,s using zn and Kn.j for
4: Get Xn, using 2h, and ka. for j <
Get kn, using X
5: Get kn. using Xa...
6: end for
6: Retain Xn,; checkpoint end for
7: Get In+1 using In and knnd
7:
8: end for
According to Eg. 7
8:
for = 8 do
9:
Get An.i using An+1 and bo., for Load checkpoint X,,s
10:
4.3 Proposed Implementation
Get b,,s using An, and Xn,
12: Discard checkpoint Xn,
The theories given in the last section were mainly
13:
end for
introduced for the numerical analysis in [41] Be-
14:
Get An using Au+1 and bn,e
cause the original expression includes recalcula-
15: Discard checkpoint In.
tions of intermediate variables, we propose the al- ternative expression in Eq. (7) to reduce the com- putational cost. The discretized adjoint system in
16:
end for
2x Eq. 7 depends on the vector-Jacobian product (VJP) AT To obtain it, the computation graph from the input Xn,; to the output f(Nnn,; + c,htn) is required. When the computation graph in the forward integration is entirely retained, the memory consumption and computational cost are of the same orders as those for the naive backpropagation algorithm. To reduce the memory consumption, we propose the following strategy as summarized in Algorithms
[ [nd
At the forward integration of neural ODE component, the pairs of system states Hn and time points in at time steps n O,.... N computation graphs are discarded, as shown in Algorithm I For M neural ODE components, the memory for checkpoints is OMMN). The backward integration is summarized in Algorithm The below steps are repeated from = N 11 to 0. From the checkpoint T.. the intermediate states Xn.; for 8 stages are obtained following the Runge-Kutta method in Eq. (5) and retained as checkpoints with memory of O(s), while all computation graphs are discarded. Then, the adjoint system is integrated from n + to n using Eq. 7 Because the computation graph of the neural network f in line i discarded, it is recalculated and the VJP 0f
are retained with a memory of ONN) as checkpoints, and all
is obtained using the backpropagation algorithm one-by-one in line II where only a single use of the neural network recalculated at time. This is why the memory consumption is proportional to the number of checkpoints MN + s plus the neural network size L. By contrast, existing methods apply the backpropagation algorithm to the computation graph of a single step composed of stages or multiple steps. The memory consumption is proportional the number of uses of the neural network between two checkpoints (s at least) times the neural network size L,, in addition to the memory for checkpoints (see Table Due the recalculation, the computational cost of the proposed strategy is O(4MNsL), whereas those of the adjoint method 2 and ACA [46] are (M 2 and ((3MN&L) respectively. However, the increase in the computation time is much less than that expected theoretically because of other bottlenecks (as demonstrated later).
Her
5 Experiments
We evaluated the performance of the proposed symplectic adjoint method and existing methods using
PyTorch 1.7.1 [35]. We implemented the proposed symplectic adjoint method by extending the
adjoint method implemented in the package torchdiffeq 0.1.1 [2]. We re-implemented ACA [46] because the interfaces of the official implementation is incompatible with torchdiffeq. In practice, the number of checkpoints for an integration can be varied; we implemented a baseline scheme
that retains only single checkpoint per neural ODE component. The source code available at
https:/ //te:///iiii ////gttuuu thub. com http://aattabbber//////occcaaaoa

---

Table 2: Results obtained for continuous normalizing flows.
MINIBOONE (M = 1)
GAS (M- = 5)
POWER (M = 5)
NLJ
mem.
time
NLL
mem.
time
NLL
mem.
time
adjoint method
[2
10.59
170
0.74
-10.53=:0255
24
4.82
-0.31+0.01
8.1
6.33
backpropagation
2
10.54
4436
0.91
-9.53
4479
12.00
-0.24:10.05
1710.9
10.64
baseline scheme
10.54
4457
1.10
-9,53=0.42
1858
5.48
-0.24410.05
515.2
4.37
ACA 46
10.57+0.30
306
0.77
-10.65
73
3.98
-0.01
29.5
5.08
proposed
10.49:0.11
95
0.84
-10.89
20
4.39
-0.31
9.2
5.73
HEPMASS (M
10)
BSDS300 (M = 2)
MNIST (M 6)
NLL
mem.
time
NLL
mem.
time
NLL
mem.
time
adjoint method
2
16.49
40
4.19
11.0.04...09
577
11.70
0.918100011
1086
10.12
backpropagation
2
17.03+#022
5254
11.82
baseline scheme
17.03++0.2
1102
4.40
ACA [46]
16.41
88
3.67
-151.27=0.77
757
6.97
0.919=0.003
4332
7.94
proposed
16.48.10.20
35
4.15
1751.17
283
8.07
0.917
1079
9.42
Negative log- ligllikelihodds (NLL), peak memory consumption [MiB], and computation time per iteration [s/itr]. See Table A2 in Appendix for standard deviations.
5.1 Continuous Normalizing Flow
Experimental Settings: We evaluated the proposed symplectic adjoint method training con- tinuous normalizing flows [12]. A normalizing flow a neural network that approximates bijective map 9 and obtains the exact likelihood of a sample by the change of variables p(g)(u) = log p(2) + log det @g(u) ). where
g(u) and p(z) denote the corresponding la- tent variable and its prior, respectively [3 6 38 A continuous normalizing flow is normalizing flow whose map is modeled by stacked neural ODE components, in particular, u = xx0) and = = xxT) for the case with M = 1. The log-determinant of the Jacobian is obtained by numerical integration together with the system state as log| det 8g(u)
So Tr 0f (ar(t,, t)dt. The trace operation Be
Tr is approximated by the Hutchinson estimator
[22. We adopted the experimental settings of the unless stated otherwise.
continuous normalizing flow, FFJORD
We examined five real tabular datasets, namely, MiniBooNE, GAS, POWER, HEPMASS, and BSDS300 datasets (33] The network architectures were the same those that achieved the best results in the original experiments; the number of neural ODE components M varied across datasets. We employed the Dormand-Prince integrator, which is fifth-order Runge- Kutta method with adaptive time-stepping, composed of seven stages (77. Note that the number of function evaluations per step is s because the last stage is reused as the first stage of the next step. We set the absolute and relative tolerances to atol = 10 and rtol = 10 10-6 respectively. The neural networks were trained using the Adam optimizer [27 with a learning rate of 10- We used batch-size of 1000 for all datasets to put a mini-batch into single NVIDIA GeForce RTX 2080Ti GPU with 11 GB of memory, while the original experiments employed a batch-size of 10 000 for the latter three datasets on multiple GPUs. When using multiple GPUs, bottlenecks such as data transfer across GPUs may affect performance, and a fair comparison becomes difficult. Nonetheless, the naive backpropagation algorithm and baseline scheme consumed the entire memory for BSDS300 dataset.
We also examined the MNIST dataset (28) using a single NVIDIA RTX A6000 GPU with 48 GB of memory. Following the original study, we employed the multi-scale architecture and set the tolerance to atol rtol 10-5 We set the learning rate 10-3 and then reduced it to 10-4 at the 250th epoch. While the original experiments used batch-size of 900, we set the batch-size to 200 following the official c The naive backpropagation algorithm and baseline scheme consumed the entire memory.
Performance:
The medians + standard deviations of three runs are summarized in Table In many cases, all methods achieved negative liglliiksssoods (NLLs) with no significant difference
httt::///ggchbbb..c///rggggee//e (MIT License)

---

because all but the adjoint method provide the exact gradients up to rounding error. and the adjoint method with a small tolerance provides a sufficiently accurate gradient. The naive backpropagation algorithm and baseline scheme obtained slightly worse results on the GAS, POWER, and HEPMASS datasets. Due to adaptive time-stepping. the numerical integrator sometimes makes the step size much smaller, and the backpropagation algorithm over time steps suffered from rounding errors. Conversely, ACA and the proposed symplectic adjoint method applied the backpropagation algorithm separately to subset of the integration, thereby becoming more robust to rounding errors (see Appendix D.I for details).
After the training procedure, we obtained the peak memory consumption during additional training iterations (mem. [MiB],, from which we subtracted the memory consumption before training (i.e.. occupied by the model parameters. loaded data, etc.). The memory consumption still includes the optimizer's states and the intermediate results of the multiply-acemmulaee operation. The results roughly agree with the theoretical orders shown in Table also Table A2 for standard deviations). The symplectic adjoint method consumed much smaller memory than the naive backpropagation algorithm and the checkpointing schemes. Owing to the optimized implementation, the symplectic adjoint method consumed smaller memory than the adjoint method in some cases (see Appendix (D22)
On the other hand, the computation time per iteration (time [s/itr]) during the additional training iterations does not agree with the theoretical orders. First, the adjoint method was slower in many cases, especially for the BSDS300 and MNIST datasets. For obtaining the gradients, the adjoint method integrates the adjoint variable X, whose size is equal to the sum of the sizes of the parameters 0 and the system state x. With more parameters, the probability that at least one parameter does not satisfy the tolerance value is increased. An accurate backward integration requires a much smaller step size than the forward integration (i.e., N much greater than N), leading to a longer computation time. Second, the naive backpropagation algorithm and baseline scheme were slower than that expected theoretically, many cases. A method with high memory consumption may have to wait for retained computation graph to be loaded or memory to be freed, leading to an additional bottleneck. The symplectic adjoint method is free from the above bottlenecks and performs faster in practice; was faster than the adjoint method for all but MiniBooNE dataset.
The symplectic adjoint method is superior (or at least competitive) to the adjoint method, naive backpropagation, and baseline scheme in terms of both memory consumption and computation time. Between the proposed symplectic adjoint method and ACA, a trade-off exists between memory consumption and computation time.
The adjoint method provides gradients with numerical errors. To evaluate the robustness against tolerance, we employed MiniBooNE dataset and varied the absolute tolerance atol while maintaining the rel- ative tolerance as rtol = 102 x atol. During the training, we obtained the computation time per iteration, summarized in the upper panel of Fig. T The computation time reduced as the tolerance increased. After training, obtained the NLLS with atol = 10 10-$, as summarized in the bottom panel of Fig. The adjoint method performed well only with atol < 10-4 With atol = 10-1, the numerical error in the backward integration was non-negligible, and the per- formance degraded. With atol >> 10-4, the adjoint method destabilized. The symplecté adjoint method performed well even with atol = 10 10-3, Even with 10 < atol < 10-2,11 it performed to certain level, while the numerical error in the forward integration was non-negligible. Because of the exact gradient, the symplectic adjoint method is robust to a large tolerance compared with the adjoint method, and thus potentially works much faster with an appropriate tolerance.
Robustness to Tolerance:
NLL time [s/itr]
adjoint method proposed
15
10
10-2
10-7
10
10-3
tolerance (atol)
Figure 1: With different tolerances.
Different Runge-Kutta Methods: The Runge Kutta family includes various integrators charac- terized by the Butcher tableau [6 17 41), such as the Heun-Euler method (a.k.a. adaptive Heun), Hormand-Prince method (a.k.a. dopri5), and eighth-order Dormand-Prince method (a.k.a. dopri8). These methods have the orders of p = 2, 3, 5, and 8 using 2, 3, 6, and 12 function evaluations, respectively. We examined these methods using
Bogacki-Shamping method (a.k.a. bosh3), fifth-order
8

---

method and Table 3: Results obtained for GAS dataset with different Runge- Kutta methods.
= 2, 5
2
3, 8
p
5, 8
p p= 8, 8
12
mem.
time
mem.
time
mem.
time
mem.
time
adjoint method [2
21: 0 247.47= 7.52
22:00 18.32+0.88
24 0
5,34
28:
9.77
backpropagation
4433:1255
11,85+1.10
baseline scheme
1858
5.82++028
4108+576
22.76=3.70
ACA
46
607 232.90=1381
69
17.72
73
4.15+0.21
138
9.36=0.55
proposed
589::44 262.99 = 5.19 43 18.59:10.5
20 0
4.78
21:
11.41
Peak memory consumption [MiB]] and computation time per iteration [s/itr].
GAS dataset, and the results are summarized Table The naive backpropagation algorithm and baseline scheme consumed the entire memory in some cases, denoted by dashes. We omit the NLLS because all methods used the same tolerance and achieved the same NLL.s.
Compared to ACA, the symplectic adjoint method suppresses the memory consumption more significantly with a higher-order method (i.e., more function evaluations s), as the theory suggests in Table ! With the Heun-Euler method, all methods were extremely slow, and all but the adjoint method consumed larger memory. lower-order method has use an extremely small step size to satisfy the tolerance, thereby increasing the number of steps N, computation time, and memory for checkpoints. This result indicates the limitations of methods that depend on lower-order integrators, such as MALI [47.. With the eighth-order Dormand-Prinee method, the adjoint method performs relatively faster. This because the backward integration easily satisfies the tolerance with higher-order method (i.e., Ñ N). Nonetheless, in terms of computation time, the fifth-order Dormand-Prince method is the best choice, for which the symplectic adjoint method greatly reduces the memory consumption and performs faster than all but ACA.
Memory for Checkpoints: To evaluate the memory consumption with varying numbers of checkpoints, we used the fifth-order Dormand -rinnce method and varied the number of steps N for MNIST by manually varying
adjoint method
backprop.
baseline scheme
ACA
proposed
the step size. We summarized the results in Fig. 2 on log-log scale. Note that, with the adaptive stepping. FFJORD needs approximately MN 200 steps for MNIST and fewer steps for other datasets. Because we set Ñ = N, but Ñ N in practice, the adjoint method is expected to require longer computation time.
memory [MiB]
10
10a
10a
The memory consumption roughly follows the theoret- ical orders summarized in Table The adjoint method needs memory of OL) for the backpropagation, and the symplectic adjoint method needs an additional mem- ory of O(MN N + s) for checkpoints. Until the number of steps M exceeds a thousand, the memory for check- points negligible compared that for the backprop- agation. Compared to the symplectic adjoint method, ACA needs a memory of (sL) for the backpropagation over 8 stages. The increase in memory is significant until the number of steps MN reaches ten thousand. For a stiff (mon-smooth) ODE for which a numerical integrator needs thousands of steps, employ higher-order integrator such the eighth-order Dormand- suppress the number of steps. For stiffer ODE, implicit integrators are commonly used, which out of the scope of this study and the related works in Table I Therefore, we conclude that the symplectic adjoint method needs the memory at the same level as the adjoint method and much smaller than others in practical ranges.
time [s [s/itr]
10
10
102
10a
10
#steps (MN)
Figure 2: Different number of steps.
possible alternative to the proposed implementation Algorithms andd 2 all intermediate states X,.; during the forward integration. Its computational cost and memory consumption are O(3MNsL) and OMMNs + L.), respectively. The memory for checkpoints can be non-negligible with practical number of steps.

---

Table 4: Results obtained for continuous-time physical systems.
KdV Equation
Cahn Hilliard System
MSF (x10-
mem.
time
MSE (x107")
mem.
time
adjoint method [2]
1.6133233
93.7=+0.0
166=16
1.68
93.7=10.0
942:124
hackpropagation
13.44
693.91000
105
4.68 1.89
3047.1:0.0
425:113
ACA 46
1.61
647.8+00.
137
5.82
648.0=0.0
484
proposed
:4.000
79.8+10.0
162
5.47:11.46
80.3
568.122
Mean-squared errors (MSEs) in long-term predictions, peak memory consumption [MiB], and computation time per iteration [ms/itr].
5.2 Continuous- Time Dynamical System
Experimental Settings: We evaluated the symplectic adjoint method on learning continuous-time dynamical systems [13 33 40). Many physical phenomena can be modeled using the gradient of system energy H as dz/dt GVH (z). where Gs coefficient matrix that determines the behaviors of the energy D We followed the experimental settings of HNN++ provided in 31 A neural network composed of one convolution layer and two fully connected layers approximated the energy function H and learned the time series by interpolating two successive samples. The deterministic convolution algorithm systems described by PDEs, namely the Korteweg-De De Vries (KdV) equation and the Cahn-Hilliard system. We used a batch-size of 100 to put mini-batch into a single NVIDIA TITAN V GPU instead of the original batch-size of 200. Moreover, we used the eighth-order Dormand-Prine composed of 13 stages, to emphasize the efficiency of the proposed method. We omitted the baseline scheme because of M = 1. We evaluated the performance using mean squared errors (MSEs) in the system energy for long-tern predictions.
enabled (see Appendix (D.3 for discussion). We employed two physical
method 177.
Performance: The medians + standard deviations of 15 runs are summarized in Table 4 Due to the accumulated error in the numerical integration, the MSEs had large variances, but all methods obtained similar MSEs. ACA consumed much more memory than the symplectic adjoint method because of the large number of stages; the symplectic adjoint method is more beneficial for physical simulations, which often require extremely higher-order methods. Due to the severe nonlinearity, the adjoint method had to employ a small step size and thus performed slower than others (i.e., N > N).
6 Conclusion
We proposed the symplectic adjoint method, which solves the adjoint system by a symplectic integrator with appropriate checkpoints and thereby provides the exact gradient. It only applies the backpropagation algorithm to each use of the neural network, and thus consumes memory much less than the backpropagation algorithm and the checkpointing schemes. Its memory consumption is competitive to that of the adjoint method because the memory consumed by checkpoints is negligible in most cases. The symplectic adjoint method provides the exact gradient with the same step size as that used for the forward integration. Therefore, in practice. it performs faster than the adjoint method, which requires small step size to suppress numerical errors.
As shown in the experiments, the best integrator and checkpointing scheme may depend on the target system and computational resources. For example, Kim et al. [26) has demonstrated that quadrature methods can reduce the computation cost of the adjoint system for stiff equation in exchange for the additional memory consumption. Practical packages provide many integrators and can choose the best ones (20 36). In the future, we will provide the proposed symplectic adjoint method as part of such packages for appropriate systems.
Acknowledgments and Disclosure of Funding
This study was partially supported by JST CREST (JPMJCR1914), JST PRESTO (JPMJPR2IC7), and JSPS KAKENHI (19K20344).
(MIT License)
10

---

References
[1] Bochev, P. B. and Scovel, C. (1994). On quadratic invariants and symplectic structure. BIT Numerical Mathematicy, 34(3)333--35.
[2] Chen, T. Q. Rubanova, Y., Bettencourt, Duvenaud, Chen, Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. (2018). Neural Ordinary Differential Equations. In Advances in Neural Information Processing Systems (NewrPPS).
[3] Chen, Z., Zhang, J., Arjovsky, M., and Bottou, L. (2020). Symplectic Recurrent Neural Networks. In International Conference on Learning Representations (ICLR).
[4] Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2018). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. arXiv.
[5] Dinh, L.,, Krueger, D., and Bengio, Y. (2014). NICE: Non-lincar Independent Components Estimation. In Workshop on International Conference on Learning Representations.
[6] Dinh, L.., Sohl-Dickstein, J., and Bengio, S. (2017). Density estimation using Real NVP. In International Conference on Learning Representations (ICLR).
[7] Dormand, J. R. and Prince, P. J. (1986). A reconsideration of some embedded Runge- Kutta formulae. Journal of Computational and Applied Mathematics, 15(2):003-11.
[8] Errico, R. M. (1997). What Is an Adjoint Model? Bulletin of the American Meteorological Society, 78(11):2577-2533
[9] Furihata. D. and Matsuo, T. (2010). Discrete Variational Derivative Method: A Structure- Preservvn Numerical Method for Partial Differential Equations. Chapman and Hal//CRC.
[10] Gholami, A., Keutzer, K., and Biros, G. (2019). ANODE: Unconditionally accurate memory-efficient gradients for neural ODEs. In International Joint Conference on Artificial Intelligence (IJCAI).
[11] Gomez, A. N., Ren, M., Urtasun, R., and Grosse, R. B. (2017). The Reversible Residual Network: Backpropagation Without Storing Activations. In Advances in Neural Information Processing Systems (NIPS).
[12] Grathwohl, W., Chen, R. T. Q., Bettencourt, J., Sutskever, I., and Duvenaud, D. (2018). FFJORD: Free- form Continuous Dynamics for Scalable Reversible Generative Models. In International Conference on Learning Representations (ICLR).
[13] Greydanus, S., Dzamba, M., and Yosinski, (2019). Hamiltonian Neural Networks. Advances in Neural Information Processing Systems (NeurIPS).
[14] Griewank, A. and Walther, A. (2000). Algorithm 799: Revolve: An implementation of checkpointing for the reverse or adjoint mode of computational differentiation. ACM Transactions on Mathematical Software, 26(1): 2691):19-44 45.
[15] Gruslys, A., Munos, R., Danihelka, I., Lanctot, M., and Graves, A. (2016). Memory--ffiiiint backpropaga- tion through time. Advances in Neural Information Processing Systems (NIPS).
[16] Hairer, E., Lubich, C., and Wanner, (2006). Geometric Numerical Integration: Structure-Preserving Algorithms for Ordinary Differential Equations, volume of Springer Series in Computational Mathematics. Springer- Verlag, Berlin/Heidelberg.
[17] Hairer, E., Nørsett, S. P., and Wanner, G. (1993). Solving Ordinary Differential Equations I: Nonstiff Problems, volume of Springer Series in Computational Mathematics. Springer Berlin Heideeberrg Berlin, Heidelberg.
[18] He, K., Zhang, X., Ren, S., and Sun, J. (2016). Deep Residual Learning for Image Recognition. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR).
[19] Herrera. C., Krach, F., and Teichmann, (2020). Neural Jump Ordinary Differential Equations: Consistent Continuous- Time Prediction and Filtering. In International Conference on Learning Representations (ICLR).
[20] Hindmarsh, A. C., Brown, P. N., Grant, K. E., Lee, S. L., Serban, R., Shumaker, D. E., and Woodward, C. S. (2005). SUNDIALS: Suite of nonlinear and differentiallalgecraic equation solvers. ACM Transactions Mathematical Software, 31(3):363-33

---

[21] Hochreiter, S. and Schmidhuber, J. (1997). Long Short-Term Memory. Neural Computation, 9(8):1735- 1780.
[22] Hutchinson, M. (1990). A stochastic estimator of the trace of the influence matrix for laplacian smoothing splines. Communications in Statistics Simulation and Computation, 19(2):433 450.
[23] Jiang, C. M., Huang, J., Tagliasacchi, A., and Guibas, L. (2020). ShapeFlow: Leamable Deformations Among 3D Shapes. In Advances in Neural Information Processing Systems (NeurlPS).
[24] Kidger, P. Morrill, J., Foster, J., and Lyons, T. (2020). Neural Controlled Differential Equations for Irregular Time Series. In Advances in Neural Information Processing Systems (NeurIPS).
[25] Kim, H., Lee, H., Kang, W. H., Cheon, S. J., Choi, B. J., and Kim, N. S. (2020). WaveNODE: A Continuous Normalizing Flow for Speech Synthesis. In ICML2020 Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihoud Models.
[26] Kim, S., Ji, W., Deng. S., Ma. Y., and Rackauckas, C. (2021). Stiff Neural Ordinary Differential Equations. arXiv.
[27] Kingma, D. P. and Ba, J. (2015). Adam: Method for Stochastic Optimization. International Conference on Learning Representations (ICLR).
[28] LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. (1998). Gradient-basse based learning applied to document recognition. In Proceedings of the IEFF, volume 86, pages 2278-2323.
[29] Li, X., Wong, T.-K. L., Chen, R. T. Q., and Duvenaud, D. (2020). Scalable Gradients for Stochastic Differential Equations. In Artificial Intelligence and Statistics (AISTATS).
[30] Lu, Y., Zhong, A., Li, Q., and Dong, B. (2018). Beyond finite layer neural networks: Bridging deep architectures and numerical differential equations. In International Conference Machine Learning (ICML).
[31] Matsubara, T., Ishikawa. A., and Yaguchi, T. (2020). Deep Energy- Based Modeling of Discrete-Time Physics. In Advances in Neural Information Processing Systems (NeurIPS).
[32] Matsuda, T. and Miyatake, Y. (2021). Generalization of partitioned Runge Kutta methods for adjoint systems. Journal of Computational and Applied Mathematics,
[33] Papamakarios, G., Pavlakou, T., and Murray, (2017). Masked Autoregressive Flow for Density Estimation. In Advances in Neural Information Processing Systems (NIPS).
[34] Pascanu, R., Mikolov, T., and Bengio, Y. (2013). On the difficulty of training Recurrent Neural Networks. In International Conference Machine Learning (ICML).
[35] Paszke, A., Chanan, G., Lin, Z., Gross, S., Yang, B., Antiga, L., and Devito, Z. (2017). Automatic differentiation in PyTorch. In Autodiff Workshop Advances Neural Information Processing Systems.
[36] Rackauckas, C., Ma, Y., Martensen, J., Warner, C., Zubov, K., Supekar, R., Skinner, D., Ramadhan, A., and Edelman, A. (2020). Universal Differential Equations for Scientific Machine Learning, arXiv.
[37] Rana, A., Li, A., Fox, D., Boots, B., Ramos, F., and Ratliff. N. (2020). Euclideanizing Flows: Diffeomorphic Reduction for Learning Stable Dynamical Systems. In Conference on Learning for Dynamics and Control (LADC).
[38] Rezende, D. J. and Mohamed, S. (2015). Variational Inference with Normalizing Flows. In International Conference on Machine Learning (ICML).
[39] Rumelhart, D. E., Hinton, G. E., and Williams, J. (1986). Learning representations by back-propagating errors. Nature, (536633))
[40] Saemundsson, S., Terenin, A., Hofmann, K., Deisenroth, M. P., Sæmundsson, S., Hofmann, K., Terenin,
A., and Deisenroth, M. P. (2020). Variational Integrator Networks for Physically Meaningful Embeddings. In Artificial Intelligence and Statistics (AISTATS).
[41] Sanz-Serna, J. M. (2016). Symplectic Runge- Kutta schemes for adjoint equations, automatic differentiation, optimal control, and more. SIAM Review, 58(1):3-33,
[42] Takeishi, N. and Kawahara, Y. (2020). Learning dynamics models with stable invariant sets. AAAI Conference Artificial Intelligence (AAAI).
12

---

[43] Teshima. T., Tojo, K., Ikeda, M., Ishikawa, I., and Oono, K. (2020). Universal Approximation Property of Neural Ordinary Differential Equations. In NeurIPS Workshop Disferential Geometry meets Deep Learning (Dif)Geo4DL).
[44] Wang. (2013). Forward and adjoint sensitivity computation of chaotic dynamical systems. Journal of Computational Physics, 235:1-13.
[45] Yang, G., Huang. X., Hao, Z., Liu, M. Y., Belongie, S., and Hariharan, B. (2019). Pointflow: 3D point cloud generation with continuous normalizing flows. International Conference on Computer Vision (ICCV).
[46] Zhuang, J., Dvomek, N., Li. X., Tatikonda. S., Papademetris, X., and Duncan, J. (2020). Adaptive Checkpoint Adjoint Method for Gradient Estimation in Neural ODE. In International Conference on Machine Learning (ICML).
[47] Zhuang, J., Dvornek, C., Tatikonda. S., Dunean, (2021). MALI: A memory efficient and reverse accurate integrator for Neural ODEs. In International Conference on Learning Representations (ICLR).
13


Reviews:
[{"review_id": "v2M_4Zs4g8", "reviewer": "Reviewer_G97n", "summary": "This paper proposed to view the state and its adjoint in Neural ODEs as a partitioned system, hence the partitioned Runge-Kutta method can be applied, hence the integrator is symplectic and is accurate in terms of gradient estimation. To my knowledge, this is the first time that the idea of the higher-order symplectic integrator is introduced to Neural ODE. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros: \\\nThe idea of a symplectic Runge-Kutta solver and view the adjoint system as a partitioned system is well-studied in the numerical analysis community, but to my knowledge, it's not systematically introduced into the deep learning community. The paper is in general well written with solid theoretical analysis.\n\nComments: \\\nI would suggest the authors remove the \"minimal memory\" in the title since the memory cost still grows linearly with integration time steps, while the original adjoint method and MALI all have a constant memory w.r.t integration time. \n\nI'm curious how does the proposed method compare with MALI which has a constant $O(1)$ memory cost and is lower than symplectic-adjoint ($O(T)$ memory). MALI is a second-order method, which hinders the numerical accuracy of high-order solvers. But from my experience, even a 2nd order solver is sufficient for most deep learning with Neural ODEs. It would be more convincing if the authors could discuss or validate in experiments.\n\nAnother suggestion is that the authors might want to expand a bit more on why ACA uses $sL$ memory while symplectic-adjoint uses $L$ memory during backward, current writing does not explain it sufficiently and it's a bit hard to understand at first reading. For an RK-step of $s$ stages, ACA performs backward through all $s$ stages at once, while symplectic-adjoint performs backward pass for each stage of the $s$ stages, and each stage can be deleted from the memory once it's computed. This depends on the result of Thm 1, however, I'm curious if Thm1 always holds for any ODE system? Does it require some properties of the ODE system such as smoothness? Can it be applied to some neural-jump ODE or stochastic ODE?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "tDqepw1eCuz", "reviewer": "Reviewer_rdkK", "summary": "In this paper, the authors propose a new adjoint method for training neural ODEs, named the symplectic adjoint method, which is an adjoint method solved by a symplectic integrator. The proposed symplectic adjoint method obtains the exact gradient with memory proportional to the number of uses plus the network size. Numerically, the authors show that the symplectic adjoint method consumes much less memory than the naive backpropagation algorithm and checkpointing schemes, and performs faster than the adjoint method, and is robust to round off errors.", "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": "Strength:\n1. Propose a new numerical algorithm for the backpropagation of NODEs.\n\n2. Runge-Kutta only (which is generally enough), so accessible to higher-order methods. It has more memory advantage for higher-order methods.\n\n3. Comprehensive numerical experiments and nice performance.\n\nWeakness:\n1. The memory consumption is proportional to the number of time steps. When the ODE is not smooth enough for higher-order methods and/or integrating for a long time period, this may cause an issue. Please address this point and I am happy to raise the rating.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No", "review_text": "", "extra_scores": {"ethical_concerns": "NA", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4 hours"}, "extra_text": {}}, {"review_id": "cl0-k37x7DO", "reviewer": "Reviewer_1e2r", "summary": "The paper proposes a new algorithm for propagating gradients through neural\nODEs, which minimizes peak memory requirements by checkpointing not only the\nindividual solver steps but also the internal stages. The considered, specific,\npartitioned Runge--Kutta method has symplectic properties and provides exact\ngradients. The experiments demonstrate the reduced peak memory cost together\nwith competitive runtimes.\n", "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": "__Originality__\nTo the best of my knowledge, the proposed method for continuous sensitivity\nanalysis appears to be novel. Previous approaches for adjoint computation of\nneural ODEs are properly cited and compared to.  \nLarge parts of the specific algorithm seem to have been developed in [36]. To\nbest judge on the differences and on the specific contribution of this paper, an\nexplicit comparison to [36] would be very helpful.\n\n__Quality__\nThe submission appears technically sound. It includes an adequate discussion of\nrelated methods for sensitivity analysis, and a thorough derivation of the\nproposed symplectic adjoint method. Proofs are included or referred to.\n\n\n__Clarity__\nThe submission is well-structured and clearly written, and seems to contain\nsufficient information for reimplementing the method. The added code further\nhelps reproducibility. More specific comments for improvement can be found\nbelow.\n\n__Significance__\nThe proposed method provides a new algorithm for computing continuous\nsensitivities with reduced memory-cost, and thereby provides a significant\ncontribution. I believe there is the potential for the proposed method to be\nused, or at least considered, in future research or application of NODEs in\nlow-memory scenarios.\n\n\n__Questions for the authors__\n- To the best of my knowledge, symplectic ODE solvers rely on fixed steps (see\n  e.g. [16]). Could you clarify how your method is able to benefit from\n  step-size adaptation while still being symplectic?\n\n  Similarly, my understanding of symplectic ODE solvers is that they are able to\n  _exactly_ preserve an _approximate_ Hamiltonian, which is not equal to the true,\n  continuous Hamiltonian. This would contradict your claim on computing the\n  _exact_ gradient. Could you clarify my possible misconception?\n- It would be nice to include two other adjoint methods into the discussion:\n  (i) the CVODES adjoint described in [2] uses interpolation of the forward\n  solution, and (ii) the quadrature-based adjoint computation described in [3]\n  employs numerical quadrature instead of an ODE solver. Both are also described\n  and used in [4]. Discussing their respective advantages and limitations would\n  provide helpful information (though I am aware that [3] and [4] are only\n  available as preprints at the current time such that their inclusion is not\n  strictly necessary).\n\n\n__Additional, minor comments and questions__\n- Specifically mentioning _\"peak\"_ memory cost could be helpful; e.g. in 2.2 I\n  was initially confused about the memory cost: M components, with O(NsL) memory\n  cost, initially suggested a total O(MNsL) cost (instead of O(M+NsL)). Being\n  reminded about \"peak\" memory usage clarified this misunderstanding for me.\n- Introduction: l. 34-35 would benefit from a citation to support the claim.\n- l. 55: \"Backprop tends to suffer from rounding errors\" - I was not able to\n  find a comment on this in [41]. However, the information provided in the\n  supplementary D.1 seems reasonable.\n- 2.1: \"Neural ODE components\" is used throughout the paper, but never\n  introduced. An introducing sentence could help the clarity\n- 2.1: Is the number of layers \"L\" really helpful for the discussion? All\n  compared methods scale the same way in \"L\", and only the discussion regarding\n  M, N, s seems relevant.\n- l. 71: It seems misleading to claim that the original study on neural ODE\n  \"introduced\" the adjoint method - they rather \"used\" this method. Multiple\n  subsequently cited papers predate [2] and reflect that statement.\n- l. 81-82: Why should the size of the adjoint variable lead to numerical\n  issues? An explanation, or a supporting citation, would be helpful.\n- Throughout the paper the \"Dormand--Prince\" method is mentioned, but to the\n  best of my knowledge there are multiple methods developed by and named after\n  Dormand and Prince - as you also mention in your experiments. Specifying the\n  specific method (i.e. \"fifth-order Dormand--Prince\", \"DOPRI5\", or possibly\n  even \"Runge--Kutta 4/5\")\n- l. 104-105: The step-size and computational cost depends on the specified\n  tolerance level - for high tolerances, low orders could be preferrable.\n- l. 105-106: High-order symplectic integrators do exist. See e.g. [1].\n- l. 127: \"automatic differentiation\" seems misleading - continuous adjoint\n  sensitivities are precisely the alternative to automatic differentiation.\n- Remark 3: Could you clarify further what \"is _automatically_ solved by the same\n  Runge--Kutta method\" means?\n- l. 261: It seems strong to state robustness to rounding errors; instead,\n  becoming \"more robust\" is already a significant improvement.\n- Experiments: Doing more than three runs would be helpful to better assess the\n  standard deviation.\n- Experiments \"Different Runge--Kutta Methods\": Please specify the used\n  tolerance levels. In this context, I would imagine that the performance of the\n  solvers strongly depends on the tolerance level; for lower tolerances, DOPRI8\n  might provide the best performance.\n- Supplementary, eq. (12), typo: t -> T\n\n\n[1] \"Composition constants for raising the orders of unconventional schemes for\nordinary differential equations\", Kahan and Li (1997).\n\n[2] \"SUNDIALS: Suite of Nonlinear and Differential/Algebraic Equation Solvers\",\nHindmarsh et al. (2005)\n\n[3] \"Stiff Neural Ordinary Differential Equations\", Kim et al. (2021)\n\n[4] \"Universal Differential Equations for Scientific Machine Learning\",\nRackauckas et al. (2020)\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors mentioned both in the checklist and in the paper that the\nmethod is only applicable to Runge--Kutta methods which is a fair and valid\nlimitation. The other mentioned limitation relates to the runtimes: Table 1\nshows the largest estimate of computational cost for the proposed method. The\nexperiments do however not reflect this limitation. The first paragraph of page\n8 mentions this observation and refers to possible memory-related issues. A\npossible additional experiment could therefore be in computational environments\nwith an abundance of memory, to observe the validity of the estimated, increased\ncomputational cost (but I do not see this as a strict requirement for this paper).\n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes the use of the symplectic adjoint method for improving the memory footprint of neural ODE models. Reviewers agree that the use of such technique in the deep learning community is novel and supported by comprehensive numerical results.
I also believe that the memory consumption issue raised by reviewer rdkK has been, at least, partially addressed in the rebuttal.
I am therefore inclined to accept this work.
Decision: Accept (Poster)