Towards Scalable and Stable Parallelization of Nonlinear RNNs

Authors: xavier  gonzalez
andrew  warrington
jimmy t h smith
scott w linderman
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 03:44:16

Abstract: Conventional nonlinear RNNs are not naturally parallelizable across the sequence length, unlike transformers and linear RNNs. Lim et. al. therefore tackle parallelized evaluation of nonlinear RNNs, posing it as a fixed point problem solved with Newton's method. By deriving and applying a parallelized form of Newton's method, they achieve large speedups over sequential evaluation. However, their approach inherits cubic computational complexity and numerical instability. We tackle these weaknesses. To reduce the computational complexity, we apply quasi-Newton approximations and show they converge comparably, use less memory, and are faster, compared to full-Newton. To stabilize Newton's method, we leverage a connection between Newton's method damped with trust regions and Kalman smoothing. This connection allows us to stabilize the iteration, per the trust region, and use efficient parallelized Kalman algorithms to retain performance. We compare these methods empirically and highlight use cases where each algorithm excels.

Main paper: 
 introduction: 
 Parallel computation has helped fuel the rise of deep learning [30] . Architectures such as transformers [73] and linear RNNs [31, 65, 53, 27, 21] are specifically designed to allow parallelization over the length of the input sequence. However, most conventional nonlinear RNNs (e.g. Elman RNNs, GRUs [15] , LSTMS [29] 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 nonlinear RNN architectures are still used widely across the scientific community [44, 62, 34] . Furthermore, recent work has suggested that linear RNNs (and transformers) are fundamentally limited in their expressivity compared to nonlinear RNNs [43] . Finally, nonlinear RNNs continue to be of significant interest in computational and theoretical neuroscience as models of neural systems [74, 69, 54, 64, 60, 19, 18, 17] . Therefore, scalable and stable parallelization methods for nonlinear RNNs offer significant benefits across many fields. Towards this goal, Lim et al. [36] proposed DEER, a method for evaluating a nonlinear RNN in parallel. DEER casts inference as finding the solution of a 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 method enjoys quadratic convergence rates [49, Chapter 11] . Lim et al. [36] also show that the inversion of the structured Jacobian matrix required by Newton's method can be cast as an associative parallel scan [7] . DEER therefore reduces the evaluation runtime over sequential evaluation by as much as factor of twenty.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. From Scalability: + Diagonal Jacobian DEER [36] Quasi-DEER ELK Quasi-ELK Stability: + Trust region + Kalman filter Figure 1 : Overview of the parallelizable methods we consider in this paper. We introduce diagonal approximations to improve complexity (quasi-DEER, Section 4.1) and link to Kalman filtering and trust regions to improve stability (ELK, Section 4.2). We combine these ideas in quasi-ELK (Section 4.2).Table 1 : Description of the relative strengths and weaknesses of the five evaluation methods we consider. We include a discussion of this in Section 7. using a parallel scan to evaluate updates from Newton's method, DEER inherits O(T D 2 ) memory complexity and O(T D 3 ) computational work [7] . These costs can be prohibitive in practical deep learning settings. The second limitation of DEER is numerical stability, inherited from Newton's method. In general, undamped Newton's method does not provide global convergence guarantees, and in practice often diverges [49] . We seek to ameliorate both these weaknesses.To do this, we leverage two techniques: quasi approximations and trust regions. Quasi approximations are a common adaptation of Newtons method, where approximate, but faster and less memory intensive updates, are used in-place of exact "full" Newton steps. Empirically, these are often observed to expedite convergence in terms of wallclock time, even though more Newton iterates are used. We apply quasi-approximations to remove the memory and compute scaling inherited by DEER, also finding accelerated convergence and reduced memory consumption. Secondly, we leverage a connection between Newton's method with a trust region and Kalman smoothing in sequential models [71] . This 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, while still being able to use parallelized Kalman smoothers [59, 12] , achieving a parallel runtime that is logarithmic in the sequence length. We refer to DEER accelerated with a quasi approximation as quasi-DEER, and DEER stabilized with trust regions as "Evaluating Levenberg-Marquardt via Kalman" (ELK). We then combine these yielding a fast and stable algorithm, which we term quasi-ELK.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: We first introduce background material, particularly focusing on DEER [36] , in Sections 2 and 3. We then present three short novel proofs: that DEER is globally convergent; that this convergence is robust to modifications of the linearized dynamics (Proposition 1); and that there is a unique solution with no local minima (Appendices A.1 and A.2). We then introduce quasi-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 (Section 4.2 and 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://github.com/lindermanlab/elk.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 in the abstract and introduction is that we improve the limitations of parallelizing the evaluation of nonlinear RNNs with Newton's method by addressing scalability and stability concerns. We develop quasi-DEER to 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 contributions made in the paper and important assumptions and limitations. A No or 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. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
problem statement: 
 We consider nonlinear Markovian state space models, with the state at time t denoted s t ∈ R D and nonlinear transition dynamics f : R D → R D . We denote the full sequence of T states as s 1:T ∈ R T ×D . Note that we will be mainly considering the transition dynamics in this paper, and so we suppress any (possibly random) input dependence of the model in the notation. Note however the algorithms in this paper extend trivially to these situations.For any collection of candidate states {s t } T t=1 and an initial state s 0 we can define the residual r(s 1:T ) := [s 1 -f (s 0 ), s 2 -f (s 1 ), s 3 -f (s 2 ), . . . , s T -f (s T -1 )] ∈ R T ×D .(This residual can be interpreted as the one-step error incurred by assuming the t th state is s t instead of f (s t-1 ). The solution of the state space model, s * 1:T , is the only trace with zero residual. Equivalently, it is the unique solution to the fixed point equation r(s * 1:T ) = 0.(EQUATIONThe conventional way of obtaining s * 1:T is to apply f sequentially T times. Sequential evaluation always yields a valid trace, but it requires O(T ) sequential operations (i.e. computational depth or span), and hence does not fully leverage the capabilities of parallel hardware. We aim to compute s * 1:T in sublinear time using parallel computation.Jacobian of the Residual For notational brevity, we overload s and r to also denote vectors in R T D , representing flattened versions of s 1:T and r 1:T . We can therefore write the Jacobian of the residual for the whole sequence, J(s), as a T D × T D matrix with block bidiagonal structure of the formEQUATION)3 DEER: Newton's Method for Parallel Evaluation of Sequential ModelsLim et al. [36] propose DEER, an algorithm using Newton's method for parallel evaluation of nonlinear sequential models, including both discrete-time nonlinear RNNs (GRUs, LSTMs, etc.) and neural ODEs [13, 32] . In this paper, we focus on the discrete-time setting, and address questions that arise from Lim et al. [36] : how to scale Newton's method, and how to make it numerically stable.In this section we introduce DEER. We begin with a simplified derivation that emphasizes the link between Newton's method on vector spaces and parallelizable linear recurrences. We then present a 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 to discuss the weaknesses of DEER, and to motivate the methods we develop in Section 4.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 to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations 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. • 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 a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in 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, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use 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. 
derivation of deer from newton's method: 
 The original derivation of DEER used Newton's method on Banach spaces and the Frécht derivative for continuous-time systems to derive the update and show convergence [36] . 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 i th Newton iterate for (2), starting at s (i) , is given byEQUATIONor equivalently,EQUATIONNote this uses the root-finding view of Newton's method, see Appendix C.2.Provided all of { ∂f /∂s} T t=2 are finite, then the Jacobian defined in (3) is invertible and all of the eigenvalues are equal to one. 1 Storing and naively inverting the Jacobian is infeasible for large D or T . However, since J(s) is block bidiagonal, we can solve for ∆s in (5) by forward substitution. This reduces to a simple recursion with the initial condition ∆s (i+1) 1 = -r 1 (s (i) ), and for t > 1, ∆sEQUATIONDEER uses the linearity of this recursion, soling it in parallel with a parallel associative scan [36, 7, 65] . Therefore, with O(T ) processors, each Newton iteration can be performed in O(log T ) time.We emphasize that the computation of the Newton step ∆s in ( 5) is being parallelized. J would, in general, be a T D × T D matrix that is prohibitive to store or invert. But by formulating this solve as an LDS in (6) , we are able to parallelize the computation of ∆s (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 O(log T ) time. We use the parallel scan from JAX [9] (see Appendix B.6).We present a proof that DEER converges globally for discrete-time RNNs to the solution s * 1:T of (2) in at most T steps. Proposition 1. Undamped Newton's method will converge to the true solution, s * 1:T , of the fixed point (2) in at most T Newton iterations, for any initial s (0) 1:T . Proof sketch. For the full proof by induction, see Appendix A.1. The structure of J(s) determines the recurrence in (6) . The update applied at time t, ∆s (i+1) t , from (6) is the summation of a linearized f applied to the update at time t -1, and the residual one-step error at time t. Therefore, if the previous timestep is correct (i.e. ∆s (i+1) t-1 = 0), then the update at time t 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 s 0 are fixed and known, the result follows that all T timesteps will have zero residual after T iterations.It is not immediately obvious from (4) or the proof given by Lim et al. [36] that DEER converges globally, but Proposition 1 shows that it in fact does, at least theoretically. This result has two crucial corollaries. First, after i Newton iterations, s (i) 1:T will have zero error for all t ≤ i. Therefore, if the iteration encounters numerical instabilities, as Newton is prone to, we can simply use a heuristic of resetting s (i) t to a finite value for all t > i. This preserves the solution for time indices t ≤ i and allows the optimization to continue, but it is equivalent to running Newton's method from scratch on s i:T . This process is repeated until the entire trace has zero residual. A second corollary is that any set of finite matrices can replace { ∂f /∂s} T t=2 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.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, ∂f /∂s is a dense D × D matrix. Therefore the parallel associative scan, which uses matrix-matrix multiplications, has O(T D 3 ) computational work and O(T D 2 ) memory complexity. Stability is limited because we often have no control over the eigenvalues of ∂f /∂s. If sufficiently many of these eigenvalues 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 convergence is dramatically slower. These weaknesses motivate the development of two new techniques for parallelized evaluation of RNNs, quasi-DEER and ELK, which we discuss in the next section.Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 crossreferenced. • 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 a 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. 
scaling and stabilizing newton's method for parallel evaluation: 
 In Section 4.1 we introduce quasi-DEER, a quasi-Newton method that addresses the intractability of DEER for large state sizes. In Section 4.2 we introduce Evaluating Levenberg-Marquardt 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 a damped Newton's method for parallel sequential evaluation that is scalable and numerically stable.As a consequence of our Proposition 1, replacing the Jacobians { ∂f /∂s} T t=2 with an arbitrary matrix will still result in global convergence of the resulting DEER-like algorithm in at most T iterations. A straightforward way to reduce the computational cost is to replace { ∂f /∂s} T t=2 with {diag ( ∂f /∂s)} T t=2 , i.e. take the diagonal entries of the Jacobians of the dynamics functions. The resulting linear recursion requires only O(DT ) memory because it only needs to store diagonal matrices, and O(DT ) work, because the parallelized associative scan only uses element-wise vector multiplies. Position-wise matrix-vector multiplies are still required to obtain the residuals, but this computation can be embarrassingly parallelized across the sequence.Quasi-Newton methods approximate the Jacobian for computational reasons, so we refer to this algorithm as 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. [36] . Quasi-DEER improves the scalability of DEER, but it does not address stability concerns. We propose a more stable solution below.Rather than treating RNN evaluation as a fixed point finding problem, let us instead consider it as an optimization problem. First, we define the merit functionEQUATIONAs in the fixed point formulation, the unique minimizer of this objective is s * . In fact, the only local minimum of the merit function ( 7) is s * , as proved in Proposition 3 in Appendix A.2. One way of minimizing this nonlinear sum of squares objective is via the Gauss-Newton algorithm [49] , which alternates between linearizing the terms in the merit function and solving the resulting linear sum-of-squares problem. The linearized objective at iteration i isEQUATION)The solution is ∆s (i+1) = -J(s (i) ) -1 r(s (i) ), which is exactly the DEER update from (5) .Formulating evaluation as nonlinear least squares also suggests more stable algorithms. The Levenberg-Marquardt algorithm [49] uses updates that solve a constrained optimization problemEQUATIONwhere D i+1 is an upper bound on the step size. We recognize this constraint as a trust region, which is often used in conjunction with Newton's method to improve numerical stability and convergence.Finally, minimizing this constrained optimization is equivalent to minimizing the LagrangianEQUATIONfor some λ i+1 ≥ 0. As noted by Särkkä and Svensson [71] , 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-Marquardt update that minimizes (10) with fixed λ i+1 is equivalent to finding the maximum a posteriori (MAP) estimate of s 1:T in a linear Gaussian state space model, which can be done in O(log T ) time on a sufficiently large parallel machine.Proof. Expanding the residual and Jacobian functions in (8) , we see that up to an additive constant, the negative Lagrangian can be rewritten as,EQUATIONwhere N (x | µ, Σ) denotes the probability density function of the multivariate normal distribution.We recognize (11) as the log joint probability of a linear Gaussian state space model (LGSSM) [70] on (s 1 , . . . , s T ). The means of the dynamics distributions are given by the linearization of f , and the emissions are the previous iteration's states, s (i) . The parameter λ i+1 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 [70] . As with the linear recursions in DEER, the Kalman smoother can be implemented as a parallel scan that scales as O(log T ) in time on a machine with O(T ) processors [59, 12] .Therefore, we can evaluate an RNN by minimizing the merit function with the Levenberg-Marquardt algorithm. Since each step of the algorithm is performed by parallel Kalman smoothing, we call this approach Evaluating Levenberg-Marquardt with Kalman (ELK). Note that DEER is a special case of ELK, where λ = 0, 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 [49, Thms. 11.7, 11.8] .As with DEER, we can substitute an approximate Jacobian into the Lagrangian to obtain the quasi-ELK algorithm. Quasi-ELK 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 to converge than ELK, each quasi-ELK iterate is faster, giving overall runtime speedups.The convergence rate of (quasi-)ELK depends on the trust region radius D i (or alternatively λ i ). While there exist methods to analytically set λ i [49, Algorithm 4.3], these approaches require factorizing ∂r /∂s, which is intractable at scale. Therefore, we treat λ as a hyperparameter set by a sweep over log-spaced values (cf. Appendix B.4).We also use Kalman filtering instead of smoothing. We do so for two main reasons: filtering requires less work and memory; and we also found it to converge in fewer Newton iterations than smoothing. We believe this is a result of Proposition 1, where the early part of the trace converges first. In Appendix A.3 we also discuss a connection that reinterprets ELK and the Kalman filter as defining a linear recurrence where the trust region attenuates the eigenvalues used in the parallel scan.Limitations The quasi-methods lose the local quadratic convergence properties of Newton (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 zeros when unstable is also motivated by Proposition 1, but does slow convergence in (quasi-)DEER methods. As a result, we develop ELK to stabilize evaluation, but, like DEER, ELK has cubic complexity in D (quasi-ELK then combats this). However, quasi-ELK adds an additional 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 wallclock 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 wallclock 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 at time of writing is most easily implemented in JAX [9] .Question: 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 conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We provide experimental details to allow for reproducibility, including hardware used, in Section 6 and Appendix B. We also provide our code at https://github. com/lindermanlab/elk Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a 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 a novel architecture, describing the architecture fully might suffice, or if the contribution is 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 is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. Justification: We have read and followd the NeurIPS code of ethics. We do not use human subjects; we use publicly available datasets; and our research is in ways to accelerate standard machine learning algorithms so their broader impacts are to allow current machine learning techniques to be more scalable and stable.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 a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
related work: 
 RNNs and Parallelism Nonlinear RNNs are a natural choice for modeling sequential data because 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 architectures that inherently admit parallelism have been proposed, including transformers [73] , deep linear RNNs [41, 31, 28, 65, 53, 27, 4, 21] , and convolutions [51, 57, 56, 55] . These obtain parallelism by developing new architectures, and do not consider parallelizing existing nonlinear architectures. DEER [36] is notable as it considers parallel evaluation and training of arbitrary nonlinear RNNs.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 [33, [1] [2] [3] and neural ODEs [13, 42] replace conventional feed forward network layers [26] with an implicit layer whose output is the root of an equation. Moreover, Song et al. [66] parallelizes the evaluation of feedforward nets using Jacobi and Gauss-Siedel iterations. In sequence modeling, parallel decoding methods [45, 58, 23] adapt ideas from Jacobi and Gauss-Siedel iterations to evaluate autoregressive sequence models in parallel. These approaches iteratively refine inputs by repeatedly feeding in previous outputs back into a parallelized sequence model. However, these methods presuppose the existence of a parallelized forward pass for the sequence model and do not leverage additional gradient information to obtain sublinear convergence.Parallelizing Dynamical Systems over Time Other work has investigated evaluating other nonlinear dynamical systems over time. ParaDIGMS [63] parallelizes sampling from diffusion models, but uses Picard iterations instead of Newton's method. In the numerical ODE and PDE communities there has been great interest in Parallel in Time methods, see Gander [24] , Ong and Schroder [50] for surveys. Vargas et al. [72] parallelized the evaluation of chaotic dynamical systems over time, but instead of casting Newton's method as a parallel scan resorts to multigrid 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 are efficient algorithms that use an approximation of the Jacobian or Hessian in Newton's method, and include approaches like BFGS [11, 22, 25, 61] and L-BFGS [37] . Other approaches use Newton's method to optimize deep nets [40] . However, these quasi-Newton algorithms do not admit efficient parallel scans. There are also conjugate gradients methods for exploiting structured Jacobians or Hessians [68] , though they often do not attain the fast convergence rates of Newton or quasi-Newton methods [49] . Methods for stabilizing and ensuring Newtons method converges globally include regularization approaches [48, 20] , backtracking line search [52] , and trust regions [16] . All these stabilization methods have strengths and weaknesses, but as noted by Nocedal and Wright [49] : "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 [6] and Bell [5] draw connections between the Gauss-Newton method and the iterated extended Kalman filter and smoother [67, 70] . Because Gauss-Newton is unstable, it is natural to use Levenberg-Marquardt [35, 39] to stabilize the filtering/smoothing problem [14, 38, 71] . These approaches start with a smoothing problem, and stabilize it using approaches from nonlinear equations; whereas we start with a nonlinear equation to solve, and make the connection with Kalman filtering to leverage parallelized algorithms [59] .We also address the practicalities of operationalizing this connection for modern deep networks. 
experiments: 
 We now experimentally examine the relative performance of these methods. Specifically, we evaluate whether quasi-DEER can provide memory savings over DEER and runtime savings over sequential evaluation, while retaining the accuracy of training and evaluation; and whether ELK and quasi-ELK can be used to stabilize evaluation in regimes where DEER is unstable. In Sections 6.1 and 6.2 we use experimental designs from Lim et al. [36] and show quasi-DEER retains 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 unstable, and that quasi-ELK is the fastest of all parallelized methods. We provide further details in Appendix B. 30K 100K 300K 1M 30K 100K 300K 1M 30K 100K 300K 1MSequence Length (T) We first use an experimental design from Lim et al. [36] . The task is to evaluate an untrained GRU across a range of hidden state sizes (D) and sequence lengths (T ) on a 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 are 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 see that quasi-DEER requires as much as an order of magnitude less memory than DEER, thus allowing the application to architectural regimes previously infeasible with DEER. In Figure 6 of Appendix B.1.1 we show that in smaller T and D regimes we observe the expected sublinear time scaling with sequence length. This experiment confirms that quasi-DEER can replicate the performance of DEER, but with a smaller memory footprint.We verify that quasi-DEER expedites training nonlinear RNN models. We replicate the third experiment from Lim et al. [36] , where a GRU is trained to classify C. elegans phenotypes from the time series of principal components of the worms' body posture [10] .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 a factor of 2.5, despite performing 2-2.5 times more Newton iterations per training step. This highlights how quasi-DEER can replace DEER when training nonlinear RNNs, bringing both time and memory savings. In our experiment, we use the quasi-approximation for the backwards pass as well, leading to gradients that are different from DEER in this setting, but we show empirically that there is negligible degradation in performance (Figure 3 , Left).DEER is prone to "spikes", where orders of magnitude more steps are required for convergence (Figure 3 , Middle). While quasi-DEER is not as susceptible to these spikes (never more than half an order of magnitude), these instabilities motivated the study of stabilizing methods. We conclude by studying an application where these numerical instabilities in DEER are critical. We use a small autoregressive GRU (hidden dimension N h = 3), where the previous sampled value is input into the GRU at the next step. Such autoregressive architectures were not examined by Lim et al. [36] , 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 s t 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. [36] ) for evaluating a 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 a 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 a noisy sine wave (see Figure 4 , Bottom). The linear recurrence relation (6) was numerically unstable in DEER and quasi-DEER. To remedy these instabilities, we take the approach described earlier of setting the unstable parts of the 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), slowing convergence (see Figure 4 , Top Right). This finding highlights how the instabilities of DEER -which are inherited from both pathologies of Newton's method and the parallel recurrence -can be crippling in even very simple scenarios. While resetting allows for convergence, the resulting convergence is very slow.We then apply ELK and quasi-ELK. We show the results in the top right and bottom panels of Figure 4 . We select the trust region size with a one-dimensional search over log-spaced values between 10 0 and 10  
conclusion: 
 In this paper we proposed methods for scalable and stable parallel evaluation of nonlinear RNNs. DEER [36] achieved speedups over sequential evaluation, but incurred quadratic memory, cubic work, and numerical instabilities. We therefore extended DEER to use quasi-Newton approximations, reduing computational complexity; and provided a novel proof that both DEER and quasi-DEER converge globally. To stabilize DEER, we leveraged a connection between the Levenberg-Marquardt method and Kalman smoothing to enable parallel evaluation of RNNs, allowing us to stabilize DEER while still leveraging fast parallel filtering. We verified empirically that quasi-DEER, ELK, and quasi-ELK 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 observed, our experiments show that quasi-DEER provides the fastest convergence in terms of wallclock time. However, if the dynamics are on the edge of stability, then of the parallelized methods, ELK offers the most stable performance, but quasi-ELK could be faster in wall-clock time and just as stable. In such settings, it is 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 of the Jacobian that still admit fast parallelism but are more faithful approximations may allow for more accurate quasi steps to be 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 further improvements to parallelized methods can prove faster than sequential evaluation for dynamical systems on the edge of stability, or whether there are fundamental limitations to the computational benefit of parallelization, are interesting questions for future work. 
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 to 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 be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 as 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, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: We have written an algorithms paper. We do not produce a pretrained language model, an image generator, or a scraped dataset. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or 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 not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets 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. [36] , whose work and code was an inspiration for this paper, and Chang et al. [12] , 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 (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.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 package should be provided. For popular datasets, paperswithcode.com/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 documentation provided alongside the assets? Answer: [Yes] Justification: We provide our code at https://github.com/lindermanlab/elk 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, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS 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. 
 
 Appendix: A.1 Proof of Proposition 1Proposition 1 Undamped Newton's method will converge to the true solution, s * , of the fixed point (2) in at most T Newton iterations, for any initial s (0) .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 s 0 and the dynamics functions f are fixed. Therefore, s * 1 = f (s 0 ) and in general s * t = f (s * t-1 ). Thus, it follows from the initial condition of the DEER recurrence relation that s(i) 1 = s * 1 for all i ≥ 1.Furthermore, we observe that if s(i) t= s * t for all t less than some t (i) , then r t (s (i) ) = 0 for all t < t (i) by the definition of the residual in (1) . Therefore, the DEER linear recurrence relation necessarily gives ∆s(i+1) t = 0 for all t < t (i) . Furthermore, because s * t (i) = f (s (i) t ), it follows that ∆s (i+1) t (i) = -r t (i) (s (i) ) = s * t (i) -s (i) t (i). Thus, it follows that after applying another Newton iteration that s(i+1) t = s *t for all t < t (i) + 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 1 of Shih et al. [63] . However, Shih et al. [63] proves a result about Picard iterations (zeroth order method). Our proof about the global convergence of Newton iterations contains the additional complication of dealing with a LDS (as a consequence of using a first order method). Note, that the global convergence comes from the zeroth order update; however, the first order gradient information can accelerate convergence.Discussion of corollaries of Proposition 1 Corollaries of this proposition include that the method will converge (in at most T iterations) to s * even if the Jacobians ∂f /∂s are replaced by arbitrary matrices, and that we can reset the values of s (i) t for t > i arbitrarily and still enjoy convergence. In more detail, the elements in the sub-block-diagonal of J := ∂r /∂s 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 ( 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 ∆s (i+1) t-1 = 0 (no update is required at position t -1), and so the update to s (i+1) t no longer depends on the Jacobian.We exploit this to develop quasi-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 (cf. Figure 7 ).Explicitly, the global convergence of quasi-DEER is a theoretical result (a corollary of Proposition 1), but the fast runtime of quasi-DEER in practice is an empirical result (cf. Figure 3 ). Proof. First, we observe that ∇L(s) = J(s) T r(s), where J(s) is defined as in (3) . Because J(s) is a lower triangular matrix with all entries on its diagonal equal to 1, it follows that all of its eigenvalues are equal to 1. Therefore, J(s) is nonsingular for all s. Thus, J(s) has trivial nullspace for all s, i.e. J(s) T r(s) = 0 ⇐⇒ r(s) = 0. But only s * satisfies r(s * ) = 0.Since there are no critical points other than s * , the merit function cannot have local minima or saddle points.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 a fixed sequence of states and outputs (note that any stochastic dynamics function can be reparameterized as deterministic by conditioning on the random inputs). Thus, s * is the only sequence with zero residual (i.e. 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 Proposition 1. Another way to see this however is that each update step ( 4) is equal to J -1 r. But, J is always invertible and so has trivial nullspace. Furthermore, the residual r can only be zero at the unique solution s * . 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.A complementary perspective on how ELK results in more stable evaluation of nonlinear RNNs is to see how the Kalman filter damps the eigenvalues of the Jacobian matrices of the transition dynamics. We first provide a high-level overview, and then provide a more detailed derivation.Overview Let A t be the Jacobians ∂f /∂s used in the linear recurrence relations and b t 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 a second linear recurrence relation where the linear operator is given by Γ t A T . Note that Γ t is a symmetric positive definite matrix with eigenvalues bounded above by 1 /1+λ. Thus, by the Spectral Theorem, it follows that the norms of the eigenvalues of Γ t A t are bounded above by the max of the norms of the eigenvalues of A t , scaled by 1 /1+λ. 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 t are the cause of the instability of DEER when evaluating unstable dynamical systems, ELK directly attenuates these large eigenvalues, explaining why the intermediate iterations using ELK remain stable.Derivation We define our dynamics used in Newton iteration (i + 1) asA t = ∂f t+1 ∂s (s (i) t ) b t = f t+1 (s (i) t ) - ∂f t+1 ∂s (s (i) t )s (i) t . Now A t ∈ R D×D and b t ∈ R D .In line with considering the system as the LDS in (11) , we set the process noise to be I D , and with the emissions governed by s(i+1) t ∼ N (s (i) t , σ 2 I D ), where σ 2 controls the size of our trust region (note that in the notation of developed in Section 4.2 we have λ = 1/σ 2 ).In the notation of Murphy [46] , we see that the predict step isµ (t+1)|t = J t µ t|t + b t Σ (t+1)|t = A t Σ t|t A T t + I D .Meanwhile, the update step isµ (t+1)|(t+1) = µ (t+1)|t + Σ (t+1)|t (Σ (t+1)|t + σ 2 I D ) -1 (y t+1 -µ (t+1)|t ) Σ (t+1)|(t+1) = Σ (t+1)|t -Σ (t+1|t) (Σ (t+1|t) + σ 2 I D )Σ T (t+1|t) .To unpack this further, we first define the attenuation matrixΓ t = σ 2 A t Σ t|t A T t + (σ 2 + 1)I D -1.Because Σ t|t is a covariance matrix, it is also symmetric positive definite, and so A t Σ t|t A T t is symmetric positive definite, and so all of its eigenvalues are greater than zero. Therefore, all the eigenvalues of A t Σ t|t A T t + (σ 2 + 1)I D are greater than σ 2 + 1. We note that Γ t is also symmetric and positive definite. Thus, by the Spectral Theorem, all eigenvalues of Γ t are positive. By the above argument, the eigenvalues of Γ t are all less than σ 2 1+σ 2 < 1. Thus, we observe that the resulting filtering is given by the recurrence relationµ (t+1)|(t+1) = linear Γ t A t µ t|t + offset Γ t b t + (A t Σ t|t A T t + I D ) A t Σ t|t A T t + (σ 2 + 1)I D -1 y t+1 Σ (t+1)|(t+1) = Γ t (A t Σ t|t A T t + I D ). Given Σ t|t T -1 t=0, we see that the filtered means (the updates for ELK) come from a linear recurrence relation with linear term Γ t A t .We therefore compare the eigenvalues of Γ t A t to eigenvalues of A t . Because Γ t is symmetric positive definite, by the Spectral Theorem we can write Γ t = QΛ t Q T , where Q is an orthogonal (and therefore unitary) matrix, and Λ t is a diagonal matrix where every entry is in (0, 1) (the entries of Λ t are the eigenvalues of Γ t , which are greater than 0 by the Spectral Theorem and less than σ 2 1+σ 2 < 1 by the argument above). Now, let's consider any arbitrary unit vector v ∈ C D , and let Λ max t denote the maximum of the norms of all eigenvalues of A t . Then ∥A t v∥ 2 ≤ Λ max t by the definition of Λ max t . However, we want to know ∥QΛ t Q T A t v∥ 2 for any arbitrary unit vector v ∈ R D . However, we know that the action of a unitary matrix cannot change the 2-norm of a vector, so∥QΛ t Q T A t v∥ 2 = ∥Λ t Q T A t v∥ 2 .Moreover, multiplying a vector by a diagonal matrix cannot increase the 2-norm of a vector by more than the absolute value of the diagonal matrix, which in the case of Λ t is bounded above byσ 2 /σ 2 +1. Thus, ∥QΛ t Q T A t v∥ 2 ≤ σ 2 σ 2 +1 ∥A t v∥, or ∥QΛ t Q T A t v∥ 2 ≤ σ 2 1 + σ 2 Λ max tfor any unit vector v ∈ C D . This highlights that we can interpret reducing σ 2 (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.In this section we elaborate on our Experiment 1, discussed in Section 6.1. We closely follow the experimental design in Section 4.1 of Lim et al. [36] , including 5 warm-up steps for all timing experiments and a batch size of 16. However, instead of running 5 seeds for 20 repetitions each, we run 20 seeds for 5 repetitions each, to get more coverage over different evaluation (as the timing for each evaluation are observed to be low variance). We also include memory profiling experiments not present in Lim et al. [36] . For these experiments we use 3 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 B.1.3.For runs with the same specifications (sequence length T , hidden state size D, and algorithm), we observe that sometimes runs with memory profiling ran out of memory whereas runs with timing profiling did not run out of memory. A difference between our time profiling runs and memory profiling runs was how we handled preallocation of memory. For time profiling, we allowed JAX to preallocate memory because that is how JAX usually runs, and so gives a better indication of wallclock time in practice. For memory profiling, we did not allow JAX to preallocate memory so we could get a more fine-grained measure of memory usage.However, JAX provides this following discussion of memory preallocation (see https://jax. readthedocs.io/en/latest/gpu_memory_allocation.html#): Difference between sequential and quasi-DEER outputs However, [not preallocating memory] is more prone to GPU memory fragmentation, meaning a 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 runs run out of memory are likely very close to the 16GB threshold, our hypothesis is that this phenomenon is a manifestation of this documented memory fragmentation.In Figure 5 we qualitatively show that for the same example used in Figure 3 of Lim et al. [36] 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 be found in Section 4.1 of Lim et al. [36] .In Figure 6 , we run the timing benchmarks of Section 6.1 on a wider range of sequence lengths T and hidden state sizes D, on a larger GPU (a V100 with 32 GB) and with a smaller batch size of 1. In doing so, we highlight that the parallel nature of 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 in the larger regimes considered in our main text and in Lim et al. [36] , we often observe linear 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 6 shows good evidence that these parallel algorithms are suffering from saturation of the GPU, and would benefit from even more optimized parallel hardwareThe 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 done 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.As we discussed in Section 4.1, quasi-DEER is O(T D) in memory while DEER is O(T D 2 ) in memory because DEER uses dense Jacobians ∂f /∂s while quasi-DEER uses a diagonal approximation, diag( ∂f /∂s). 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; however, such an approach would still be O(T D 2 ) in memory required. There are two implementation 1K 3K 10K 30K 100K 300K 1M 1K 3K 10K 30K 100K 300K 1M 1K 3K 10K 30K 100K 300K 1M 1K 3K 10K 30K 100K 300K 1M 1K 3K 10K 30K 100K 300K 1M 1K 3K 10K 30K 100K 300K 1MSequence Length (T)Figure 6 : Evaluating an untrained GRU. Sublinear and linear timing regimes for parallelized algorithms. The above experiments were run on a 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 of parallel algorithms. However, when the GPU becomes saturated, the benefits of parallelization are reduced and we begin to see linear scaling in wall-clock time with T .workarounds. One is to loop over computing partial derivatives, effectively trading time for memory. The second is simply derive the diagonal entries of the Jacobian for the architecture of interest.For the purpose of showcasing the O(T D) 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 we JIT compile a lax.scan in JAX, and we 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 O(T ). While there may be more memory efficient ways to sequentially evaluate an RNN, we keep the same benchmarking structure as Lim et al. [36] for to make comparison easier.Here we discuss the experimental details for Experiment 2 in Section 6.2. We follow the same experimental set up as in Section 4.3 and Appendix B.3 of Lim et al. [36] . 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 et al. [36] , using a 16GB V100 SXM2. However, we note that if 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 a 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.Here we discuss our experimental details for our Experiment 3, discussed in Section 6.3.The architecture is a GRU with hidden states h t ∈ R 3 and scalar inputs x t ∈ R. However, at every point in the sequence t, we readout the hidden state h t ∈ R 3 and use it to parameterize a mean µ t+1 ∈ R and a variance σ 2 t+1 ∈ R + . We then sample x t+1 according to x t+1 ∼ N (µ t+1 , σ 2 t+1 ); this output x t+1 is then fed into as the input to the AR GRU at time step t + 1 to make the new hidden step h t+1 . This AR GRU is trained using standard sequential evaluation and backpropagation-through-time to produce a noisy sine wave of length 10,000. We train the AR GRU on 1024 traces x 1:T 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 x 1:T .Once the AR GRU has been trained, it can generate its own trace x1:T given an initial hidden state h 0 and noises ϵ 1:T .We note that such a system is Markovian with dimension D = dim(h) + dim(x), as together the hidden state h t and output x t+1 determine the next hidden state h t+1 and output x t+2 . Thus, in the notation of Section 2, a hidden state s t of the Markovian state space model is s t = (x t+1 , h t ). Therefore, we can apply fixed point methods to try to find the correct trace s * in a parallelized manner instead of autoregressively.We note that one distinction of this set-up with respect to the notation developed in Section 2 is that the dynamics functions f are effectively time-varying because the way in which x t+2 is generated from (x t+1 , h t ) depends on the noise ϵ t+2 , 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 f t .The timing experiments were carried out as follows on an Nvidia A100 GPU. 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 ϵ 1:T and therefore different nonlinear dynamics), and timed each for a 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 as wallclock 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 standard 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 2 is 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 to convergence to be quite stable across random seeds (see Figure 7 ).Table 2 : Time to evaluate a length T = 10, 000 trained AR GRU using sequential vs parallelized methods. We note the dynamax package [12] we used for the parallel Kalman filter implementation in ELK is not optimized for speed, and hence these run times could be further improved. 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 a linear recurrence relation, or equivalently an undamped Newton step, instead of 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 are replace O(D 3 ) matrix-matrix multiplication with O(D) diagonal matrix multiplication. The O(D 3 ) work required by a parallel scan on a dense linear recurrence likely saturates the GPU). We see in Table 2 that individual steps in the dense DEER/ELK are (approximately) a 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.Thus, we find that our fastest parallelized method on wall-clock time is quasi-ELK, but even so it is approximately two times slower than sequentially evaluating this AR GRU. Therefore, an interesting direction for future work would be to characterize regimes where parallel methods can outperform sequential methods, and to investigate whether this autoregressive setting is such a 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.We provide more details on how to set the hyperparameters for ELK. Figure 7 shows how to set the hyperparameter for ELK in the context of the evaluating the AR GRU that generates a 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 of Figure 7 ). We see a 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 a well-behaved dependence that can be optimized on a 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 7 ). 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 stabilized ELK and quasi-ELK 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.These experiments were run on a single Nvidia A100 with 80GB of onboard memory. We include an extra experiment where we tackle the parallel evaluation of the classic non-linear 5dimensional 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 as our nonlinear dynamics function f , i.e. the architecture/time evolution is the Lorenz-96 ODE system. The state is the five-dimensional Lorenz system state. The input is therefore the initial condition of the ODE; and the outputs are the T × 5 subsequent system states.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 is a special case of ELK for λ → 0). DEER is faster (in terms of wallclock time) because of the extra work done per ELK iteration. However, ELK has stabilized convergence, whereas DEER relies heavily on resetting. Interestingly we see that quasi is slower by all metrics, suggesting that the chaotic dynamics may require the more accurate updates. Quasi methods can be implemented to consume notably lower memory, however, and so may be preferable in certain circumstances.In Figure 8 , 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 the entire state sequence. We set hyperparameters using 10 different evaluations of the Lorenz96 (i.e. starting from 10 different initial points).These experiments were run on a single Nvidia A100 with 80GB of onboard memory. For a more detailed reference on parallel scans, the interested reader should refer to Appendix H of Smith et al. [65] or to Blelloch [7] .In our codebase, we leverage jax.lax.associative_scan with the correct binary associative operator. The binary associative operator for DEER and quasi-DEER is simply the composition of affine maps, while the binary associative operation for Kalman filtering can be found in Särkkä and García-Fernández [59] and in dynamax [12] .In this appendix, we provide additional background on Newton's method, and why it is of use for parallelizing nonlinear RNNs.Newton'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, Newton's method can exhibit fast convergence in practice. However, a major motivation for this paper is that globally, Newton's method can be unstable and converge slowly. This instability is a major motivation for our development of ELK.A core insight from Lim et al. [36] is that in the setting of evaluating RNNs, Newton'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 a very powerful tool across a wide variety of domains (e.g. Taylor expansions), this linear approximation can be a good approximation, leading to rapid convergence. For example, if we were dealing with linear RNNs, DEER would converge in one Newton iteration. In this paper, we are instead dealing with nonlinear RNNs, so more Newton iterations are required.We provide a brief discussion of Newton's method for root finding. A great resource for further study is Nocedal and Wright [49] .Let's say we want to find the solution s * to the nonlinear equation r(s) = 0, and we have a guess s (i) at iteration i. Newton's method linearizes r(s) at the guess s (i) , i.e. r(s) := r(s (i) ) + ∂r ∂s(s (i) )(s -s (i) ).Thus, we get our new guess s (i+1) as the solution to r(s) = 0. Therefore, s (i+1) satisfiess (i+1) -s (i) = -J -1 r(s (i) ),where we define J := ∂r ∂s .C.2 Newton, Gauss-Newton, Root-Finding, and OptimizationIn this paper, we seek to find the root of a nonlinear equation r(s) = 0. In Appendix C.1 we discuss how to use Newton's method for root finding to obtain the updates (i+1) ← s (i) -J -1 r.However, another approach is to consider minimizing the merit function L(s) := ∥r(s)∥ 2 2 /2. The root s * of r will also minimize L(s), so the goal of root-finding to solve r(s) = 0 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 L(s) (see Boyd and Vandenberghe [8] for a great reference on Newton's method and optimization), the update obtained is actuallys (i+1) ← s (i) -H(s (i) ) -1 g(s (i) ),where H is the Hessian of L and g = J T r is the gradient of L with respect to s. The Gauss-Newton method approximates this optimization update for minimizing the merit function by making the approximation H ≈ J T J, and so the Gauss-Newton update for minimizing the merit function ends up being the same as the Newton update for the finding the root of r.Newton's method only converges within a suitable basin [47, §1.2.4, p. 37], but establishing best practices for initialization is an open problem. For instance, Yap provides a bound on the norm of the basin [75, 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 [47, 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 [47, §1.2.4, p. 39].To see this fact, note that the characteristic polynomial of J(s) in (3) is (λ -1) T D . 


Decision: 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.

Reviews: review: paper_summary: paper_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. 

Given 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. 

The 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. 

I had a difficult time delineating the contributions, but I believe the following to be the main contributions: 

1. 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. 
2.  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.
strength_weakness: strength_weakness: I prefer to package most of my perceived weaknesses in the questions section, as they are mostly minor. Here are my main concerns:

* 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. 
* 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.
questions: 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? 
* 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? 
* 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. 
* 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
* To me it seems that the NN layer updates must be deterministic for this approach to work. Is this impression correct? 
    * 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? 
* “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? 
* I find the following poorly stated/supported as of my reading so far: 
    * 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. 
    * Replacing J with any **arbitrary matrix**? really? Why? 
    * For a similar reason: 		A 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. 
    * Why does this work? Is this a theoretical result, or an empirically valid approach? 
* Question on Figure 6: is the later linear growth of the wall clock because of the limited parallel compute on the GPU? 
* 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? 
* Figure 4: I’m not sure I understand top-right for some traces, 
    * why the DEER (red) traces goes down at such a late iteration? I thought it was unstable and would never converge. 
    * Also, the quasi-DEER (orange), goes up and then down, what does that mean? \
    * Why is the IKE going down faster first, but then quasi-IKE converges faster?
limitations: limitations: I am a little concerned about the way that limitations are currently addressed. 
- 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. 
- As also mentioned in questions, I would appreciate if authors make some comments on architectural limitations of this work.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
questions: 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: limitations: Yes.
paper_summary: paper_summary: This work extends the parallel RNN (DEER) method to improve its efficiency and stability. The authors make two main modifications:
1. Replace the full Jacobian matrix inverse with its diagonal, allowing for linear-time inversion.
2. Introduce damping via Kalman filter to enhance the stability of Newton methods.
strength_weakness: strength_weakness: Overall, the paper have solid contributions. But there are several weakness
 
1. The potential drawbacks of these modifications are not fully explored
   - Efficiency gains may come at the cost of slower convergence
   - Damping could potentially compromise accuracy
   
2. Limited experimental scope
   - The experiments may not reveal issues that could arise in harder, larger-scale cases
   - For example: does the model work on chaotic systems such as Lorenz 96?

3. it will be interesting to compare with popular models like transformers and xLSTM
questions: questions: Harder examples: how does the model work on chaotic system such as Lorenz 96 or Weather forecast?
More benchmark, how does the model compared to transformers and xLSTM?
limitations: limitations: The limitations are not sufficiently addressed.
paper_summary: paper_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).
The authors corroborate their improvements with some theoretical results and experiments
strength_weakness: strength_weakness: I will put here weaknesses and questions at the same time:

1) 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.

2) 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.

3) 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.
questions: questions: above.
limitations: limitations: Well discussed by authors
paper_summary: paper_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.

Strengths

1. The paper is well written and easy to follow.
2. The proposed methods are very well justified.
2. The proposed technique is reasonably designed.

Weaknesses

1. The experiments are thin.
2. Although the proposed methods are novel, the adopted techniques are already well-known.

Overall, I recommend a weak acceptance of this paper.
strength_weakness: strength_weakness: 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.
2. 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.
questions: questions: See the weaknesses.
limitations: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
questions: 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.
- More experimental details are required such as how the parallelism is implemented and what kind of parallel computing is applied in the experiments. 
- 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: limitations: The authors have not adequately addressed the limitations.

Comments: title: Thanks!
comment: > What is T?


Yes! $T=1000$

> Is it correct?

Yes, on the Lorenz task, DEER converges faster on wallclock time. However, DEER is a special case of IKE with $\lambda=0$, and the rebuttal pdf shows how to set this hyperparameter robustly.

We discuss this point in more detail in our original rebuttal to you, in the bullet point about the Lorenz experiment, beginning with “We really liked your suggestion of using our algorithms to perform inference in a Lorenz96 system.”

Please let us know if we can provide additional clarification!
Thank you for advocating to accept this paper!
Thank you for your response. I remain my stance to accept this paper.
Thanks for the clarification. Two more questions:
- what is T in this case? Can these algorithms work for large T such as 1000 or 10000 (which are common in chaotic system).
- Is it correct that the DEER algorithm, while unstable at the first 100 iterations, converges the fastest in wall time? Does it imply DEER is more efficient compared to the proposed methods?
Thank you for asking follow-up questions. We were very tight on characters so we couldn’t include much information initially! We can certainly add some more clarity:


> If the model is not auto-regressive, how does it model the time evolution?

We tackle the parallel evaluation of the classic non-linear 5-dimensional Lorenz-96 system, with F=8 which results in chaotic dynamics.  The setup is the same for (q-)DEER and (q-)IKE. We directly use the Lorenz-96 dynamics as our nonlinear dynamics function $f$, i.e. the architecture/time evolution _is_ the Lorenz-96 ODE system.  The state is the five-dimensional Lorenz system state (instead of using an RNN approximator with its own hidden state predicting the next timestep of the true system). 
We chose this in part to explore the parallel evaluation of other architectures; but also to allow us to focus on the parallel evaluation of the Lorenz-96 system itself – as opposed to the capability of a function approximator to approximate the system.

> Specifically, what are the input and output of the model?

The input is therefore the initial condition of the ODE; and the outputs are the $T \times 5$ subsequent system states (plotted in Review PDF Figure 2, bottom row).

> What error metric is being used?

Our solvers [(q)-DEER and (q)-IKE] are optimizing the merit function defined in eq. (7), which is the sum of squares of the one-step prediction error (defined in eq. (1)). However, this loss expresses the pairwise consistency of consecutive inferred states under the specified dynamics;  as opposed to the more conventional notion of “accuracy”: the difference between the inferred sequence and the true sequence. We therefore plot this as the 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 the entire state sequence.

We hope this clears any points of confusion. Please let us know if you have any further questions!

-- Submission 13253 Authors.
Thank you for including the Lorenz experiment. However, I’m a bit confused about how the task is defined. Typically, the task involves predicting the next state auto-regressively given the history and the iteration is the time stepping. Since the Lorenz system is ergodic, the state should not converge to a stationary state. However it seems the iteration is something different here. Could the authors provide more details about the experiment? 
- Specifically, what are the input and output of the model?
- if the model is not auto-regressive, how does it modeling the time evolution? 
- what error metric is being used?
This is very good to hear. Yes, we will certainly flesh out these intuitions in any camera ready version, and agree that they will make the paper even better.
Thank you. This does indeed give some intuition. I'm not sure if this intuitions were already in the paper. If not, adding them it would help the readers who aren't intimately familiar with Newton's method, to understand where is the main source of speedup.
Thank you for your helpful comments. We agree that they have made the paper better! 

Thank you as well for upgrading your score and for advocating for acceptance of this paper!
Thank you so much for advocating for the acceptance of this paper!
Thank you for reading our rebuttal, for your high score, and for advocating for acceptance of this paper. We are glad that our clarifications were helpful. We agree that your commentary has made the paper better!

> The upper bound is...

The intuition for why far fewer steps than $T$ are often needed in practice comes from the fact that DEER is equivalent to Newton's method. As we discuss in the rebuttal to reviewer jU7g, Newton'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, Newton's method can exhibit fast convergence in practice. However, a major motivation for this paper is that globally, Newton's method can be unstable and converge slowly. This instability is a major motivation for our development of IKE. 

Another intuition is available from the scan (dynamical system) perspective. At each "Newton iteration," DEER linearizes the nonlinear dynamics of the RNN it is evaluating. To the extent that linear approximations are a very powerful tool across a wide variety of domains (think Taylor expansions), this linear approximation can be a decent approximation, and so lead to rapid convergence. For example, if we were dealing with linear RNNs, DEER would converge in one Newton iteration. In this paper, we are instead dealing with nonlinear RNNs, so more Newton iterations are required.

Please let us know if there is anything else we can clarify!
I would like to thank the authors for clarifying the questions I asked, and addressing my concerns. I believe adding several of the responses would to be added to the main manuscript and make it more readable. 

> The upper bound is ... Each Newton step requires more work than a sequential step; but fewer steps are required and the work is parallelizable.

I can see that the theory is worst-case and in practice "fewer steps but more work" makes some sense on a very high level. But  I was hoping the authors can provide some intuition as to why there are fewer steps than $T$ needed in practice? 

Given that I had already a high score, I would keep my overall score as is but increase contribution score (from 2 to 3)
Thank you for your helpful clarification and I would like to upgrade my score. I suggest the authors add these discussions in the revised version.
Thanks for the rebuttal. Very interesting - especially your view on what is the method doing to eigenvalues in practice. Keeping my accept score. Good luck!