Towards Scalable and Stable Parallelization of Nonlinear RNNs

Firstly, we thank the reviewer for taking the time to review our submission and for their positive comments! We were particularly pleased by the comments on the foundation, clarity and presentation of our work.

Weaknesses

We discuss these in global review (“empirical validation” in Section 2 and “implementation” in Section 1.1). Broadly, regarding “empirical validation” we appreciated your perspective to bring our robust experimental validation to other datasets, and so took the suggestion of Reviewer VSf9 and demonstrated our methods on the chaotic Lorenz-96 system. Regarding “implementation,” we note that our method can be dropped into any stateful architecture.

As I am…

This is a very interesting point, that we probably didn’t discuss in enough detail in the paper. The key insight of DEER was remapping the matrix inversion (a prohibitively expensive operation) required by the Newton iteration into a parallel scan. This allows the “inverse” to be solved efficiently and while exploiting parallelism. However, the parallel scan incurs $O(TD^3)$ work, and hence does not scale well (and is prone to instabilities due to the unboundedness of the individual Jacobians). Our insight was to introduce quasi-approximations into this, to reduce this complexity to $O(TD)$, while proving this retains global convergence (c.f. Proposition 1). (This suggested the “resetting” heuristic for combating instabilities; and also led us towards the connection with trust region stabilisation.)

The parallel scan, given sufficiently many processors, scales as $O(\log T)$. As we show in Figure 6, we see this speedup at low model sizes and sequence lengths. Once the processors are saturated, we see a linear increase in the runtime (since the amount of work we do is linear), but it is making much more effective use of the GPU, resulting in a constant factor speedup over sequential application at larger model sizes/sequence lengths.

We have added an expanded discussion of these points, thank you for highlighting them.

More experimental details…

The JAX library includes an implementation of parallel scan which we use off-the-shelf. As a result the implementation (with regard to parallelism) is very straightforward. We included code in the original submission and are also preparing a full public code release. We have also added more tabulated hyperparameters, configs, runtimes etc; and added a paragraph explaining parallel scans in detail and how we leverage them.

Will the word…

We believe “parallel” is the correct word. Crucially, we pose the inherently sequential operation as an optimization over the length of the sequence, where each optimization step is parallelizable (sub-linear parallel application time in the sequence length, $O(\log T)$, to be exact). We have reinforced this point in the introduction. If we have misunderstood the reviewers query, then please let us know!

The authors have…

We refer the reviewer to the general response in regard to discussion of limitations. We have added much more explicit discussion of the limitations in a single place to correct this shortfall. If the reviewer has additional points they would like us to include, then please let us know, we are more than happy to include them!

We again thank the reviewer for taking the time to review our submission and for raising some really interesting points of confusion that we had not foreseen. Fixing these has definitely made the paper better. Please let us know if you have any further questions!

— Submission 13253 Authors

Firstly we thank the reviewer for taking the time to review our submission and for their positive comments! We were particularly pleased by the comments on the foundation, clarity and presentation of our work.

The central reservation of the reviewer is that our experimental results are “thin”. In regard to this, we refer the reviewer to the general response and the supplemental review PDF, where we have included additional experiments.

The reviewer also comments that many of the components we discuss exist in the literature. We ultimately agree with this statement; but argue that the real merit in our paper is assembling these components and exploring (theoretically and empirically) their relative strengths and weaknesses, and stratagems to operationalizing these components in the context of deep learning and non-linear systems. As such, we believe the ties to existing work highlights the timeliness and utility of our paper – as we clearly build directly and pragmatically on existing literature that is of general interest to many in the community.

We again thank the reviewer for their evaluation of our paper, their kind words, and for pointing to avenues to improve the paper. Please let us know if you have any further questions!

— Submission 13253 Authors

We thank the reviewer for taking the time to review our paper and providing some helpful feedback!

Limited experimental scope

We point the reviewer to the global response for additional experiments we have added:

• In Review PDF Figure 1 we study the sensitivity to the $\lambda$ parameter (related to Reviewer S8Ch’s feedback), showing that $\lambda$ can be effectively set with a 1-d hyperparameter sweep, and that actualized speed-ups for a selected $\lambda$ are robust for a given data modality.

• We also add additional experiments exploring the wallclock time/iteration/accuracy tradeoff of the methods (r.e. Your comment “The potential drawbacks…”). In Review PDF Figure 1, we show for the AR GRU that while full methods take fewer steps, quasi methods converge faster in wallclock time across a broad range of desired accuracy levels.

• We really liked your suggestion of using our algorithms to perform inference in a Lorenz96 system. We include results for this experiment in the Review PDF Figure 2. We see that DEER and IKE methods converge in a comparable number of steps (this makes sense as DEER is a special case of IKE for $\lambda \to 0$). DEER is faster (in terms of wallclock time) because of the extra work done per IKE iteration. However, we do see that IKE has stabilised convergence, whereas DEER relies heavily on the resetting heuristic. Interestingly we see that quasi is slower by all metrics, suggesting that the chaotic dynamics may require the more accurate updates. Quasi methods still consume notably lower memory, however, and so may be preferable in certain circumstances.

It will be interesting to…

The comparison to alternative architectures, such as the xLSTM and Transformers, is a potential point of comparison. However, we stress that this paper is about accelerating and operationalizing the parallel evaluation of sequential, non-linear models. Papers developing or benchmarking natively parallelizable architectures (e.g. Transformers) often state that sequential architectures such as (e.g.) GRUs cannot be parallelized, and are often therefore eliminated from consideration or score poorly on speed tests. However, DEER raised the point that this is not entirely accurate, and that alternative black-box methods may be within reach for efficiently evaluating these oft-maligned architectures. Our work focuses on extending and operationalizing DEER, without changing the architecture or its asymptotic performance. We have clarified this in the main text.

The limitations are…

As mentioned in the global review, we have further highlighted and discussed the limitations in separate paragraphs in the conclusion (while also retaining the “inline” discussion to retain flow). If the reviewer has further limitations we have not addressed, then please let us know, we are more than happy to include them!

We again thank the reviewer for evaluation, and for raising many points that have improved the paper. Please let us know if you have any further questions!

— Submission 13253 Authors

We thank the reviewer for taking the time to review our paper, for their positive feedback, and for providing some very interesting discussion points!

Approximation quality:

This is a great observation. We have added Review PDF Figure 1 probing the compute or computational budget required to achieve different levels of accuracy in the AR GRU experiment. There are two major findings here: (a) the speed ups are fairly robust across elements within a dataset (i.e. once tuned the speedups are fairly low variance); and (b) quasi-methods took more steps, but were faster as each optimization step is faster. We also find that IKE methods converged faster for the AR GRU. However, we note this relationship is reversed in the additional Lorenz96 experiments (proposed by reviewer VSf9; see global review), where IKE and DEER converge in a similar number of steps, and quasi methods require notably more. This reinforces a point made in the discussion of the original submission that different methods may be faster or more stable in different settings, but that the memory savings and ultimate convergence are consistent. We have expanded discussion and guidance in the paper.

Speed:

Yes, dense matrices are much slower to process, and this matches our experience (the parallel scan on a dense linear recurrence requires $O(D^3)$ work which saturates the GPU). Table 2 in the paper shows the average time per step and average number of steps required to converge. The key take-away is that individual steps in DEER/IKE are (approximately) a factor of between 3.5 and 30 times slower per step than their quasi variants (which roughly tallies with your experience). However, they take (on average in Figure 4) a factor of between 2 and 10 fewer iterations. We have added some additional discussion of this tradeoff.

What is this doing?

This is a super interesting question. We have two notable observations about (re-)interpreting the methods we present. Below we include a condensed proof on how IKE can be re-interpreted as stabilising the parallel scan by attenuating the eigenvalues used in the parallel scan, and hence stabilising their cumulative product. We will add the full proof to the main paper.

To your question on quasi: we don’t have a particularly satisfactory answer. Very early in our investigations when exploring quasi approximations, we noticed that the off-diagonal elements were often much smaller in magnitude than the diagonal elements (which makes intuitive sense for many recurrent models). This gave us some confidence to explore quasi approximations further, but we were unable to derive any deeper insight, and so validated it empirically. It is a really interesting avenue of future work looking into this, or other families of approximations and their relative merits and drawbacks.

Summary: We again thank the reviewer for positive feedback, and for raising some really unique discussion points. Please let us know if you have any further questions!

— Submission 13253 Authors

—

Condensed additional proof that IKE attenuates the largest eigenvalues in linear recurrence: (We include a full proof in revised paper)

Let $\{J_t\}$ be the Jacobians used in the linear recurrence relations and $\{b_t\}$ be the offsets. Then the prediction step of the Kalman filter (IKE) is the same as DEER. However, after applying the update step in IKE (which imposes the trust region), we obtain a second linear recurrence relation where the linear operator is given by $\Gamma_t J_t$. It can be shown that $\Gamma_t$ is a symmetric positive definite matrix with eigenvalues bounded above by $\frac{1}{1 + \lambda}$. Thus, making using of the Spectral Theorem, it follows that the norms of the eigenvalues of $\Gamma_t J_t$ are bounded above by the max of the norms of the eigenvalues of $J_t$, scaled by $\frac{1}{1 + \lambda}$. Note that larger $\lambda$ corresponds to more regularization/smaller trust region; and therefore correspondingly results in smaller effective eigenvalues in the scan. We recover DEER exactly if $\lambda = 0$. Thus, while large eigenvalues in $J_t$ are the cause of the instability of DEER when evaluating unstable dynamical systems, IKE directly attenuates these large eigenvalues, explaining why the intermediate iterations using IKE remain stable.

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!

We thank the reviewer for highlighting some interesting clarification points.

The result proving…

This is a really interesting line of thought that we will clarify in the paper. There are lots of points that your comment touches on, so we will try and answer them sequentially:

Re: Generality of Prop 1 with regard to eigenvalues: We state in Prop. 1 that we assume the Jacobians must be finite. Given this assumption, the proof by induction in Appendix A.1 is rigorous and general for real numbers. Thanks to your feedback, we will include an explicit pointer to the full proof.

The finiteness of the Jacobians is a core requirement for DEER. Thus, Prop. 1 (a statement about the DEER algorithm) is applicable anywhere that DEER is applicable. We also argue that the finiteness of the Jacobians is a weak assumption, as it is a requirement for training models with backpropagation and is satisfied by many commonly used architectures.

However, it is absolutely true that the eigenvalues can be greater than one, and this leads to the product of the Jacobians computed by the scan blowing up in magnitude. In fact, the sensitivity of DEER to blowing up is directly linked to the eigenvalues of the Jacobians. It is this explosion that causes overflow on finite-precision architectures.

The resetting heuristic combats this, and is motivated by a result derived as part of Proposition 1: that the first $t$ steps have zero residual after $t$ Newton steps (see ~Line 111 or 447), e.g., that we can restart the optimization with the correct $t$-length initialization. We can view this heuristic as allowing the Newton steps to resume when overflow is encountered, as opposed to failing or falling back to sequential evaluation.

IKE directly combats eigenvalues greater than one: A related point is that IKE was motivated to prevent overflow by limiting the step size in the latent state using a Kalman filter (KF). We have since made a connection from the Kalman filter back to the original scan: the trust region/KF can be interpreted as attenuating the eigenvalues by a factor of at least $\frac{1}{\lambda+1}$, stabilising the scan without the use of the heuristic (see Proof that IKE attenuates the largest eigenvalues in linear recurrence in response to 5uzK). IKE combats exploding cumulative products without the use of the heuristic.

Re: discussion of the re-initialization heuristic: we respectfully disagree that we “sweep it under the rug”. We discuss the heuristic in several places, and explicitly highlight the heuristic in Figure 4. The heuristic is a necessary intervention — not originally discussed by Lim et al. — to combat instabilities of undamped Newton in floating-point representations. We see the heuristic as a (albeit small) novel contribution, and a practical corollary of Prop. 1. We have expanded discussion of the limitations and opportunities posed by the reinitialization heuristic.

In general…

Re: connections back to floating point: Our analysis was a theoretical exploration to develop methods we would eventually test empirically, and so we didn’t make explicit theoretical connections to floating point representations. Floating point precision can introduce instabilities into Newton through overflow and limited numerical precision in functions with large gradients. The finite precision also sets a floor on the fidelity of the convergence. These limitations are further motivation for stabilisation techniques, such as IKE, to combat the instabilities inherent to Newton’s method when executed in finite-precision.

Re: Convergence rates of Newton’s Method: We spent a lot of time thinking about this. Newton’s method only converges within a suitable basin [Nesterov, 2018, Lectures on Convex Optimization, §1.2.4, p37] — but establishing best practices for initialization is an open problem. For instance, Yap provides a bound on the norm of the basin [Yap, Fundamental Problems in Algorithmic Algebra, 1993, Lecture IV, §10, p. 174]. However, this definition requires bounding the derivative of the objective function, which is harder than the original problem. Nesterov derives a basin for quadratic convergence around the true solution [Theorem 1.2.5; Nesterov, 2018, Lectures on Convex Optimization, §1.2.4, p39], but does not provide information on how to locate this basin apriori. Indeed, Nesterov defaults to taking standard gradient steps early in optimization until you assume you are in the basin, and then using Newton steps [Nesterov, 2018, Lectures on Convex Optimization, §1.2.4, p39]. We experimented with this initially, but found it to be worse (by any metric we conceived).

Quasi-methods do not, in general, enjoy quadratic convergence (but showing the global convergence of quasi-DEER was a major motivation behind Proposition 1), but are often observed to converge in practice. We have not been able to prove results about the convergence rates of the quasi methods we propose. However, our experiments show that their rates are acceptable in practice. Thus, studying quasi-DEER from a theoretical perspective is an interesting open problem. As a result of your points, we have included in the conclusion a list of open theoretical problems inspired by your questions. We have also added a brief survey on the convergence properties of Newton variants.

Limitations: Yes

Is the reviewer saying there are limitations that we do not discuss? Or are that we do satisfactorily discuss limitations? If the former, please see the Global Response for information on how we have updated the paper.

Thank you again for providing some really interesting and unique insights. The extra discussion has definitely improved the paper. Please let us know if you have any further questions!

Thank you for your thorough and positive review!

We refer the reviewer to the global response for discussion of an ablation study on $\lambda$ (r.e. “As authors list”) and limitations (r.e. “In responses to” and “As also mentioned”). We now respond to your other comments:

Contributions: The manuscript…

We will explicitly list our contributions in the introduction. Specifically we:

• Show DEER is globally convergent
• Introduce quasi-approximations to improve efficiency
• Introduce trust regions (IKE) to stabilise DEER
• Empirically validate the methods, and provide guidance on how to select and tune them
• Have since explicitly shown there is a unique global solution to DEER (see “I was missing…” response) and provide a reinterpretation of how trust regions stabilise the original scan (see 5uzK response).

The proposition 1… and From seeing the…

You are absolutely correct. Proposition 1 tackles the theoretical global convergence of the algorithm, showing we can expect convergence. The upper bound is impractical as it requires $T$ iterations; but this is a worst case convergence and we often require much fewer than $T$ iterations (e.g. Figure 4a). Each Newton step requires more work than a sequential step; but fewer steps are required and the work is parallelizable. Proposition 1 was also important in developing the resetting heuristic for (q-)DEER. We have added clarification on this.

In my initial…

Thank you for raising this. While DEER was inspired by the root finding perspective, IKE was inspired by the optimization perspective (see Section 4.2). We have further clarified the relationship between perspectives and their implications.

I was missing…

We add a proof of the uniqueness of the solution, and global convergence to this solution, as we did not make this connection clear enough. To summarise:

For a deterministic forward function and fixed inputs, there is a fixed sequence of latents and outputs. This is therefore the only sequence with zero residual (i.e. there is a unique sequence $s_{1:T}^*$ generated by the deterministic dynamics). Furthermore, DEER cannot get stuck at any point that is not this sequence. We prove this in Proposition 1. Another way to see this however is that each update step (eq. (5)) is equal to $J^{-1} r$. But, we have established (see “Provided all of” below) that $J$ is always invertible and so has trivial nullspace. Furthermore, the residual $r$ can only be zero at the unique solution $s_{1:T}^*$. Thus $J^{-1} r$ is nonzero everywhere except at the true solution, where it is zero. Thus, DEER cannot get stuck en route to finding the true and unique solution.

To me it seems…

This is an interesting point: DEER and IKE can handle stochastic models – because stochastic models become deterministic when the stochasticity is fixed i.e. has been reparameterized as in Figure 4. In fact, DEER and IKE have the same requirements as backpropagation, and so can be applied to any stateful recurrent model that admits backpropagation – a huge and practical class of models.

Provided all of…

The eigenvalues of a lower-triangular matrix are equal to its main diagonal. The Jacobian in eq. (3) is the derivative of eq. (1). Each term in the vector is of the form $s_t - f(s_{t-1})$. Therefore, the Jacobian with respect to $s_t$ is the identity matrix; hence the leading diagonal is the identity; and all its eigenvalues equal one and the Jacobian is invertible (no zero eigenvalues).

Efficiently solving this inverse is not tractable (see Line 92). However, we avoid computing the inverse by exploiting the structure in $J$ to instead pose eq. (5) as the recursion in eq. (6). This recursion is linear and can be solved in parallel with a scan.

I find the…

We will make this discussion more clear in the paper: The elements in the sub-block-diagonal of $J$ can be replaced with arbitrary values – but the main block diagonal must remain as the identity and all other entries must be zero. Retaining convergence under modifications to the sub-block-diagonal portion is a corollary of Proposition 1, and can be seen from eq. (6): If all the states up to and including position $t-1$ at the $(i)$th Newton iteration are correct, then the update in eq. (6) at Newton iteration $(i+1)$ for position $t$ will use $\Delta s_{t-1}^{(i+1)} = 0$ (no update is required at position $t-1$), and so the update to $s_{t}^{(i+1)}$ no longer depends on the Jacobian. We will explicitly state this, as it motivates q-DEER and is, as you point out, a surprising result.

We exploit this to develop q-DEER, retaining only the diagonal of the Jacobians. This reduces the parallel scan from $O(D^3)$ to $O(D)$ making each iteration faster (while still admitting global convergence as above), but needs more Newton iterations to converge due to approximate updates. We find that this trade-off often yields a faster wallclock time (see Review PDF Figure 1).

Explicitly, the global convergence of q-DEER is a theoretical result (a corollary of Proposition 1), but the fast runtime of q-DEER in practice is an empirical result.

Question on Figure 6…

Yes, that's a great point! We will more heavily emphasise this.

Section 6.1: why…

Yes! The trained model does have larger eigenvalues, both at the true latent sequence, and at the sequences visited during the Newton iteration.

Figure 4: I’m…

We use “unstable” to mirror Newton’s method parlance. Proposition 1 shows that $t \leq (i)$ have zero error, and so the instability applies only to $t>(i)$. Unstabilized methods can be arbitrarily bad in this regime until global convergence, but may converge before then. (q-)IKE removes the instabilities, leading to faster overall convergence. We have added clarification on this around Figure 4.

Thank you for your thorough analysis, and for raising many points that have improved the paper. Please let us know if you have any further questions!

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)

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.