Adaptive higher order reversible integrators for memory efficient deep learning

## Abstract

Abstract. The depth of networks plays a crucial role in the effectiveness of deep learning. However, the memory requirement for backpropagation scales linearly with the number of layers, which leads to memory bottlenecks during training. Moreover, deep networks are often unable to handle time-series appearing at irregular intervals. These issues can be resolved by considering continuous-depth networks based on the neural ODE framework in combination with reversible integration methods that allow for variable time-steps. Reversibility of the method ensures that the memory requirement for training is independent of network depth, while variable time-steps are required for assimilating time-series data on irregular intervals. However, at present, there are no known higher-order reversible methods with this property. High-order methods are especially important when a high level of accuracy in learning is required or when small time-steps are necessary due to large errors in time integration of neural ODEs, for instance in context of complex dynamical systems such as Kepler systems and molecular dynamics. The requirement of small time-steps when using a low-order method can significantly increase the computational cost of training as well as inference. In this work, we present an approach for constructing high-order reversible methods that allow adaptive time-stepping. Our numerical tests show both the advantages in computational speed and improved training accuracy of the new networks when applied to the task of learning dynamical systems.

## 1. Introduction

Deep neural networks are widely used across various learning tasks [, ], and their depth often plays a crucial role in the effectiveness of learning. These networks have also been shown to be particularly useful in the tasks of learning models of dynamical systems []. It was recently shown that the use of numerical methods for neural network architectures can provide impressive results with theoretical guarantees [, , , ].
In this work we use the theory of symmetric numerical methods for the construction of a new class of reversible neural networks. The new class allows memory efficient computations of gradients in training and reduced computational costs in learning models of dynamical systems.

The high memory costs of computing the gradient of very deep neural networks using the backpropagation algorithms poses a significant bottleneck in their training, hindering their scalability and efficiency.
To address this, a neural ODE approach combined with the adjoint method has been proposed for gradient computations [], which avoids storing intermediate states during forward propagation, potentially making the cost of gradient computation independent of network depth. However, it was quickly realized [] that using this approach with arbitrary discretization methods leads to incorrect gradients.

The solution of this problem is to use reversible integrators, which ensure accurate gradient computations by reconstructing intermediate states precisely during backward integration. However,
symplectic reversible integrators [], which are a large class of well studied integrators, do not allow adaptivity in the step size and require a particular structure in the neural ODE. This makes them unsuitable for use in time-series applications where data appears at irregular intervals, or when time-step needs to be decreased adaptively to achieve a prescribed accuracy in the learning of dynamical systems.

The existing reversible integrators that allow the use of variable time-steps are asynchronous leapfrog (ALF), which is used in MALI network [], and reversible Heun []. These methods are based on operating on an augmented space: a neural ODE in the original variable $z$
is extended to a larger space and replaced by a neural ODE in $(z,v)$. Both methods are known to be of order of accuracy $(2,1)$ in $(z,v)$, which makes them computationally costly in learning tasks that require high accuracy in the integration of the neural ODE, in particular, in learning of dynamical systems.
To be able to reach high accuracy, the lower order methods are forced to use small step sizes and, as a result, higher computational costs. This phenomenon was highlighted in the examples in []. Therefore, there is a need for higher order reversible methods, which we address in this paper.

The problem of learning of models of dynamical systems is fundamental for understanding complex phenomena observed in nature. In addition, it is crucial in engineering for the development of digital twins. In all these cases, there are complex high dimensional and nonlinear systems involved. Deep learning has already proven to be useful in this context [, , , ] and there are different approaches which already exist. For example, if the training data consists of time-series data that corresponds to a constant time-step, neural ODEs can be trained with a low order method and time-series can be predicted with high accuracy provided that the trained neural ODE is integrated with the same integrator that was used during training [, , , ]. However, in realistic examples, snapshots of trajectories with variable time-steps need to be processed. Moreover, a discretization-independent prediction of the system’s evolution is desired. In this case, the trained neural ODE is required to represent the continuous dynamical system to high accuracy and higher order methods with adaptive time stepping can be employed to achieve this efficiently.

The main contribution of this paper is the development of a methodology to construct reversible neural networks based on higher order numerical methods. First, we prove that ALF
is of order 2 in both $(z,v)$ at even time steps. Using this property and the theory of composition methods, we construct a class of reversible networks of any even order.
A particular architecture based on 4th order networks is compared to the already known ALF method
on examples of learning dynamical systems. The obtained advantages are twofold. On the one hand, the comparison of the two methods with the adaptive time stepping shows that the higher order method is more computationally efficient. On the other hand, comparing the methods with fixed step size we observe that the higher order method is able to reach a higher training accuracy in the learned dynamics and does not require the use of very small time steps to reach it.

## 2. Background

### 2.1. Neural ODE

Assume that an unknown function $F:X\rightarrow Y$ is approximated by a neural network based on training data $\{x_{i},y_{i}=F(x_{i})\}^{n}_{i=1}$. The neural network architecture is given as a composition of parameterized nonlinear activation functions $\sigma$ and the propagation in a feature space $\mathbb{R}^{n}$ is defined by

$$ (2.1) $\displaystyle z_{j+1}=\sigma(z_{j},\theta_{j}),\quad j=0,\dots,N-1,$ $\displaystyle z_{0}=z,$ $$

where $z_{j}$ is the value of the feature variable at the $j$th layer with the total number of layers $N$ and $z=(x_{1},\dots,x_{n})$.
Such kind of networks can be constructed based on discretized differential equations. In this case the differential equation is called neural ordinary differential equation (neural ODE). Then, $z_{0},\dots,z_{N}$ satisfying (2.3) approximate a flow of a neural ODE given by

$$ (2.2) $\dot{z}(t)=f(z(t),\theta(t)),$ $$

for some vector field $f$. In addition, there is a natural new parameter $h$ appearing in the discretized neural ODE, which stands for the time step size. The resulting network parameterized by $h$ is as follows

$$ (2.3) $\displaystyle z_{j+1}=\sigma_{h}(z_{j},\theta_{j}),\quad j=0,\dots,N-1,$ $\displaystyle z_{0}=z.$ $$

The learning problem is to find parameters $\{\theta_{j}\}_{j=0}^{N-1}$ which lead to the best approximation of $F$. The parameters are usually found as a solution of the following optimization problem.

$$ (2.4) $\displaystyle\min_{\{\theta_{j}\}}J=L(z_{N},y)$ $\displaystyle z_{j+1}=\sigma_{{\color[rgb]{0,0,0}h}}(z_{j},\theta_{j}),\quad j =0,\dots,N-1,$ $\displaystyle z_{0}=z,$ $$

where $L(\cdot,\cdot)$ is a loss function which measures the distance between the output of the network and $y=(y_{1},\dots,y_{N})$ from the training data.

Because of the connection between network (2.3) and the corresponding neural ODE (2.2), there exists a continuous counterpart of (2.4) which makes (2.4) an approximation of an optimal control problem of the form

$$ (2.5) $\displaystyle\min_{\theta(t)}J=L(z(T),y)$ $\displaystyle\dot{z}(t)=f(z(t),\theta(t)),\quad t\in[0,T],$ $\displaystyle z(0)=z.$ $$

Solutions of (2.4) are usually found using methods based on gradient descent. Such methods require computations of the gradients of the loss function with respect to all parameters $\theta_{j}$ for $j=1,\dotsm N-1$.
By the chain rule

$$ $\frac{\partial}{\partial\theta_{j}}J=\nabla_{z}L(z_{N},y){{}^{\top}}\frac{ \partial z_{N}}{\partial z_{N-1}}\cdots\frac{\partial z_{j+2}}{\partial z_{j+1 }}\frac{\partial z_{j+1}}{\partial\theta_{j}}.$ $$

To avoid computationally expensive multiplication of large matrices, the formula is evaluated from the left to the right (backpropagation). This requires that the intermediate values $z_{j}$ ($j=1,\ldots,N$) are available. Their size corresponds to the width of the layer and their number to the network’s depth $N$.

### 2.2. Methods of gradient computations

There are two main approaches to calculate the gradients in deep learning, namely, the approach called discretise-then-optimise and the approach called optimise-then-discretise. The main difference is that in the first approach, the gradients are calculated with respect to the discrete flow (2.4) and in the second approach the gradients
relate to the continuous problem (2.5).
The first approach is the default approach in deep learning. Its problem is in the memory requirement, which is resolved in the second approach at the cost of restrictions on the network structure.

##### Discretize then optimize.

The calculation of the gradients in this approach is done in two steps. First step is to propagate forward through the network $z_{j}\rightarrow z_{j+1}$ using the current values of parameters $\theta_{j}$ determined in the previous iteration of the optimization algorithm or by initialisation. All the values of $z_{j}$ and the computational graph are saved to be accessed next in the backpropagation step. In this step the gradient $\nabla_{\theta_{j}}J$ is calculated using the reverse mode of automatic differentiation (AD). By construction, all the gradients are computed exactly. Therefore, the obtained gradients are exact for the discrete optimization problem.

##### Optimise then discretise

This approach is based on the formula for the continuous gradients via adjoint variables $p(t)$. The adjoint variables are the solution of

$$ (2.6) $\dot{p}=-\frac{\partial}{\partial z}f(z(t),\theta(t))^{\top}p,\qquad p(T)= \nabla L(z(T),y),$ $$

on time interval $[0,T]$ and the differential of $J(\theta)$ with respect to $\theta(t)$ is calculated as follows

$$ (2.7) $DJ(\theta(t))=p^{\top}(t)\frac{\partial}{\partial\theta}f(z(t),\theta(t)).$ $$

This method is particularly interesting when the available memory is limited. In this case, the forward propagation is done to obtain the value of $z(T)$ and $\nabla L(z(T),y)$, there is no need to save the intermediate values $z_{j}$ and the computational graph, because the backpropagation is realized by integrating numerically backward in time the equations for $(z(t),p(t))$. The correct discretisation of the state-adjoint equations leads to the same expression for the gradient as the one obtained in the discretise-then-optimise approach. Still, the values of $z_{j}$ obtained by the numerical integration backward do not always coincide with the values obtained in the forward pass. This is why this approach usually leads to inexact gradients. This issue can be solved by considering reversible networks.

### 2.3. Reversible neural network

Reversible network is a network with the property that there exists an explicit formula for backward propagation $z_{j+1}\mapsto z_{j}$ that exactly inverts a forward pass $z_{j}\mapsto z_{j+1}$, i.e., there exists a map $\tilde{\sigma}(\cdot)$, such that $z_{j}=\tilde{\sigma}(z_{j+1})$ and $z_{j+1}=\sigma_{h}(\tilde{\sigma}(z_{j+1}),\theta_{j})$. It requires that the time steps $t_{0},\ldots,t_{N}$ have been stored when the forward propagation was computed but it does not require storage of the (potentially very high-dimensional) intermediate values $z_{j}$. The discretizations of neural ODE (2.3), which admit this property are called reversible methods. As for now, there are only two known reversible methods allowing for an adaptive choice of the step size, namely, asynchronous leapfrog [] and reversible Heun [].

This notion of reversibility for neural networks needs to be contrasted with the notion of time-reversibility or symmetry for numerical integrators.
In the context of neural networks, reversibility means that there exists an explicit, efficient formula to invert the forward pass. In numerical integration theory, a time-reversible or symmetric numerical integrator is a formula to advance the solution of an ordinary differential equation by time $h$ such that its inverse is obtained by substituting $h$ by $-h$ [, II.3]. In case of the dynamical system $f(z(t),\theta(t))$ from (2.3), if the discretization by a numerical method $z_{j+1}=\sigma_{h}(z_{j},\theta(t_{j}))$ is symmetric, then it implies $z_{j}=\sigma_{-h}(z_{j+1},\theta(t_{j+1}))$, or equivalently, $z_{j}=\sigma_{-h}(z_{j+1},\theta_{j+1})$. Therefore, we can set $\tilde{\sigma}(\cdot)=\sigma_{-h}(\cdot,\theta_{j+1})$, which implies that the method is reversible. The symmetry of integrators is beneficial in the context of the article as inverses of the methods required for backpropagation take simple forms and efficient classical techniques to construct higher order methods [, II.4] apply.

##### Asynchronous Leapfrog (ALF) method

As the optimization parameter $\theta(t)$ in the dynamics $f(z(t),\theta(t))$ depends on time, it can be seen as a part of $f$ and written simply $f_{\theta}(z(t),t)$. The ALF method requires the augmentation of the pair of state and time $(z,t)$ with the velocity $v$ which approximates $f_{\theta}(z(t),t)$. We denote a step forward of ALF method with the step size $h$ by $\Psi^{ALF}_{h}$. Given a triple $(z_{j},v_{j},t_{j})$ and a step size $h$, the algorithm generates in the forward pass the next values $(z_{j+1},v_{j+1},t_{j+1})$ as follows

$$ (2.8) $\begin{pmatrix}z_{j+1}\[8.5359pt] v_{j+1}\end{pmatrix}={\color[rgb]{0,0,0}\Psi^{ALF}_{h}(z_{j},v_{j},t_{j})=} \begin{pmatrix}z_{j}+hf_{\theta}(z_{j}+\frac{h}{2}v_{j},t_{j}+\frac{h}{2})\[8 .5359pt] 2f_{\theta}(z_{j}+\frac{h}{2}v_{j},t_{j}+\frac{h}{2})-v_{j}\end{pmatrix},\quad t _{j+1}=t_{j}+h.$ $$

The step backward calculates $(z_{j},v_{j},t_{j})$ from $(z_{j+1},v_{j+1},t_{j+1})$ as follows

$$ (2.9) $\begin{pmatrix}z_{j}\[8.5359pt] v_{j}\end{pmatrix}={\color[rgb]{0,0,0}\Psi^{ALF}_{-h}(z_{j+1},v_{j+1},t_{j+1}) },\quad t_{j}=t_{j+1}-h.$ $$

If the method is initialized at $(z_{0},f(z_{0},t_{0}),t_{0})$, then ALF is a second order method in $z$ and first order method in $v$, as it was shown in []. The order of accuracy is the order in step size $h$ of the error of the numerical flow with respect to the exact flow of the ODE []. The ODE in this case is $\dot{z}(t)=f_{\theta}(z(t),t)$ and its flow is extended to $(z(t),v(t)=f_{\theta}(z(t),t))$. Notice that ALF is symmetric by definition.

##### Reversible Heun method

Another reversible method based on state-space augmentation
was introduced in []. The state is augmented to $(v,z,\hat{z},t)$ with initialization at $(f_{\theta}(z_{0},t_{0}),z_{0},z_{0},t_{0})$. The reversible Heun method was shown to be also of second order in $z$ and first order in $v$. In addition, it is a symmetric method.

In the following part of the paper we will concentrate on the construction of higher order methods based on ALF, but the same can be also applied to the reversible Heun method.

## 3. New reversible architectures

A general approach in numerical analysis to construct higher order symmetric methods is by composition []. In this case one can start with a lower order numerical method and construct a new method by composition of the lower order method with a particular choice of step sizes. This construction leads to a method of higher order of accuracy.

###### Example 3.1 .

Consider a simple example of a differential equation given by

$$ (3.1) $\dot{z}=z^{2}+t+\sin(zt)+\frac{1}{z^{2}+1}.$ $$

Solving the equation numerically using ALF method leads to second order in the error with respect to a high accuracy solver in both $z$ and $v$ as displayed in Figure 1. This is surprising as the order of consistency of ALF in $v$ is only 1 []. Indeed, as we show below, a method consisting of a composition of two steps of ALF has order of consistency 2 in $(z,v)$, which explains the convergence behaviour. Since there is no theory for composition of mixed order methods, this observation also proves crucial in the application of composition methods to (two steps of) ALF.

Figure: Figure 1. Log-log plot of the global error of trajectories $(z(t),v(t))$ of (3.1) defined on time interval $[0,1.0]$ and obtained by ALF with $h$ ranging from $0.5$ to $10^{-3}$.
Refer to caption: x1.png

###### Theorem 3.2 .

Composition of two steps of ALF methods, i.e. $\Psi^{ALF}_{h/2}\circ\Psi^{ALF}_{h/2}$, applied to $\dot{z}=f(z,t)$ provides second order accurate approximations of position $z$ and velocity $v=\dot{z}$.

The proof of the theorem is based on comparing the terms in the Taylor series of the exact flow of a differential equation and the numerical flow obtained by composition of two steps of the ALF method. We refer to the appendix for the computations.
The composition of two steps of the ALF method, each with time step $\frac{h}{2}$, will be called ALF2 and denoted by $\Psi^{ALF2}_{h}$. Now we are in a classical situation, with ALF2 a one step reversible method of even order and we can apply the composition methods to construct higher order methods. In this work we consider the Yoshida approach []. Yoshida composition permits to construct methods of a higher accuracy by composing numerical methods of order $2k$ for some integer $k$. It is defined by a symmetric composition of the same method $\Psi^{\mathrm{2k}}$ ($2k$ stands for the order) with different step sizes

$$ $\Psi^{\mathrm{Y}}_{h}=\Psi^{\mathrm{2k}}_{ah}\circ\Psi^{\mathrm{2k}}_{bh}\circ \Psi^{\mathrm{2k}}_{ah}$ $$

with time steps defined by

$$ $a=\frac{1}{2-2^{\frac{1}{2k+1}}},\qquad b=1-2a.$ $$

###### Theorem 3.3 ( [ 22 ] ) .

Yoshida composition of a reversible method $\Psi^{\mathrm{2k}}$ of order $2k$ has order $2k+2$ and is reversible.

Yoshida is not the only composition method that can be used, another possible approach is Suzuki composition []. Several approaches are reviewed in [, ].

###### Remark 3.4 .

The approach based on Yoshida composition might require checkpoints in case of certain neural ODEs, e.g., when the learning task is to learn a dispersive partial differential equation such as heat equation. This is because the Yoshida composition forces the use of negative time-steps, which can be a problem in dissipative cases, where it can lead to instability.

In the following, we will denote by $\Psi^{Y,2k}_{h}$ the higher order methods obtained by Yoshida composition of ALF2, where $2k$ is the order of the method. The constructed higher order reversible method can be used for the construction of a reversible network. In this case, the step forward and the step backward are defined recursively based on the steps forward and backward of a lower order method $\Psi^{Y,2k-2}_{h}$. The starting method of order $2$ is the ALF2 method, i.e. $\Psi^{Y,2}_{h}=\Psi^{ALF2}_{h}$ by abuse of notation.

###### Example 3.1 .

###### Theorem 3.2 .

###### Theorem 3.3 ( [ 22 ] ) .

###### Remark 3.4 .

##### Adaptive stepping

One of the main advantages in the construction of reversible methods based on ALF is that they allow for adaptive step sizes []. This can be done in the same manner as for ALF [], where the main idea is to delete the computational graph and all the variables needed for the step size computations and only the value of the accepted new step size $h_{j}$ is saved. As a result, values $h_{1},\dots,h_{N}$ are saved and then accessed in the integration backward needed for gradient computations. This can be done in exactly the same manner for our Yoshida-based methods. The resulting steps forward and backward are summarized in Algorithm 1 and Algorithm 2.

Figure: Algorithm 1 Step forward of $2k$-th order Yoshida

Figure: Algorithm 2 Step backward of $2k$-th order Yoshida

###### Remark 3.5 .

Notice that even though the reversible Heun method was proven to be of order $(2,1)$ in $(z,v)$ in [], it was noted in the same paper that it gains the second order in both variables at even steps. This implies that the composition approach can be used in this case as well.

###### Remark 3.5 .

##### Gradient computations

The augmentation of the feature space leads to the new variable which we denote by $\phi=(z,v)$. Then, the learning problem is formulated as follows with $P_{z}(\phi)$ projection of $\phi$ to $z$

$$ (3.2) $\displaystyle\min_{\{\theta_{j}\}}J=L(P_{z}(\phi_{N}),y)$ $\displaystyle\phi_{j+1}=\Psi^{Y,2k}_{h}(\phi_{j},\theta_{j}),\quad j=0,\dots,N -1,$ $\displaystyle\phi_{0}=(z,f(z,\theta_{0})).$ $$

Following [], the discrete version of (2.6) associated with (3.2) is given by

$$ (3.3) $\left(\lambda_{N}\right)^{\top}=\nabla L(\phi_{N}),\quad\lambda_{j}=\left( \frac{\partial\phi_{j+1}}{\partial\phi_{j}}\right)^{\top}\lambda_{j+1},$ $$

and the gradients are computed by

$$ (3.4) $\frac{\partial J(\theta)}{\partial\theta_{j}}=\lambda_{j+1}^{\top}\frac{ \partial\phi_{j+1}}{\partial\theta_{j}}.$ $$

The adjoint method for the gradient computation as in the MALI network [] and the reversible Heun network [] is based on the propagation $\phi_{j}\rightarrow\phi_{j+1}$ and automatic differentiation for the computation of the step backward of the adjoint variable following (3.3). The resulting method of gradient computation is summarized in Algorithm 3.

Figure: Algorithm 3 Computation of gradients

## 4. Experiments

### 4.1. Parameter identification in dynamical systems

We consider the identification problem of unknown parameters of a dynamical system. The structure of the differential equations is assumed to be known, but some parameters in the equations are unknown. The training data is given by a set of a time series of trajectories points $\{x_{l}(t_{i})\}_{i,l}$ with $l=1,\dots L,\ i=0,\dots,I$. The goal is to learn the parameters from the given trajectories. This class of problems can be naturally treated using the neural ODE approach. The vector field $f$ in (2.2) is given by the known differential equation and $\theta=\theta_{1},\dots,\theta_{s}$ is the set of unknown parameters. In this case, the learning problem can be stated in the form of (3.2), where the same $\theta_{1},\dots,\theta_{s}$ appear all at each layer. The training data $z_{0}=(x_{1}(t_{0}),\dots,x_{L}(t_{0}))$ stands for the initial points and $y$ includes all the other points of the given trajectories. We denote by $y(t_{i})$ the points in $y$ corresponding to trajectories at time $t_{i}$ for $i=1,\dots,I$. The loss have a special structure in this case as it depends on the intermediate states obtained during the integration of neural ODE, namely, it depends on $(z_{N_{1}},\dots,z_{N_{I}})$, to measure the distance with the given trajectories points $(y(t_{1}),\dots,y(t_{I}))$. As a result, it takes the form $L=\sum_{i=1}^{I}L_{i}(z_{N_{i}},y(t_{i}))$. Because of the additive form of the loss, the gradients can be computed as a sum of the corresponding gradients of $L_{1},\dots,L_{I}$ as follows

$$ $\frac{\partial L}{\partial\theta_{i}}=\frac{\partial L_{1}}{\partial\theta_{i} }+\cdots+\frac{\partial L_{I}}{\partial\theta_{i}},\qquad i=1,\dots,s,$ $$

where each of the terms in the sum is computed using Algorithm 3. The memory efficiency is still important in this case, because we don’t store all the intermediate states at the propagation forward, but only the states which approximate the trajectories at the desired times $t_{1},\dots,t_{I}$.

#### 4.1.1. Kepler problem

We consider the Kepler problem, where the dynamics describes the evolution of the position $q$ and velocity $v$ of a mass point moving around a much heavier body. It is modeled on the 4-dimensional space $x=(q,v)\in\mathbb{R}^{2}\times\mathbb{R}^{2}$.
The equations are defined on time interval $[0,1]$ as follows

$$ (4.1) $\dot{q}=v,\quad\dot{v}=-\frac{\alpha}{\|q\|^{3}}\,q,$ $$

with an unknown parameter $\alpha\in\mathbb{R}$. The training set is given by the initial condition $x(t_{0})$ and $q$-coordinate of $5$ points on a trajectory of (4.1) generated with $\alpha=\pi/4\approx 0.785$, i.e. $\{q(t_{i})\}^{5}_{i=1}$. The task is to learn $\alpha$ as accurately as possible. From the training set we form $z_{0}=x(t_{0})$ and the corresponding $y(t_{i})=q(t_{i})$ for $i=1,\dots,5$. We set up a learning problem in the form of (3.2) with the loss defined by $L=\sum_{i=1}^{5}\|q_{N_{i}}-q(t_{i})\|^{2}$ with $q_{N_{i}}$ projection of $z_{N_{i}}$ to $q$-coordinate and $N_{i}$ the number of time steps used in the integration from $t_{i-1}$ to $t_{i}$. We compare two algorithms for performing the numerical integration during training, namely, ALF and the Yoshida composition of ALF2 of order 4. We write Y4 for the Yoshida composition method for shortness. First, we test the wall-clock time required to reach loss accuracy $10^{-8}$ using adaptive methods. The tests are run for different initializations of $\alpha$ in optimization. The results can be seen in Table 1. It can be observed that in all tests, Y4 is at least four times faster than ALF.

**Table 1. Time to reach accuracy $10^{-8}$ using adaptive methods in the Kepler problem**
| Computation time |  |  |
| --- | --- | --- |
| Initial value of $\alpha$ | adaptive ALF | adaptive Y4 |
| 0.1 | 7.68 sec | 2.42 sec |
| 0.7 | 4.07 sec | 1.02 sec |
| 0.75 | 3.26 sec | 0.803 sec |
| 0.8 | 2.5 sec | 0.44 sec |
| 1.3 | 8.39 sec | 3.85 sec |

Next, we compare the accuracy in learning the parameter using methods with fixed step-sizes.
The accuracy is analyzed using the error landscape for ALF and Y4 for fixed step size $h=0.01$. The results in Figure 2 show that the minimizer of Y4 method is closer to the true value of the parameter than the minimizer of ALF, which means that Y4 will provide more accurate results in the optimization.

Figure: Figure 2. Error landscape of ALF and Y4 methods for Kepler problem showing the loss computed for the parameters in a neighbourhood of the true value of $\alpha$ displayed by a vertical line.
Refer to caption: extracted/5921894/Y_ALF_error_landscape.png

#### 4.1.2. Nonlinear harmonic oscillator

In the second example we consider a system of coupled Duffing oscillators, which describes the movements of a coupled system of mass points attached with springs with nonlinear elastic forces. The dynamics of $N$ mass points is given by the following equations

$$ (4.2) $\dot{q}_{i}=v_{i},\quad\dot{v}_{i}=-a_{i}q_{i}-b_{i}q_{i}^{3}-\sum_{j=1}^{N}e_ {i,j}(q_{i}-q_{j}),\qquad i=1,\dots,N,\\$ $$

with the condition $e_{i,j}=e_{j,i}$. Positions of $N$ mass points are given by $q=(q_{1},\cdots,q_{N})\in\mathbb{R}^{N}$ and velocities by $v=(v_{1},\cdots,v_{N})\in\mathbb{R}^{N}$. We set $x=(q,v)\in\mathbb{R}^{2N}$. In the numerical experiments we fix $N=10$ and assume that parameters $a_{i},b_{i},e_{i,j}\in\mathbb{R}$ for $i,j=1,\dots,10$ are unknown. As a result, (4.2) has dimension $20$ with $65$ unknown parameters. The training set consists of initial and final positions of 200 trajectories, that is $z_{0}=(x_{1}(t_{0}),\dots,x_{200}(t_{0}))$ and $y=y(t_{1})=(x_{1}(t_{1}),\dots,x_{200}(t_{1}))$. In this setting, we compare the computational time to reach a certain training accuracy of ALF and Y4 with adaptive time stepping and the training accuracy of ALF and Y4 with fixed step-size. We present the wall-clock times to reach the training accuracy $10^{-4}$ in Table 2. The time required by Y4 to reach accuracy $10^{-4}$ is almost two times smaller which
illustrates the lower computational costs of the method. Next, we test ALF and Y4 with the step size fixed to $h=0.1$ and the training until either the training accuracy reaches $10^{-4}$ or the number of epochs reaches 500. Figure 3 shows that ALF is stuck at the training accuracy $10^{-2}$ and the training stops because of reaching 500 epochs, while Y4 converges to accuracy $10^{-4}$ with $181$ epochs. The same behaviour is observed for different parameter initialization. Decreasing the step size to $h=0.01$ permits ALF to reach accuracy $10^{-4}$.

Figure: Figure 3. Training loss is displayed in logarithmic scale for the parameter identification in case of coupled oscillation for ALF and Y4 with fixed step size $h=0.1$.
Refer to caption: x2.png

**Table 2. Time to reach accuracy $10^{-4}$ using adaptive methods in nonlinear oscillators problem**
| Computation time |  |  |
| --- | --- | --- |
| Mean parameter error at initialization | adaptive ALF | adaptive Y4 |
| $0.28897009$ | $81608$ sec | $51720$ sec |
| $0.29821727$ | $68645$ sec | $40646$ sec |
| $0.30549358$ | 56764 sec | 29990 sec |
| $0.30289593$ | 96301 sec | 46524 sec |
| $0.29106813$ | 22161 sec | 13790 sec |

### 4.2. Learning of dynamical systems parameterized by neural network

We consider a problem, where a part of the structure of the differential equations is known and the unknown part is approximated using a neural network. Our goal is to find the neural network parameterization such that the resulting trajectories of the system are as close as possible to given trajectories from the training set. As before, the problem can be treated by the neural ODE approach. In this case the vector field $f$ in (2.2) is given by a neural network. The learning problem (3.2) is then to find the parameters of the network. As before, we need to integrate $f$ by a reversible method. The training is done as in the problem of identification of parameters. The only difference is that the true vector field is only approximated by a network and we learn the best possible approximation.

#### 4.2.1. Nonlinear harmonic oscillator

We consider the problem of approximating the potential function of a physical system. In this case, the vector field of the physical system is approximated by a neural network and then integrated to obtain the final time points of the trajectories. The result is compared with the true trajectories. We consider again the Duffing oscillators, this time with two mass points, where the equations can also be equivalently written as

$$ (4.3) $\dot{q}=v,\quad\dot{v}=-\nabla V(q),$ $$

with $q=(q_{1},q_{2}),\ v=(v_{1},v_{2})$ and $V(q)$ stands for the potential energy of the system given by

$$ (4.4) $V(q)=\frac{a_{1}}{2}q_{1}^{2}+\frac{a_{2}}{2}q_{2}^{2}+\frac{e}{2}(q_{2}-q_{1} )^{2}+\frac{b_{1}}{4}q_{1}^{4}+\frac{b_{2}}{4}q_{2}^{4}.$ $$

The learning task is to learn the unknown potential. The gradient of the potential is approximated using a neural network with 51500 parameters. The potential is obtained by integrating the neural network numerically. We compare the computational time of ALF and Y4 to reduce the value of the loss function below $10^{-2}$. The results presented in Table 3 show the better performance of the higher order method.

**Table 3. Time to get accuracy $10^{-2}$ by adaptive methods in oscillators problem parameterized by NN**
| Computation time |  |  |
| --- | --- | --- |
| Random initialization of parameters in NN | adaptive ALF | adaptive Y4 |
| Initialization 1 | 2504 sec | 1974 sec |
| Initialization 2 | 2961 sec | 1857 sec |
| Initialization 3 | 4185 sec | 2627 sec |
| Initialization 4 | 3542 sec | 2125 sec |
| Initialization 5 | 3396 sec | 2616 sec |

## 5. Conclusion

In this work, we construct higher order reversible methods. These constitute explicit numerical integrators which are compatible with adaptive step-size selection strategies. The methods are employed to train deep neural networks that are based on neural ODEs. Thanks to the reversibility property, we avoid high memory requirements for backpropagation in the optimization procedure of the network parameters.
We showed the advantages of the newly constructed networks on the example of a network based on a 4th order method. The advantages are twofold. On the one hand, the higher order method leads to lower computational costs which
results in faster training.
Moreover, in parameter identification tasks, higher accuracy is reached using higher order method at a given step size.

Such a construction of a neural network can be used for approximation properties of general functions. We only identified the advantages on the tasks of learning of dynamical systems. The advantages of the suggested network in case of other learning tasks has to be further investigated. One direction for future research are higher order reversible techniques to train neural SDEs as investigated in [].

## Appendix A Error analysis of ALF2

In this appendix we show the proof of Theorem 3.2. Namely, we show that the local error of ALF2 in $(z,v)$ is of order $O(h^{3})$.

Let us consider the Taylor expansion of the exact flow $(z(t),v(t))$ around $(z(t_{0}),v(t_{0})=f(z_{0},t_{0}))$.

$$ $z(t_{0}+h)=z_{0}+hf(z_{0},t_{0})+\frac{h^{2}}{2}\left(\frac{\partial f}{ \partial z}(z_{0},t_{0})\circ f(z_{0},t_{0})+\frac{\partial f}{\partial t}(z_{ 0},t_{0})\right)+O(h^{3}),$ $$

$$ $v(t_{0}+h)=f(z_{0},t_{0})+h\left(\frac{\partial f}{\partial z}(z_{0},t_{0}) \circ f(z_{0},t_{0})+\frac{\partial f}{\partial t}(z_{0},t_{0})\right)+\\ \frac{h^{2}}{2}(\frac{\partial^{2}f}{\partial z^{2}}(z_{0},t_{0})(f(z_{0},t_{0 }),f(z_{0},t_{0}))+\frac{\partial f}{\partial z}(z_{0},t_{0})\circ\frac{ \partial f}{\partial z}(z_{0},t_{0})\circ f(z_{0},t_{0})\\ +2\frac{\partial^{2}f}{\partial z\partial t}(z_{0},t_{0})\circ f(z_{0},t_{0})+ \frac{\partial f}{\partial z}(z_{0},t_{0})\circ\frac{\partial f}{\partial t}(z _{0},t_{0})+\frac{\partial^{2}f}{\partial t^{2}}(z_{0},t_{0}))+O(h^{3}).$ $$

Now we consider the same for the numerical flow obtained with ALF2, that is composition of two steps of ALF each with the step size $\frac{h}{2}$. One step of ALF2 from $(z_{0},v_{0})$ leads to $(z_{1},v_{1})$ of the form

$$ $z_{1}(h)=z_{0}+\frac{h}{2}(f(z_{0}+\frac{h}{4}f(z_{0},t_{0}),t_{0}+\frac{h}{4} )+f(z_{0}+hf(z_{0}+\frac{h}{4}f(z_{0},t_{0}),t_{0}+\frac{h}{4})\\ -\frac{h}{4}f(z_{0},t_{0}),t_{0}+\frac{3h}{4})),$ $$

and

$$ $v_{1}(h)=v_{0}+2(f(z_{0}+hf(z_{0}+\frac{h}{4}f(z_{0},t_{0}),t_{0}+\frac{h}{4}) -\frac{h}{4}v_{0},t_{0}+\frac{3h}{4})\\ -f(z_{0}+\frac{h}{4}f(z_{0},t_{0}),t_{0}+\frac{h}{4})).$ $$

Writing down the Taylor expansion in $h$ for $(z_{1}(h),v_{1}(h))$ we find exactly the same terms as in $(z(t_{0}+h),v(t_{0}+h))$ up to terms of the third order $O(h^{3})$. The computations to obtain the Taylor expansion of $(z_{1},v_{1})$ were done using Maple software. This implies that the local error of ALF2 is of the 3rd order, and therefore, the global error is of order 2. This finishes the proof.

## Appendix B Details of numerical experiments

In all the numerical experiments, our implementation of the Yoshida composition method uses the code of the MALI network []. We use the steps forward and backward of the ALF method as composition steps to compute ALF2 and its Yoshida composition.

### B.1. Kepler problem

The training data for the comparison of the computational time in Table 1 is given by a trajectory $x$ of (4.1) with initial condition $x_{0}=(0.75,0,0,\frac{0.9\pi}{4}\sqrt{\frac{5}{3}})$ on time interval $[0,T]=[0,1]$, which is an elliptic orbit. The trajectory is obtained by numerical integration using sci.integrate.odeint with relative and absolute tolerances $10^{-7}$ and $10^{-8}$ respectively and maximum step size $10^{-5}$. The optimizer used in the training is SGD from PyTorch with initial learning rate 0.1 and scaled by 0.95 for each epoch.

The error landscape in Figure 2 is obtained by considering 81 trajectories obtained using the same integration method as for the time comparison explained above for 81 different initial conditions $(x_{0})_{i}$ in a neighborhood of $x_{0}$, given by a 4-dimensional box of diameter 0.4 around $x_{0}$. The points $(x_{0})_{i}$ are chosen on a grid with a step size $0.1$, which includes $x_{0}$ as its point. The loss visualized in Figure 2 as a function of $\alpha$ is $L(\alpha_{k})=\frac{1}{81}\sum_{i=1}^{81}\sum_{j=1}^{5}\sum\|(q_{N_{j}}({
\alpha_{k}}))_{i}-q(t_{j},(x_{0})_{i})\|^{2}$ with $q_{N_{j}}$ projection of $z_{N_{j}}$ to $q$-coordinate and $\alpha_{k}$ taking 300 values in $[\frac{\pi}{4}-10^{-4},\frac{\pi}{4}+10^{-4}]$.

### B.2. Nonlinear harmonic oscillator

There are two settings considered for the learning of the dynamics (4.2). In the first setting, the learning problem is the parameter identification as presented in Section 4.1.2. In the second case we consider the parametrization of the potential by a neural network as described in Section 4.2.1. Here we give more details on both problems.

#### B.2.1. Identification of parameters

In the experiments for the time comparison shown in Table 2 we consider a set of $200$ trajectories in the training data with the initial conditions generated by the Halton sequence in a 20-dimensional box around zero vector $x_{0}$ with diameter 2.0. The trajectories are obtained by numerical integration using sci.integrate.odeint with relative and absolute tolerances $10^{-13}$ and $10^{-14}$ respectively. In the training we use AdamW optimizer from PyTorch with learning rate scheduler ExponentialLR. The results shown in Table 2 are obtained with different learning rates, namely, the first two with the initial learning rate $10^{-2}$ and $\gamma=0.995$, the last three with the initial learning rate $10^{-1}$ and $\gamma=0.998,0.997,0.99$ for the tree results respectively. At each epoch we consider all 200 trajectories, so that the loss is $L=\frac{1}{200}\sum_{i=1}^{200}\|(z_{N})_{i}-x(T,(x_{0})_{i})\|^{2}$ with $T=0.5$.

#### B.2.2. Neural network parametrization

The goal is to find the unknown potential governing (4.3). For this we assume a particular form of the potential, namely,

$$ $V(q)=\sum_{i=1}^{s}\sum_{j=1}^{n}c_{i,j}\sigma_{i}(q_{j})+\sum_{i=1}^{d}\sum_{ j=1}^{n}\sum_{k=j+1}^{n}C_{i,j,k}\Sigma_{i}(\|q_{j}-q_{k}\|),$ $$

where $\sigma_{i}$ stand for different single particle potentials and $\Sigma_{i}$ for double particle potentials. In the case considered in (4.4), we have

$$ $\displaystyle c_{1,1}=\frac{a_{1}}{2},\ c_{1,2}=\frac{a_{2}}{2},\ \sigma_{1}(q )=q^{2},$ $\displaystyle c_{2,1}=\frac{b_{1}}{4},\ c_{2,2}=\frac{b_{2}}{4},\ \sigma_{2}(q )=q^{4},$ $\displaystyle C_{1,1,2}=\frac{e}{2},\ \Sigma_{1}(x)=x^{2}.$ $$

In the learning problem, we assume that functions $\sigma_{1},\sigma_{2}$ and $\Sigma_{1}$ are unknown as well as parameters $a_{1},a_{2},b_{1},b_{2},e$. We parameterize the derivatives $\frac{a_{1}}{2}\sigma_{1}^{\prime},\frac{a_{2}}{2}\sigma_{1}^{\prime},\frac{b_
{1}}{4}\sigma_{2}^{\prime},\frac{b_{2}}{4}\sigma_{2}^{\prime}$ and $\frac{e}{2}\Sigma_{1}^{\prime}$ by neural networks each and use them to model the dynamics in (4.3). We use 5 neural networks, which we denote by $\xi_{1},\xi_{2},\xi_{3},\xi_{4},\xi_{5}$. All of them have the same architecture $q\mapsto W_{1}\tanh{\left(W_{2}\tanh{(W_{3}q)}\right)}$, where $W_{1}$ is a matrix of parameters of size $1\times 100$, matrix $W_{2}$ is of size $100\times 100$ and $W_{3}$ is of size $100\times 1$. The resulting dynamics is defined by

$$ $\displaystyle\dot{q}_{1}=v_{1},\quad\dot{v}_{1}=-\xi_{1}(q_{1})-\xi_{3}(q_{1}) -\xi_{5}(q_{1}-q_{2}),$ $\displaystyle\dot{q}_{2}=v_{2},\quad\dot{v}_{2}=-\xi_{2}(q_{2})-\xi_{4}(q_{2}) -\xi_{5}(q_{2}-q_{1}).$ $$

The equations parameterized by neural networks are then integrated using ALF or Yoshida composition of ALF2 at each epoch in the training. The training data is set to be a set of $1000$ trajectories with the initial conditions generated by the Halton sequence in a 4-dimensional box around $x_{0}=(0.8,-0.4,0.0,0.0)$ with diameter $2.0$. The optimizer is AdamW with initial learning rate $10^{-3}$ and scheduler ExponentialLR with $\gamma=0.995$. In addition, we consider batches of 300 trajectories at each epoch with the resulting loss function of the same form as in the case of the parameter identification problem.

## Peer Reviews

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: Are step sizes treated as fixed on the backward pass, or do we somehow take gradients w.r.t step size calculations? If the former, how large are the resulting errors compared to the errors in backward propagated states as computed with a non-time-reversible integrator of the same order?

Line 505: "The potential is obtained by integrating the neural network numerically" -- please explain.

The description of Heun's method (line 213) is both very terse (what is it?) and very detailed (initilaization values etc.)
soundness: 2
strengths: The idea of using higher-order reversible integrators to solve ODEs is inherently appealing, as the potential to reduce memory costs and improve accuracy is good. The ideas are mostly clearly explained, and the examples, if somewhat simplistic, are at least relevant to the stated goals.
summary: This paper uses time-reversible ODE integration to solve neural ODEs. It focuses on a specific class of higher-order reversible integrators, and uses the time reversibility property to avoid inconsistency between the system states $z$ on the forward vs. backward passes. This allows a backward propagation of gradient information with memory costs independent of the sequence length.

The authors demonstrate this strategy for tuning parameters of dynamical systems or learning ODE terms parametrized by neural networks.
weaknesses: Overall the paper doesn't provide significant theoretical results, and the experimental results are fairly weak (details below):

Literature review and citations are inadequate. There is a considerable and active literature on learning dynamical systems with ML, but this hasn't really been acknowledged. As a start, foundational papers such as Chen (2018) and Rackauckas (2020) should be cited, and many of the references in Ghadami (2022) will be relevant. There is also a considerable literature on physics-constrained (e.g. McGreivy 2023) and Hamiltonian neural networks (Greydanus 2019 and later articles) that has not been adequately cited. The first paragraph of the introduction makes fairly broad claims ("learning models of dynamical systems") with loose terminology ("numerical methods") but the cited references are all highly specific, a broad overview of the field should be given with major relevant works cited.

The assessments of improvements to speed/accuracy tradeoffs arising from the proposed approached are quite weak:
* In both tables 2 and 3 we never compare to alternative methods or even weak baselines such as a non-reversible integrator with adaptive step size (e.g. Forward Euler or RK4), but rather only between low- and high-order version of reversible integrators. We are missing a clear picture of how the proposed technique improves on alternative methods. It can' simply be that higher-order integrators are more efficient on these particular problems, that's simply not of sufficient general interest
* For parameter estimation tasks, we should consider accuracy with regard to the reference (ground truth) parameters, or prediction errors on held-out initial conditions. The training loss doesn't necessarily indicate which method is actually better in these senses.
* A fixed step size was used in table 2, providing an unrealistic picture of how this would be used and losing sight of the fundamental reasoning behind the technique
* Fig. 2 seems to exaggerate the benefits of Y4 over ALF, since in a realistic scenario the step size would be automatically set by each method. The loss landscape should (also) be shown for adaptive step sizes scenario, and for the same 10^-8 accuracy tolerance as used for computing table 1.
* The problems are fairly easy. The system state was fairly low dimensional in all examples. The only example where a neural network was trained involved learning a low-dimensional potential with strong inductive biases from physics. There are many examples of challenging dynamical system learning tasks in the literature, including some benchmark studies. Ideally we should see an improvement of cost-accuracy tradeoffs on some established task.

One of the advantages of backward mode differentiation over forward mode is that we don't need to compute input-output gradients $d\phi_{j+1}/d\phi_j$ which are of size $n^2$. But here we do, at least in the current formulation, calculate and store the matrix explicitly. Is this strictly necessary, or is there some way around it? What happens when $d\phi_{j+1}/d\phi_j$ is sparse, for example when representing local interactions in a discretized PDE? The authors claim that their time integrators could work for "potentially very high-dimensional $z$" (line 173) but this doesn't seem to scale.

Checkpointing for gradient calculation should be mentioned as an alternative to the present approach for reducing memory costs, and the relative merits of each should be discussed.

For the proposed tasks and dataset sizes, many other approaches are possible. Examples include simulation-based inference (Cranmer 2020) that treats the forward process as a black box, or simply training an attention-based or recurrent network to impute missing in continuous time. These should at least be discussed, with references.

The notation in 2.1 is fairly confused. $\sigma$ is supposed to be an activation function but then turns out to have learnable parameters. The role of $h$ in eq. 3 isn't clearly stated. $F$ seems to be used once for $f$, and $T$ is used without definition. At the same time, there is probably too much detail here, as we don't really need an explanation of backprop.

The manuscript could benefit from some careful proofreading, e.g. for grammatical errors (line 369: "have a special structure").

References:

Chen RT, Rubanova Y, Bettencourt J, Duvenaud DK. Neural ordinary differential equations. Advances in neural information processing systems. 2018;31.

Rackauckas C, Ma Y, Martensen J, Warner C, Zubov K, Supekar R, Skinner D, Ramadhan A, Edelman A. Universal differential equations for scientific machine learning. arXiv preprint arXiv:2001.04385. 2020 Jan 13.

Ghadami A, Epureanu BI. Data-driven prediction in dynamical systems: recent developments. Philosophical Transactions of the Royal Society A. 2022 Aug 8;380(2229):20210213.

Greydanus S, Dzamba M, Yosinski J. Hamiltonian neural networks. Advances in neural information processing systems. 2019;32.

McGreivy N, Hakim A. Invariant preservation in machine learned PDE solvers via error correction. arXiv preprint arXiv:2303.16110. 2023 Mar 28.

Cranmer K, Brehmer J, Louppe G. The frontier of simulation-based inference. Proceedings of the National Academy of Sciences. 2020 Dec 1;117(48):30055-62.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
questions: See weakness.

Minor comments:
- References are cited without parentheses, which hinders readability. Proper citation using \citep tags with parentheses is recommended.
- The composition method could construct more complex architectures, and such an exploration could add value to the study.
soundness: 3
strengths: The focus of the study is reasonable, with a thorough survey of related works and sound theoretical foundations. If this paper were a contribution to applied mathematics, it would deserve high praise.
summary: "Improving neural network architectures inspired by ODEs" and "using neural networks to learn ODEs" have been investigated alternately. This paper proposes a novel architecture based on a composition method, which improves model design and achieves memory efficiency during training thanks to its reversibility.
weaknesses: However, the validation falls short. The experiments are limited to simple datasets, which are insufficient to demonstrate the general applicability of the proposed architecture. In what settings is this architecture genuinely effective? For instance, while reversible dynamical systems, like the Kepler problem used in the study, are a clear application, this alone is somewhat narrow, as one could achieve similar outcomes by applying reversible integrators within Neural ODE frameworks. On the other hand, could applying this architecture to general tasks, like image processing, compromise performance due to the reversibility constraint? The study would benefit from a clarified scope of applicability and a more comprehensive set of experiments.

While the theoretical memory efficiency is convincing, there is no experimental validation. Although ALF is presumably the baseline, only computation time is compared. An experimental comparison of memory consumption is also necessary, along with theoretical orders for computational and memory costs.

Considering ICLR's focus on machine learning, this paper falls short of expected standards due to the ambiguity of its contributions and insufficient empirical support.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: I don't quite understand the purpose of section 4.2. "LEARNING OF DYNAMICAL SYSTEMS PARAMETERIZED BY NEURAL NETWORK". Could the authors comment on the novel insights or advantages gained from this section compared to 4.1, or to clarify how this section extends or complements the results from 4.1?
soundness: 2
strengths: - I find the idea novel and original. The paper extends the capabilities of existing reversible integrators like ALF and reversible Heun, which were previously limited to low orders of accuracy. By demonstrating the feasibility of higher-order reversible integrators, the authors open new possibilities for improving the efficiency and accuracy of neural ODE-based deep learning.

- The paper is presented with exceptional clarity, guiding readers seamlessly through key concepts. It begins with a thorough setup of neural ODEs, followed by a well-balanced discussion of the advantages and limitations of two gradient computation methods, effectively motivating the notion of reversibility. The structured flow continues with an in-depth explanation of numerical methods, culminating in the core contribution: the construction of higher-order methods using ALF or Heun’s method. The integration of these methods within the neural ODE training framework (i.e. leveraging state-adjoint equations), is explained with precision, enhancing comprehension and practical relevance.

- The significance of this work is moderate. While it can be seen as a direct follow-up to the contributions of Zhuang et al. (2021) and Kidger et al. (2021), it still brings meaningful value, particularly in the context of neural network training for dynamical systems with multiple time scales.
summary: In summary, the paper presents a novel approach to construct higher-order reversible integrators for memory-efficient deep learning, with the main contribution being the development of higher-order reversible integrators based on the asynchronous leapfrog (ALF) method. The proposed methods demonstrate improvements in training speed and accuracy, particularly in learning dynamical systems. It opens up new avenues for training neural ODEs with restricted memory budgets.
weaknesses: - My primary concern lies with the soundness of the conclusion. Indeed, we see smaller errors and computation time of adaptive Y4 when comparing with adaptive ALF in both Kepler and nonlinear harmonic oscillator under some settings. It's not clear to me:
   -  Where the savings of the computation time come from? Is it because adaptive Y^4 in general uses larger step sizes? Is there a limit on $k$ for adaptive $Y^{2k}$ to outperform adaptive ALF on computation time? Why or why not there is a limit? I would suggest the authors to add more comparisons by varying parameter $k$ and report the results.
   -  Why the final loss of Y4 is lower? (When solving differential equations, it's easier to understand that higher order methods often outperform lower order methods w/ the same step size. However, when training neural networks, it's less obvious - the training is more or less "agnostic" to the computation schemes you are incorporating ... eg. why Runge Kutta is better than Euler?) What if we increase the number of epochs? Will it help adaptive ALF to achieve lower loss? Could the authors provide some justifications and experimental data to support them?

- The author may consider adding more contexts for the "adaptive stepping" presented in Hairer et al. (2006) in the Background section.

- The authors may consider adding a broader set of literatures for the readers. For examples,
  - known reversible methods other than asynchronous leapfrog and Heun:
     - Loup Verlet. Computer” experiments” on classical fluids. i. Thermodynamical properties of Lennard-Jones molecules. Physical review, 159(1):98, 1967.
  - for composition of mixed-order methods, although it is less developed than the theory for composing same-order integrators, there are some frameworks addressing this topic, especially in contexts like symplectic integrators and splitting methods
     - "Splitting methods for differential equations" by Sergio Blanes, Fernando Casas, Ander Murua
     - “Geometric Numerical Integration” by Hairer, Lubich, and Wanner
  - some other works on deep learning of dynamical systems
     - Rudy SH, Kutz JN, Brunton SL. 2019 Deep learning of dynamics and signal-noise decomposition with time-stepping constraints. J. Comput. Phys. 396, 483–506. (doi:10.1016/j.jcp.2019.06.056)
     - Y. Liu, J. N. Kutz, and S. L. Brunton, “Hierarchical deep learning of multiscale diﬀerential equation time-steppers,” arXiv:2008.09768
(2019)
     - Raissi M, Perdikaris P, Karniadakis GE. 2018 Multistep neural networks for data-driven discovery of nonlinear dynamical systems. Preprint. (https://arxiv.org/abs/1801.01236)

Meta-review: The paper introduces a higher order reversible time integrator for computing the gradients in the neural ODE setup to train neural networks.   They show results demonstrating the computational savings for the higher order integration with adaptive time stepping compared to ALF steps (prior work for lower order time integration). 
The main strengths of this paper are in the positive results for higher order integration that can have significant impact for complex dynamical systems and good explanation of methods that allows for easy understanding from the reader.
The main weakness of this paper is limited contributions - though very promising, reviewers agreed that experimental evidence falls short due to the focus on simplistic systems.

Decision: Reject