Towards Scalable and Stable Parallelization of Nonlinear RNNs
Authors: Xavier Gonzalez
Andrew Warrington
Jimmy T.H. Smith
Scott Linderman
Venue: NeurIPS 2024 poster (2024)
Keywords: RNNs
Newton's method
Parallel algorithms
Scalability
Numerical Stability
Primary area: optimization
TL;DR: We introduce methods for parallelizing the evaluation of RNNs that are scalable and numerically stable with quasi-Newton methods and trust regions.

Abstract: Transformers and linear state space models can be evaluated in parallel on modern hardware, but evaluating nonlinear RNNs appears to be an inherently sequential problem. Recently, however, Lim et al. '24 developed an approach called DEER, which evaluates nonlinear RNNs in parallel by posing the states as the solution to a fixed-point problem. They derived a parallel form of Newton's method to solve the fixed-point problem and achieved significant speedups over sequential evaluation. However, the computational complexity of DEER is cubic in the state size, and the algorithm can suffer from numerical instability. We address these limitations with two novel contributions. To reduce the computational complexity, we apply quasi-Newton approximations and show they converge comparably to Newton, use less memory, and are faster. To stabilize DEER, we leverage a connection between the Levenberg-Marquardt algorithm and Kalman smoothing, which we call ELK. This connection allows us to stabilize Newton's method while using efficient parallelized Kalman smoothing algorithms to retain performance. Through several experiments, we show that these innovations allow for parallel evaluation of nonlinear RNNs at larger scales and with greater stability.

Towards Scalable and Stable Parallelization of Nonlinear RNNs
Xavier Gonzalez! Andrew Warrington 12.3 Jimmy T.H. Smith Scott W. Linderman
'Departennt of Statistics, Stanford University.
Wu Tsai Neurosciences Institute, Stanford University,
GE Healthcare. "ICME, Stanford University. 'Liquid Al.
(xavier18, scott (anddrrman@ssaanfoooo edu
Abstract
Transformers and linear state space models can be evaluated in parallel on modern hardware, but evaluating nonlinear RNNs appears to be an inherently sequential problem. Recently, however, Lim al. [1] developed an approach called DEER, which evaluates nonlinear RNNs in parallel by posing the states as the solution fixed-point problem. They derived a parallel form of Newton's method solve the fixed-point problem and achieved significant speedups over sequential evaluation. However, the computational complexity of DEER is cubic in the state size, and the algorithm can suffer from numerical instability. We address these limitations with two novel contributions. To reduce the computational complex- ity, we apply quasi-Newton approximations and show they converge comparably to Newton, use less memory, and are faster. To stabilize DEER, we leverage a connection between the Levenberg-Marquae
algorithm and Kalman smoothing, which we call ELK. This connection allows us to stabilize Newton's 's method while using efficient parallelized Kalman smoothing algorithms to retain performance. Through several experiments, we show that these innovations allow for parallel evaluation of nonlinear RNNs larger scales and with greater stability.
Introduction
Parallel computation has helped fuel the rise of deep learning [2]. Architectures such transform- ers [3] and linear RNNs [4-8] are specifically designed to allow parallelization over the length of the input sequence. However, most conventional nonlinear RNNs (e.g. Elman RNNs, GRUs [9], LSTMs [10] etc.) are not readily parallelizable over the sequence length due to their sequential architecture. Thus, they do not benefit as much from parallel hardware. Nonetheless, these nonlin- ear RNN architectures are still used widely across the scientific community [1 [11-13]. Furthermore, recent work has suggested that linear RNNs (and transformers) are fundamentally limited in their expressivity compared to nonlinear RNNs [14]. Finally, nonlinear RNNs continue to be of signif- icant interest in computational and thcoretical neuroscience as models of neural systems [15-22]. Therefore, scalable and stable parallelization methods for nonlinear RNNs offer significant benefits across many fields.
Towards this goal, Lim et al. [1] proposed DEER, method for evaluating nonlinear RNN in parallel. DEER casts inference as finding the solution of fixed-point equation designed specifically to capture the nonlinear dynamics of the RNN. Newton's method is used to solve the resulting fixed- point equation. With good initialization, Newton's 's method enjoys quadratic convergence rates [23, Chapter I]. Lim et al. [1] also show that the inversion of the structured Jacobian matrix required by Newton's 's method can cast as an associative parallel scan [24]. DEER therefore reduces the evaluation runtime over sequential evaluation by as much as factor of twenty.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Scalablity:
Diagosal Jacobian
Table 1:
Description of the relative strengths and weaknesses of the five evaluation methods we consider. We discuss further in Section 7.
Stability: -Trust region Kalman filter
DEER [1]
Quasi-DEER
ELK
Quasi- ELK
Desiderata
Method
Figure 1: Overview of the paralleliz- able methods we consider in this paper. We introduce diagonal approximations to improve complexity (quasi-DEER, Section 4.1) and link to Kalman filter- ing and trust regions to improve stability (ELK, Section 4.2). We combine these ideas in quasi-ELK (Section 4.2).
Parallel
Work
Memory
Stability
Sequential
No
O
(())
Very high
DEER [1]
Yes
O(TD*) O(TD*)
Low
Quasi DEER
Yes
O(TD) O(TD)
Low
ELK
Yes
O(TD*) O(TD*)
High
Quasi-ELK
Yes
O(TD) O(TD)
Moderate
However, DEER also inherits the weaknesses of Newton's method and parallel scans. The first weakness is scalability. Let D denote the state dimension and T denote sequence length. Using a parallel scan to evaluate updates from Newton's 's method, DEER inherits O(TD²) memory com- plexity and (TD))) computational work [24]. These costs can be prohibitive in practical deep learning settings. The second limitation of DEER is numerical stability, inherited from Newton's 's method. In general, undamped Newton's 's method does not provide global convergence guarantees and in practice often diverges [23]. We seek to ameliorate both of these weaknesses.
To do this, we leverage two techniques: quasi-Newton approximations and trust regions. Quasi- Newton approximations are a common adaptation of Newton's method, where approximate, but faster and less memory intensive updates are used instead exact Newton steps. Empirically, these methods often expedite convergence in terms of wall-clock time, even though more iterations are performed. We apply quasi-Newtor approximations to reduce the memory and compute required by DEER, and we find accelerated convergence and reduced memory consumption. Secondly, we leverage a connection between Newton's 's method with trust region and Kalman smoothing in se- quential models [25]. This connection allows us to stabilize the Newton iteration by limiting the step size to the radius of the trust region, preventing large and numerically unstable steps. The update can be computed with a parallel Kalman smoother [26, 27], achieving a runtime that is log- arithmic in the sequence length. We refer to DEER accelerated with a quasi-Newton approximation as quasi-DEER, and DEER stabilized with trust regions as "Evaluating gevenberg-Marquand via Kalman" (ELK). We then combine these approaches to yield a fast and stable algorithm, which we call quasi- -LKK.
Crucially, DEER, ELK, and their "quasi" variants are algorithms for parallelizing any discrete- time nonlinear dynamical system, including stateful architectures such as RNNs, that may or may not include stochasticity. We use "parallel" to refer to the fact that each iteration of our iterative algorithm operates on the entire T-length sequence (and not on each sequence element one at a time).
We outline the key contributions and organization of the paper here. First, we introduce background material, particularly focusing on DEER [1], in Sections and 3. We then present three short novel proofs: that DEER is globally convergent; that this convergence is robust to modifications of the linearized dynamies (Proposition 1): and that there is unique solution with no local minima (Appendices A. and A.2), We then introduce quasi-Newton approximations to DEER to improve efficiency (quasi- DEER, Section 4.1), and trust regions to stabilize DEER (ELK, Section 4.2) We also provide an interpretation of how trust regions stabilize the dynamics by damping the eigenvalues of the Jacobians (Appendix A.3). We show empirically that quasi-DEER remains accurate, with reduced runtime and memory consumption (Section 6). In regimes where DEER is numerically unstable or convergences slowly, we show ELK and quasi-ELK can enjoy fast, numerically stable convergence. We conclude by discussing the relative strengths and weaknesses of each method, providing guidance on how to select and tune them, and highlighting avenues for future research (Section 7). We provide our code at https: // thub. com/ 1 aadeerrmmlabbeek..
2

---

2 Problem Statement
We consider nonlinear Markovian state space models, with the state at time denoted St RDD and nonlinear transition dynamics f _2 + RR We denote the full sequence of 'I' states as $1.T E RR''D Note that we primarily focus on the transition dynamics, so we suppress any (possibly random) input dependence of the model in the notation. However, the algorithms in this paper extend easily these situations.
For any collection of candidate states {sr]T= and an initial state So we can define the residual
r($1.r) = [s: f(s)), $2 f(s:), $3 fis ...., f(sT-1)] RTxD
(1)
This residual can be interpreted as the one-step error incurred by assuming the thh state is St in- stead of (i) The solution of the state space model, si.r, is the only trace with zero residual. Equivalently, it is the unique solution to the fixed-point equation
r(sirr) -
(2)
The conventional way to obtain si. is to apply f sequentially T times. Sequential evaluation always yields a valid trace, but it requires ((T) sequential operations (i.e. computational depth or span), and hence does not fully leverage the capabilities of parallel hardware. We aim to compute si.r in sublinear time using parallel computation.
Jacobian of the Residual For notational brevity, we overload s and r to also denote vectors in R representing flattened versions of S1.T and r1:7. We can therefore write the Jacobian of the residual for the whole sequence, [s), as TD x TD matrix with block bidiagonal structure of the form
1D 8s
0
0
(81)
0
dr Os
J(s) = (s)
(3)
1p
0s
(ST-1)
ID/
0s
3 DEER: Newton's Method for Parallel Evaluation of Sequential Models
Lim et al. [1] propose DEER, an algorithm using Newton's method for parallel evaluation of non- linear sequential models, including both discrete-time nonlinear RNNs (GRUs, LSTMs, etc.) and neural ODEs [28, 29]. In this paper, we focus on the discrete-time setting. and address questions that arise from Lim et al. [1]: how to scale Newton's 's method, and how to make it numerically stable.
In this section introduce DEER. We begin with simplified derivation that emphasizes the link between Newton's method on vector spaces and parallelizable linear recurrences. We then present new proof that DEER theoretically always converges globally. This proof also highlights why global convergence can be numerically unstable and/or slow in practice. We conclude by using these insights discuss the weaknesses of DEER, and motivate the methods we develop Section 4.
3.1 Derivation of DEER from Newton's Method
The original derivation of DEER used Newton's method on Banach spaces and the Fréchet derivative for continuous-time systems derive the update [1]. We specialize to the setting of discrete- time RNNs, and present a streamlined derivation that more directly connects the structure of the Jacobian in (3) to the linear recurrence relation in (6). This connection highlights why DEER incurs cubic work in D and may encounter numerical instabilities. We will also use this form to prove DEER's global convergence in Section 3.2.
The ,th Newton iterate for (2), starting at s(0), is givee by s s(++1) + - -(
(4)
or equivalently,
As(1+1) = = (i+1) - (t) = -
(5)
3

---

Note this uses the root-finding view of Newton's 's method, see Appendix C.2.
The Jacobian defined in (3) is invertible and all of the eigenvalues equal one.' Storing and naively inverting the Jacobian is infeasible for large D or T. However, since J(s) is block bidiago- nal, we can solve for As in (5) by forward substitution. This reduces simple recursion with the
initial condition As} (1+1) -r1(s(() and for > 1,
af
11+1)) As,
(8) 5 (5/21 As
(++1) rr(
(6)
DEER uses the linearity of this recursion, solving it in parallel with parallel associative scan [1, 24]. Therefore, with ((T) processors, each Newton iteration can be performed (log 7') time.
We emphasize that the computation of the Newton step As in (5) is being parallelized. J would, in general, be a TD TD matrix that is prohibitive to store invert. But by formulating this solve as an LDS in (6), we are able to parallelize the computation of As (which consists of T state updates, each of dimension D) over the sequence length. With sufficient processors, each update in (5) can be computed in (log T) time. Our implementation uses the parallel associative scan from JAX [30] (see Appendix B.6).
3.2 Global Convergence of DEER
We present a proof that DEER converges globally for discrete-tims RNNs to the solution si. of (2) in at most T steps.
Proposition 1. Undamped Newton's 's method will converge to the true solution, Sir of the fixed-
point equation (2) in at most T Newton iterations, for any initial $1-1-
Appendix A.1. The structure of (() determines the recurrence in (6). The update applied at time t, As! +1) from (6) is the summation of a linearized applied to the update time t 1, and the residual one-step error at time t. Therefore, if the previous timestep is correct (i.e. As = O), then the update at time is just the one-step residual, which is defined exactly as the error. Therefore, if the previous value is correct, the updated current value will be correct. Given that f and 5o are fixed and known, the result follows that all T timesteps will have zero residual after T iterations.
Proof sketch. For the full proof by induction,
not immediately obvious from (4) or past work Lim et al. [1] that DEER converges globally, but Proposition 1 shows that it does, at least theoretically. This result has two crucial corollaries. First, after i Newton iterations, $ will have zero error for all S Therefore, if the iteration encounters numerical instabilities, as Newton is prone to, we can simply use heuristic of resetting ( finite value for all
This preserves the solution for time indices & S and allows the optimization to continue, but it is equivalent to running Newton's 's method from scratch on 8,r- This process repeated until the entire trace has zero residual. A second corollary that any set of finite matrices can replace {af/ae) in (3) or (6), and the resulting quasi-Newton method will still converge globally in at most T iterations. This preservation of global convergence provides further motivation for exploring quasi-Newton methods, as we discuss in the next section.
3.3 Weaknesses of DEER
Despite the theoretical convergence of DEER, its formulation as a linear recurrence relation in (6) highlights its limited scalability and stability. Scalability is limited because, in general, 0f/8s is dense D x D matrix. Therefore the parallel associative sean, which uses matrix-matrix multipli- cations, has (TD computational work and O(TD²) memory complexity. Stability is limited because we often have no control over the cigenvalues of 0f/os. If sufficiently many of these eigen- values over the sequence length are larger in magnitude than one, then the linear recurrence relation will be numerically unstable. The heuristic approach of resetting unstable values is sufficient to ensure global convergence, but as we show in Section 6.3, it comes at the cost of runtime, as conver- gence is dramatically slower. These weaknesses motivate the development of two new techniques for parallel evaluation of RNNs: quasi-DEER and ELK, which we discuss in the next section.
"o see this fact, note that the characteristic polynomial of (() in (3) is (A- 1)TD.

---

Scaling and Stabilizing Newton's Method for Parallel Evaluation
In Section 4.1 we introduce quasi-DEER, quasi-Newton method that addresses the intractabilityy of DEER for large state sizes. In Section 4.2 introduce Evaluating Levenberg-Manquar with Kalman (ELK), a damped Newton method for numerically stable, parallel evaluation of nonlinear RNNs. We also introduce quasi-ELK, which combines quasi-DEER and ELK to create damped Newton's method for parallel sequential evaluation that is scalable and numerically stable.
4.1 Quasi-DEER: Scaling DEER with Diagonal Jacobian Approximations
As consequence of our Proposition 1, replacing the Jacobians {8f/0>] with arbitrary ma- trix will still result in global convergence of the resulting DEER-like algorithm in at most T it- erations. A straightforward way to reduce the computational cost is to replace {ar/a>)/=2 _ with (diag (0f/0=)) i.e. take the diagonal entries of the Jacobians of the dynamics functions. The resulting linear recursion requires only O(TD) memory because it only needs to store diagonal ma- trices, and (TD)) work, because the parallelized associative scan only uses element- -ise wise vector multiplies. Fosition-wise matrix-vector multiplies are still required to obtain the residuals, but this computation can be naively parallelized across the sequence.
Quasi-Newton methods approximate the Jacobian for computational reasons, so we refer to this algorithm quasi-DEER. In Section 6, we show that quasi-DEER outperforms DEER on wall- clock time and memory usage on the tests from Lim et al. [1]. Quasi-DEER improves the scalability of DEER, but it does not address stability concerns. We propose more stable solution below.
4.2 ELK: Stabilizing DEER with Trust Regions
Rather than treating RNN evaluation fixed-point finding problem, let instead consider it as an optimization problem. First, we define the function
((s) = ||r(s)]
(7)
As in the fixed-point formulation, the unique minimizer of this objective is s*. In fact, the only
local minimum of the merit function (7) s*, as proved Proposition way of minimizing this nonlinear sum of squares objective is via the Gauss-Newton algorithm [23], which alternates between linearizing the terms in the merit function and solving the resulting linear sum-of- squares problem. The linearized objective iteration is
Appendix A.2. One
Egto/ (As) 1
||((
J(s J(s As
(8)
The solution is (i(++)) - (ss r()()), which is exactly the DEER update from (5).
Formulating evaluation as nonlinear least squares also suggests more stable algorithms. The algorithm [23] uses updates that solve a constrained optimization problem
min Zgto (As) subject to | < Di+1. As
where D is upper bound on the step size. We recognize this constraint as trust region, which is often used in conjunction with Newton's 's method to improve numerical stability and convergence.
Finally, minimizing this constrained optimization is equivalent to minimizing the Lagrangian
E(As, = Egto (As) + Ai- |
(10)
for some Ac+1 2 O. noted by Särkkä and Svensson [25], the minimizer of this Lagrangian can be obtained by a Kalman smoother. We emphasize this connection in the following proposition.
Proposition 2. Solving for the Levenberg- Levenberg-Murquard update that minimizes (10) with fixed X is equivalent to finding the maximum a posteriori (MAP) estimate of Suz a linear Gaussian state space model, which can be done in O I') time on sufficiently large parallel machine.
5

---

Proof. Expanding the residual and Jacobian functions in (8), we see that up to an additive constant, the negative Lagrangian can be rewritten as,
[(s,, \et1) = log (s1 ((s0)) ID) +
Σ log 5/
Sts Ait1
ID
T
0f Üs
log.V V
St f(s
f(s
+
(s (s(21)
(sL-1 (51-1- (1I)
where N/x μr, Σ) denotes the probability density function of the multivariate normal distribution.
We recognize (11) as the log joint probability of linear Gaussian state space model (LGSSM) [31] on (81,..., 8r). The means of the dynamics distributions are given by the linearization of ʃ. and the emissions are the previous iteration's s states, s()) The parameter x sets the precision of the
emissions, governing how far the posterior mode deviates from the previous states.
The minimizer of (10) is the posterior mode of the LGSSM (11), and can be obtained by Kalman smoothing [31]. As with the linear recursions DEER, the Kalman smoother can be implemented as parallel scan that scales as Oloo 7') in time on a machine with O("") processors [26, 27].
Therefore, we can evaluate an RNN by minimizing the merit function with the Levenberg-Marquarrd algorithm. Since each step of the algorithm is performed by parallel Kalman smoothing, we call this approach Evaluating
with Kalman (ELK). Note that DEER is special case of ELK, where À = O, which can be seen as minimizing the unpenalized linearized objective (8), or, alternatively, taking a Newton step with an infinitely large trust region. Moreover, under certain conditions, ELK also enjoys global convergence guarantees [23, Thms. 11.7, 11.8).
Sevenberr-Marquanne
Quasi-ELK: Scalability and Stability As with DEER, we can substitute an approximate Jacobian into the Lagrangian to obtain the quasi-ELK ELK algorithm. Quasi-Ela enjoys the compute and memory scaling of quasi-DEER, as well as stability from the trust region damping from ELK. We show empirically in Section 6.3 that while quasi-ELK takes more iterates converge than ELK, each quasi-ELK iterate is faster, giving overall runtime speedups.
Implementation Details The convergence rate of (quasi-) ELLK depends on the trust region radius D; (or alternatively A²). Although there exist methods to analytically set X; [23, Algorithm 4.3], these approaches require factorizing 0r/os, which intractable at scale. Therefore, we treat as a hyperparameter set by sweep over log-spaced values (cf. Appendix B.4).
We also use Kalman filtering instead of smoothing. We do for two main reasons: filtering requires less work and memory; and we also found it to converge in fewer Newton iterations than smoothing. We hypothesize that this follows from Proposition 1, where the early part of the trace converges first. In Appendix A.3 we also discuss an alternative interpretation of ELK and the Kalman filter as defining a linear recurrence where the trust region attenuates the eigenvalues used in the parallel scan.
The quasi-Newton methods lose the local quadratic convergence properties of New- but remain globally convergent (cf. Proposition 1). Our implementation of quasi-DEER for training uses approximate gradients (cf. Section 6.2). The heuristic of resetting to the states to zero when they become unstable is also motivated by Proposition 1, but it slows convergence in (quasi-)DEFR methods. As a result, we develop EI.K to stabilize evaluation. Like DEER, ELK has cubic complexity D. which we combat with quasi-ELK. However, quasi-ELK adds an ad- ditional hyperparameter. Note that all four parallelized methods discussed in this paper, as well as sequential evaluation of RNNs, have different regimes where they are fastest. For example, in our evaluation of autoregressive RNNs (Section 6.3), the ELK methods are faster than the DEER methods on wall-clock time, but they are slower than sequential. In our evaluation of the Lorenz96 system (Appendix B.5), ELK is more stable than DEER, but DEER is faster on wall-clock time. An area for future research is characterizing the properties of dynamical systems and hardware where each method is fastest. Finally, at the core of the implementation of the parallelized methods is the parallel associative scan (cf. Appendix B.6), which currently available JAX [30] but not in the standard Pyrorch package.
Limitations

---

Related Work
Nonlinear RNNs are a natural choice for modeling sequential data be- cause of their inductive biases and memory efficiency. However, most nonlinear RNNs are not parallelizable over the sequence length, and architectures that can exploit parallel computational hardware have been core to the success of deep learning. Therefore, a range of sequence architec- tures that inherently admit parallelism have been proposed, including transformers [3], deep linear RNNs [4-8, 32-35], and convolutions [36-39]. These methods obtain parallelism by developing new architectures, and do not consider parallelizing existing nonlinear architectures. DEER [1] is notable as it considers parallel evaluation and training of arbitrary nonlincar RNNs.
RNNs and Parallelism
Root Finding in Deep Learning Beyond DEER, there has been broad interest in using root finders/fixed- point iterations in deep learning and sequence modeling. Deep implicit layers [40-43] and neural ODEs [28, 44] replace conventional feed forward network layers [45] with an implicit layer whose output is the root of an equation. Moreover, Song al. [46] parallelize the evaluation of feedforward nets using Jacobi and Gauss-Siedel iterations. In sequence modeling, parallel decod- ing methods [47-49] adapt ideas from Jacobi and Gauss-Siedel iterations evaluate autoregressive sequence models in parallel. These approaches iteratively refine inputs by repeatedly feeding in previous outputs back into parallelized sequence model. However, these methods presuppose the existence of parallelized forward pass for the sequence model and do not leverage additional gra- dient information to obtain sublinear convergence.
Parallelizing Dynamical Systems over Time Other work has investigated evaluating other non- linear dynamical systems over time. ParaDIGMS [50] parallelizes sampling from diffusion models but uses Picard iterations instead of Newton's 's method, while Selvam et al. [51] use Parareal itera- tions. In the numerical ODE and PDE communities there has been great interest in Parallel in Time methods; see Gander [52]. Ong and Schroder [53] for surveys. Vargas et al. [54] parallelized the evaluation of chaotic dynamical systems over time, but instead of casting Newton's method as a parallel scan, they resort to multi-grid methods to evaluate at different hierarchies. Moreover, these methods have not yet been applied to parallelizing RNNs.
Scaling and Stabilizing Newton's Method Quasi-Newton methods efficient algorithms that use an approximation of the Jacobian or Hessian in Newton's method, and include approaches like BFGS [55-58] and L-BFGS [59]. Other approaches use Newton's method optimize deep nets [60]. However, these quasi-Newton algorithms do not admit efficient parallel scans. There are also conjugate gradients methods for exploiting structured Jacobians Hessians [61]. though they often do not attain the fast convergence rates of Newton or quasi-Newton methods [23]. Meth- ods for stabilizing and ensuring Newton's method converges globally include regularization ap- proaches [62, 63], backtracking line search [64], and trust regions [65]. All these stabilization methods have strengths and weaknesses, but as noted by Nocedal and Wright [23]: "the trust-region
Newton method has proved to be highly effective in practice," leading us to apply it to evaluating RNNs.
Nonlinear Least Squares and Kalman Smoothing Bell and Cathey [66] and Bell [67] draw connections between the Gauss-Newton method and the iterated extended Kalman filter and smoother [68, 31J. Because Gauss-Newton is unstable, it is natural to use Levenberg- Levenberg-Marquarrt [69, 70] to stabilize the filtering/smoothing problem [25, 71, 72]. These approaches start with a smooth- ing problem and stabilize it using approaches from nonlinear equations, whereas we start with nonlinear equation to solve and make the connection with Kalman filtering to leverage parallelized algorithms [26]. We also discuss the practicalities of this connection for modern deep networks.
6 Experiments
We now experimentally examine the relative performance of these methods. Specifically, we eval- uate: 1) whether quasi-DEER can provide memory savings over DEER and runtime savings over sequential evaluation, while retaining the accuracy of training and evaluation; and 2) whether ELK and quasi-EIK can be used to stabilize evaluation in regimes where DEER is unstable. In Sec- tions 6.1 and 6.2 we use experimental designs from Lim et al. [1] and show that quasi-DEER retains
7

---

... Sequential
DEER
Quasi-DEER
D
10a
Wallelock
10
Memory (GB)
10a
100
30K
100K
300K
1M 30K
100K 300K
1M 30K
100K 300K
1M 30K
100K
300K
1M
Sequence Length (T)
Figure 2: Evaluating an untrained GRU. Relative performance of sequential, DEER and quasi- DEER for evaluating randomly initialized (and untrained) GRU on (Top Row) wall-clock time, averaged over 20 random seeds and (Bottom Row) memory, averaged over 3 random seeds. All experiments use 16GB V100 SMX2 (memory capacity indicated by the black dashed line) and Newton methods were run to convergence. Missing points in each series indicate the GPU ran out of memory. Quasi-DEER has a runtime commensurate with DEER, but with lower memory consumption, allowing quasi-DEER to work at scales where DEER cannot. The accuracy of the final converged solution is similar for all methods (see Figure 5 in Appendix B.1).
the fast runtime and accuracy of DEER and can reduce memory consumption by up to an order of magnitude. In Section 6.3 we show that (quasi-)ELK remains stable when DEER becomes un- stable, and that quasi-ELK is the fastest of all parallelized methods. We provide further details in Appendix B.
6.1 Quasi-DEER for Evaluation
We first use an experimental design from Lim et al. [1]. The task is to evaluate an untrained GRU across a range of hidden state sizes (D) and sequence lengths (T) 16GB V100 GPU; the inputs to the RNN also have dimension D. We compare the wall-clock time and memory usage of three methods: sequential evaluation, DEER, and quasi-DEER. Results shown Figure 2.
Both DEER and quasi-DEER are up to twenty times faster than sequential evaluation. The runtimes are similar between DEER and quasi-DEER for small networks, because although quasi-DEER steps are faster, quasi-DEER takes more iterations to converge. For larger networks, the difference in runtime is more pronounced. We also that quasi- DEER requires much nitude less memory than DEER, thus allowing the application to architectural regimes previously infeasible with DEER. In Figure 6 of Appendix 1. show that in smaller T and D regimes observe the expected sublinear time scaling with sequence length. This experiment confirms that quasi-DEER replicate the performance of DEER, but with smaller memory footprint.
order of mag-
6.2 Quasi-DEER for Training
We verify that quasi-DEER expedites training nonlinear RNN models. We replicate the third exper- iment from Lim et al. [1], where GRU is trained to classify C. elegans phenotypes from the time series of principal components of the worms' body posture [73].
We show results in Figure 3. We see that the training dynamics under quasi-DEER leads to the similar validation accuracy trajectories. However, every quasi-DEER training step is faster by factor of 2.5, despite performing around 2 times more Newton iterations per training step. This finding highlights how quasi-DEER can replace DEER when training nonlinear RNNs, yielding both time and memory savings. In our experiment, we use the quasi-DEER approximation for the backward pass well, leading to gradients that are different from DEER in this setting. However, we find that there is negligible degradation in performance (Figure 3, left).

---

100
Wallelock time
16
Accuracy (%)
Newton iters per update
Validation
per param update (8)
0.4
Sequential
10
50
DEER
0.2
Quasi-DEER
0.0
0
50K
100K
50K
100K
0
50K
100K
Training Step
Training Step
Training Step
Figure 3: Training a GRU with DEER. Comparison of DEER and quasi-DEER during GRU
training for the C. elegans time-series classification task (Section 6.2). Each time series has length T = 17.984. We show the median, and 5-95% interval across a rolling window of 20 training steps. (Left) DEER and quasi-DEER have the similar validation accuracy trajectories, indicating similar training dynamics. The sequential trace shown is for 24 hours of training (compared to 11 and 4 hours for the whole DEER and quasi-DEER traces). (Center) Each quasi training iteration is 2.5 times faster than each DEER training iteration. Sequential training steps took more than 6 seconds each (not pictured). (Right) Each quasi training iteration requires (approximately) times more Newton iterations to converge, indicating that each quasi Newton step is approximately 5 times faster than the corresponding DEER Newton step.
DEER is prone to "spikes", where orders of magnitude more steps are required for convergence (Figure 3, right). While quasi-DEER is not susceptible to these spikes (never more than half an order of magnitude), these instabilities motivated the study of stabilizing methods.
6.3 ELK and Quasi-ELK for Evaluating Autoregressive RNNs
We conclude by studying an application where these numerical instabilities in DEER are critical. We use small autoregressive GRU (hidden dimension N₂, 3), where the previous sampled value is input into the GRU at the next step. Such autoregressive architectures were not examined by Lim al. [1]. but are an important class of models. We describe the precise details of the AR GRU we use in Appendix B.3. Crucially, the Markovian state 8, used by all four parallelized methods must be expanded to include the current sampled output value, as well as the current GRU state.
Initialized AR GRU We first repeat the analysis in Section 6.1 (and similar to the evaluation in Lim et al. [1] for evaluating randomly initialized autoregressive GRU. We see in the top left panel of Figure 4 that all four parallelized methods converge rapidly and stably to the correct trace, indicated by low mean absolute discrepancy (MAD) between the true trace and the generated trace.
Trained AR GRU We then study a pre-trained GRU that generates noisy sine wave (see Figure 4, bottom). The linear recurrence relation (6) was numerically unstable DEER and quasi-DEER. To remedy these instabilities, trace to a fixed value (here zero). Doing so ensures convergence, but at the cost of "resetting" the optimization for large swathes of the trace (Figure 4, bottom) and slowing convergence (see Figure 4, top right). This finding highlights how the instabilities of DEER which are inherited from both pathologies of Newton's 's method and the parallel recurrence can be crippling in very simple scenarios. While resetting allows for convergence, the resulting convergence is very slow.
take the approach described earlier of setting the unstable parts of the
We then apply ELK and quasi-FLKK K. We show the results in the top right and bottom panels of Fig- ure 4. We select the trust region size with a one-dimensional search over log-spaced values between 10" and 10". We see ELK has stabilized convergence, with the evaluation never incurring numerical instabilities or requiring heuristics. Crucially, by taking more stable steps (and not needing stabiliz- ing heuristics) ELK and quasi-ELK converge faster than DEER and quasi-DEER. ELK can stabilize and expedite the convergence DEER, with BLKK faster still (by wall-clock time).
However, on this task, all parallelized methods (including DEER) are slower than sequential gen- eration. Quasi-ELK is the fastest parallel method, taking 221 milliseconds, compared to sequential evaluation, taking 96 milliseconds. For comparison, DEER took 1,225 milliseconds. Quasi-ELK therefore still represents large improvement in runtime over previous parallel methods. We pro- vide timing details and further discussion in Appendix B.3.2.

---

DEER
Quas--DEER
ELK
Quasi-ELA -LK
Untrained ARGIUU MAD
Trained ARGIRU MAD
10a
10
10a
10-7
10-3
10-10
10-6
5
10
15
2000
4000
6000
8000
10000
Nowton Itcrations
Newton Iterations
20
0
series aftrr
Newton iteration:
-20
20
100
1000
0
-20
2000
20
-20
0
2000
4000
6000
8000
10000
Timestep,
Figure 4: ELK stabilizes parallel evaluation of an AR GRU. (Top Left) The mean absolute differ- (MAD) evaluated the outputs converges rapidly for all four methods a sequence gener- ated by an untrained AR GRU. (Top Right) The MAD for evaluating a trained AR GRU. Undamped DEER variants unstable and converge slowly (using the reset heuristic). ELK stabilizes and celerates convergence, (Bottom) The output after 1, 100, 1000, and 2000 Newton iterations. The black dotted line the true trace. ELK and quasi-ELK converge rapidly, but DEER and quasi-DEER are unstable. The lines where DEER and quasi-DEER are zero depict the zeroing heuristic.
7 Conclusion
In this paper proposed methods for scalable and stable parallel evaluation of nonlinear RNNs. DEER [I] achieved speedups over sequential evaluation, but incurred quadratic memory, cubic work, and numerical instabilities. We therefore extended DEER to use quasi-Newton approximations that reduced computational complexity, and we provided a novel proof that both DEER and quasi-DEEs converge globally. To stabilize DEER, leveraged connection between the Levenberg- Marrguadar method and Kalman smoothing to enable parallel evaluation of RNNs, allowing us to stabilize DEER while still leveraging fast parallel filtering. We empirically verified that quasi-DEER, ELK, and quasi-FI.A K improve convergence across a range of metrics and examples. This result allows parallel evaluation of nonlinear RNNs to be scaled beyond what is possible with DEER.
When selecting an approach, we offer the following advice: If rapid convergence is reliably ob- served, our experiments show that quasi-DEER provides the fastest convergence in terms of wall- clock time. However, the dynamics are on the edge of stability, then ELK offers the most stable performance of these parallelized methods, but quasi-ELK could be faster in wall-clock time and just as stable. In such settings, it worth sweeping the hyperparameter to choose the best version of ELK (note that for À 0, ELK specializes to DEER). However, in the setting of chaotic dynamics, standard sequential evaluation may ultimately be faster.
Our experiments and these observations also highlight avenues for future research. While we found success with a diagonal approximation, structured approximations the Jacobian that still admit fast parallelism but are more faithful approximations may allow for more accurate quasi-Newton steps to taken. Secondly, quantifying the convergence rates of quasi-ELK would allow us to provide tighter bounds than those derived in Proposition 1. Finally, theoretically investigating whether fur- ther improvements to parallelized methods can prove faster than sequential evaluation for dynamical systems the edge of stability, whether there are fundamental limitations to the computational benefit of parallclization, are interesting questions for future work.
10

---

Acknowledgments
We thank John Duchi, David Zoltowski, and the members of the Linderman Lab for their help- ful feedback. thank Amber Hu for her work on preliminary versions of this project, and Leo Kozachkov for pointing out that the conditions we established for the merit function are equivalent to invexity. We thank the anonymous NeurIPS reviewers whose feedback improved this paper.
This work was supported by grants from the NIH BRAIN Initiative (U19NS)))) 13201, MUINS131987 & RFIMHH33778), the NSE/NIH CRCNS Program (R01 N8137899,. X.G. would also like to ac- knowledge support from the Walter Byers Graduate Scholarship from the NCAA. S.W.L. is sup- ported by fellowships from the Simons Collaboration on the Global Brain, the Alfred P. Sloan Foun- dation, and the McKnight Foundation. The authors have no competing interests to declare.
Some of the experiments were performed the Sherlock cluster. We thank Stanford University and the Stanford Research Computing Center for providing computational resources and support that contributed to these research results.
References
[1] Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non-linear sequential models over the sequence length. In International Conference Learn- ing Representations, 2024.
[2] Sarah Hooker. The hardware lottery. Communications of the ACM, 64(12):58 65, 2021.
[3] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, tukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Proceedings of the 31 International Conference on Neural Information Processing Systems, pages 6000-6010, 2017.
[4] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and Francois Flcuret. Transformers are RNNs: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning, pages 5156-5165. PMLR, 2020.
[5] Jimmy T.H. Smith, Andrew Warrington, and Scott W. Linderman. Simplified state space layers for sequence modeling. In International Conference on Learning Representations (ICLR), 2023.
[6] Antonio Orvieto, Samuel Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences, In Inter- national Conference on Machine Learning, pages 26670-26698, PMLR, 2023.
[7] Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. In First Conference on Language Modeling, 2024.
[8] Leo Feng, Frederick Tung, Mohamed Osama Ahmed, Yoshua Bengio, and Hossein Hajimir- sadeghi. Were mns all we needed? arXiv, 2024. httsttss://aoi org/ /00.000.. 48550/ 2410.01201
[9] Kyunghyun Cho, Bart van Merriënboer, Çaglar Gülçehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder- decoder for statistical machine translation. arXiv preprini arXiv: 1078, 2014.
[10] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, (8):1735-1700, 1997.
[11] Ibomoiye Domor Mienye, Theo G. Swart, and George Obaido. Recurrent neural networks: A comprehensive review of architectures, variants, and applications. Information, 15:517, 2024. doi: 10.3390/infol 5090517. URI httpp::/tddooo /tt///ooo/////////o///oo/ooooel/ 3390/1nfo155900117, Academic Editor: María N. Moreno García.
[12] Oleksandr Shchur, Marin Bilos, and Stephan Gunnemann. Intensity-free learning of temporal point processes. In Internatinal Conference on Learnng Representationss 2020.
11

---

[13] Dieterich Lawson, Allan Raventós, Andrew Warrington, and Scott Linderman. SIXO: Smooth- ing inference with twisted objectives. In Advances Neural Information Processing Systems, volume 35, 2022.
[14] William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models. In Forty-first International Conference on Machine Learning, 2024.
[15] Saurabh Vyas, Matthew D. Golub, David Sussillo, and Krishna Shenoy. Computa- tion through neural population dynamics. Annual Review of Neuroscience, 43:249-275, 2020. doi: 10.1 1 www.o-n-2n---------e--- annurev- newro-eero--ee www1o0900o0-00-1-----0
URL https httpe://dooooo//iiiieee org/10.1 1146/
[16] David Sussillo and Omri Barak. Opening the black box: Low-dimensional dynamics in high- dimensional recurrent neural networks. Neural Computation, 25(3):626-64) 2013.
[17] Chethan Pandarinath, Daniel J. O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D. Stavisky, Jonathan C. Kao, Eric M. Trautmann, Matthew T. Kaufman, Stephen I. Ryu, Leigh R. Hochberg, Jaimie M. Henderson, Krishna V. Shenoy, L. F. Abbott, and David Sussillo. Infer- ring single-trial neural population dynamics using sequential auto-encoders. Nature Methods, 15:805 815, 2018.
[18] Jimmy Smith, Scott Linderman, and David Sussillo. Reverse engineering recurrent neural networks with Jacobian switching linear dynamical systems. Advances in Neural Information Processing Systems, 34:16700-16713, 2021.
[19] Marine Schimel, Ta-Chu Kao, Kristopher T. Jensen, and Guillaume Hennequin. II-QR-AE: Control-Based Learning of Input-Driven Dynamics with Applications to Neural Data. Pro- ceedings of the International Conference Learning Representations (ICLR). ICLR, 2022.
[20] Fatih Dine, Adam Shai, Mark Schnitzer, and Hidenori Tanaka. CORNN: Convex optimization of recurrent neural networks for rapid inference of neural dynamics. Adv. Neural Inf. Process. Syst., November 2023.
[21] Carina Curto and Katherine Morrison. Graph rules for recurrent neural network dynamics. Notices of the American Mathematical Society, 70(4), April 2023. doi: 10. 10.1009//ot22661. URL https://doi: org/oog//// 10. www1000000ni2660 not.1661........
[22] C. Costacurta, S. Bhandarkar, D. M. Zoltowski, and S. W. Linderman. Structured flexibility in recurrent neural networks via neuromodulation. NeurIPS, 2024. doi: 10.1101/2024.07.26 605315.
[23] Jorge Nocedal and Stephen J. Wright. Numerical Optimization. Springer, edition, 2006.
[24] Guy Blelloch. Prefix sums and their applications. Technical Report CMU-CS-90-
Carnegie Mellon University, School of Computer Science, 1990.
[25] Simo Särkkä and Lennart Svensson. smoothers. In ICASSP 2020 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages 5875-5879. IEEE, 2020. doi: 10.1 10.1109//CSSP400776 2020.9054764.
Marquurrd and line-search extended Kalman
[26] Simo Särkkä and Ángel F. Garefa-Fermández. Temporal parallelization of bayesian smoothers. IEEE Transactions on Automatic Control, 66(1):299-306, 2021. doi: 10.1109/TAC.202. 2976316.
[27] Peter Chang, Giles Harper-Donnelly, Aleyna Kara, Xinglong Li, Scott Linderman, and Kevin Murphy, Dynamax: State space models library in JAX, 2023. URL https://githhbbb com/ probma/dynamax.
[28] Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31, pages 6571-6583 2018.
[29] Patrick Kidger. On Neural Differential Equations. PhD thesis, University of Oxford, 2021.
12

---

[30] James Bradbury, Roy Frostig. Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, and Skye Wanderman- Milnee. JAX: composable transformations of Python+NumPy programs, 2018. http: thub. com/ com/google/ jax.
[31]
Simo Särkkä. Bayesian Filtering and Smoothing. Cambridge University Press, Cambridge, UK, 2013. ISBN 978-1-107-03385-3.
[32] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with struc- tured state spaces. In International Conference on Learning Representations (ICLR), 2021.
[33] Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length. In International Conference on Learning Representations, 2018.
[34] Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prud- nikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xl- stm: Extended long short-term memory. In Thirty-eighth Conference on Neural Information Processing Systems, 2024.
aattps://arxx: wwwg wwtp/405..ee.ee
[35] Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transformers with the delta rule sequence length. In Proceedings of NewrIPS, 2024.
[36] Aaron van den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals, Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu. Wavenet: A generative model for raw audio. arXiv preprint arXiv: 1609 03499, 2016.
[37] David W. Romero, Anna Kuzina, Erik J Bekkers, Jakub Mikolaj Tomezak, and Mark Hoogen- doorn. CKConv: Continuous kernel convolution for sequential data. In International Confer- ence on Learning Representations, 2022.
[38] Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models. In International Conference on Machine Learning, pages 28043-28078.
PMLR, 2023.
[39] Rom Parnichkun, Stefano Massaroli, Alessandro Moro, Jimmy T.H. Smith, Ramin Hasani, Mathias Lechner, Qi An, Christopher Re, Hajime Asama, Stefano Ermon, Taiji Suzuki, Michael Poli, and Atsushi Yamashita. State-free inference of state-space models: The transfer function approach. In Forty-first International Conference on Machine Learning. 2024.
[40] Zico Kolter, David Duvenaud, and Matt Johnson. Deep implicit layers Neural ODEs, Deep Equilibrium Models, and Beyond, 2020. NeurIPS 2020 Tutorial.
[41] Shaojie Bai, J Zico Kolter, and Vladlen Koltun. Deep equilibrium models. In Advances in Neural Information Processing Systems, volume 32, 2019.
[42] Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Multiscale deep equilibrium models. In Advances in Neural Information Processing Systems, volume 33, pages 5238-5250, 2020.
[43] Shaojie Bai, Vladlen Koltun, and J Zico Kolter. Neural deep equilibrium solvers. In Interna- tional Conference on Learning Representations, 2021.
[44] Stefano Massaroli, Michael Poli, Sho Sonoda, Taiji Suzuki, Jinkyoo Park, Atsushi Yamashita, and Hajime Asama. Differentiable multiple shooting layers. Advances in Neural Information Processing Systems, 14:12-53211 2021.
[45] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep Learning. MIT Press, 2016. ISBN 978-0262035613.
[46] Yang Song, Chenlin Meng, Renjie Liao, and Stefano Ermon. Accelerating feedforward compu- tation via parallel nonlinear equation solving. In International Conference on Machine Learn- ing, 2021.
13

---

[47] Tsvetomila Mihaylova and André F. T. Martins. Scheduled sampling for transformers. In Fer- nando Alva-Manchego, Eunsol Choi, and Daniel Khashabi, editors, Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics: Student Research Work- shop, pages 351 1-3666, Florence, Italy, July 2019. Association for Computational Linguistics. doi: 10.18653/11/111-2000
[48] Andrea Santilli, Silvio Severino, Emilian Postolache, Valentino Maiorca, Michele Mancusi, Riccardo Marin, and Emanuele Rodolà. Accelerating transformer inference for translation via parallel decoding. In Proceedings of the 61st Annual Meeting of the Association for Computa- tional Linguistics (Volume 1: Long Papers), pages 12336-12355, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10. (0.18653/v1122 10005//11/222..0 http:/222...a-looo.6..
[49] Yichao Fu, Peter Bailis, Ion Stoica, and Hao Zhang. Break the sequential dependency of LLM inference using lookahead decoding. In Forty-first International Conference on Machine Learning, 2024.
[50] Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sam- pling of diffusion models. 37th Conference on Neural Information Processing Systems, 2023. URL https: org/10. 00e//004855000 a20.v.2305.............. 16317. 37th Conference on Neural Infor- mation Processing Systems.
[51] Nikil Roashan Selvam, Amil Merchant, and Stefano Ermon. Self-refining diffusion samplers: Enabling parallclization via parareal iterations. In Advances in Neural Information Processing Systems (NeurIPS 2024), 2024.
[52] Martin J. Gander. 50 years of time parallel time integration. In Thomas Carraro, Markus Geiger, Sebastian Körkel, and Rolf Rannacher, editors, Multiple Shooting and Time Domain Decomposition Methods, volume 9 of Contributions in Mathematical and Computational Sci- ences, pages 69-113. Springer, Cham, 2015. doi: 10.10077978-33319-2232111 https://doi. org/10.00
URL
[53] B. W. Ong and J. B. Schroder. Applications of time parallelization. Computing and Vi- sualization in Science, 23:1-10, 2020. doi: 10. 10.10077:00991-00000000 //doi.org/10
URL https:
[54] D. A. Vargas, D. Falgout, S. Günther, and Schroder. Multigrid reduction in time for chaotic dynamical systems. SIAM Journal on Scientific Computing, A2019-22099 2023. doi: 10.113772211519605
[55] C.G. Broyden. The convergence of a class of double-rank minimization algorithms. IMA Journal of Applied Mathematics, 6(1):76-90, 1970.
[56] R. Fletcher. A new approach to variable metric algorithms. The Computer Journal, 13(3): 317-322, 1970.
[57] D. Goldfarb. A family of variable-metric methods derived by variational means. Mathematics of Computation, 24(109):2--26, 1970.
[58] D.F. Shanno. Conditioning of quasi-newton methods for function minimization. Mathematics of Computation, 24(111):647-656, 1970.
[59] Dong C. Liu and Jorge Nocedal. On the limited memory BFGS method for large scale opti- mization. Mathematical PracicaallPPrrr Programming, 45(1-3):503-528 1989.
[60] James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored ap- proximate curvature. In International conference machine learning, pages 2408-2417. PMLR, 2015.
[61] T. Steihaug. The conjugate gradient method and trust regions in large scale optimization. SIAM Journal on Numerical Analysis, 20(3):626 637, 1983.
[62] Yurii Nesterov and Boris T Polyak. Cubie regularization of newton method and its global performance. Mathematical
ProneattiiaaPooPrrrreeeeea Programming, 177(1): 2006.
14

---

[63] Nikita Doikov and Yurii Nesterov. Gradient regularization of Newton method with Bregman distances. Mathematical Programming, pages 1-25, 2023.
[64] James Ortega and Werner Rheinboldt. Iterative Solution of Nonlinear Equations in Sev- eral Variables. SIAM, 2000.
[65] A.R. Conn, N.LMM. Gould, and Ph.L. Toint. Trust Region Methods, volume Society for Industrial and Applied Mathematics, 2000.
[66] S. M. Bell and F. W. Cathey. The iterated Kalman filter update as a Gauss-Newton method. IEEE Transactions on Automatic Control, 38(2):294-297, 1993.
[67] B. M. Bell. The iterated Kalman smoother as Gauss-Newton method. SIAM Journal on Optimization, 4(3):626-634) 1994. doi: 10.1 10.1137/8804035
[68] H. W. Sorenson. Kalman filtering techniques. In H. W. Sorenson, editor, Kalman Filtering: Theory and Application, page 90. IEEE Press, New York, 1966.
[69] Kenneth Levenberg. method for the solution of certain non-linear problems in least squares. Quarterly of Applied Mathematics, 2: 168, 1944.
[70] Donald W. Marquardt. An algorithm for least-squares estimation of nonlinear parameters. Journal of the Society for Industrial and Applied Mathematics, 11(2):431-41, 1963.
[71] Y. Chen and D. S. Oliver. Levenberg-Marquardt forms of the iterative ensemble smoother for efficient history matching and uncertainty quantification. Computational Geosciences, 17(4): 689-703, 2013. doi: 10. (007//10599-003-93335
[72] J. Mandel, E. Bergou, S. Gürol, S. Gratton, and I. Kasanický. Hybrid Levenberg-Marquardt and weak-constraint ensemble Kalman smoother method. Nonlinear Processes in Geophysics, 23(2):59-73, 2016.
[73] André EX Brown, Eviatar I Yemini, Laura J Grundy, Tadas Jucikas, and William R Schafer. A dictionary of behavioral motifs reveals clusters of genes affecting caenorhabditis elegans locomotion.
Proceedings of the National Academy of Sciences, 110(2):991-76. 2013.
[74] B. D. Craven and B. M. Glover. Invex functions and duality. Journal of the Australian Mathe- matical Sociery, 39(1):1-20, 1985. ISSN 155.-6115. doi: 10. N0.1017/S144400 140177/114467788700022226
[75] A. Ben-Israel and B. Mond. What is invexity? The ANZIAM Journal, 28(1):1-9, 1986. ISSN 1839-4078. doi: 10.1077503327000055142
[76] Kevin Murphy. Probabilistic Machine Learning. Cambridge, 2022.
[77] Stephen Boyd and Lieven Vandenberghe. Convex Optimization. Cambridge University Press, Cambridge, UK, 2004. ISBN 9780521833783.
[78] Yurii Nesterov. Lectures on Conver Optimization, volume 137 of Springer Optimization and Its Applications. Springer, 2nd edition, 2018. ISBN 978-3-319-91577-44
doi: 10.1007/ 778-3-319-91578-1 URL https: /ttp://///.aao.rr///....//////// 9-99-999---111-1-----1-
www.oogg///oo 1007
[79] Chee Yap. Fundamental Problems in Algorithmic Algebra. Oxford University Press, Ox- ford, UK, 1993. ISBN 978-0-19-512537-5
15

---

A Theoretical Results
A.1 Proof of Proposition
Proposition 1 Undamped Newton's 's method will converge to the true solution, s", of the fixed-point equation(2) in at most T Newton iterations, for any initial s())
Proof. We prove this result by induction on the sequence length.
In general, the guess s(0) need not equal the solution s" anywhere. However, the initial state So and the dynamics functions/ are fixed. Therefore, s f(o) and in general s? f( 1). Thus, it follows from the initial condition of the DEER recurrence relation that ss = for all i 2 1.
Furthermore, observe that if s 57 for all less than some {(0) then (s 0 for all 1 (00 by the definition of the residual in (1). Therefore, the DEER linear recurrence relation necessarily gives As, (({11) = 0 for all t < ((0) Furthermore, because (/ = f(s("). ) follows that (6)
(s((- -( (s()) k/00
((0)) Thus, it follows that after applying another Newton iteration
that st (i+1) s for all < (0) + 1. The global convergence result and bound on Newton iterates follows by induction.
We note that this proof technique (induction) is very similar to that used to prove Proposition of Shih et al. [50]. However, Shih et al. [50] proves a result about Picard iterations (zeroth or- der method). Our proof about the global convergence of Newton iterations contains the additional complication of dealing with a LDS (as consequence of using first order method). Note, that the global convergence comes from the zeroth order update; however, the first order gradient in- formation can accelerate convergence. The fact that both Picard and Newton iterations have very similar proofs of global convergence suggests that they are closely related, and we will expand on this connection in future work.
Discussion of corollaries of Proposition 1 Corollaries of Proposition 1 include that the fixed- point iterations will converge (in most T iterations) to s* even if the Jacobians 8f/as are replaced (i) by arbitrary matrices, and that we can reset the values of St for 1 > i arbitrarily and still enjoy convergence.
In more detail, the elements in the sub-block-diagona of J= 0r/as can be replaced with arbitrary values but the main block diagonal must remain the identity and all other entries must be zero. Retaining convergence under modifications to the sub-block-diagone portion is a corollary of Proposition 1, and can be seen from (6): If all the states up to and including position t 1 at the (i)th Newton iteration are correct, then the update in (6) at Newton iteration (i + 1) for position t
will use As 0 (no update is required at position 1), and so the update s (i+1) no longer depends on the Jacobian.
We exploit this to develop quasi-DEER retaining only the diagonal of the Jacobians. Doing so reduces the parallel sean from ((D3) work ((D) work, making each iteration faster (while still admitting global convergence as above), but needs more Newton iterations to converge due to approximate updates. We find that this trade-off often yields faster wall-clock time (cf. Figure 8).
Explicitly, the global convergence of quasi-DEER is a theoretical result (a corollary of Proposi- tion 1). but the fast runtime of quasi- in practice is an empirical result (cf. Figure 3).
A.2 The Merit Function Has No Local Minima or Saddle Points
Proposition 3. The merit function L(s) defined in (7) has global minimum at the trace s*, satisfying L(s*) = 0. It has no other critical points, i.e. no such that L((s) = other than at the unique s* for which r(s*) = O.
Proof. First, we observe that VL(s) = (ss))r(s) where J(s) defined in (3). Because J(s) lower triangular matrix with all entries on its diagonal equal to 1, it follows that all of its eigenvalues
16

---

are equal to 1. Therefore, (s) is nonsingular for all s. Thus, J(s) has trivial nullspace for all i.e. J(s)" r(s) =0 r(s) O. But only s* satisfies r(s")
Since there are no critical points other than s", the merit function cannot have local minima or saddle points. Since we have shown that every stationary point is a global minimum, follows that the merit function [(s) is invex (cf. [74, 75)).
We also discuss further the uniqueness of the global minimizer s* of the merit function L.
For a deterministic forward function f and fixed inputs there is fixed sequence of states and outputs (note that any stochastic dynamics function can be reparameterized deterministic by conditioning on the random inputs). Thus, s* the only sequence with zero residual (i.c. there is a unique sequence generated by the deterministic dynamics).
Furthermore, DEER cannot get stuck at any point that is not this sequence. We prove this in Propo- sition 1. Another way to see this however is that each update step (4) is equal to Jr'r. But, J is always invertible and so has trivial nullspace. Furthermore, the residual can only be zero at the unique solution s". Thus J-ir is nonzero everywhere except at the true solution, where it is zero Thus, DEER cannot get stuck en route to finding the true and unique solution.
A.3 Kalman Filtering Damps the Eigenvalues of the Dynamics Matrices
complementary perspective on how ELK results in more stable evaluation nonlinear RNNs is to how the Kalman filter damps the eigenvalues of the Jacobian matrices of the transition dynamics. We first provide high-level overview, and then provide a more detailed derivation.
Overview Let A, be the Jacobians af/as used in the linear recurrence relations and b, be the offsets. Then the prediction step of the Kalman filter (ELK) is the same as DEER. However, after applying the update step in ELK (which imposes the trust region), we obtain second linear recur- rence relation where the linear operator is given by r,A Note that I, is a symmetric positive definite matrix with eigenvalues bounded above by 1/1+A. Thus, by the Spectral Theorem, it fol- lows that the norms of the eigenvalues A, are bounded above by the max of the norms of the eigenvalues of A,, scaled by 1/1+x. Note that larger À corresponds to more regularization/smaller trust region; and therefore correspondingly results in smaller effective eigenvalues in the scan. We recover DEER exactly if À 0. Thus, while large eigenvalues in A, are the cause of the insta- bility of DEER when evaluating unstable dynamical systems, ELK directly attenuates these large cigenvalues. explaining why the intermediate iterations using ELK remain stable.
Derivation
We define our dynamics used in Newton iteration (i + 1) as
A. fttt Os
(s)
b. fs+1(s
0fc+1
(s (=))()))))) (i)
Ös
Now A, E DxD and b, E P.
In line with considering the system as the LDS in (11), we set the process noise to be ID, and with the emissions governed by
(++1)
(ss (6) 0 Ip),
s
where 2² controls the size of our trust region (note that in the notation of developed in Section 4.2 have X = 1/0²).
In the notation of Murphy [76], we see that the predict step
M(tti)jt J.prje + b E E(+tije- = A,E AT + ID-
Meanwhile, the update step is
Ma+1))++) Paenje + E(ennjeE + -'In)~ (N411 "( M(+1)))
Ea+1)444 = E(+D)
+ +02ID) E(411(E(441)))
1T TT
17

---

To unpack this further, we first define the attenuation matrix
o²=0² A + (o² (
Because Efe is a covariance matrix, also symmetric positive definite, and so A is symmetric positive definite, and so all of its cigenvalues are greater than zero. Therefore, all the cigenvalues of A, - (o² + 1)Ip are greater than o² +
We note that F, is also symmetric and positive definite. Thus, by the Spectral Theorem, all eigen- values of Tn are positive. By the above argument, the eigenvalues of T are all less than o² I402 < 1.
Thus, observe that the resulting filtering is given by the recurrence relation
offset
linear
M(+4)(((4)) P.A + T,b, + (A,E + +In) (A + +o2 + +1JIn)
yr+1
((+1))((+1) T, EuA Ip). T-1
3-1 we see that the filtered means (the updates for ELK) come from a linear recur- ==
Given Exjt]
rence relation with linear term P,
We therefore compare the eigenvalues of A, to eigenvalues of A,. Because F, is symmetric positive definite, by the Spectral Theorem we can write r, QA.Q², where Q an orthogonal (and therefore unitary) matrix, and A is diagonal matrix where every entry is in (0, 1) (the entries of Ar are the eigenvalues of Tt, which are greater than by the Spectral Theorem and less than 14-02 by the argument above).
Now, let's consider any arbitrary unit vector v E CD, and let Amax denote the maximum of the norms of all eigenvalues of A, Then ||, S A by the definition of A max However, we want to know IQAQ" A, for any arbitrary unit vector E RD, However, we know that the action of a unitary matrix cannot change the 2-norm of a vector, so QA,Q" A = Q/A.Q" A,v Moreover, multiplying a vector by diagonal matrix cannot increase the 2-norm of vector by more than the absolute value of the diagonal matrix, which in the case of A, is bounded above by o /o"t1.
Thus, QQ A <
क. A or
o²
0²
Q A,V 2 1+
A.
for any unit vector v E CD. This highlights that we can interpret reducing 6² (reducing the size of the trust region and increasing stabilization) as directly attenuating the eigenvalues in the linear recurrence, helping to combat eigenvalues with large magnitude.
A.4 Scale-ELK
Motivated by our derivation Appendix A.3, which shows that ELK reduces the magnitude of the eigenvalues of the Jacobian matrices in the transition dynamics, recommend more lightweight version of ELK which we call scale-ELK.
Scale-ELK uses hyperparameter uses linear dynamical system just like DEER, with the dynamics defined as
[o. 1] (as opposed to À E oo) used by ELK). Scale-ELK
0st+1 (s (0)
A. = (1 k)
ds
b, f(41(s - (1 - 0f+1 (0)
(s )s
Os
Thus, setting = 0 recovers DEER, while setting A = recovers a (computationally expensive form of) sequential evaluation. Ideally. is chosen to keep the magnitude of the eigenvalues of {A.TT below 1. By Proposition 1, scale-ELK also enjoys global convergence.
Scale-ELK enjoys two primary benefits over ELK. First, an evaluation of scale-ELK uses fewer FLOPs than ELI,, as scale-FIKK is just parallelizing an LDS with ELK uses parallelized Kalman filter. Second, the Kalman filter involves inverses which run the risk of introducing numerical insta- bility, while scale-ELK avoids these complications.
18

---

GRU outputs for last 200 indices, DEER vs Sequential
Difference between sequential and DEER outputs
M. AMM
M MMM
0.0
DEER Sequential
-0.5
5800
5820 9840 s860 9880 9500 9520 9940 5960 9980
2000
4000
6000
8,000
10000
GRU outputs for last 200 Indices, quasi-DEER vs Sequential Difference between sequential and quasi-DEER outputs
1--5
AMM Ques DEER Sequential
VM
0.0
-0.5
5800
0875
9900
9925
9950
10000
2000
4000
6000
8000
10000
Sequence index (a)
Sequence index (b)
Figure 5: The accuracy evaluating with parallelized methods (DEER and quasi-DEER) as op- posed sequential evaluation. The parallelized methods converge the correct within nu- merical precision. The hidden state size D = 4 and the sequence length = 10, 000.
B Experimental Details
B.1 Quasi-DEER for Evaluation
In this section elaborate the experimental design in Section 4.1 of Lim et al. [1], including 5 warm-up steps for all timing experiments and batch size of 16. However, instead of running 5 seeds for 20 repetitions each, run 20 seeds for repetitions each, to get more coverage over different evaluation (as the timing for each evaluation are observed be low variance). We also include memory profiling experiments not present in Lim et al. [1]. For these experiments we use random seeds and only one repetition, because the memory usage is extremely stable in between runs. We discus the memory profiling experiments in more detail in Section 1.1.3.
Experiment 1, discussed in Section 6.1. We closely follow
For runs with the same specifications (sequence length T, hidden state size D, and algorithm), observe that sometimes runs with memory profiling ran out of memory whereas runs with timing profiling did not profiling runs was how we handled preallocation of memory. For time profiling, we allowed JAX preallocate memory because that is how JAX usually runs, and gives better indication of wall-clock time in practice. For memory profiling, we did not allow JAX to preallocate memory so could get more fine-grained measure of memory usage.
of memory. A difference between
time profiling runs and memory
However, JAX provides this following discussion of memory preallocation (see https://jaat readthedocs. 1o/ het.//nnnaattppppe/eeeeee http///guummmmrrrrrrr gaap__mmmrr_llloaatiiooo aww.cation..nnniooo html#):
However, [not preallocating memory] is more prone to GPU memory fragmenta- tion, meaning JAX program that uses most of the available GPU memory may OOM with preallocation disabled.
Because the specifications where the time profile runs stay within memory but the memory out of memory are likely very close to the 16GB threshold, our hypothesis that this phenomenon manifestation of this documented memory fragmentation.
B.1.1 Numerical Precision of DEER and Quasi-DEER
In Figure we qualitatively show that for the same example used Figure 3 of Lim et al. [1] that quasi-DEER converges within numerical precision to the correct trace in the untrained GRU benchmarking task discussed in Section 6.1. Similar results for DEER can found in Section 4.1 of Lim et al. [1].
19

---

Sequential
DEER
Quasi-DEER
D
D
16
D 32
D 64
10P
Wallelock (8)
Memory (GB)
10P
1K
3K
30K
100k
d00k
1K
30
30X
100K
300K
IM
IK
20K
1OnK
IM
IK
3BC
30K
10ll
301K
1M
IK
10K
30K
100K
1M
1K
10K
300
100K
1M
18
10K
30K
100K
300K
IM
Sequence Length (T)
Figure 6: Evaluating an untrained GRU. Sublinear and linear timing regimes for parallelized algorithms. The above experiments were run on 32 GB V100 with a batch size of 1. As in Figure 2, we use 20 seeds for timing. 3 seeds for memory, and the dashed black line indicates the memory capacity of the GPU (32 GB). We observe that in smaller regimes in D and T that the wall-clock time shows sublinear scaling indicative of the use parallel algorithms. However, when the GPU becomes saturated, the benefits of parallelization are reduced and we begin to see linear scaling wall-clock time with
B.1.2 Different Scaling Regimes Depending GPU Saturation
In Figure we run the timing benchmarks of Section 6.1 on wider range of sequence lengths T and hidden state sizes D, on larger GPU (a V100 with 32 GB) and with smaller batch size of doing so, we highlight that parallel nature DEER and quasi-DEER, as their wall-clock time scales sublinearly in the sequence length T in smaller (D, T) regimes. However, we note that the larger regimes considered in our main text and in Lim et al. [1], we often observe lincar scaling in the sequence length T for the wall-clock time of DEER and quasi-DEER, even though these algorithms are still faster than sequential evaluation. Figure shows good evidence that these parallel algorithms are suffering from saturation of the GPU, and would benefit from even more optimized parallel hardware
The parallel scan, given sufficiently many processors, scales O T). As we show in Figure 6, we see this speedup at low model sizes and sequence lengths. Once the processors are saturated, we see a linear increase in the runtime (since the amount of work done is linear), but it is making much more effective use the GPU, resulting in constant factor speedup over sequential application at larger model sizes/sequence lengths.
20

---

B.1.3 Memory Profiling Details
As we discussed in Section 4.1, quasi-DEER is O(TD) in memory while DEER is O(TD²) in memory because DEER uses dense Jacobians af/os while quasi-DEER uses a diagonal approxi- mation, diag({f/as). However, to implement quasi-DEER with automatic differentiation, the most standard approach would be to compute the dense Jacobian, and then to take the diagonal; how- ever, such an approach would still be OTTD²) in memory required. There two implementation workarounds. One is to loop over computing partial derivatives, effectively trading time for mem- ory. The second is simply derive the diagonal entries of the Jacobian for the architecture of interest. For the purpose of showcasing the ((TD) memory usage of quasi-DEER in Section 6.1, we take this second approach, deriving the diagonal entries of the Jacobian of the GRU nonlinear dynamics and implementing them in JAX. However, for our other experiments, where memory capacity is not a problem, we simply use the less memory efficient version of quasi-DEER.
We also see linear memory scaling in evaluating the RNN sequentially. This behavior occurs because JIT compile lax. scan in JAX, and track the maximum memory used on the GPU at any point in the computation. Because the inputs and the hidden states of the RNN scales are both of length T. the memory usage of OTT). While there may be more memory efficient ways to sequentially evaluate an RNN, we keep the same benchmarking structure as Lim et al. [1] for to make comparison easier.
B.2 Quasi-DEER for Training
Here we discuss the experimental details for Experiment 2 in Section 6.2. We follow the same experimental setup as in Section 4.3 and Appendix B.3 of Lim et al. [1]. As an aside, we note that the choice of hardware can impact behavior of the algorithms dramatically. For replicability, we run on the same hardware as Lim al. [1], using a 16GB V100 SXM2. However, we note that we try to run these same experiments on A100, DEER struggles to converge numerically, although quasi-DEER shows no such difficulty. If we run on CPU, both DEER and quasi-DEER converge numerically. On balance, on the eigenworms time series classification task, both DEER and quasi- DEER are numerically stable for the most part; the numerical instabilities we have observed for DEER on an A100 are likely specific to some particular detail of JAX/hardware interaction.
In our implementation of quasi- DEER for this experiment, for ease of implementation we do not use the memory efficient version, i.c. we instantiate the full Jacobian and then take the diagonal of it. Nonetheless, we still demonstrate speed gains. Directly implementing the diagonal derivative would likely lead to further speed and memory gains.
The RNN used in this experiment is 5 layer GRU. When we evaluate this architecture in parallel, we evaluate each layer in parallel using (Quas,)-DEER, In Figure 3 (right), we report the number of (quasi)-DEER iterations averaged over all layers and batches. In Figure 7, we report the number of iterations needed for convergence in the last layer only.
B.3 ELK and Quasi-ELK for Evaluating Autoregressive RNNs
Here we discuss our experimental details for our Experiment 3, discussed in Section 6.3.
B.3.1 AR GRU Architecture
The architecture is GRU with hidden states h, R and scalar inputs x, E lR. However, every point in the sequence 1, we readout the hidden state hy E IR3 and use it to parameterize a mean M++1 E R and variance of+1 R4. We then sample T++1 according to xc+1 ~N(u41 of+1): this output Iz+1 is then fed into as the input to the AR GRU at time step t + 1 to make the new hidden step h
This AR GRU is trained using standard sequential evaluation and backpropagatioo-ttrougggh produce noisy sine wave of length 10,000. We train the AR GRU on 1024 traces X1:7' generated from a sine wave with amplitude 10 and white noise applied to each time step, and the training objective is to minimize the the negative log probability of the X1.7.
to
Once the AR GRU has been trained, can generate its own trace X1. given an initial hidden state hg and noises ETT
21

---

20.0
17.5
Newton iters (deepest layer)
15.0
per update
12.5
10.0
7.3
5.0
DEER
2.5
Quasi-Dras DEER
0.0
50K
100K
Training Step
Figure 7: The number of Newton iterations needed for the deepest layer of the layer GRU used for the time-series classification task in Section 6.2.
We note that such system is Markovian with dimension = dim(h) + dim(x), as together the hidden state h, and output T2++1 determine the next hidden state h4+1 and output T++2- Thus, in the notation of Section 2, a hidden state S, of the Markovian state space model is s, = (I,-1. 1 h,). Therefore,
apply fixed-point methods to try find the correct trace s* in parallelized manner instead of autoregressively.
We note that distinction of this set-up with respect the notation developed in Section 2 is that the dynamics functions f are effectively time-varying because the way which Z++2 is generated from (I++1, h₂) depends the noise cr42, the value of which varies across the sequence. However, all the results in the paper still apply after subsuming the input dependence into a time-varying dynamics function fi.
B.3.2 Wall-clock Time Benchmark
The timing experiments were carried out as follows on an Nvidia A100 GPU with 80 GB of GPU memory, using python 3.9 and jax 0.4.11. Exact wall-clock times will vary depending hardware, software, and implementation.
We ran sequential evaluation of the trained AR GRU to produce noisy sine waves of length T = 10,000, as well as the four parallelized method we consider in this paper (DEER, quasi-DEER, ELK, and quasi-ELK).
We ran 20 different random seeds (which lead to different values of the Ci. and therefore different nonlinear dynamics), and timed each for total of 4 repetitions (i.e. 80 timing runs per method). We record the wall-clock time needed to evaluate the length T sequence sequentially, as well wall- clock time, divided by T needed to run T Newton iterations of each of the parallelized methods (thus, we obtain the time per Newton iteration for each of the parallelized methods).
Over these 80 timing runs, the sequential evaluation took an average of 96 milliseconds, with stan- dard deviation of 1.55 ms. We report the average time per Newton iteration, the total number of iterations needed for convergence, and the total wall-clock time to convergence in Table 2. Note that the third column of Table the product of the first two columns.
We effectively read the number of Newton iteration to convergence off of the graphs in Figure 4, but find the number of Newton iterations convergence to be quite stable across random seeds (see Figure 8).
These timing results are illustrative of multiple themes of our paper. We see that while the undamped Newton steps are individually faster because they are carrying out fewer computations (they are just computing linear recurrence relation, equivalently an undamped Newton step, instead of
22

---

Table 2:
Time to evaluate length T 10, 000 trained AR GRU using sequential vs parallelized
methods. We note the dynamax package [27] we used for the parallel Kalman filter implementation ELK is not optimized for speed, and hence these run times could further improved.
Algorithm
Time per New- step (ms,
Newton steps to convergence
Total time to convergence
mean std)
(ms)
Sequential Evaluation
Sequential
N/A
N/A
96
Parallelized Methods
DEER
0.282 + 0.0005
4449
1255
Quasi-DEER
0.087 0.0002
7383
642
ELK
3.600 + 0.0670
172
619
Quasi-ELK
0.141 + 0.0004
1566
221
computing a filtering pass, or equivalently solving a trust region problem). However, because the undamped Newton methods are numerically unstable, they take dramatically more Newton steps to convergence.
Similarly, we see that the quasi methods are dramatically faster than their dense counterparts as they replace ((D3) ((D)) work required by parallel scan on a dense linear recurrence likely saturates the GPU). We in Table 2 that individual steps in the dense DEER/ELK are (approximately) factor of between 3.5 and 30 times slower per step than their quasi (diagonal) variants. However, they take a factor of between 2 and 10 fewer iterations.
multiplication with O(D) diagonal matrix multiplication. The
Thus, we find that our fastest parallelized method on wall-clock time is quasi-ELK, but even so is approximately two times slower than sequentially evaluating this AR GRU. Therefore, interest- ing direction for future work would be to characterize regimes where parallel methods can outper- form sequential methods, and to investigate whether this autoregressive setting is such regime, or whether parallelized methods can benefit from further speed-ups by leveraging adaptive trust region sizes, clever initialization strategies, or even more modern parallelized hardware.
B.4 Setting the Hyperparameter for the AR GRU
We provide more details on how to set the hyperparameters for ELK. Figure 8 shows how to set the hyperparameter for ELK in the context of the evaluating the AR GRU that generates noisy sine wave (Figure 4).
We sweep over the hyperparamter for 15 different input sequences, and plot the median and quartiles of the cost to convergence in terms of Newton iterates and runtime (left column Figure 8). We bathtub curve: large À takes needlessly small steps, slowing progress; small À results in many resets, slowing convergence. Crucially, we see there is little variance across individual sequences. These results show that there is well-behaved dependence that can be optimized validation set with a simple 1-d grid search.
We also chart the approximation error against cost for the AR GRU (center and right column of Figure 8). We see that the approximation error reduces in fewer Newton steps with full DEER as opposed to quasi-DEER, but, crucially, the wall-clock time (the more important of the two metrics) notably lower across all accuracies for quasi-DEER. This indicates that our more efficient but approximate- quasi-DEER is broadly preferable to the more expensive but exact DEER updates. Furthermore, the stabilized ELK and quasi-ELK are better still. We also show the steps/time to convergence for range of accuracy thresholds, and see that our methods outperform DEER across the full range of thresholds and metrics.
These experiments were run on single Nvidia A A100 with 80GB of onboard memory.
23

---

DEER
q- Q-DEER
ELK
q-ELK
Newton steps
for MAD < 0.001
Newton steps
104
MAD
1043
10a
for MAD
10a
109
10a
10a
5000
10000
109
10
10
10
Wallclock time (8)
Newton steps
Wallelock time
for MAD < 0.001
for MAD <e (s)
1.0
MAD
1025
100
109
0.1
10-3
10
0.0
0.5
1.0
1.5
10a
10
10
10
Wallelock time (s)
Figure Experiment to show how the hyperparameters for (quasi)-ELK on the AR GRU pre- trained to generate noisy sine wave (Figure in the main text). Top row plots Newton steps; bottom row plots wall-clock time. Lower is better for all plots. (Left) median steps/time to convergence over À over 15 sequences. Quartiles are shaded but are very small. DEER methods are independent of X. (Center) Updated version of Figure instead plotting MAD a function of wall-clock time. (Right) Time to convergence is robust as function of convergence threshold Median and quartiles across 15 sequences are shown. DEER methods are nearly constant at the thresholds considered (very slight positive slope). Note we plot for increasing À corresponding to smaller trust region, and reducing e corresponding to a tighter convergence threshold.
B.5 Additional Experiment: Evaluating Chaotic Lorenz96 Systems
We include an extra experiment where we tackle the parallel evaluation of the classic non-linear 5- dimensional Lorenz-96 system, with F = 8 which results in chaotic dynamics. We seek to evaluate this system (for T = 1000 timesteps) using (quasi)-DEER and (quasi)-ELK. We directly use the Lorenz-96 dynamics the Lorenz-96 ODE system. The state the five-dimensional Lorenz system state. The input is therefore the initial condition of the ODE; and the outputs are the T subsequent system states.
nonlinear dynamics function f. i.e. the architecture/time evolution is
We demonstrate that all the parallelized methods converge to the correct trace, but that (quasi)-ELK is dramatically more stable at intermediate Newton iterations prior to convergence. We see that DEER and ELK methods converge in a comparable number of steps (this makes sense as DEER special case of ELK for the extra work done per ELK iteration. However, ELK has stabilized convergence, whereas DEER that quasi is slower by all metrics, suggesting that
0). DEER is faster (in terms of wall-clock time) because of
relies heavily resetting. Interestingly the chaotic dynamics may require the more accurate updates. Quasi methods can be implemented to consume notably lower memory, however, and may be preferable in certain circumstances.
Figure 9, we report mean absolute deviation (MAD) of the time series at Newton iteration (i) against the true state sequence. "Iteration" then refers to the number of Newton iterations, i.e. the number of updates applied to entire state sequence. evaluations of the Lorenz96 (i.e. starting from 10 different initial points).
set hyperparameters using 10 different
These experiments were run on single Nvidia A100 with 80GB of onboard memory.
B.6 Background Parallel Scans
For more detailed reference on parallel scans, the interested reader should refer to Appendix H of Smith et al. [5] or Blelloch [24].
In our codebase, we leverage jax. associative vcan with the correct binary associative operator. The binary associative operator for DEER and quasi-DEER is simply the composition of

---

DEER
q G-DEER
ELK
q-ELK
Newton steps
MAD 0..1
10a
Newton steps
1000
MAD
1025
MAD <
109
10a
100
10
10
10
10
500
1000
10a
10
10
Wallelock time (+)
Newton steps
Wallelook time
for MAD < E (s)
MAD
1025
1.0
for MAD <
109
10a
0.1
10
10
10
10
0.0
0.5
1.0
1.5
10
10-1
10
Wallclock time (s)
Iteration 50
Iteration 100
Iteration 200
Iteration 500
True trace
DEER
DEDE
ELK
E-ELK
Figure 9: Evaluating the Lorenz96 system in parallel. (Top two rows): Same format as Figure 8. (Bottom row): Plot of Lorenz96 trajectory during optimization. DEER methods are noticeably more unstable than ELK methods.
affine maps, while the binary associative operation for Kalman filtering be found in Särkkä and Gareía-Fernánder [26] and in dynamax [27].
C Additional Background on Newton's Method
In this appendix, we provide additional background on Newton's method, and why it of use for parallelizing nonlinear RNNs.
Newton's 's method provably enjoys quadratic (very fast) convergence in a basin near the true solution. Moreover, as exhibited by the widespread usage of Newton's method across many domains, New- ton's s method can cxhibit fast convergence in practice. However, a major motivation for this paper is that globally, Newton's 's method be unstable and converge slowly. This instability is major motivation for our development of ELK.
A core insight from Lim et al. [1] is that in the setting of evaluating RNNs, Newton's 's method can be cast as a parallel scan (called DEER). At each "Newton iteration." DEER linearizes the nonlinear dynamics of the RNN it is evaluating. To the extent that linear approximations are very powerful tool across wide variety of domains (e.g. Taylor expansions), this linear approximation can be good approximation, leading to rapid convergence. For example, if we were dealing with linear RNNs, DEER would converge one Newton iteration. In this paper, we are instead dealing with nonlinear RNNs, more Newton iterations are required.
25

---

C.1 Newton's Method for Root-Finding
We provide brief discussion of Newton's method for root finding. A great resource for further study is Nocedal and Wright [23].
Let's say want to find the solution s* to the nonlinear equation r(s) 0, and we have guess s(t) at iteration Newton's method linearizes r(s) the guess s(t), i.e.
dr (s
r(s) == r(s! +
()()))(s ())) Ös
Thus, we get our new guess ((++)) as the solution to f(s) = 0. Therefore, s(++1) satisfies
s ((+1)) (1) =JJ --]-r(s())),
where we define J == Ös
C.2 Newton, Gauss-Newton, Root-Finding, and Optimization
In this paper, we seek find the root of nonlinear equation r(s) 0. Appendix C.I discuss how to use Newton's 's method for root finding to obtain the update
(1-1) + (() -r'r.
However, another approach is to consider minimizing the merit function C(s) := ||r())1//2. The root s" of r will also minimize [(s). so the goal of root-finding to solve r(s) = is the same as trying to find the minimize of L(s). However, if one applies Newton's method for optimization to try to minimize C(s) (see Boyd and Vandenberghe [77] for great reference on Newton's 's method and optimization), the update obtained is actually
((+ (() H(s()))) g(s
where H is the Hessian of L and g J is the gradient of C with respect The Gauss-Newton method approximates this optimization update for minimizing the merit function by making the approximation H J and so the Gauss-Newtor update for minimizing the merit function ends up being the same as the Newton update for the finding the root
C.3 Convergence of Newton's Method
Newton's method only converges within suitable basin [78, $1.2.4, p. 37], but establishing best practices for initialization is an open problem. For instance, Yap provides a bound on the norm the basin [79, Lecture IV, $10, p. 174]. However, this definition requires bounding the derivative of the objective function, which is harder than the original problem. Nesterov derives a basin for quadratic convergence around the true solution [78, Thm 1.2.5, $1.2.4. p. 39], but does not provide information on how to locate this basin a priori. Indeed, Nesterov defaults to taking standard gradient steps early in optimization until you assume you are in the basin, and then using Newton steps [78, 11.2.4, p. 39].
D Algorithm Block
For the reader's convenience, we provide an algorithm block depicting "the ungulates"? (parallelized RNN algorithms, i.c. DEER, ELK, and the quasi-variants)
? An ungulate large hooved mammal, and affectionate term for DEER, ELK, and the quasi-variants
26

---

Algorithm ParallelizeRNN
1: procedure PARALLELIZERNN(S, 8n, init_guess, tol, method, quasi)
2: diff <- 00
3: states +- init_guess
4: while diff tol do
5: shifted states statee + [so, states|: -1]]
6: fs "(shified_statee
7:
Js GETJACOBIANS(F, shifted_statee
8:
if quasi then
9:
Js < DIAG(Js)
bs < fs Js seifted_statee
11:
if method = 'deer' then
PARALLELSCAN(Jss,bb.00
13:
else method = 'elk' then
new sta_ttates - PARALLELKKALMANPIIRRRR.SS states, so)
15:
diff + ||ttates - new_ stw_states|so
16:
states new_states
17:
return states
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The main claim the abstract and introduction is that we improve the lim- itations of parallelizing the evaluation of nonlinear RNNs with Newton's 's method by ad- dressing scalability and stability concerns. We develop quasi-DEER address scalability concerns and ELK to address stability concerns. We combine them in quasi-ELK.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the con- tributions made the paper and important assumptions and limitations. No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: See limitations.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged create a separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to vi- olations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
27

---

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We prove Proposition 1.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide experimental details to allow for reproducibility. including hard- ware used, in Section and Appendix B. We also provide our code at https: //github. com/1 com/linderranlab/eee
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed
28

---

instructions for how to replicate the results, access to hosted model (e.g.. in the case of large language model). releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
model (e.g., large language model), then there should
(d) We recognize that reproducibility may be tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility. In the case dlosed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We provide our code at
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
public/guides/ Seaiddds/CooeSSumiiiiiiiooo sinttoiiissoooocicc for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is acceptable answer. Papers cannot be rejected simply for not including code, unless this central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / /twpps.cpppblll..// httpcccpcccc////e/.. gquudes/CooeSiin
wwssssssooololo cy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new pro- posed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the seript and why.
At submission time, to preserve anonymity, the authors should release anonymized ver-
sions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide experimental details to allow for reproducibility, including hard- ware used, Section and Appendix B. We provide out code lindermanlab/elk
Guidelines:
The answer NA means that the paper does not include experiments.
29

---

The experimental setting should be presented in the core of the paper to a level of detail that necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Yes, we provide standard deviations in Table 2. Sometimes we don't provide error bars when the runs are extremely similar (low variance).
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful ures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
show in tables fig-
If eтror bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yss]
Justification: Yes, we are very careful to specify the type of hardware, including memory capacity.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: http::///eeuippeeeeee httpppulllccccicccueeeeeeeeeeeee hetiiic//hhhiiiiiitie
Answer: [Yes]
Justification: We have read and followd the NeurlPS code of ethics. We do not use human subjects; we use publicly available datasets; and our research is in ways to accelerate stan- dard machine learning algorithms so their broader impacts are to allow current machine learning techniques to be more scalable and stable.
30

---

Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require
deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consider- ation due to laws regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: We discuss how we make machine learning more scalable. Such scalability can lead more energy efficient usage. Any negative impacts would occur from more efficient machine learning being used for pernicious ends.
Guidelines:
The answer NA means that there is no societal impact of the work performed. If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mecha- nisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We have written an algorithms paper. We do not produce pretrained lan- guage model, an image generator, or a scraped dataset.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions access the model implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
31

---

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite Lim et al. [1], whose work and code was an inspiration for this paper, and Chang et al. [27], which we used to scaffold our implementation of ELK.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include a URL.
The name of the license (c.g.. CC-BY 4.0) should be included for each asset.
For scraped data from particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paper saperswithcode com con/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification:
We provide our code at httt:://ggiuuu..co///deeeeaaaaa
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limi- tations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved data collection, curation, or other labor should be paid at least the minimum wage in the country of the data col- lector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
32

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
33


Reviews:
[{"review_id": "dt0QZxxMzT", "reviewer": "Reviewer_S8Ch", "summary": "This paper addresses the problem of parallel computation in RNNs, to be able to tap into the full potential highly efficient parallel machines (GPUs) that are available today: A naive inference of a recurrent model would apply each layer on the current hidden state, thereby requiring the length of the sequence $T$ sequential computations. The approach taken by this paper, which is introduced in the framework of Newton's method of root finding, proposes finding the solution to a chain of equations that, if satisfied, will verify that the found activations are equal to the sequential inference one. \n\nGiven sequence states $s_1,\\dots , s_T$, the general approach (based on earlier method DEER() is to define a sequence of $T$ residuals $$r = [s_1-f(s_0),\\dots, s_T-f(s_{T-1})]$$ which capture the discrepancies between successive activations and the update layer corresponding to that layer. The update rule is based on the gradient of these residuals wrt to the layer activations $$\\Delta s_t^{(i+1)} = \\left[\\frac{df}{ds_{t-1}} (s_{t-1}^{(i)}) \\right] \\Delta s_{t-1}^{(i+1)} - r_t(s^{(i)})$$. Since at each step we can compute activations in parallel, we can execute this update for all $t=1,\\dots, T$ in parallel. The proposition 1 states that at most $T$ updates will be needed to arrive at the correct solution. This is already introduced in DEER, but this paper goes on to address two main problems with DEER. \n\nThe linear recurrent update suffers from scalability due to the $D\\times D$ size of the full Jacobians that impact both memory and time complexity, and numerical instability dueo to the singular values larger than one in Jacobian which can lead to numerical instabilities. \n\nI had a difficult time delineating the contributions, but I believe the following to be the main contributions: \n\n1. Quasi-DEER: using a simplified form for the Jacobian where only diagonal terms of the Jacobian are considered, leading to a linear memory and time complexity per each layer. However, this model might still suffer from the numerical instabilities. \n2.  Iterative Kalman Evaluation (IKE), where the updates to the activations are \"damped\" by adding an L2 regularization term to the update rules, which can effectively prevent instabilities.", "questions": "* The proposition 1 seems to suggest that the convergence stops after T steps. but that is not any speedup over the trivial sequential computation of T steps of the RNN, so it is of little value it seems. Or am I missing something? Is this T-step iterative method still better than a simple sequential RNN computation? perhaps that’s just theory, If the practical results are better, they better be stated here? \n* From seeing the figure 2, it seems like the Deer & quasi-DEER are faster, in light of my previous questions, can authors elaborate why is that the case? Is this because the actual number of steps required to stop is much less than T, or is it some hardware-related efficiency of the (quasi-)DEER implementation, vs the naive implementation? \n* in my initial reading understanding, I confused the Newton's method for root finding with Newton's method for optimisation, which would involve Hessian of loss for finding the minima of a function, while the formula involving Jacobian is for finding the “root” of a function. I'm not sure if other readers would have a similar confusion, but perhaps a small note would help. \n* I was missing a bit the high level intuition behind the residual sequence and how is it leading to the same sequential solution. For example, I wasn't sure if/how the equations would lead to a unique solution all the time or not, and why is that so. I think I would have understood it better\n* To me it seems that the NN layer updates must be deterministic for this approach to work. Is this impression correct? \n    * Also, perhaps it is beneficial for readers if authors elaborate on what types of sequential RNN architectures would be compatible with their approach for parallel computation? \n* “Provided all of {df/dst−1} T t=2 are finite, then the Jacobian defined in eq. (3) is invertible and all of the eigenvalues are equal to one.” I don't understand why this is the case? Why are eigenvalues equal to one? Matrix J has off-diagonals, so it’s not quite clear why J inverse is so trivial to compute? \n* I find the following poorly stated/supported as of my reading so far: \n    * 137  As a consequence of our Proposition 1, replacing the Jacobians {df/dst−1}Tt=2 with an arbitrary  matrix will still result in global convergence of the resulting DEER-like algorithm in at most T  iterations. \n    * Replacing J with any **arbitrary matrix**? really? Why? \n    * For a similar reason: \t\tA straightforward way to reduce the computational cost is to replace {df/dst−1}Tt=2   with {diag (df/dst−1 )}Tt=2 ; that is, to take the diagonal entries of the Jacobians of the dynamics ,  functions. \n    * Why does this work? Is this a theoretical result, or an empirically valid approach? \n* Question on Figure 6: is the later linear growth of the wall clock because of the limited parallel compute on the GPU? \n* section 6.1: why is convergence of a random AR stable, while convergence of a trained AR model is more unstable? Can you elaborate? Is this setup leading to the large top eigenvalue problem mentioned earlier? \n* Figure 4: I’m not sure I understand top-right for some traces, \n    * why the DEER (red) traces goes down at such a late iteration? I thought it was unstable and would never converge. \n    * Also, the quasi-DEER (orange), goes up and then down, what does that mean? \\\n    * Why is the IKE going down faster first, but then quasi-IKE converges faster?", "limitations": "I am a little concerned about the way that limitations are currently addressed. \n- In responses to the questions, authors have listed section 6.3 as a limitation, but I do not see where this section is exposing their method limitations. \n- As also mentioned in questions, I would appreciate if authors make some comments on architectural limitations of this work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper consideres an important problem, ie, parlallel computation of RNNs, which is of high practical importance. \n- The key contributions of the paper in quaski-DEER and (quasi-)IKE are novel and address issues of scalability and numerical instability.\n- To the best of my knowledge, the proposed contributions are novel.", "weaknesses": "I prefer to package most of my perceived weaknesses in the questions section, as they are mostly minor. Here are my main concerns:\n\n* As authors list lambda's for IKE as a hyper parameter, it warrants a more comprehensive discussion on its choice and impact on the numerical results. Namely, an abalation study that shows how sensitive the choice of lambda is. Also, whether it is task and data-dependent it could be, or it is fairly robust across these. \n* Contributions:  The manuscript would benefit from a more clear distinction on their contributions. Since the paper draws heavily from previous work that introduced DEER, it was a bit challenging for me at the end to delineate what in the paper was novel and the current paper’s contribution. While IKE and quasi-IKE and the propositions were clearly novel, other points were not as clear to me.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough and positive review! \n\nWe refer the reviewer to the global response for discussion of an ablation study on $\\lambda$ (r.e. “As authors list”) and limitations (r.e. “In responses to” and “As also mentioned”). We now respond to your other comments:\n\n> Contributions: The manuscript… \n\nWe will explicitly list our contributions in the introduction.  Specifically we:\n- Show DEER is globally convergent\n- Introduce quasi-approximations to improve efficiency\n- Introduce trust regions (IKE) to stabilise DEER\n- Empirically validate the methods, and provide guidance on how to select and tune them\n- Have since explicitly shown there is a unique global solution to DEER (see “I was missing…” response) and provide a reinterpretation of how trust regions stabilise the original scan (see 5uzK response). \n\n> The proposition 1… **and** From seeing the…\n\nYou are absolutely correct. Proposition 1 tackles the theoretical global convergence of the algorithm, showing we can expect convergence. The upper bound is impractical as it requires $T$ iterations; but this is a _worst case_ convergence and we often require much fewer than $T$ iterations (e.g. Figure 4a). Each Newton step requires more work than a sequential step; but fewer steps are required and the work is parallelizable. Proposition 1 was also important in developing the resetting heuristic for (q-)DEER. We have added clarification on this. \n\n> In my initial…\n\nThank you for raising this. While DEER was inspired by the root finding perspective, IKE was inspired by the optimization perspective (see Section 4.2). We have further clarified the relationship between perspectives and their implications.\n\n> I was missing…\n\nWe add a proof of the uniqueness of the solution, and global convergence to this solution, as we did not make this connection clear enough. To summarise:\n\nFor a deterministic forward function and fixed inputs, there is a fixed sequence of latents and outputs. This is therefore the **only** sequence with zero residual (i.e. there is a unique sequence $s_{1:T}^*$ generated by the deterministic dynamics). Furthermore, DEER cannot get stuck at any point that is not this sequence. We prove this in Proposition 1. Another way to see this however is that each update step (eq. (5)) is equal to $J^{-1} r$. But, we have established (see “Provided all of” below) that $J$ is always invertible and so has trivial nullspace. Furthermore, the residual $r$ can only be zero at the unique solution $s_{1:T}^*$. Thus $J^{-1} r$ is nonzero everywhere except at the true solution, where it is zero. Thus, DEER cannot get stuck en route to finding the true and unique solution.\n\n> To me it seems…\n\nThis is an interesting point: DEER and IKE  *can* handle stochastic models – because stochastic models become deterministic when the stochasticity is fixed i.e. has been reparameterized as in Figure 4. In fact, DEER and IKE have the same requirements as backpropagation, and so can be applied to any stateful recurrent model that admits backpropagation – a huge and practical class of models. \n\n> Provided all of…\n\nThe eigenvalues of a lower-triangular matrix are equal to its main diagonal. The Jacobian in eq. (3) is the derivative of eq. (1). Each term in the vector is of the form $s_t - f(s_{t-1})$. Therefore, the Jacobian with respect to $s_t$ is the identity matrix; hence the leading diagonal is the identity; and all its eigenvalues equal one and the Jacobian is invertible (no zero eigenvalues). \n\nEfficiently solving this inverse is not tractable (see Line 92). However, we avoid computing the inverse by exploiting the structure in $J$ to instead pose eq. (5) as the recursion in eq. (6). This recursion is linear and can be solved in parallel with a scan. \n\n> I find the…\n\nWe will make this discussion more clear in the paper: The elements in the sub-block-diagonal of $J$ can be replaced with arbitrary values – but the main block diagonal must remain as the identity and all other entries must be zero. Retaining convergence under modifications to the sub-block-diagonal portion is a corollary of Proposition 1, and can be seen from eq. (6): If all the states up to and including position $t-1$ at the $(i)$th Newton iteration are correct, then the update in eq. (6) at Newton iteration $(i+1)$ for position $t$ will use $\\Delta s_{t-1}^{(i+1)} = 0$ (no update is required at position $t-1$), and so the update to $s_{t}^{(i+1)}$ no longer depends on the Jacobian. We will explicitly state this, as it motivates q-DEER and is, as you point out, a surprising result. \n\nWe exploit this to develop q-DEER, retaining only the diagonal of the Jacobians. This reduces the parallel scan from $O(D^3)$ to $O(D)$ making each iteration faster (while still admitting global convergence as above), but needs more Newton iterations to converge due to approximate updates. We find that this trade-off often yields a faster wallclock time (see Review PDF Figure 1).\n\nExplicitly, the global convergence of q-DEER is a theoretical result (a corollary of Proposition 1), but the fast runtime of q-DEER in practice is an empirical result.\n\n> Question on Figure 6…\n\nYes, that's a great point! We will more heavily emphasise this.\n\n> Section 6.1: why…\n\nYes! The trained model does have larger eigenvalues, both at the true latent sequence, and at the sequences visited during the Newton iteration.  \n\n> Figure 4: I’m…\n\nWe use “unstable” to mirror Newton’s method parlance. Proposition 1 shows that $t \\leq (i)$ have zero error, and so the instability applies only to $t>(i)$. Unstabilized methods can be arbitrarily bad in this regime until global convergence, but *may* converge before then. (q-)IKE removes the instabilities, leading to faster overall convergence. We have added clarification on this around Figure 4.\n\nThank you for your thorough analysis, and for raising many points that have improved the paper. Please let us know if you have any further questions!", "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": "1w9UQy0wzX", "reviewer": "Reviewer_jU7g", "summary": "This paper presents a stable and scalable method to parallelize evaluation of an RNN across the sequence length. Furthermore, the paper presents a sketch of when the more recently proposed DEER algorithm converges. Alongside these contributions, the paper presents experiments showing the efficacy of the proposed method.", "questions": "- In general, getting to exactly zero error is impractical (floating point errors among other reasons) so how can the theorem be generalized to considering number of iterations to achieving \\epsilon error? As a side consequence, can the authors clarify when DEER can achieve quadratic convergence rate of Newton method? How sensitive DEER is to blowing up errors when propagating information across sequence length?", "limitations": "Yes.", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "+ Identifies a shortcoming with the DEER algorithm\n+ Presents a computationally efficient and scalable alternative to DEER", "weaknesses": "- The result proving the convergence of DEER isn't quite rigorous/general -- it relies on the Jacobians being finite which is quite a strong assumption which isn't true in practice, particularly in situations when the eigenvalues are > 1. The paper sweeps this under the rug through re-initialization, which suggests it isn't quite a general statement.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for highlighting some interesting clarification points. \n\n> The result proving…\n\nThis is a really interesting line of thought that we will clarify in the paper.  There are lots of points that your comment touches on, so we will try and answer them sequentially:\n\n**Re: Generality of Prop 1 with regard to eigenvalues**: We state in Prop. 1 that we assume the Jacobians must be *finite.* Given this assumption, the proof by induction in Appendix A.1 is rigorous and general for real numbers. Thanks to your feedback, we will include an explicit pointer to the full proof.\n\nThe finiteness of the Jacobians is a core requirement for DEER. Thus, Prop. 1 (a statement about the DEER algorithm) is applicable anywhere that DEER is applicable. We also argue that the finiteness of the Jacobians is a weak assumption, as it is a requirement for training models with backpropagation and is satisfied by many commonly used architectures.\n\nHowever, it is absolutely true that the eigenvalues can be greater than one, and this leads to the *product* of the Jacobians computed by the scan blowing up in magnitude. In fact, the sensitivity of DEER to blowing up is directly linked to the eigenvalues of the Jacobians. It is this explosion that causes overflow on finite-precision architectures. \n\nThe resetting heuristic combats this, and is motivated by a result derived as part of Proposition 1: that the first $t$ steps have zero residual after $t$ Newton steps (see ~Line 111 or 447), e.g., that we can restart the optimization with the correct $t$-length initialization. We can view this heuristic as allowing the Newton steps to resume when overflow is encountered, as opposed to failing or falling back to sequential evaluation. \n\n**IKE directly combats eigenvalues greater than one**:  A related point is that IKE was motivated to prevent overflow by limiting the step size in the latent state using a Kalman filter (KF). We have since made a connection from the Kalman filter back to the original scan:  the trust region/KF can be interpreted as attenuating the eigenvalues by a factor of at least $\\frac{1}{\\lambda+1}$, stabilising the scan _without_ the use of the heuristic (see **Proof that IKE attenuates the largest eigenvalues in linear recurrence** in response to 5uzK).  IKE combats exploding cumulative products without the use of the heuristic.  \n\n**Re: discussion of the re-initialization heuristic**: we respectfully disagree that we “sweep it under the rug”. We discuss the heuristic in several places, and explicitly highlight the heuristic in Figure 4. The heuristic is a necessary intervention — not originally discussed by Lim et al. — to combat instabilities of undamped Newton in floating-point representations. We see the heuristic as a (albeit small) novel contribution, and a practical corollary of Prop. 1. We have expanded discussion of the limitations and opportunities posed by the reinitialization heuristic.\n\n> In general…\n\n**Re: connections back to floating point**: Our analysis was a theoretical exploration to develop methods we would eventually test empirically, and so we didn’t make explicit theoretical connections to floating point representations. Floating point precision can introduce instabilities into Newton through overflow and limited numerical precision in functions with large gradients. The finite precision also sets a floor on the fidelity of the convergence. These limitations are further motivation for stabilisation techniques, such as IKE, to combat the instabilities inherent to Newton’s method when executed in finite-precision.\n\n**Re: Convergence rates of Newton’s Method**: We spent a lot of time thinking about this. Newton’s method only converges within a suitable basin [Nesterov, 2018, Lectures on Convex Optimization, §1.2.4, p37] — but establishing best practices for initialization is an open problem. For instance, Yap provides a bound on the norm of the basin [Yap, Fundamental Problems in Algorithmic Algebra, 1993, Lecture IV, §10, p. 174]. However, this definition requires bounding the derivative of the objective function, which is harder than the original problem. Nesterov derives a basin for quadratic convergence around the true solution [Theorem 1.2.5; Nesterov, 2018, Lectures on Convex Optimization, §1.2.4, p39], but does not provide information on how to locate this basin apriori. Indeed, Nesterov defaults to taking standard gradient steps early in optimization until you assume you are in the basin, and _then_ using Newton steps [Nesterov, 2018, Lectures on Convex Optimization, §1.2.4, p39].  We experimented with this initially, but found it to be worse (by any metric we conceived).\n\nQuasi-methods do not, in general, enjoy quadratic convergence (but showing the global convergence of quasi-DEER was a major motivation behind Proposition 1), but are often observed to converge in practice. We have not been able to prove results about the convergence rates of the quasi methods we propose. However, our experiments show that their rates are acceptable in practice. Thus, studying quasi-DEER from a theoretical perspective is an interesting open problem. As a result of your points, we have included in the conclusion a list of open theoretical problems inspired by your questions. We have also added a brief survey on the convergence properties of Newton variants. \n\n> Limitations: Yes\n\nIs the reviewer saying there are limitations that we do not discuss? Or are that we do satisfactorily discuss limitations? If the former, please see the Global Response for information on how we have updated the paper. \n\nThank you again for providing some really interesting and unique insights. The extra discussion has definitely improved the paper. Please let us know if you have any further questions!", "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": "c1eBOWJJFH", "reviewer": "Reviewer_VSf9", "summary": "This work extends the parallel RNN (DEER) method to improve its efficiency and stability. The authors make two main modifications:\n1. Replace the full Jacobian matrix inverse with its diagonal, allowing for linear-time inversion.\n2. Introduce damping via Kalman filter to enhance the stability of Newton methods.", "questions": "Harder examples: how does the model work on chaotic system such as Lorenz 96 or Weather forecast?\nMore benchmark, how does the model compared to transformers and xLSTM?", "limitations": "The limitations are not sufficiently addressed.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Solid contributions to improving the DEER method\n- Potential for increased efficiency and stability in parallel RNN implementations", "weaknesses": "Overall, the paper have solid contributions. But there are several weakness\n \n1. The potential drawbacks of these modifications are not fully explored\n   - Efficiency gains may come at the cost of slower convergence\n   - Damping could potentially compromise accuracy\n   \n2. Limited experimental scope\n   - The experiments may not reveal issues that could arise in harder, larger-scale cases\n   - For example: does the model work on chaotic systems such as Lorenz 96?\n\n3. it will be interesting to compare with popular models like transformers and xLSTM", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our paper and providing some helpful feedback! \n\n> Limited experimental scope\n\nWe point the reviewer to the global response for additional experiments we have added: \n\n- In Review PDF Figure 1 we study the sensitivity to the $\\lambda$ parameter (related to Reviewer S8Ch’s feedback), showing that $\\lambda$ can be effectively set with a 1-d hyperparameter sweep, and that actualized speed-ups for a selected $\\lambda$ are robust for a given data modality. \n\n- We also add additional experiments exploring the wallclock time/iteration/accuracy tradeoff of the methods (r.e. Your comment “The potential drawbacks…”). In Review PDF Figure 1, we show for the AR GRU that while full methods take fewer steps, quasi methods converge faster in wallclock time across a broad range of desired accuracy levels.\n\n- We really liked your suggestion of using our algorithms to perform inference in a Lorenz96 system. We include results for this experiment in the Review PDF Figure 2. We see that DEER and IKE methods converge in a comparable number of steps (this makes sense as DEER is a special case of IKE for $\\lambda \\to 0$). DEER is faster (in terms of wallclock time) because of the extra work done per IKE iteration.  However, we do see that IKE has stabilised convergence, whereas DEER relies heavily on the resetting heuristic.  Interestingly we see that quasi is slower by all metrics, suggesting that the chaotic dynamics may require the more accurate updates.  Quasi methods still consume notably lower memory, however, and so may be preferable in certain circumstances.  \n\n> It will be interesting to…\n\nThe comparison to alternative architectures, such as the xLSTM and Transformers, is a potential point of comparison. However, we stress that this paper is about accelerating and operationalizing the parallel evaluation of sequential, non-linear models. Papers developing or benchmarking natively parallelizable architectures (e.g. Transformers) often state that sequential architectures such as (e.g.) GRUs cannot be parallelized, and are often therefore eliminated from consideration or score poorly on speed tests. However, DEER raised the point that this is not entirely accurate, and that alternative black-box methods may be within reach for efficiently evaluating these oft-maligned architectures. Our work focuses on extending and operationalizing DEER, without changing the architecture or its asymptotic performance. We have clarified this in the main text.\n\n> The limitations are…\n\nAs mentioned in the global review, we have further highlighted and discussed the limitations in separate paragraphs in the conclusion (while also retaining the “inline” discussion to retain flow). If the reviewer has further limitations we have not addressed, then please let us know, we are more than happy to include them! \n\nWe again thank the reviewer for evaluation, and for raising many points that have improved the paper. Please let us know if you have any further questions!\n\n— Submission 13253 Authors", "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": "41iAIJgbvL", "reviewer": "Reviewer_5uzK", "summary": "This paper presents improvements over the approach of Lim et al. last year for the parallelization of nonlinear recurrences in the sequence length. The paper jumps right into the problem after a quick introduction, and presents a few improvements on Deer: Quasi-Deer (an approximate method), IKE (stable but expensive), and Quasi-IKE (moderate stability, moderate memory and work).\nThe authors corroborate their improvements with some theoretical results and experiments", "questions": "above.", "limitations": "Well discussed by authors", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper flows very well and jumps right into the math: the authors have lots to show and I must say I was impressed by the quality and flow of the paper. While I have some questions and think this line of work does not end with this paper, I believe this contribution is worth publishing. I think the author correctly shows their improvements, and I particularly liked the quasi-deer approach. I also really like how the paper is hones on limitations and clearly provides a gateway to future research.", "weaknesses": "I will put here weaknesses and questions at the same time:\n\n1) Approximation quality: One missing plot is the approximation error: even in a toy setting and with random input, it is interesting to understand how well the methods can approximate the nonlinear RNN and at which cost. You have something similar in Figure 3, but this is not enough: I would be very curious to see, e.g., the evolution of the hidden state norm.\n\n2) Speed: My experience with non-diagonal RNNs in Jax using parallel scan is that associative scans on dense recurrences are slower than linear processing - in contrast to the diagonal setting where I usually observe something like a factor 10. You do comment a bit on this, but I am curious about how much it costs, in your framework, to do a parallel scan on dense matrices. This is of course a lower bound on your compute, right? I am not a hardware person, so this counts more as a curiosity, which I think should be better developed in the paper.\n\n3) What is this doing?: I think that, philosophically maybe, approximations sometimes work not because they are close to the object of interest, but for some other reason. This was, for instance, the case of S4 - where the ZOH discretization introduces some structure which ends up being the reason for its success. So I ask: can you present an example where you provide a precise computation of what your recurrence is doing? I think this would be super interesting, and perhaps connect your method to some other literature on SSMs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for taking the time to review our paper, for their positive feedback, and for providing some very interesting discussion points! \n\n> Approximation quality:\n\nThis is a great observation. We have added Review PDF Figure 1 probing the compute or computational  budget required to achieve different levels of accuracy in the AR GRU experiment. There are two major findings here: (a) the speed ups are fairly robust across elements *within* a dataset (i.e. once tuned the speedups are fairly low variance); and (b) quasi-methods took more steps, but were faster as each optimization step is faster. We also find that IKE methods converged faster for the AR GRU. However, we note this relationship is reversed in the additional Lorenz96 experiments (proposed by reviewer VSf9; see global review), where IKE and DEER converge in a similar number of steps, and quasi methods require notably more. This reinforces a point made in the discussion of the original submission that different methods may be faster or more stable in different settings, but that the memory savings and ultimate convergence are consistent. We have expanded discussion and guidance in the paper. \n\n> Speed:\n\nYes, dense matrices are much slower to process, and this matches our experience (the parallel scan on a dense linear recurrence requires $O(D^3)$ work which saturates the GPU). Table 2 in the paper shows the average time per step and average number of steps required to converge.  The key take-away is that individual steps in DEER/IKE are (approximately) a factor of between 3.5 and 30 times slower _per step_ than their quasi variants (which roughly tallies with your experience). However, they take (on average in Figure 4) a factor of between 2 and 10 fewer iterations. We have added some additional discussion of this tradeoff. \n\n> What is this doing?\n\nThis is a super interesting question. We have two notable observations about (re-)interpreting the methods we present.  Below we include a condensed proof on how IKE can be re-interpreted as stabilising the parallel scan by attenuating the eigenvalues used in the parallel scan, and hence stabilising their cumulative product.  We will add the full proof to the main paper.\n\nTo your question on quasi: we don’t have a particularly satisfactory answer.  Very early in our investigations when exploring quasi approximations, we noticed that the off-diagonal elements were often much smaller in magnitude than the diagonal elements (which makes intuitive sense for many recurrent models). This gave us some confidence to explore quasi approximations further, but we were unable to derive any deeper insight, and so validated it empirically. It is a really interesting avenue of future work looking into this, or other families of approximations and their relative merits and drawbacks.  \n\n**Summary**:  We again thank the reviewer for positive feedback, and for raising some really unique discussion points. Please let us know if you have any further questions!\n\n— Submission 13253 Authors\n\n\n—\n\n\n**Condensed additional proof that IKE attenuates the largest eigenvalues in linear recurrence**:  (We include a full proof in revised paper)\n\nLet $\\{J_t\\}$ be the Jacobians used in the linear recurrence relations and $\\{b_t\\}$ be the offsets. Then the prediction step of the Kalman filter (IKE) is the same as DEER. However, after applying the update step in IKE (which imposes the trust region), we obtain a second linear recurrence relation where the linear operator is given by $\\Gamma_t J_t$.  It can be shown that $\\Gamma_t$ is a symmetric positive definite matrix with eigenvalues bounded above by $\\frac{1}{1 + \\lambda}$. Thus, making using of the Spectral Theorem, it follows that the norms of the eigenvalues of $\\Gamma_t J_t$ are bounded above by the max of the norms of the eigenvalues of $J_t$, scaled by $\\frac{1}{1 + \\lambda}$.  Note that larger $\\lambda$ corresponds to more regularization/smaller trust region; and therefore correspondingly results in smaller effective eigenvalues in the scan.  We recover DEER exactly if $\\lambda = 0$.  Thus, while large eigenvalues in $J_t$ are the cause of the instability of DEER when evaluating unstable dynamical systems, IKE directly attenuates these large eigenvalues, explaining why the intermediate iterations using IKE remain stable.", "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": "eUMsJakKh2", "reviewer": "Reviewer_KXRw", "summary": "In this paper, the authors propose to improve DEER, a previous method that evaluates non-linear RNN in parallel by viewing it as a fixed point problem. Specifically, instead of using Newton's method to solve the fixed point problem, the authors leverage quasi-Newton methods to approximate the procedure. This alleviates the intractability of DEER. Moreover, they stabilize DEER by connecting it to Kalman smoothing technique. Empirical results show that the improved methods demonstrate superior performance in both speed and memory usage.\n\nStrengths\n\n1. The paper is well written and easy to follow.\n2. The proposed methods are very well justified.\n2. The proposed technique is reasonably designed.\n\nWeaknesses\n\n1. The experiments are thin.\n2. Although the proposed methods are novel, the adopted techniques are already well-known.\n\nOverall, I recommend a weak acceptance of this paper.", "questions": "See the weaknesses.", "limitations": "The authors have adequately addressed the limitations.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "1. The paper is well written and easy to follow.\n2. The proposed methods are very well justified.\n2. The proposed technique is reasonably designed.", "weaknesses": "1. The experiments are thin. In this paper, the authors evaluate the quality of different methods on only one dataset with a fixed model. It is hard to predict how the proposed methods will behave on other datasets with various RNN architectures.\n2. Although the proposed methods are novel, the adopted techniques are already well-known. For example, quasi-Newton methods and the Levenberg-Marquardt update are well-studied before. This weakens the novelty of this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Firstly we thank the reviewer for taking the time to review our submission and for their positive comments! We were particularly pleased by the comments on the foundation, clarity and presentation of our work. \n\nThe central reservation of the reviewer is that our experimental results are “thin”. In regard to this, we refer the reviewer to the general response and the supplemental review PDF, where we have included additional experiments.\n\nThe reviewer also comments that many of the components we discuss exist in the literature.  We ultimately agree with this statement; but argue that the real merit in our paper is assembling these components and exploring (theoretically and empirically) their relative strengths and weaknesses, and stratagems to operationalizing these components in the context of deep learning and non-linear systems.  As such, we believe the ties to existing work highlights the timeliness and utility of our paper – as we clearly build directly and pragmatically on existing literature that is of general interest to many in the community.  \n\nWe again thank the reviewer for their evaluation of our paper, their kind words, and for pointing to avenues to improve the paper. Please let us know if you have any further questions!\n\n— Submission 13253 Authors", "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": "gzAlE37pLt", "reviewer": "Reviewer_4jAU", "summary": "This paper aims to address the challenge of parallelizing the evaluation of nonlinear Recurrent Neural Networks (RNNs). Key contributions include the introduction of quasi-Newton approximations and trust region-based methods to reduce computational complexity and improve numerical stability, respectively. These enhancements enable reductions in memory usage and computational time, while maintaining the accuracy of the model's evaluations. The techniques presented are empirically validated, demonstrating their efficacy and potential application across various domains where RNNs are employed.", "questions": "- As I am not very familiar with this topic, I am wondering which equation will be speeding up using parallel computing. And how the method scales when parallelism is increasing or decreasing.\n- More experimental details are required such as how the parallelism is implemented and what kind of parallel computing is applied in the experiments. \n- Will the word \"parallel\" be proper for this type of method? Initially, I thought the parallelism was w.r.t the sequential length from the title. However, it seems that parallelism is applied for the quasi-newton method, which to me is equal to speeding up the computation in optimization with parallel computing. How it is related to parallelizing RNN in terms of the sequential length?", "limitations": "The authors have not adequately addressed the limitations.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The integration of quasi-Newton methods provides a fresh perspective on overcoming the challenges in parallelizing the evolution of nonlinear RNNs. The theoretical insight extends the scalability of RNNs soundly.\n\n- By employing quasi-Newton approximations, the paper effectively reduces the computational complexity from cubic to nearly linear with respect to the sequence length. This reduction is critical for deploying RNNs in resource-constrained environments.\n\n- The introduction of trust regions and their integration with Kalman smoothing techniques stabilizes the Newton iterations. This is a notable advancement over traditional methods, which often suffer from numerical instabilities due to undamped Newton updates.\n\n- The methods are not only theoretically sound but are also empirically validated to demonstrate their practical effectiveness. This includes significant reductions in memory usage and computational time, crucial metrics for real-world applications.", "weaknesses": "- The paper's empirical validation, though robust, is limited to specific architectures and scenarios. Expanding the testing to a wider range of nonlinear RNN types and more diverse datasets could provide a more comprehensive understanding of the methods’ applicability and limitations.\n- The implementation of quasi-Newton and trust-region methods might be complex and could require significant modifications to existing neural network training pipelines. This complexity could hinder widespread adoption without additional tools or simplified frameworks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Firstly, we thank the reviewer for taking the time to review our submission and for their positive comments! We were particularly pleased by the comments on the foundation, clarity and presentation of our work. \n\n> Weaknesses\n\nWe discuss these in global review (“empirical validation” in Section 2 and “implementation” in Section 1.1). Broadly, regarding “empirical validation” we appreciated your perspective to bring our robust experimental validation to other datasets, and so took the suggestion of Reviewer VSf9 and demonstrated our methods on the chaotic Lorenz-96 system. Regarding “implementation,” we note that our method can be dropped into any stateful architecture.\n\n> As I am…\n\nThis is a very interesting point, that we probably didn’t discuss in enough detail in the paper. The key insight of DEER was remapping the matrix inversion (a prohibitively expensive operation) required by the Newton iteration into a parallel scan. This allows the “inverse” to be solved efficiently and while exploiting parallelism. However, the parallel scan incurs $O(TD^3)$ work, and hence does not scale well (and is prone to instabilities due to the unboundedness of the individual Jacobians). Our insight was to introduce quasi-approximations into this, to reduce this complexity to $O(TD)$, while proving this retains global convergence (c.f. Proposition 1).  (This suggested the “resetting” heuristic for combating instabilities; and also led us towards the connection with trust region stabilisation.) \n\nThe parallel scan, given sufficiently many processors, scales as $O(\\log T)$. As we show in Figure 6, we see this speedup at low model sizes and sequence lengths. Once the processors are saturated, we see a linear increase in the runtime (since the amount of work we do is linear), but it is making much more effective use of the GPU, resulting in a constant factor speedup over sequential application at larger model sizes/sequence lengths. \n\nWe have added an expanded discussion of these points, thank you for highlighting them.\n\n> More experimental details…\n\nThe JAX library includes an implementation of parallel scan which we use off-the-shelf. As a result the implementation (with regard to parallelism) is very straightforward. We included code in the original submission and are also preparing a full public code release. We have also added more tabulated hyperparameters, configs, runtimes etc; and added a paragraph explaining parallel scans in detail and how we leverage them. \n\n> Will the word…\n\nWe believe “parallel” is the correct word. Crucially, we pose the inherently sequential operation as an optimization over the length of the sequence, where each optimization step is parallelizable (sub-linear parallel application time in the sequence length, $O(\\log T)$, to be exact). We have reinforced this point in the introduction. If we have misunderstood the reviewers query, then please let us know! \n\n> The authors have…  \n\nWe refer the reviewer to the general response in regard to discussion of limitations. We have added much more explicit discussion of the limitations in a single place to correct this shortfall. If the reviewer has additional points they would like us to include, then please let us know, we are more than happy to include them! \n\nWe again thank the reviewer for taking the time to review our submission and for raising some really interesting points of confusion that we had not foreseen. Fixing these has definitely made the paper better. Please let us know if you have any further questions!\n\n— Submission 13253 Authors", "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:
Firstly, we thank all six reviewers for their positive feedback and insightful comments. 

We present methods for parallelizing the evaluation of non-linear RNNs, building on a recent method, DEER, from Lim et al. We first proved the global convergence of DEER. We then ameliorate DEER’s two major weaknesses: cubic computational complexity and instability. We achieve this by introducing quasi-approximations and trust regions (IKE), respectively. We evaluate all methods and find that quasi and IKE variants match or outperform DEER across a range of metrics, including accuracy, wall clock time, memory and iterations.

We first comment on the positives highlighted by the reviewers: 
- **Well Received**: We feel all reviewers understood and commended the intention and merits of our submission.
- **Foundations and Connections**: Many reviewers noted how well-founded our approach is and how it dovetails nicely with recent work. 
- **Quality of Presentation**: We were especially pleased by the consistent praise for the quality of our submissions presentation.

There were some critical themes shared in some reviews. We comment on these here:
 
1.  **Discussion of Limitations**: Despite several comments praising our discussion of limitations, some reviewers commented we did not adequately address the limitations of the methods we present. In a bid to create better flow we attempted to “inline” the limitations, addressing them as they arose. However, we clearly missed the mark here. We have added an additional paragraph to the discussion explicitly repeating those limitations (we do note that we did discuss these limitations in the original submission, but appreciate the reviewers emphasis on highlighting limitations):

    a. Quasi methods lose the quadratic convergence properties of Newton (but we show empirically that convergence is *faster in terms of wallclock time* as a result of the efficient approximation). 

    b. Although motivated by Proposition 1, the quasi approximation can be significant. 

    c. IKE stabilises evaluation, but, like DEER, has cubic complexity in the state dimension (quasi-IKE then combats this).

    d. The heuristic of resetting to zeros when unstable is motivated by Proposition 1, but does slow convergence in (q-)DEER methods. 
  
    e. (q-)IKE adds an additional hyperparameter ($\lambda$). 

    f. Extended discussion on when each technique is likely to yield the best results.

    g. Architectural/implementation limitations (see more below).

1.1 **On Architectural/implementation limitations**: (q-)DEER and (q-)IKE can be applied to _any_ stateful architecture with state in $\mathbb{R}^D$ with _minimal-to-no_ modifications to the overall pipeline. All methods effectively define a different `forward` method for the model with identical interfaces. Sequential/(q-)DEER/(q-)IKE can be switched between simply by changing a flag in the model class. (q-)DEER and (q-)IKE are then implemented in a model-agnostic manner. 

The main architectural/implementation limitation, at least at the time of writing, is that Torch does not have an implementation of parallel scan. This is a major reason we wrote our code in JAX. However, with the explosion of parallel-scan-powered methods, it is our understanding that the Torch team is actively developing a parallel scan implementation. 

2. **Experiments**: Several reviewers commented on the experimental evaluation – with some excellent suggestions of ways to strengthen them. Off the back of these suggestions, we have included some further experimental evaluation:

- **Review PDF Figure 1**: In response to the comments of S8CH, VSf9 and 5uzK, we explore $\lambda$ in (quasi)-IKE for the AR-GRU (Left Column). We sweep over $\lambda$ for 15 different input sequences, and plot the median and quartiles of the cost to convergence in terms of Newton iterates and runtime. We see a bathtub curve (large $\lambda$ takes needlessly small steps, slowing progress; small $\lambda$ results in many resets, slowing convergence).  Crucially, we see there is little variance across individual sequences. These results show that there is a well-behaved $\lambda$ dependence that can be optimised on a validation set with a simple 1-d grid search.

Furthermore, In response to the comments of S8CH, VSf9 and 5uzK, we chart the approximation error vs cost for the AR GRU (Center and Right Column). We see that the approximation error reduces in fewer Newton steps with full DEER as opposed to quasi-DEER, but, crucially, the wallclock time (the more important of the two metrics!) is notably lower across all accuracies for quasi-DEER. This indicates that our more efficient – but approximate – quasi-DEER is broadly preferable to the more expensive – but exact – DEER updates. Furthermore, the stabilised IKE and quasi-IKE are better still. We also show the steps/time to convergence for a range of accuracy thresholds, and see that our methods outperform DEER across the full range of thresholds and metrics.  

- **Review PDF Figure 2**: In response to the comments from VSf9, KXRw and 4jAU, we include an extra experiment applying (q-)DEER and (q-)IKE to approximate the chaotic Lorenz-96 system. We demonstrate that all the parallelized methods converge to the correct trace, but that (q)-IKE are dramatically more stable at intermediate Newton iterations prior to convergence. 

> Explicit pointer to full proof of Prop1 in Appendix A.1

We noticed that we forgot to include a pointer in the main text to the full proof of Prop1, which can be found in Appendix A.1 of the original submission. We now include this explicit pointer.

**Summary**: We hope that the expanded and clarified discussion of limitations and these additional experiments allay the reviewers' concerns, and help further elucidate the relative benefits of the methods we introduce. Thank you again, and please do not hesitate to reply if there are further clarifications we can make!

— Submission 13253 Authors

Meta-review: Parallelized RNNs require solving a fixed-point equation. A previous study proposed solving this equation using Newton's method. This paper improves the computational cost and stability of Newton's method for the parallelization of RNNs. To reduce computational cost, the paper proposes a quasi-Newton method with significantly lower per-iteration complexity. To enhance stability, a trust region-based method combined with Kalman smoothing is utilized. According to reviews, the paper is well-written, and the contribution is sound. To address concerns about the experiments, the authors provided additional experimental results that elaborate on the stability enhancements. While there are concerns about comparing the proposed method with DEER, it is difficult to establish a good convergence rate for quasi-Newton methods, which makes the convergence-cost tradeoff challenging. According to reviewers, the contributions are sufficient for acceptance.
Decision: Accept (poster)