Parallelizing non-linear sequential models over the sequence length
Authors: Yi Heng Lim
Qi Zhu
Joshua Selfridge
Muhammad Firmansyah Kasim
Venue: ICLR 2024 poster (2024)
Keywords: parallel algorithm
recurrent neural networks
neural ordinary differential equations
sequential models
Primary area: general machine learning (i.e., none of the above)
TL;DR: We present a parallel algorithm to evaluate and training sequential models (e.g., RNN & NeuralODE) despite their inherent sequential nature

Abstract: Sequential models, such as Recurrent Neural Networks and Neural Ordinary Differential Equations, have long suffered from slow training due to their inherent sequential nature.
For many years this bottleneck has persisted, as many thought sequential models could not be parallelized.
We challenge this long-held belief with our parallel algorithm that accelerates GPU evaluation of sequential models by up to 3 orders of magnitude faster without compromising output accuracy.
The algorithm does not need any special structure in the sequential models' architecture, making it applicable to a wide range of architectures.
Using our method, training sequential models can be more than 10 times faster than the common sequential method without any meaningful difference in the training results.
Leveraging this accelerated training, we discovered the efficacy of the Gated Recurrent Unit in a long time series classification problem with 17k time samples.
By overcoming the training bottleneck, our work serves as the first step to unlock the potential of non-linear sequential models for long sequence problems.

Published as a conference paper at ICLR 2024
PARALLELIZING NON-LINEAR SEQUENTIAL MODELS OVER THE SEQUENCE LENGTH
Y. H. Lim', Q. Zhu J. Seltridge" M. F. Kasim Machine Discovery Ltd., UK, University of Oxford, UK (yi. heng, qi zzu, muhammadJ&macce meeiinediie dee-diisoovrry com joshua selfr
ABSTRACT
Sequential models, such Recurrent Neural Networks and Neural Ordinary Dif- ferential Equations, have long suffered from slow training due to their inherent sequential nature. For many years this bottleneck has persisted, many thought sequential models could not be parallelized. We challenge this long-held belief with our parallel algorithm that accelerates GPU evaluation of sequential models by up to 3 orders of magnitude faster without compromising output accuracy. The algorithm does not need any special structure in the sequential models' architec- ture, making it applicable to wide range of architectures. Using our method, training sequential models can be more than 10 times faster than the common se- quential method without any meaningful difference in the training results. Lever- aging this accelerated training. we discovered the efficacy of the Gated Recurrent Unit in long time series classification problem with 17k time samples. By over- coming the training bottleneck, our work serves as the first step to unlock the potential of non-linear sequential models for long sequence problems.
INTRODUCTION
Parallelization is arguably main workhorse in driv- ing the rapid progress in deep learning over the past decade. Through specialized hardware acceler- ators such as GPU and TPU, matrix multiplications which are prevalent in deep learning can be eval- uated swiftly, enabling rapid trial-and-erre in re- search. Despite the widespread use of parallelization in deep learning, sequential models such as Recur- rent Neural Networks (RNN) (Hochreiter & Schmid- huber 1997 Cho et al.) 2014) and Neural Ordinary Differential Equations (NeuralODE) (Chen et al. 2018 Kidger et al.) 2020) have not fully benefited from it due to their inherent need for serial evalua- tions over sequence lengths.
(0)
000
(9)
Zo
Zz
(3)
C30
63
(3)
Zo
Zz
Iterate until convergence
Z2
Zg
42
(a)
(b)
Figure 1: Evaluating sequential models us- ing (a) sequential method and (b) iterative method that is parallelizable
Serial evaluations have become the bottleneck in training sequential deep learning models. This bot- tleneck might have diverted research away from se- quential models. For example, attention mechanism (Bahdanau al. 2014) and transformers
Vaswani et al.) 2017) have dominated language modelling over RNN in recent years, partly due their ability to be trained in parallel (Hooker) 202T). Continuous Normalizing Flows (CNF) CChn et al. 2018 Grathwohl et al.) 2018), which used to utilize NeuralODE as their models, has moving towards the direction where the training does not involve simulating the ODE (Lipman
et al. 2022
Rozen et al. 202
Ben-Hamu et al. 2022). More recent works (Orvieto et al. 2023
WWokk done during internship at Machine Discovery 1 'Corresonning author

---

Published as a conference paper at ICLR 2024
Huang et al. 2022) have attempted to resuirect the sequential RNN, but they focus on linear recur-
rent layers that can be evaluated in parallel with prefix scan (Blelloch 1990; Martin & Cundy) 2018
Smith et al. 2022 leaving non- -inear recurrent layers unparalleizable over their sequence length.
In this paper, we present an algorithm that can parallelize the evaluation and training of non-linear sequential models like RNN and NeuralODE without changing the output of the models beyond rea- sonable numerical precision. We do this by introducing a general framework to solve non-linear dif- ferential equations by restating them as fixed-point iteration problems with quadratic convergence, equivalent to Newton's 's method for root finding. The fixed-point iteration involves parallelizable operations and an inverse linear operator that can be evaluated in parallel even for sequential models like RNN and ODE (see Figure []) As the convergence is quadratic, the number of fixed-point iter- ations can be quite small especially when the initial starting point is close to the converged solution. Most importantly, the proposed algorithm does not need special structure of sequential models, removing the need to change models' architecture to be able take advantage of the parallelization.
2 RELATED WORKS
There have been some attempts to parallelize the evaluation and training of sequential models, es- pecially for RNN. However, most (if not all) of them require special structures of recurrent layers. Lei et al. (2017) changed the matrix multiplication involving the states into element-wise multipli- cation to enable parallelization. Luo et al. (2020) segments the sequence into several groups to be evaluated by RNN in parallel and their groups interdependency are learned by a higher-level RNN. Huang et al. (022)) Orvieto et al. (2023). and Martin & Cundy (2018) use linear recurrent layer where it can be evaluated in parallel using prefix scan (Blelloch) 1990).
On the NeuralODE side, the parallelization effort mainly comes from the past works in parallelizing the ODE solver. One of the main idea the multiple shooting method (Kiehl) 1994 Gander & Vandewalle, 2007) Chartier& & Philippe [1993 [ellen & Zennaro 1989 Lions et al. |001 where the
time sequence is split into several segments then multiple ODE solvers are executed for each segment in parallel. The process is repeated iteratively until the solutions from all segments are matched. Although multiple ODE solvers can be executed in parallel, each ODE solver itself still requires sequential operation. The multiple shooting method has been adapted in training NeuralODE in Massaroli et al. (2021). The multi-grid idea is also applied to parallelize a special RNN unit and ResNet by converting them an ODE (Gunther et al.) 2020 Moon & Cyr 2022). Another work, neural rough ODE (Morrill et al.) 2021 has shown that it is possible to train NeuralODE for very long sequence by computing log-signature of the signals over large steps and enable the ODE solver to take big step. The calculation of the log-signature can be done in parallel. However, solving the ODE still requires sequential operation even though the time step is larger than the original.
Over the past few years, there have been increasing interests in using fixed points finders (e.g., root- finder) in deep learning. DEQ (Bai et al. 2019: 202T) employs root-finding in evaluating neural networks with infinite identical layers. The concept of infinite layers paired with root-finders has been applied to various problems and has yielded impressive results (Bai et al. 2020 [iu ct al. 2022 Huang et al. 2021). Nonetheless, these works do not talk about parallelizing sequential models.
Several recent works have also looked into parallelizing sequential models using fixed-point itera- tions or root finders. In the context of stochastic generative modelling, Shih al. (2023) divides the time span into several regions and then employs Picard iteration that are parallelizable in solving the ODE. Wang & Ragni (202T) recast RNN evaluation a fixed- fixed-point iteration and solve it by only do- ing a small number of iterations without checking the convergence. As a result, this approach might produce different results than evaluating RNN sequentially. [Song al. (20TT) evaluates feedforward computations by solving non- non-linear equations using Jacobian or Tauss-Seidee iterations. Notably, the works mentioned above only use zeroth order fixed-point iterations that converge slower than our method and might not be able to converge at all if the mapping is not contracting.
DEER FRAMEWORK
We will present the DEER framework: "non-linear Differential Equation as fixed point itERation" with quadratic convergence and show its relation to Newton's 's method. This framework can be ap-
2

---

Published as a conference paper at ICLR 2024
plied to ID differential equations, ie., ODEs, well differential equations in higher dimensions, ie. Partial Differential Equations (PDEs). The same framework can also be adopted to discrete difference equations to achieve the same convergence rate, which can be applied to RNN. With the framework, we can devise a parallel algorithm to evaluate RNN and ODE without significant changes to the results.
3.1 DEER FRAMEWORK
Consider output signal of interest (r) E Ih which consists of n signals d-dimensional space, where the coordinate is denoted as r TRd. The output signal, y(r), might depend on an input signal, x(r), via some non-linear delayed differential equation (DE),
L[y(r) = f (y - yi), $2), yr Sp), x(r),
(I)
where T[:]: Rd -> RR the linear operator of the DE, and f the non-linear function that depends on values of at P different locations, external inputs and parameters 0. This form is general enough to capture various continuous differential equations such as ODE (with L|-1 = d/dt and t), partial differential equations (PDEs), even discrete difference equations for RNN.
Now let's add (r)r)((r - sp) terms on the left and right hand side, where (r R"x" is an n-by-n matrix that depends the location r. The values of Gp will be determined later. Equation I now becomes
P
L(y(r)) CGp(r)y(r - sp) f (y(r x()),...)(()))))
Σ (r(rv((r- sp)
(2)
p=1
P
y(r) LG f - si)., +
Σ Gp(r)y(r sp) p=1
(3)
The left hand side of equation 2 is linear equation with respect to which can be solved more easily than solving the non-linear equation in most cases. equation we introduce the notation LG [.] as linear operator that solves the linear operator the left-hand side of equation 2 with some given boundary conditions.
Equation can be seen as a fixed-point iteration problem, i.e., given an initial guess (0) (r). we can
iteratively compute the right hand side of the equation until it converges. To analyze the convergence near the true solution, let's denote the value of y at i-th iteration as (+)(()) = yy(r) + 5y(() (r) with
y²(r) (r) as the true solution that satisfies equation 3 Putting (()) into equation 3 to get y (i+1), and performing Taylor expansion up to the first order, we obtain
P Σ
(++1)) (r) = La Eiar
(r 0
where 0p is the Jacobian matrix of f with respect to its p-th parameters. From the equation above, the first order term of (b(++1) can be made 0 by choosing
(r) = -pf (y(r - $1), .... y(r - sp), x(r),
(5)
It shows that the fastest convergence around the solution can be achieved by choosing the matrix G, according to the equation above. It can also be shown that the iteration in equation 3 and equation equivalent to the realization of Newton's method in Banach space, therefore offering a quadratic convergence. The details of the relationship can be found in Appendix A. and the proof of quadratic convergence can be seen in Appendix A..
Iterative process in equation 3 involves evaluating the function f. its Jacobian, and matrix multipli- cations that can be parallelized in modern accelerators (such as GPU and TPU). If solving the linear equation can be done in parallel, then the whole iteration process can take advantage of parallel computing. Another advantage of solving non- non-linear differential equations as fixed point iteration problem in the deep learning context is that the solution from the previous training step be used as the initial guess for the next training step if fits the memory. Better initial guess can lower the number of iterations required to find the solution of the non-linear DE.
3

---

Published as a conference paper at ICLR 2024
3.1.1 DERIVATIVES
To utilize the framework above in deep learning context, we need to know how calculate the forward and backward derivatives. For the forward derivative, we would like to know how much y perturbed (denoted as ōy) if the parameter is slightly perturbed by 66. By applying the Taylor series expansion to the first order to equation | and using Gr(r) as in equation 5 we obtain
Of
Py =
(y(r si,, yy. y(r sp). x()), 8)68
00
The equation above means that to compute dy, one can compute the forward derivative of the outputs of from d0 (denoted as of), then execute the inverse linear operator LG on öf.
For the backward derivative, the objective is to obtain the gradient of a loss function L with respect to the parameter, 0L/00, given the gradient of the loss function to the output signal, aC/Oy. To obtain the backward derivative, write
ac
OL ōy By 30
((
0C Oy
Of (y(r (i)r-s)), -i),..,, (S)-ss)))x))))))) -(.))))((())))) (p))) .)
LG
(7)
00
30
Note that the order of evaluation in computing the backward gradient should follow the brackets in the equation above. In contrast to equations and 6 the linear operator L is operated to the left in the innermost bracket in equation 7 This is known the dual operator of L and in practice can be evaluated by applying the vector-Jccoban product of the linear operator [-] The results of the dual operator in the innermost bracket are then the followed by vector-Jacobian product of the function f. The forward and backward derivatives for x are similar to the expressions for just substitute the differential with respect to 0 into the differential with respect to x.
The forward and backward gradient computations involve only one operation of T in contrast to the forward evaluation that requires multiple iterations of LG evaluations. This allows the gradient computations to be evaluated more quickly than the forward evaluation. Moreover, the trade-off between memory and speed can be made in the gradient computations. If one wants to gain the speed, the matrix G from the forward evaluation can be saved be used in the gradient computation. Otherwise, the matrix G can be recomputed in the gradient computation to save memory.
The gradient equations above apply even if the forward evaluation does not follow the algorithm in the previous subsection. For example, if equation 3 does not converge and forward evaluation is done differently (e.g., sequentially for RNN), the backward gradient can still be computed in parallel This might still provide some acceleration during the training.
according to equation 7
3.2 PRACTICAL IMPLEMENTATION
Equation I has a very general form that can capture ordinary differential equations (ODEs), most partial differential equations (PDEs), and even discrete difference equations. To apply DEER frame- work in equation 3 to a problem, there are several steps that need to be followed. The first step is recast the problem into equation| [ to define the variable y, the linear operator L[-], and the non- linear function f(.). The second step is to implement what we call the shifter function. The shifter function takes the whole discretized values of (r) and returns list of the values of y at the shifted positions, i.c. y(r sp) for p {l....,P) P). The shifter function might need some additional infor- mation such as the initial boundary conditions. The output of the shifter function will be the input to the non-linear function. The next step, and usually the hardest step, is to implement the inverse operator I [h] given the list of matrices G, and the vector values discretized at some points. The inverse operator LG [h] might also need the information on the boundary conditions.
Once everything is detined, iterating equation 3an be implemented following the code in appendix B.I The Jacobian matrix in equation 5 can be calculated using automatic differentiation packages Paszke et al. 2017 Frostig et al. 2018).
DEER framework can be applied to any differential or difference equations as long as the require- ments above can be provided. This includes ordinary differential equations, discrete sequential mod- els, and even partial differential equations (see Appendix (A44). To keep focused, we only present the application of DEER in parallelizing ODE and discrete sequential models.
4

---

Published as a conference paper at ICLR 2024
3.3 PARALLELIZING ORDINARY DIFFERENTIAL EQUATIONS (ODE)
An ODE typically takes the form of dy /dl - ((y x(1, 0) where the initial condition y(O) is given. The ODE form above can be represented by equation | with r = 1, L = d/dl, P = 1, and
$1 O. This means that the operator T.g in ODE is equivalent solving the linear equation below
given the initial condition yO),
dy
dt (1) + G(UY()) = = y(1)= (((()
(8)
Assuming that Gtt) and ztt) are constants between t = to and = 4441 as G, and z, respectively, we can write the relations between yi+1 =y(ti-1) and yi ytt)) as
yi+1 G.y; +2
A, 1++1 +,. I the identity matrix, and exp(.) the matrix exponential. Equation can be evaluated using the parallel prefix scan algorithm as described in Blelloch (190)) and Smith al. (2022). Specifically, we define pair of variables c++1 (G.(2,) for every discrete time point 44 the initial values co- = =[(y)), and an associative operator,
with G; = exp
C++1 CJ+1 G,G, 22))
(10)
Given the initial value of co and the associative operator above, we can run the associative scan in parallel to get the cumulative value of the operator above. The solution y; can be taken from the second element of the results of the parallel scan operator.
On the implementation note, we can reduce the computational error by taking the G; and 2, as the mid-point value, i.e. G; [G((t) G(+++1)] and z; 1 [z((t) + z(t++1)] By taking the mid-point value, we can obtain the third order local truncation error ((A,) instead of second order
error using either the left or the right value, with only a small amount of additional computational
expenses. The details can be seen Appendix A.5
The materialization of DEER framework on ODE can be seen as direct multiple shooting method (Chartier & Philippe 1993 Massaroli et al., 202T) that splits the time horizon as multiple regions. However, in our case each region is infinitesimally small, making the Newton step follows equation and parallelizable. The details of the relationship between our method with direct multiple shooting method can be seen in Appendix
3.4 PARALLELIZING RNN
Recurrent Neural Network (RNN) can be seen as a discretization of ODE. Having the input signal index i x; and the previous states y:-1, the current states can be written as y, f(y;-1, 1, x,, 0) This form can capture the common RNN units, such as LSTM (Hochreiter & Schmidhuber 1997) and GRU (Cho al., 2014). Also, the form can be written as equation [ with = i, L y. P 1 and 8 = 1. This means that the inverse linear operator can be calculated by solving the equation below, given the initial states yo.
y: G.y =
Y1..T LG [z..rl
(11)
Solving the equation above is equivalent to solving equation 9 from the previous subsection. It means that it can be parallelized by using the parallel prefix sean with the defined associative oper- ator in equation 10
3.5 COMPLEXITY AND LIMITATIONS
The algorithms for solving non-linear ODE and RNN in parallel are simply repeatedly evaluating equation from an initial guess. However, the equation requires the explicit Jacobian matrix for each sequence element for each shift. If y has n elements, L sampling points, and P shifted arguments, storing the Jacobian matrices requires LPn)LP) memory. Also, the prefix sean for solving the linear differential equation requires matrix multiplication of matrices G, at different sampling points. This would introduce LLP) time complexity. As the algorithm has (nn²) memory complexity and (n²)) time complexity, this algorithm would offer significant acceleration for small number of n.
5

---

Published as a conference paper at ICLR 2024
Speed up DEER GRU sequential GRU (forward) with batch
6m
10a
dm
Speed
0Rh 16
dem--32
10
61
1K
3k
10k
JUK
100K
300k
1M
Sequence length
Speed of DEER GRU over sequential GRU (forward grad) with batch size 16
Speed up
10
30k
Sequence length
Figure 2: The speed of GRU calculated using DEER method (this paper) vs commonly-used sequential method on V100 GPU for (top) forward and (bottom) forward + gradient calculations. The missing data for large number of dimensions and sequence lengths is due to insufficient memory in the DEER method. The bar height represents the mean speed up over different random seeds.
As the proposed method is realization of Newton's method in Banach space, it has the same limitations as Newton's 's method. the starting point is sufficiently far from the solution, the iteration might not converge. This problem might be addressed by using modified Newton's method with convergence guarantee such as Nesterov & Polyak (206)) and Doikov & Nesterov (2023). However, leave the use of globally-convergel Newton's s method as the future work.
There is only one additional hyperparameter in our method: the tolerance for convergence. As our method has quadratic convergence, the tolerance value does not have big effect on the number of iterations required for convergence, as long as it is not too close to the numerical precision value and not too big. In fact, using respectively 10-4 and 10-7 for single- and double-precision floating point precisions are enough for our experiments.
Parallelizing ODE and RNN both require parallelizing the evaluation of equations and [1 Those equations can be parallelized using parallel prefix scan of a custom associative operation in equation 10 This is relatively straightforward using JAX's jax lax associate iettiivescaal (Frostig al.) 2018) or TensorFlow's Ep. math. scan ivve (Abadi) 2016). However, of the time of writing this paper, this operation cannot be implemented casily using Py PyTorch (Paszke et al.) 2017). For this reason, used JAX for experiments in this paper well flax and equinox (Kidger Garcia 2021 for the neural networks.
EXPERIMENTS
4.1 PERFORMANCE BENCHMARKING
The first test to compare the speed on evaluating an RNN using the presented DEER method against the common sequential method. Specifically, we're using untrained Gated Recurrent Unit (GRU) cell from flax. linen (a JAX neural network framework) using 32-bits floating points random inputs with 16 batch size, various number of dimensions (i.e., n)), and various number of sequence lengths. The initial guess for DEER is all zeros for all the benchmark runs. The speed up on a V100 GPU obtained by the method presented in this paper is shown in figure 2

---

Published as a conference paper at ICLR 2024
Differecce between sequential and DEER outputs
GRU ousputs the last 200 Indices
MMM | MA
0.2
0.0
-0.5
Sequential
CEER
-1.0
9500
9825
9050
9675
3925 Sequence index
9900
9950
3975 10000
2000
4000
6000
8000
10000
Sequence index (b)
(b)
Figure 3: (a) The comparison between the outputs of GRU evaluated with sequential method vs DEER method. The line for sequential method output is almost not visible because overlaid by the output of DEER method. Only the last 200 indices are shown for clarity. (b) The difference between the outputs of sequential and DEER method for the whole 10k sample length.
The figure shows that the largest speed up is achieved on long sequence lengths with small number of dimensions. With IM sequence length. 16 batch size, and 1, the sequential evaluation required 8.7 s while DEER method only took 15 ms, which translates to a speed up of over 500. However, the speed up decreases as the number of dimensions increases, where the speed up is only about 25% with 64 dimensions (64 hidden elements in GRU). This is due to the explicit computation of Jacobian matrix and the matrix multiplication that scales to O(n)).
The speed up for forward + gradient calculations is even greater than the speed up for forward evaluations only. With the same set up, the speed up for 1M sequence length with 1 dimension could be more than 1000 times faster. This is because the backward gradient calculations require only one evaluation of LG in equation as opposed to multiple evaluations in forward calculations.
Table in Appendix C.2 presents the tables of speed up for Figure aas well speed up when using smaller batch sizes. Generally, the speed up increases with smaller batch size, where speed up of above 2600 can be achieved with batch size 2. This means that with more devices, greater speed up can be achieved by having a distributed data parallel, effectively reducing the batch size per device.
Figure 3 shows the comparison between the output of an untrained GRU of 32 hidden elements evaluated using sequential vs DEER method. The input to the GRU is a Gaussian-randon tensor with 10k sequence length and 32 dimensions. From figure we see that the output of GRU evaluated using DEER method is almost identical to the output obtained using the sequential method. The small error in figure b)) is due numerical precision limits of the single precision floating point.
4.2 LEARNING PHYSICAL SYSTEMS WITH NEURALODE
We test the capability of DEER method in training a deep neural network model using a simple case from physics. Given the positions and velocities as a function of time of two-body system interacting with gravitational force, we trained Hamiltonian Neural Networks (HNN) (Crreydanus) jet al. 2019) to match the data. There are n = states in this case. Note that this set up is different from the original HNN paper where they train the network to match the velocity + acceleration, given position + velocity at every point in time. In this case, only positions and velocities as function of time are given and the network needs to solve the ODE to make a prediction. We use 10,000 time points sampled uniformly, whereas other works in this area typically only use less than 1,000 sampling points for the training (Matsubara & Yaguchi 2022 Chen et 2019). The details of the setup can be seen in Appendix B.2
The losses during the training using DEER method vs using RK45 (Atkinson) 1991) are shown in figure/aa, b). We use the RK45 algorithm from JAX's experimental feature. From the figure, it can be seen that the training can be 1l times faster when using DEER method presented in this paper than using the ordinary ODE solver without significant difference in the validation losses between the two methods. To achieve 55k training steps, DEER method only spent day + hours, while training with RK45 required about 2 weeks of training. The small difference of the validation losses between the two methods potentially comes from different methods in solving the ODE as well as from the numerical precision issue as shown in Figure 3).
7

---

Published as a conference paper at ICLR 2024
RX455 method
Valisattio
20
10:1
faster
10
10
200
150
200
250
300
0
23
30
40
50
Hours
Training steps (x1u))) (b)
Mr momm hwwnm//ww//// M Vn
Mm wmw R
Vallcitt accuracy
0.9
M
0.8
0.8
07-
0.-
DEER methad Sequertil
0.5
05
Hours
Training suepp (x10P)
(a)
Figure 4: (Top) The validation losses of HNN with NeuralODE training using DEER method (shown in blue) RK45 method (in orange) function of (a) training hours and (b) training steps. (Bottom) The validation accuracy of RNN training using DEER method (blue) vs the sequential method (orange) as function of (c) training hours and (d) training steps.
4.3 TIME-SERIES CLASSIFICATION WITH RECURRENT NEURAL NETWORK (RNN)
Faster training enabled by DEER
method allows us use classical RNN architectures, such Gated Recurrent Unit (GRU, Cho et al.] (2014)), for problems with long time series. In this subsection, we trained a neural network that consists of GRUs to classify Figen- Worms dataset (Brown et al.) 2013) from UEA (Bagnall al. 2018) with 17,984 time samples for each entry. The details of the training and architecture can be found in Appendix B.3
Model
Accuracy (%)
ODE-RNN (folded), step: 128
47.9 = 5.3
NCDE, step:
66.7 11.8
NRDE (depth 3), step: 32
75.2 = 3.0
NRDE (depth 2), step: 128
76. + 5.9
NRDE (depth 2), step: 4
83.8 + 3.0
UnICORNN (2 layers)
90.3 - 3.0
LEM
92.3
GRU (from this paper) LEM (our reproducibility attempt)
88.0 05.0444 4.4
92.3 + 2.1
Figure 4c)-(d) shows the compari- son of validation accuracy during the training of the GRU network using DEER method vs the common sequen- tial method. From the figure, we can see that the validation accuracy plot when using DEER is similar the valida- tion plot obtained with the sequential method. The difference the validation accuracy might be due to numerical pre- cision issue as shown in Figure that accumulated at long sequence. However, the training using DEER is up 22 times faster than the training using the sequential method. What would have taken more than day to train using the common sequential method, only takes hour using DEER.
Table 1:
The classification accuracy of ForenWorms dataset for various methods, including folded ODE-RNN (Rubanova et al.) 2019). Neural CDE (Kidger al. 2020), Neural RDE (Morrill et al.) 2021). UnICORNN (Rusch & Mishra 2021). LEM (Rusch al. 2021), and GRU. The mean and standard deviations of the accu- racy were obtained from times repetition with different seeds. The numbers of non-GRU methods were obtained from Morrill al. (2021) and Rusch et al. (202T).
The classification results the test dataset from Eigen Worms for various methods are shown in Table As we from the table, the network with GRUs can be competitive more
8

---

Published as a conference paper at ICLR 2024
modern methods (Morrill al., 2021: Rusch & Mishra 2021 Rusch al. 2021 for this long time series dataset. The long training time of GRUs using the sequential method could be the main factor hindering the trial-and-error process of exploring GRU architectures for this dataset with long sequences. However, our DEER method enables much faster training of GRUs, facilitating iterative experimentation identify optimal GRU architectures and hyperparameters.
To demonstrate the applicability of DEER to other architecture, we performed an experiment to reproduce the results from LEM (Rusch et al. 2021). The reimplementation of LEM from PyTorch to JAX makes the reproducibility attempt not straightforward. Moreover, as the Eigen EigenWorms Worms dataset very small, slight changes could change the results quite significantly. With DEER, we can run more experiments than using sequential methods, allowing us to tune hyperparameters faster. Our attempt has produced similar results to the original attempt at [Rusch al. (202T). With DEER, one epoch can be done in 18 seconds while the sequential method could take about 116 seconds.
4.4 SEQUENCE CLASSTFICATION WITH MULTIPLE HEADS RNN
One way to avoid unfavorable scaling with re- spect to n in DEER is to use multi-head recur- rent unit. For example, instead of having unit with n 256 channels, we split the unit into 32 heads with each head has n 8 channels. Each head can also have different strides with expo- nentially increasing values, for example, the first head has stride 29, second head has stride 2 and so on. The use of multiple heads is inline with the idea from state-spaces (Gu et al.) 2021bb a 2022 Smith et al. 2022) where states are grouped and each group has different length scale.
Model
Accuracy
State- Mate-space or linear recurrent
LSSL (Gu et al. (0221b)
84.65%
S4 (Gu el al. 2001aa)
91.80%
Liquid-S4 (Hasani et al. 2022
92.02%
LRU (Orvicto et al. 2023)
89.0%
Convolution
TrellisNet (Bai et al. 2018)
73.42%
FlexConv (Romero et al. MultiresNet (Shi et al. 2023)
2021
80.82%
93.15%
Non-linear recurrent r-LSTM (Trinh et al. UR-GRU (Gu et al. 2020)
We performed an experiment using the multi- head GRU on a standard sequential CIFAR-I0 dataset. The details of the experimental setup can be found in Appendix B.4 Running an experi- ment with DEER using multi-head GRU is about times faster than using sequential method. The speed up is not as big in the previous sub- sections because of shorter effective sequence lengths due to larger strides. However, the speed up still allows us to run multiple experiments to find a good configuration.
2018
72.2%
74.4%
Multi-head GRU (ours)
90.25%
The classification test dataset accu- racy of sequential CIFAR- 10 dataset for various methods, grouped according to their architecture classes. Underlined scores are the best score for each architecture class.
Table 2:
Table 2 shows that multi-head GRU can achieve 89.35% test accuracy. Although the result is not the state-of-the-art, this is the best result on sequential CIFAR-10 using non-linear recurrent units. We show that simple modification of GRU and the capability to run more experiments enabled by DEER has uncover the potential of classical recurrent units such as GRU in competing with more modern architectures. Please note that our aim for this paper is not to focus on new architectures, but focus on a new computational method that enables faster experimentation on non-linear sequential models. Finding new non-linear sequential models enabled by DEER is for the future work.
5 CONCLUSION
We introduced a method to parallelize the evaluation and training of inherently sequential models, such as ODE and RNN. Evaluations using our approach can up to orders magnitude faster than traditional sequential methods when dealing with long sequences. When training sequential models with reasonable settings, our method can achieve over a 10-fold speed increase without significantly altering the results. By having technique accelerate the training and evaluation of sequential models, we anticipate a hastened pace of research in this domain, potentially catalyzing the emergence of novel and interesting non-linear sequential models in the future.

---

Published as a conference paper at ICLR 2024
REPRODUCIBILITY STATEMENT
The code required for reproducing the algorithm and results in this paper can be found in http://sstabbboommmmmccciiiiiler
ACKNOWLEDGEMENT
We acknowledge the support from Google Cloud start up program for their cloud computing credit that being used in large part of this research.
REFERENCES
Martín Abadi. Tensorflow: learning functions at scale. In Proceedings of the 21st ACM SIGPLAN international conference on functional programming, pp. 1-1, 2016.
Kendall Atkinson. An introduction numerical analysis. John wiley & sons, 1991.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXfv:1607.06450, 2016.
Anthony Bagnall, Hoang Anh Dau, Jason Lines, Michael Flynn, James Large, Aaron Bostrom, Paul Southam, and Eamonn Keogh. The uea multivariate time series classification archive, 2018. arXiv preprint 1rXiv: 0e0/////1.0005. 2018.
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint arXiv:1409.0473. 2014.
Shaojie Bai, Zico Kolter, and Vladlen Koltun. Trellis networks for sequence modeling. arXiv preprint arXiv:1810.06663 2018
Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. Advances in Neural Information Processing Systems, 32, 2019.
Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Multiscale deep equilibrium models. Advances in Neural Information Processing Systems, 33:5238-5250, 2020.
Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Neural deep equilibrium solvers. In International Conference on Learning Representations, 2021.
Alfredo Bellen and Marino Zennaro. Parallel algorithms for initial-value problems for difference and differential equations. Journal of Computational and applied mathematics, 25(3):341-350, 1989.
Heli Han-Hamu,, Samuel Cohen, Joey Bose, Brandon Amos, Maximillian Nickel, Aditya Grover, Ricky TQ Chen, and Yaron Lipman. Matching normalizing flows and probability paths on mani- folds. In International Conference on Machine Learning, pp. 1749-1763. PMLR, 2022.
Guy E Blelloch. Prefix sums and their applications. 1990.
André EX Brown, Eviatar Yemini, Laura Grundy, Tadas Jucikas, and William Schafer. dic- tionary of behavioral motifs reveals clusters of genes affecting caenorhabditis elegans locomotion. Proceedings of the National Academy of Sciences, 110())791-796, 2013.
P Chartier and B Philippe. A parallel shooting technique for solving dissipative ode's. Computing, 51(3-4):209-233 1993.
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David K Duvenaud. Neural ordinary differential equations. Advances in neural information processing systems, 31, 2018.
Zhengdao Chen, Jianyu Zhang, Martin Arjovsky, and Léon Bottou. Symplectic recurrent neural networks. In International Conference Learning Representations, 2019.
10

---

Published as a conference paper at ICLR 2024
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint arXiv: 1et..14999.999 2014.
Yann N Dauphin, Angela Fan, Michael Auli, and David Grangier. Language modeling with gated convolutional networks. In International conference on machine learning, pp. 933-941. PMLR, 2017.
Nikita Doikov and Yurii Nesterov. Gradient regularization of newton method with bregman dis- tances. Mathematical Programming, pp. 1-25, 2023.
Roy Frostig, Matthew James Johnson, and Chris Leary. Compiling machine learning programs via high-level tracing. Systems for Machine Learning. 4(9), 2018.
Martin J Gander and Stefan Vandewalle. Analysis of the parareal time-pane one-parallll time-integn lime-integration method. SIAM Journal on Scientific Computing, 29(2):55--578, 2007.
Will Grathwohl, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. Ffjord: Free-form continuous dynamics for scalable reversible generative models. In International Con- ference on Learning Representations, 2018.
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advances in neural information processing systems, 32, 2019.
Albert Gu, Caglar Gulcehre, Thomas Paine, Matt Hoffman, and Razvan Pascanu. Improving the gat- ing mechanism recurrent neural networks, In International Conference on Machine Learning, pp. 1800-3809 PMLR, 2020.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.00300 2021a.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Com- bining recurrent, convolutional, and continuous-time models with linear state space layers. Ad- vances neural information processing systems, 34:572-585, 2021b.
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971- 35983, 2022.
Stefanie Gunther, Lars Ruthotto, Jacob B Schroder, Eric C Cyr, and Nicolas R Gauger. Layer- parallel training of deep residual neural networks. SIAM Journal on Mathematics of Data Science, 2(1):1-23, 2020.
Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state- space models. arXiv preprint 2tt.:22009.1955//. 2022.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()):
1735-1780, 1997.
Sara Hooker. The hardware lottery. Communications of the ACM, 64(12):58-66 65,, 2021.
Feiqing Huang, Kexin Lu, CAI Yuxi, Zhen Qin, Yanwen Fang, Guangjian Tian, and Guodong Li. Encoding recurrence into transformers. In The Eleventh International Conference Learning Representations, 2022.
Zhichun Huang, Shaojie Bai, and J Zico Kolter. Implicit": Implicit layers for implicit representa- tions. Advances Neural Information Processing Systems, 34:9639-9999 9650, 2021.
Carl T Kelley. Iterative methods for linear and nonlinear equations. SIAM. 1995.
Patrick Kidger and Cristian Garcia. Equinox: neural networks in JAX via callable PyTrees and filtered transformations. Differentiable Programming workshop at Neural Information Processing
Systems 2,21, 2021.

---

Published as a conference paper at ICLR 2024
Patrick Kidger, James Morrill, James Foster, and Terry Lyons. Neural controlled differential equa- tions for irregular time series. Advances in Neural Information Processing Systems, 33:6696- 6707, 2020.
Martin Kiehl. Parallel multiple shooting for the solution of initial value problems. Parallel comput- ing, 20(3):275-295, 1994.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 1412. 2014.
Tao Lei, Yu Zhang, Sida Wang, Hui Dai, and Yoav Artzi, Simple recurrent units for highly paral- lelizable recurrence. arXiv preprint arXiv:1709.0275.. 2017.
Jacques-Louis Lions, Yvon Maday, and Gabriel Turinici. Résolution d' edp par un schéma temps "pararéel". Comptes Rendus de I' d'Académie des Sciences-Series I-Mathematics, 332(():661166)) 2001.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow match- ing for generative modeling. In The Eleventh International Conference on Learning Representa- tions, 2022.
Juncheng Liu, Bryan Hooi, Kenji Kawaguchi, and Xiaokui Xiao. Mgnni: Multiscale graph neural networks with implicit layers. Advances in Neural Information Processing Systems, 35:21358- 21370, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv: 2017.
Yi Luo, Zhuo Chen, and Takuya Yoshioka. Dual-path mn: efficient long sequence modeling for time-domain single-channel speech separation. In ICASSP 2020-2020 IEEE International Con- ference on Acoustics, Speech and Signal Processing (ICASSP), pp. 46-50. IEEE, 2020.
Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations, 2018.
Stefano Massaroli, Michael Poli, Sho Sonoda, Taiji Suzuki, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Differentiable multiple shooting layers. Advances in Neural Information Pro- cessing Systems, 34:16532-16544 2021.
Takashi Matsubara and Takaharu Yaguchi. Finde: Neural differential equations for finding and preserving invariant quantities. In The Eleventh International Conference on Learning Represen- tations, 2022.
Gordon Euhyun Moon and Eric C Cyr. Parallel training of gru networks with a multi-griin solver for long sequences. arXiv preprint 22X.422003.04738 2022.
James Morrill, Cristopher Salvi, Patrick Kidger, and James Foster. Neural rough differential equa- tions for long time series. In International Conference on Machine Learning, pp. 7829-7838.
PMLR, 2021.
Yurii Nesterov and Boris T Polyak. Cubie regularization of newton method and its global perfor- mance. Mathematical Programming, 108(1):177-205, 2006.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. arXiv preprint arXiv:2303.063443 2023.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang. Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
David W Romero, Robert-Jan Bruintjes, Jakub M Tomezak, Erik J Bekkers, Mark Hoogendoom, and Jan C van Gemert. Flexconv: Continuous kernel convolutions with differentiable kernel sizes. arXiv preprint /0.0050, 2021.
12

---

Published as a conference paper at ICLR 2024
Noam Rozen, Adiyya Grover, Maximilian Nickel, and Yaron Lipman. Moser flow: Divergence- based generative modeling on manifolds. Advances in Neural Information Processing Systems, 14.17669-17680. 2021.
Yulia Rubanova, Ricky TQ Chen, and David K Duvenaud. Latent ordinary differential equations for irreuulally-smppled time series. Advances in neural information processing systems, 32, 2019.
T Konstantin Rusch and Siddhartha Mishra. Unicornn: A recurrent model for learning very long time dependencies. International Conference on Machine Learning, pp. 2168-9176 PMLR, 2021.
expressive memory for sequence modeling. arXiv preprint vrxiv2211000744: 2021. T Konstantin Rusch, Siddhartha Mishra, N Benjamin Erichson, and Michael W Mahoney. Long
Jiaxin Shi, Ke Alexander Wang, and Emily Fox. Sequence modeling with multiresolution convo- lutional memory. International Conference on Machine Learning, pp. 31312-31327. PMLR, 2023.
Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffusion models. arXiv preprint arXiv:2305. 16317, 2023.
Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.04933, 2022.
Yang Song, Chenlin Meng, Renjie Liao, and Stefano Ermon. Accelerating feedforward computation via parallel nonlinear equation solving. In International Conference on Machine Learning, pp. 9791-9800. PMLR, 2021.
Trieu Trinh, Andrew Dai, Thang Luong, and Quoc Le. Learning longer-term dependencies in mns with auxiliary losses. In International Conference Machine Learning, pp. 4965-4974, PMLRR 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Zhengxiong Wang and Anton Ragni. Approximate fixed-points in recurrent neural networks. arXiv preprint arXiv:2 106:22006.000 024/7, 2021.
13

---

Published as a conference paper at ICLR 2024
A THEORETICAL RESULTS
A.1 RELATION TO NEWTON'S 'S METHOD
The main method presented in this paper is basically realization of Newton's method in Banach space. Finding the root of functional q(y) = 0 can be done by performing the Newton's method iteration (f++1) y(() [Dq(y())] - where Dq is the Fréchet derivative of q. In equation the function q(y) is given by (()(r)) [[y(r) f(y(r- 81),.., x(r,, 0). Therefore, the
((y(())
Fréchet derivative of functional q operated on a function h(r) is given by
(a(r) [.(((r)]
0, - sp)-
(12)
p=1
Therefore, the Newton's 's method iteration to find the root of functional q(y) 0 is
yl++1) y - (Dq(x)) -1 [
f(y sı), x(r) (r).
(13)
The inverse Fréchet derivative terms can be computed by solving the equation below for y y(t+1,
Da(y) (x() - -((tt) - f(y
(14)
Expanding Dq(y) from equation 12 on (+), we obtain
P Σ 0,ty (r Sp)
Da(y)ly (1+1)) = f(y(0 (r (1-s)))-.)))))))))))(((()))) s.)...,... x(r),
(15)
p=1
Solving for (+++)) the equation above will produce the equation
3 showing that equation is
just a realization of Newton's 's method in solving equation I
A.2 RELATION TO DIRECT MULTIPLE SHOOTING
One popular way to parallelize an ODE solver is to perform direct multiple shooting (MS) method
Chartier & Philippe 1993 Massaroli et al. 2021). ). The following parts will follow Massaroli et al. 2021) very closely. Consider initial value problem with dy /dl f(y,t) for 1 E [0, 7] with y(0)= = Yo. With MS, we first split the time horizon into N multiple regions with 0 to < t1 < to ... < tN = T, have a guess on initial values each region, (t,) b,, then apply the ODE solver for each region. After applying the ODE, the guessed values b, can then be updated to match the results of ODE from the previous region. Denote a function Φ(b, 4, f441) as the value of y at 4 after solving the initial value problem with initial value y(t,) b; Evaluating this function & requires solving an ODE from 6; to 4.41. With the function above, we can write the constraints to be solved, b (bb,, 4, 441) with by Yo.
The works by Chartier & Philippe (193) and Massaroli et al. (202T) update by by applying Newton method. Specifically, at k-th iteration with values of b (R) are available for all i-s, the values at the next iteration can be determined by (Massaroli et al.) 2021)
dΦ (8) (K) les (1)4+++++++ bb (b. l,, 411 (b((+1)) b (1) by (k+1) = yo- (16)
b = <bb!
The term b (k+1) makes the equation above needs to be evaluated sequentially, in addition to com- puting the function & that requires solving an ODE.
If we assume there are large number of regions, N method becomes equivalent to the DEER realization on solving the ODE In this limit, we can write
and A4;
4++1 l - O, then the MS
Φ(b;, 4, 4-1) b, + f(b,, 4) and OΦ
Of (b,, 4),
(17)
ab
with At, = - - to and I an identity matrix. Substituting equation 17 to equation 16
we obtain
b (l+1)
b+1 I+ af (b. (k) A))(,) b (A+11) F(b
(R) ..)
of (b. (8) 1.)b
(8)
At,
(18)
ab
14

---

Published as a conference paper at ICLR 2024
By re-arranging the equation above, get
bi+1*
(k+1)
b (k+1)
ab (b. (8) bbbD (k+1) = =fbb,", .4)
of
ab (b, (k) ,b)bb
Of
(k)
(19)
At
With the limit A4;
+ O and N + oc, the term b, can be written as (44) and (b (b (b+1) -
equals to dy (y(k)1) /dt. By rewriting the equation above in y, obtain
dy(++1) dt ay (y (t)y(k+1) f(y(")
öf
of
(y (k) t)y
(20)
With a little bit of algebra, one can show that the equation above is equivalent to equations 3 and By having large number of regions, we can approximate Φ as in equation 17 that can be evaluated very quickly.
A.3 PROOF OF QUADRATIC CONVERGENCE OF DEER ITERATION
This proof is inspired by standard results on the order of convergence of Newton's method and follows the argument of Kelley (1995). $$4 and 5.1. Throughout, we will make repeated use of the Sauchy-Schwarz (CS) and triangle (T) inequalities.
We employ the standard Euclidean nomn on R".
y|| vy y V of + va y -yV1....,, Bn)² ER"
(21)
and the induced operator norm on x"
A|l == sup Ay ||>|| 1, E ER y YER")
(22)
Note that the Fréchet derivative Dq(y) E R R"X
We denote the value of y at the ith iteration as ( = y (8) equation 3 Following from A.I we assume that Dq(y) is Lipschitz continuous with Lipschitz constant y, iee.:
(0 (r), where y exactly satisfies
Dq(y) - Dq(z))l S 7 z||
Vy,zER".
(23)
Consider the operators A Dq(y) and B = Dq(y"). We assume that both invertible and that -I-B I-BB-Al| <:
A B = -I- [I (1-B-'A)
(24)
(--B-A)
(Neumann series for matrix inverse).
m=0
Post-multiplying by B-1 and taking norms:
|--1||
Σ
I B
B-)
B-1
h=0
(CS)
Σ [ B B-'A)"
B-1
IA3
Σ
[I - B ---A--)-- -1)1)) B-
m=0
(25)
(CS)
Σ
| - B B AA""
|B-1|
B |B-1|
B- 3-1A A
||[Dq(y") 1
[[[Qq()]]
-)
15

---

Published as a conference paper at ICLR 2024
We denote the ball of radius as B(c) {y | E}, where 8y y y"
I [Dq(y")] Dg(y)
= [[ (Dq(y') Da(y)) ||
(CS)
(/Da(y")1=1 Da(y") - Dq(y)
(26)
23
"1 -yll
|| [|[Pa(y")]
Let 0 << (27 [Dq(y))]- Then y E B(e):
I- [-1a(y"")~ - ا [[Qq(>)]]
- >
(27)
1
This means that
[[ s 2 "-
(28)
The fundamental theorem of calculus can be expressed as:
q(y) - q(y") = S'' Da + +A - y")) (y- yy) da
(29)
q(y) = [/¹ Daly" + 16y) öydh.
Therefore,
8y 55((0) Daly -1a0(((1))))))))))))))
-8y(0 (a(( S¹" Da (y+ + )oy(1) 8y
(30)
= [Dq( L'' (Da(s - Da(y + (oy by 1
Let Let B(e). Then:
| ( L/' (Dq(si)) Da (s(" + (oy(!)) b8
(CS) [Dq 111 L (Da(S Da Dy
KP By
2||[Dq(y")) )1 / //'''' (Da(>(0) -DG)))DD((
Dq (y Da
+ +dy
(orr (
2 "[ /.') /| (Da(v Da (s* + 1or ا0-r0
19 [[[[D(>>)]=11) /." Pa(s Da (y" +1oy
(31)
2|[1a((>))) "11 L/'² 7/1 211x00 - (y- + (1
1,,r rP
( ) /
2/(0 - 2) 1 da |
(CS)
< 2|l[Da(>")] 11
1| >all> >111
2
16

---

Published as a conference paper at ICLR 2024
So | (1+))) is bounded by [] I with constant >|[[Dqyy")] For quadratic convergence.
the norms successive errors must decrease. If we shrink such that y (P < 1, then:
| - >||[Dq(y"" 66
< <>e|[Da(>>"] |88
(32)
< <|16y|01
and the sequence therefore converges quadratically to y*
A.4 APPLICATION TO PARTIAL DIFFERENTIAL EQUATIONS (PDES)
We have mentioned that equation I can be applied Partial Differential Equations (PDEs). As this paper is concerning about Ordinary Differential Equations (ODEs) and discrete difference equation, we only present in this appendix the example of equation | in representing simple PDE. Let's rewrite the equation I as reminder,
[[((r) f (y(r - ...,..r y(r xp)x))))
(33)
where [[:] the linear operator, is the non-linear function, is the signal of interest, is the external signal, is the coordinate, 0 is the parameters of f, and s; is the shifted location.
We will take the Burgers' equation an example. Burgers' equation can be written
du (u2))
8u =0 Ox
v
(34)
Ot 2 Or
where a is function of and I, u(r, 1). we define w u², we can write the Burgers' equation to be
8
8
03
w
u2
(35)
From the equation above, we can define the signal of interest = (u, w) the coordinates r = (a, tT", the non-linear function f = (0, u²)T, and the linear operator
0
0 D3
(36)
Given the relations between the Burgers' equation and equation G matrix and the argument of L.
we can analytically compute the
Of
0
0
G dy
=
- 22
0)
(37)
0
h =f + Gy
(38)
Therefore, the linear operation that needs to be solved, LG [h],
is
Du
1 &w Su
+ 2 dr
v dr²
=0
(39)
g21 (z. t)u + w -
(40)
given h) 2 and 921 (I, can still be further simplified into
-2 2u with from the previous iteration. The equations above
i
[z. t)
1 1ggz
0²
1 0h2 (x. f).
(xr, 1)
(41)
0t
2
dr
2 dr
2 dr
The equation above be repeatedly evaluated until the convergence is achieved.
17

---

Published as a conference paper at ICLR 2024
A.5 ERROR BOUND ON TAKING THE MIDPOINTS IN ODE
The exact analytical solution the ODE
dy dt
1 + G()x(t) = z(t)
(42)
y(t) = e -M()) y(0)+ + 1/" (M(T)((T)dT))
(43)
where
M(t) 1, G(A)da.
(44)
Equation 43 is the form that equation takes in the ODE case. The approximation we wish to take involves inserting G(t) = ] (t.) + (t++1)) and z(t) = 1/2 (z(t,) )) + z(t++1)) into which yields
equation | In order to find the local truncation error when applying this method, we need to find
and compare the Taylor expansions of equations 43 and
We detine
d
so that, for instance, Go G(0), G
4G dt t=
G
L dx
and
de
t=
1==0
G6
d. de
We can then write the Taylor expansion of M(t) as:
1=0
M(I) / G(A)dA
//'' Go GOA+ + 1 + +0(29) da
(45)
Got + - Gb + 6
Ggi + O
2
whence we immediately get
M
_ (n-1)
(46)
expected from the Leibniz integral rule. Therefore:
MM = M(1) + M(1) + M(1) + ((M()))
1
1
= Got (G) + G6) _ + 12
(2G) + 3G. G. 3G + 2G%) & + O(t))
I= Got 2 + B O(t))
(47)
where for ease of notation have introduced the new variables A 1/2 (G) + G)) 12 and B
12 (2G + G/ + 3G' Go + 2G²). As the ODE in section is originally cast in the form dy = x(y()),x(t), 0), the right hand side of is
(t) ((y(t),((t))))))))))) G(t)x(t),
(48)
which can be expressed using the Taylor expansions of f, G, and y as:
t+ z(t) = +o+ Goyo Goyó + Goyo) ++ (f' + Gayd + 2G/y/ + Giyo) 2 Gf +
(49)
18

---

Published as a conference paper at ICLR 2024
Then,
S. (( f( (I+ GL+ (f++ Goyo + (f) + Goyo + Goro)r
Gnyo
1
(60 + Goy' 22/ + G(y)) + +((7)) dr
+ 2
= (fo + Govol t + 2
[ + Goyo + G/vo + Gofo + G ts
1
+
6 (5 + Goyo + 2 + Coyo 2Go + 20 + Gyo
+G + Goyo - Gi + Gy 12 0(th)
Ct Dr + Et O(t)),
(50)
where we have introduced C, D, and as the coefficients of 1, 12 and 18 respectively. Then,
e e-(( S(" (alr)((T)d
II Gol + A + ( (CL+DI2 - E EIE 0(14)
1 = (fo + Goyo) (+ +
2 (G) + Goy' + Goyo- - Gofo - Gūyo) 122
(6 + Goy" 2G + G%yo Goff Gāv'
GoGo + G + Giye 20-o 2G Goya) 8 0(4)
= Ft + Hr + Jt + O(t),,
(51)
where again we have introduced new variables for the coefficients: F.H. and J. We define D, = (t,,), and take 0 and 1 This means that Co = D and substitute with A; Therefore, treating equation 43 as fixed-point iteration problem gives:
(A)) = MM(( y(0)+ + L/= Mi)) (T) + G(T)) dr
= -I- + A - -B-BB FA, + HA + JAY + 0(4,)
1 (f) + - -
=y.+fA
f." GGy +2 - G.f:' Giy!
2
+ G.C.y. - 2G - 2
C. +
(52)
As 2t f(y(t), x(t), 0) + (t)) equation 9 is:
Yit1 = 6 G.A. y; + G (I
G.S.) (fe + Gaye)
1 G/A?
1 1 ( y,+ + G.1 (G.A, / G.A?+ + =
I- G. (I-GAA,+
=
2
6
x (fa + Gay.) + 0(A))
- y.)) A, - GG.f. + Gy. -v:) A? + 1 GG G. (y. - y.))
= =y+ (f. + GG (vo
4/+ 1+00 +(A)),
(53)
with f. = (f. + 11), G. (G, G.t1), and y. (y: + y:4) for 4 S 1 < 4:11.
19

---

Published as a conference paper at ICLR 2024
The Taylor expansions of the midpoint approximations are:
f. = f. + 1 GA + 1
2
+ +
G. G. G - G/'4 +
(54)
2 2 " + -
Ye y:
Recall that y (A,) is the (i + 1)th iterated guess of the form of y, given y(0) and the previous guess (t), evaluated 4 A, Meanwhile, Yi+1 is the value of the approximation to y(i+1) at t- +++1 A., given y0) and the previous guess (t). This gives the local truncation emror as:
LTE = y (141) (A,) -i+t
1 1 [ G.y G. + G.f.+ - yi)] A
= [f fe - -G - -y.)
1 M + +.y + 2G'y, G.f. - Gly - G/ 2G
c. G.ft G. (y. v.) A +7+0(4
f f. - - - ! 2 G. G./ G A
1
f + G - GGf. + G.f, + 2 c.rA ! G 1 c 47
+ 2
1
K" G.y 2 G.f, Gly + 1
+ G.G.y, + Gf G.f- 22
6
G.G.y, GGt A
1
f - G.y.' + Glyl + G.f, + Giyl +
- GG 4 + +(A
12
(55)
Recalling from section that G(t) af ōy*
define p(
gm+nf Jy سx
and expand the total
t=
derivatives:
r = -((2 y: + P..01) p(0.1)x
f"" = p(110)y"" y" + (F.200) y:) N2 + (1.1.)ly:)) x3 (F! Ni + (1.222 x, + p((01) x
(56)
G. = - - (1.0)
G² (22.00 y/
1 1.1)
The local truncation error can therefore also be written:
2 ((220) yy) y! - (r! x 2 (f.41)))) 8 - (f.022)2) X -p. ((011)
LTE 12
-((00 p(o. x + f( (1.00 (F!2.0 y, + 1.1100 (F! y: + (5!2.0 f;
(57)
+1(11x2) f. y() p(1.0) (1(2))2) p(2.00 y: a O(A7).
Accordingly, taking the midpoint values for the approximation enables us to obtain a cubic local truncation error, LTE K KA, + O(A)).
20

---

Published as a conference paper at ICLR 2024
Interpolation
Error
Expressions
Midpoint
O(A3)
Yi+1 = Gay: + ā with G; Gxp(- and z, G
G,)z, where G; (G(4) + G(t++1))/2 and z; (z(t;) +
2(4-1))/2
Left value
0(43)
Yi+1 = G.y. + & with G; = exp(- and ž; = G, G G.)z,, where G; (l+) and 2; = z(4;)
Right value
0(43)
M1+1 G.y, 2, with G; exp( - and z; G (I
G.)z, where G. G(I++1) and 2, z(li+1)
Linear
O(A3)
V(+1 e e-M ys + flitt e((r)((T) dr| where M(t)
G(t,)(t- - 44) +
(((:11) -
(t t.)² and z(t) z(ti) + 211441-4
2(t+))-o(tt(( (t (t:) 411-4
Quadratic
0(A9)
Y441 e-M()) ys flit1 M()) ((T) dr where M(I)
Sr G(T) dr with G(T) and 2(T) are interpolated quadratically us-
ing the values at 4--1, 4, and 1,+1.
Table 3: Error bounds and expressions be evaluated in ODE for various types of interpolations.
A.6 OTHER KINDS INTERPOLATIONS
The choice on midpoint interpolation was made based on the consideration of both error bound and the evaluation simplicity. Table 3 provide the error bounds and expression be evaluated for various interpolations. From the table, it is clear that midpoint interpolation provides a good balance between error and ease-of-come compute expression.
B EXPERIMENTAL DETAILS
B. CODE IN JAX
def deer iveriiteraiion
invlin: Callable| [List
inp. hmp.rddrray, Anyl, idarray,)
func: Callable| [ [List (tpp.jnnarrr.....e. ndarrayl, Any, Anyl, jnp. ndarray),
shifter_fune: Callablel [inp. ndarray, Anyl, List (jnp. indrrayy)
p_num: int,
params: Any,
gradable
xinput: Any, # gradable
invlin_params: Any, & gradable
shifter_fung params: Any, gradable yinit Sinit_guess:: jnp.ndarray,
10
11
max iter: int 100,
12
-> pippleljnn.ddaee
Callable]:
I obtain the functions to compute the jacobians and the function
14
jacfunc vax. (ja..jaaaaaa. jacfwd (func, argnums=0), in_axes= (0, D, None)
15
func2 jax. vmap (fune, in_axes= 0, None)
16
dtype Sinit_guess.dtype
set the tolerance
19
1c-7 if dtype jnp. f4 else 1e-4
20
def iter_ (iter_inp)) err, yt, liter iter_inp yt: (nsamples, ndims)
21

---

Published as a conference paper at ICLR 2024
ytparams shifter func (yt, shifter func params)
25
[p_num] insamples, ndims, ndins)
gts = l-gt for gt in jacfunc (ytparams, xinput, parans) 26
FUNCEVAL
27
rhe: (nsamples, ndims)
rhs func2 (ytparams, xinput, params) # FUNCEVAL
rhs sum cum((jnn (ij](]..iissum (" .....ii........... .........->.....................
gt, ytp)
30
for gt, ytp in zip(gts, ytparams) GTMULT # yt_next: insamples, ndims)
31
32
yt_next = invlinggts, rhs, inviin params) + INVLIN
err jap (wa.axxxj.p...... ads nyt_next yt)) checking convergence return err, yt_next, iiter
34
35
36
def cond_func
-> bool:
37
_ iiter iter inp jnp. logicalllllc..................ee and and (err tcl, iiter max iter)
39
40
very high initial error gt = jnp. jep.zeros (yinit_ggessss shape [0], yinit_guess. [-1],
err = array (1c10, dtype=dtype)
41
yinit_guess www.oesss.hap...oeeee.c shape dtype=dtype)
gts [gt] P_num
44
iiter jnp.array (0, jonee-jnp.nn..3. www.lax.wwieee...
45
err, yt, iiter
46
cona func, fune, (err, suiit_ iiter)) return yt
B.2 HAMILTONIAN NEURAL NETWORK TRAINING WITH NEURALODE
Let's denote as the states of the system. In the two-body case, the states are positions and velocities of the two masses, (#1, /1, Uz1, Uyl, 22, Y/, Ur2, Uy2)" For typical works in learning physical systems, a neural network is designed to take the states as the input and produces the states dynam- ies as the output. those works, set of states and the states dynamics are given as training data, {.... (s.. .... so the neural network can just be trained without solving an ODE.
In our set up, the training data is only given as the states as a function of time (t)....())....) To train the model with the given dataset, need to roll out the states using NeuralODE. our case, the dataset generated by generating the initial condition randomly and the dynamics governed by the gravitational force. The initial conditions orbits of the two-body system does not diverge and the orbits are still close to a circle, to make the simulation numerically stable. The states rolled out from time. Within the time period, the system makes about 2-4 rounds of orbit. There are 1000 rows of dataset generated and split into 800 for the training, 100 for the validation, and 100 for the test.
function of time s(t)
chosen the
0 to / 10 with 10,000 sampled
In order to speed up the training, we start from 20 time points at the beginning of the training and increase the number of time points by 20 every 50 training steps until reaches 10k time points. We performed the training using ADAM optimizer (Kingma & Ba 2014).
For the model, 6 linear layers with softplus activation except on the last linear layer. The hidden layers have 64 elements each. The input to the network has elements (corresponding to the states s) and the output of the network has only element that corresponds to the Hamiltonian value.
using Hamiltonian neural network (Greydanus et al. 2019) that consists of
For every training step during the training with DEER method, we save the predicted trajectory for every row of the dataset. The saved trajectory will be used as the initial guess of the DEER method for the next training step. The training was performed using ADAM optimizer with 10-3 learning rate. To speed up the training, we start from 20 time points at the beginning of the training and increase the number time points by 20 every epoch (50 training steps) until it reaches the maximum 10k time points. The loss function for the training simple mean squared error.
B.3 TIME SERIES CLASSIFICATION WITH GRU
The dataset used in this case is Eigenworms (Bronn et al. 2013). There are 259 worm samples with very long sequences (length N 17984) in this dataset, which then divided into train, validation
22

---

Published as a conference paper at ICLR 2024
output
Sequence Mean
Decoder
Add & Norm
MI.P
Add & Norm
GRU
Encoder
Input
Figure 5: Architecture for the Eigen Worms experiments.
and test sets using 70%, 15%, 15% split following Morrill et al. (202T). Each sample classified as either wild-type one of four mutant types based on the representation using six base shapes.
The model used shown in Figure It consists of four main components: encoder, GRU, MLP and decoder. The encoder would project the input features to higher dimension of 24, and the encoded input would then go through layers of GRU and MLP pair, in that order. The number hidden states for each GRU and MLP is set to 24. Then, the output of the GRU and MLP pair would go through a decoder to project it down to 5 classes. We then take the mean over the sequence length to obtain the final output. Residual connection followed by LayerNorm is applicd each GRU and MLP sublayer. Both the encoder and decoder are also simple MLPs, and all MLPs the architecture have depth of with ReL.U activation function in the hidden layer.
23

---

Published as a conference paper at ICLR 2024
The training is done using cross entropy loss with the patience for early stopping set to 200 epochs per validation accuracy. The optimization algorithm we used is ADAM optimizer with x 10 learning rate and the gradient is clipped at 1.0 per global norm. The initial guesses for the DEER method are always initialized at zeros.
B.4 SEQUENTIAL CLASSIFICATION WITH MULTI-HEAD GRU
The dataset we are using for this subsection is CIFAR-10 from torchvision. Each channel is nor- malized with mean (0.4914, 0.4822, 0.4465) and standard deviation (0.2023, 0.1994, 0.2010) and applying random horizontal flipping before transforming it to a sequence with channels by sim- ply reshaping the signal. Although the dataset from torchvision is already shuffled, reshuffle the training dataset and split training and validation with 90% and 10% portions respectively, Our results that we present in this paper is the accuracy on the test dataset from torchvision using the checkpoint that gives the best validation accuracy.
For the architecture, we first use linear layer convert 3 channels into 256 channels, then we apply M = 4 composite layers of multi-head GRUs, and finally we apply another linear layer to convert from 256 channels into 10 channels for classification. Each composite layer of multi-head GRU consists of (1) multi-head GRU with 32 heads, each with 8 channels with strides {2 22, ..... 2?, uniformly assigned each head, (2) a linear layer from 256 channels 512 channels, (3) gated linear unit (GLU) (Dauphin et al.] 2017) to convert the channels back to 256, and (4) skip con- nection followed by (5) layer norm (Ba et al. 2016). Dropout with 0. drop probability was also inserted to the model. The total number of trainable parameters of the model is 1,347,082.
During the training, we use the cosine annealing with linear warmup for 100,000 training steps. The linear warmup stage took 10,000 steps to increase the learning rate from 10-7 to 2 10 while the cosine annealing took the remaining 90,000 steps to get down to 10-7, In each training step, we applied gradient clipping by global norm equals to 1.0. The training was performed using ADAMW optimizer (Loshchilov & Hutter 2017) with 0.01 weight decay.
C
ADDITIONAL RESULTS
C.I EFFECT OF TOLERANCE LEVEL ON CONVERGENCE
The only hyperparameter of DEER is the tolerance level for the convergence criteria. In our im- plementation, if the maximum absolute deviation of the iterated values is below the tolerance level, then it is marked as "converge". The number of iterations reach the convergence for float32 and float64 precision for various tolerance values are shown in Figure 6 From the figure, we can see that setting the tolerance level of 10-4 for float32 has the same number of convergence if setting the tolerance level of 3 10-7. In both cases, either 10- or 3 x 10-7 tolerance levels, the max- imum absolute error of the DEER output with respect to the sequential method are both relatively low: 1.788 x 10-7 This shows the insensitivity of the hyperparameter of DEER method to the convergence of the method.
C.2 BENCHMARKS
Table 4
shows the speed up for batch size 16 and 8. It can be seen that the smaller batch size can achieve higher speed up average.
Another interesting finding that we found can be seen in Figure 7 The figure shows the speed up comparison when using V100 GPU and when using A100 GPU. Although for smaller number of dimensions A 100 can achieve larger speed up than V100, the speed up when using 32 dimensions on A100 suddenly drops to below 1. This effect is not as severe in V100.
Furthermore, we performed profiling to see where the bottleneck is. On the code in appendix B.I have denoted three lines of interest with FUNCEVAL, GTMULT, and INVLIN. Table 4 shows the run time using GRU with various number of dimensions with batch size of 16. From the table, can that the bottleneck is in solving 1
24

---

Published as a conference paper at ICLR 2024
Precision: float32
Precision: float64
Mean lterations
Mean Iterations
201
10a
10
10
10
10
10
10-15
10-13
10-11
10
10
10.5
10
Tolerance
Tolerance
Figure 6: The mean number iterations required to achieve the convergence for float32 and float64 precision. The case is an untrained GRU with 2 hidden size and 10000 sequence length. The mean and std of the plots were obtained from 16 data points each. The vertical black dashed lines show the 'eps" value of float32 and float64 (1.1927 10-7 and 2.22 x 10-16 respectively).
Speed up GRU over sequential GRU (forward) with batch size
dim
307
8m
Speed up
4m
dim
1k
3k
10k
30k
100K
300k
1M
Sequence length
Speed up of DCCR GRU over sequential GRU with batch size 16
dim
dim
Speed up
302
dim
303
10a
Sequence length
Figure 7: Speed up of DEER GRU over sequential method achieved in (top) V100 and (bottom) A100.
25

---

Published as a conference paper at ICLR 2024
Table 4: The speed up of GRU calculated using DEER method (this paper) vs commonly-used sequential method on V100 GPU with batch size of 16, 8, 4, and 2. The missing data for large number dimensions and sequence lengths are due to insufficient memory in computing the DEER method. The numbers in the table represent the mean speed up over 5 different random seeds.
Sequence lengths (batch size
16)
#dims
1k
3k
10k
30k
100k
300k
1M
15.7
43.5
114
182
309
394
516
2
24.4
47.6
75.0
78.4
72.4
77.1
74.6
22.1
35.4
46.5
44.4
52.1
54.2
51.3
8
14.8
20.0
21.9
24.0
25.2
25.1
7.11
8.37
8.89
9.23
9.46
32
3.74
3.76
4.16
4.24
64
1.29
1.26
1.27
Sequence lengths (batch size 8)
#dims
lk
3k
10k
30k
100k
300k
1M
16.9
44.9
131
264
521
604
945
22.5
58.0
113
149
119
141
138
4
24.1
50.7
79.0
95.9
101
103
104
21.4
31.9
42.3
45.1
47.9
49.6
49.6
16
11.7
15.3
16.8
18.5
18.6
18.7
32
6.32
7.72
8.02
8.32
8.30
64
2.64
2.67
2.78
2.54
Sequence lengths (batch size = 4)
#dims
1k
3k
10k
30k
100k
300k
1M
I
16.6
45.1
130
316
809
1110
1530
2
23.1
62.7
162
256
301
294
306
4
24
63.1
131
174
171
199
207
8
23.1
48.1
73.9
89.3
95.7
98.4
100
16
18.2
28.6
32.7
35.1
37.4
39.1
38
32
11.8
14.9
15.9
17
17.3
15.2
64
5.14
5.7
5.62
5.22
Sequence lengths (batch size = 2)
#dims
lk
3k
10k
30k
100k
300k
1M
1
17.2
43.1
128
357
1030
1860
2660
2
24.8
66.7
187
382
589
516
610
4
25.2
64
172
280
372
416
433
8
24.4
60.5
114
162
178
181
202
16
22.5
45.5
62.8
68.5
75.3
78.1
77.7
32
16.7
25.7
32
31.4
34.3
35.1
64
9.18
10.4
11.1
10.4
10.7
Table The run time of one iteration of a GRU cell in V100 GPU nanoseconds. Numbers that are less than 30 yas are not present in our profiling tool.
Number of dimensions
Line label FUNCEVAL
2
4
16
32
31,430
432,184
911,824
1,194,010
3,419,370
5,249.474
GTMULT INVLIN
88,256
248,562
428,030
1.357,965
4,724,395
147,669
1,336,836
2,015,883
4,460,318
9,786,933
19,248,959
26

---

Published as a conference paper at ICLR 2024
Table 6: GPU memory consumption in MiB for evaluating GRU using DEER method with batch
size = 16 and for various number of dimensions.
Number of dimensions
8
16
32 5038.08
GPU memory (MiB)
18.32
73.25
161.14
380.87
1351.68
accuracy
2x faster
0.8
0.7
0.7
0.6
0.6-
0.5
DEER methad
Sequertial method
0.00
0.03
0.30
3.55
0.0 62
08
0.2
1.9
Hours
Training steps (x10*) (b)
Figure 8: (a) The wall-clock comparison and (b) training step comparison of training LEM with DEER and with sequential method using the same amount of memory by increasing the batch size for sequential method.
We also did benchmark the GPU memory consumption as function of number of dimensions with batch size equals in the GRU case. The results can be seen in Table From the table, it can be seen that the GPU memory consumption approximately grows quadratically, as expected from storing the matrices G explicitly.
C.3 DEER VS SEQUENTIAL WITH THE SAME MEMORY CONSUMPTION
To make another fair comparison,
experiment comparing DEER sequential method using the same memory consumption. In this case, we use LEM architecture (Rusch et al. [021] on Eigen Worms (Brown al. 2013) dataset. All of the hyperparameters are to be the same, except the batch size. We increase the batch size for sequential method to match the memory consumption of DEER. With DEER, we set the batch size equals to 3 while for sequential method, we set 70. The GPU memory usage for both cases is about 2.6 GB GPU.
The results can be seen in From the figure, can see that although DEER uses smaller batch size and requires more steps converge, it can achieve desirable results faster than using sequential method in terms of wall-clock time.
27


Reviews:
[{"review_id": "IscNAb6uzz", "reviewer": "Reviewer_QySF", "summary": "This paper introduces a method for parallel evaluation of non-linear sequence models (neural differential equations, GRUs) as a limit case of differentiable multiple shooting. The key idea is to linearize the multiple shooting update equation and use exact parallel methods (e.g., parallel scans) to evaluate each step of the root-finding procedure. The method is evaluated on time series classifiction, hamiltonian neural ODE training, and benchmarked for efficiency.", "questions": "* How is the recurrent step baseline implemented here? Since you use JAX, have you tried to otimize the recurrence (jit-compile, or fuse)?\n* Although the scaling of the proposed method in with state dimension is quite unfavourable, many existing state-of-the-art sequence models (S4, H3, Hyena, RWKV) take a specific \"multi-SISO\" (single-input, single-output) form, where each channel of the model (in width) is assigned its own state. Although these models typically use linear recurrences that can simply apply the convolution theorem or a parallel-scan, the same structure can be used in the nonlinear case. Since these models are usually trained with small state dimensions (8 - 64), this would correspond to integrating many more systems in parallel (model width times batch size, rather than batch size as in the GRU case). Can the authors commend on whether DEER would be applicable in this case (and ideally show some preliminary result on a simple task?)  \n* How many steps do you need to converge with DEER? How much higher is the latency of one DEER evaluation vs one step of a corresponding GRU?\n* Can you comment on the different rates of approximation of this linearized version of multiple shooting, and regular multiple shooting?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "* The method is clearly presented. The authors do a good job contextualizing the approach with respect to direct multiple shooting, including recent work on adapting it to Neural ODEs (Appendix A.1, A.2)\n* The efficiency evaluation is quite thorough, investigating the effect of batch size, state dimension and sequence length.", "weaknesses": "* It is not clear whether this approach improves over direct multiple shooting (which is also applicable to GRUs). Is there some drawback to the linearization you introduce? Do you require more steps to converge?\n* The benchmarking is quite limited, both tasks showcased are small scale. The tasks do a good job at showing relative performance (end-to-end time) improvements, but they do not provide any insight on the method itself. Are there important hyperparameters, methods that could impact the final results and convergence? You mention mid-point in passing in Sec 3.3, could you elaborate on the choice of other interpolation methods that you have tried?", "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": "", "extra_scores": {}}, {"review_id": "ZxyJtAOLys", "reviewer": "Reviewer_EMpw", "summary": "This paper proposes a novel parallelization scheme for sequential models. The method works by recasting the model as a fixed point problem and using Newton's method to solve it in parallel. Involves computing the Jacobian and Hessian of the model explicitly to enable quadratic convergence. An adjoint method is developed to compute the gradient in parallel. Experimental results are shown in speeding up RNNs and neural ODEs, showing significant speedups especially for long sequence lengths. Extensive proofs are given.", "questions": "+ Can you provide comparisons to the sequential method when the throughput of the GPU is held constant (e.g. the memory is approximately fully utilized)?\n+ How do RNNs with smaller hidden sizes compare to those with larger hidden sizes? Is it the case that the increased speed of DEER is able to offset for a reduction in hidden size that might be needed to use DEER?\n+ Would you be able to use an approximation for the Jacobian (low-rank, etc) to reduce the memory usage?\n+ How many iterations are typically required to converge? How does this compare against zero-order methods that don't use a Jacobian?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 2, "strengths": "+ The method clearly has large speedups in certain training regimes, particularly for long sequences and small batch sizes\n+ The theory of the method is clearly described, with proofs provided in the appendix.\n+ The proofs in the appendix are clearly described. \n+ The method is generally applicable to any sequential method, unlike previous methods which require specific architectures or structural assumptions.\n+ Convergence is quicker than previous methods which did not incorporate Jacobians\n+ Potentially, the method could also be applied at inference time to speed up sequence generation, not just at training time.", "weaknesses": "+ The practical importance of this method is somewhat unclear.\n  From a practical point of view, the DEER method is a mechanism to use more memory in order to speed up the forward and backwards pass.\n  Therefore, many of the experiments, particularly figure 2, are not really a fair comparison, as it's very common to increase the batch size until the memory is fully utilized.\n  In other words, the fairer comparison would be to fix the throughput (i.e. FLOP/s or memory usage) of DEER and the sequential method the same and to compare the two methods. I think it's quite plausible that for a given throughput, using DEER to do more minibatch steps\n  with smaller batch size can result in faster convergence. However, I don't believe this argument is made in the paper. Figure 4 gets close, but I don't think this is normalized in terms of\n  throughput. \n+ The increased memory usage of DEER at higher hidden dimensions seems a big stumbling block. As far as I am aware, a hidden dimension of 8 would be considered very small and impose a limit on the expressiveness\n  of a recurrent model. This issue is not really addressed in the work.", "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": "", "extra_scores": {}}, {"review_id": "VheTUFyeNe", "reviewer": "Reviewer_gw76", "summary": "This paper proposes parallelizing evaluation and training of nonlinear sequential models using fixed point iteration methods. The paper proposes a method that restates nonlinear differential equations as fixed-point iterations with quadratic convergence, as in Newton's root finding method. The result is the DEER method. Speed and performance results are presented for a NeuralODE and GRU that are parallelized using DEER. Theoretical equivalences and convergence rates are established in the appendix.", "questions": "1. Can you somehow quantify the claim that the DEER method can improve the performance of the GRU since it allows more hyperparameter tuning? Perhaps report the results of a single sequential GRU run, and then show the improvement that can be achieved by hyperparameter tuning the GRU for multiple runs (allowed by the parallelization of the DEER method).\n\n2. Can you apply the DEER method to LEM and/or UniCORNN and report the results and compare the runtimes? \n\n3. Can we add more tasks? Is it possible to run a GRU or other nonlinear RNN on Long Range Arena? Or does the cubic complexity of DEER prevent running large enough models for some of these tasks? Can we at least add other common RNN benchmarks such as sequential CIFAR?  \n\n4. Can you add a PDE example (discussed in the appendix)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The proposed DEER method is well-motivated and presented clearly\n\n- The theoretical results seem sound, though I only skimmed the proofs to follow the arguments, rather than check every detail line-by-line\n\n- The method is general and can be applicable to a host of nonlinear differential equation methods such as neural ODEs and any nonlinear RNN (e.g. LSTM, GRU, etc) of broad interest to the sequence modeling community.\n\n- The method has the advantage of theoretically having quadratic convergence. Experimental results support the claim that this method should lead to nearly equivalent results (up to numerical precision) as evaluating the nonlinear ODE sequentially\n\n- For small hidden sizes, the empirical speed ups are substantial", "weaknesses": "- The biggest weakness of the paper is the empirical results, particularly related to performance. Only two tasks are considered, a synthetic physics system where HNNs are trained and the EigenWorms task where a GRU is trained. \n  - For the EigenWorms task, the GRU does not significantly perform that much better than several of the baselines and is outperformed by others.  It is claimed that the DEER method enables faster training and thus experimentation to identify optimal GRU architectures. But is this the best the GRU could do? \n  -  It would have been interesting to see the DEER method applied to some of the other methods such as LEM or UnICORNN since they could presumably be trained with the DEER method also. Would they maintain the same performance? Could better performance be achieved with the DEER method since they could potentially be trained faster and thus be hyperparameter tuned more? An exploration of this would strengthen the paper and broaden the impact.\n  - A greater variety of tasks would increase the impact of the paper and potentially broaden the audience. E.g. there are many other common sequential \"long sequence\" tasks ranging from sequential MNIST, sequential CIFAR, to the LRA benchmarks. It would have been interesting to see how a nonlinear RNN such as GRU performed on LRA. \n\n- There are many interesting theoretical results pushed to the Appendix. The authors might potentially consider packaging some of the formal statements as propositions in the main paper, to signpost the results and guide the reader to understand the rigor behind the method as well as point to the location of the proofs more explicitly in the Appendix.\n\n- A weakness of the proposed method is the cubic complexity in the hidden size that the DEER method incurs. This limits scalability and the applicability of the method to larger scale systems. This is mentioned, but perhaps a further discussion of potential paths to address this could be interesting and helpful.", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)