Transformers and linear state space models can be evaluated in parallel on modern hardware, but evaluating nonlinear RNNs appears to be an inherently sequential problem. Recently, however, Lim et al. '24 developed an approach called DEER, which evaluates nonlinear RNNs in parallel by posing the states as the solution to a fixed-point problem. They derived a parallel form of Newton's method to solve the fixed-point problem and achieved significant speedups over sequential evaluation. However, the computational complexity of DEER is cubic in the state size, and the algorithm can suffer from numerical instability. We address these limitations with two novel contributions. To reduce the computational complexity, we apply quasi-Newton approximations and show they converge comparably to Newton, use less memory, and are faster. To stabilize DEER, we leverage a connection between the Levenberg-Marquardt algorithm and Kalman smoothing, which we call ELK. This connection allows us to stabilize Newton's method while using efficient parallelized Kalman smoothing algorithms to retain performance. Through several experiments, we show that these innovations allow for parallel evaluation of nonlinear RNNs at larger scales and with greater stability.
Paper
Similar papers
Peer review
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.
Strengths
- The integration of quasi-Newton methods provides a fresh perspective on overcoming the challenges in parallelizing the evolution of nonlinear RNNs. The theoretical insight extends the scalability of RNNs soundly. - By employing quasi-Newton approximations, the paper effectively reduces the computational complexity from cubic to nearly linear with respect to the sequence length. This reduction is critical for deploying RNNs in resource-constrained environments. - The introduction of trust regions and their integration with Kalman smoothing techniques stabilizes the Newton iterations. This is a notable advancement over traditional methods, which often suffer from numerical instabilities due to undamped Newton updates. - The methods are not only theoretically sound but are also empirically validated to demonstrate their practical effectiveness. This includes significant reductions in memory usage and computational time, crucial metrics for real-world applications.
Weaknesses
- The paper's empirical validation, though robust, is limited to specific architectures and scenarios. Expanding the testing to a wider range of nonlinear RNN types and more diverse datasets could provide a more comprehensive understanding of the methods’ applicability and limitations. - 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
- 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?
Rating
6
Confidence
2
Soundness
3
Presentation
3
Contribution
2
Limitations
The authors have not adequately addressed the limitations.
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.
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. 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
See the weaknesses.
Rating
6
Confidence
3
Soundness
3
Presentation
4
Contribution
3
Limitations
The authors have adequately addressed the limitations.
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
Strengths
The paper flows very well and jumps right into the math: the authors have lots to show and I must say I was impressed by the quality and flow of the paper. While I have some questions and think this line of work does not end with this paper, I believe this contribution is worth publishing. I think the author correctly shows their improvements, and I particularly liked the quasi-deer approach. I also really like how the paper is hones on limitations and clearly provides a gateway to future research.
Weaknesses
I will put here weaknesses and questions at the same time: 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
above.
Rating
7
Confidence
4
Soundness
4
Presentation
4
Contribution
3
Limitations
Well discussed by authors
Thanks!
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!
Thank you so much for advocating for the acceptance of this 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.
Strengths
- Solid contributions to improving the DEER method - Potential for increased efficiency and stability in parallel RNN implementations
Weaknesses
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
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?
Rating
5
Confidence
3
Soundness
3
Presentation
3
Contribution
2
Limitations
The limitations are not sufficiently addressed.
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.
Strengths
+ Identifies a shortcoming with the DEER algorithm + Presents a computationally efficient and scalable alternative to DEER
Weaknesses
- The result proving the convergence of DEER isn't quite rigorous/general -- it relies on the Jacobians being finite which is quite a strong assumption which isn't true in practice, particularly in situations when the eigenvalues are > 1. The paper sweeps this under the rug through re-initialization, which suggests it isn't quite a general statement.
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?
Rating
3
Confidence
4
Soundness
2
Presentation
3
Contribution
2
Limitations
Yes.
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.
Strengths
- The paper consideres an important problem, ie, parlallel computation of RNNs, which is of high practical importance. - The key contributions of the paper in quaski-DEER and (quasi-)IKE are novel and address issues of scalability and numerical instability. - To the best of my knowledge, the proposed contributions are novel.
Weaknesses
I prefer to package most of my perceived weaknesses in the questions section, as they are mostly minor. Here are my main concerns: * 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
* 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?
Rating
6
Confidence
3
Soundness
3
Presentation
3
Contribution
3
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.
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 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!
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.
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 for your helpful clarification and I would like to upgrade my score. I suggest the authors add these discussions in the revised version.
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 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?
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.
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?
> 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!
Decision
Accept (poster)