ParaRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language Models
Authors: Federico Danieli
Pau Rodriguez
Miguel Sarabia
Xavier Suau
Luca Zappella
Venue: ICLR 2026 Oral (2026)
Keywords: RNN
Mamba
SSM
Transformers
Parallelization
Parallel scan
Nonlinear
Primary area: foundation or frontier models, including LLMs
TL;DR: We break the sequential bottleneck of nonlinear RNNs, enabling training of billion-scale LSTM/GRU models, competitive with modern architectures

Abstract: Recurrent Neural Networks (RNNs) laid the foundation for sequence modeling, but their intrinsic sequential nature restricts parallel computation, creating a fundamental barrier to scaling. This has led to the dominance of parallelizable architectures like Transformers and, more recently, State Space Models (SSMs). While SSMs achieve efficient parallelization through structured linear recurrences, this linearity constraint limits their expressive power and precludes modeling complex, nonlinear sequence-wise dependencies.
To address this, we present ParaRNN, a framework that breaks the sequence-parallelization barrier for nonlinear RNNs. Building on prior work, we cast the sequence of nonlinear recurrence relationships as a single system of equations, which we solve in parallel using Newton's iterations combined with custom parallel reductions. Our implementation achieves speedups of up to $665\times$ over na\"ive sequential application, allowing training nonlinear RNNs at unprecedented scales. To showcase this, we apply ParaRNN to adaptations of LSTM and GRU architectures, successfully training models of 7B parameters that attain perplexity comparable to similarly-sized Transformers and Mamba2 architectures.
To accelerate research in efficient sequence modeling, we release the ParaRNN codebase as an open-source framework for automatic training-parallelization of nonlinear RNNs, enabling researchers and practitioners to explore new nonlinear RNN models at scale.

Published as a conference paper at ICLR 2026
PARARNN: UNLOCKING PARALLEL TRAINING OF NONLINEAR RNNS FOR LARGE LANGUAGE MODELS
Federico Danieli
Pau Rodríguez
Miguel Sarabia
Xavier Suau Luca Zappella
Apple
danieli@apples
ABSTRACT
Recurrent Neural Networks (RNNs) laid the foundation for sequence modeling, but their intrinsic sequential nature restricts parallel computation, creating a fun- damental barrier scaling. This has led to the dominance of parallelizable archi- tectures like Transformers and, more recently, State Space Models (SSMs). While SSMs achieve efficient parallelization through structured linear recurrences, this linearity constraint limits their expressive power and precludes modeling complex, nonlinear sequence-wis dependencies. To address this, we present ParaRNN, a framework that breaks the esquence-parallelizane Building on prior work, we cast the sequence of nonlinear recurrence relationships single system of equations, which solve in parallel using Newton's 's iter- ations combined with custom parallel reductions. Our implementation achieves naïve sequential application, allowing training non-
barrier for nonlinear RNNs.
speedups of up 665x linear RNNs at unprecedented scales. To showcase this, we apply ParaRNN to adaptations of LSTM and GRU architectures, successfully training models of 7B parameters that attain perplexity comparable to similarly-sized Transformers and Mamba2 architectures. To accelerate research efficient sequence modeling. we release the ParaRNN codebase as an open-source framework for automatic training-paan parallelization of nonlinear RNNs, enabling researchers and practitioners to explore new nonlinear RNN models at scale.
INTRODUCTION
Since its introduction by Vaswani et al. (2017), the Transformer architecture has quickly imposed itself as the de-facto choice for sequence modeling, surpassing previous ofate-of-hhe-rt, RNN-based models such as GRU and LSTM (Cho et al., 2014; Hochreiter & Schmidhuber, 1997). One key
RNN Call Application Timings [ms)]
10a
Model Size
Marhh?
PaLLSTM
5.52
PenLSTM
10
10
Pauratt
Speeddin wpp.upp.... 447%
Sompssa - 0050
10a
2.9B
10:99
10-1
18
15:21
11:03
10
12:5
22
23
24
25
28
27
28
29
0
2
6
10
12
14
Sequence length
Perplexity
Figure 1: Our ParaRNN framework makes possible to apply classical RNNs in parallel, dramati- cally speeding up their training, and allowing them to be used competitively for language modeling.
Codebase is available at https::///// http:p::/ggthbbb.c///a////////// htrra/////ppppe/eeprrrrrrrr.../.

---

Published as a conference paper at ICLR 2026
reason behind the rapid adoption of Transformers lies in the efficiency of their application training time: their core sequence mixer, the attention mechanism, can in fact be applied in parallel along the length of the input sequence. This effectively overcomes one main limitation of classical RNNs, whose application must be unrolled sequentially along the input sequence.
In more recent times, however, interest in RNNs has been rekindled, largely due to their reduced memory footprint and improved efficiency at inference time. In particular, recent advancements in State Space Models (SSMs) such as Mamba (Gu & Dao, 2023; Dao & Gu, 2024) have started gaining popularity and are imposing themselves as a potential alternative to Transformers, at least for moall-to-mid-sizeer models (Zuo et al., 2024). To grant parallelization during training (and hence a comparable performance to attention), SSMs simplify the recurrence relationship at their core to one that purely linear in the hidden state. This simplification enables leveraging associativity and parallel reduction operations to quickly compute the output of the application of an SSM to whole input sequence in parallel along its length. Despite the success of modern SSMs, the linearity constraint remains a limitation which hinders their expressive power (Merrill et al., 2025; Cirone et al., 2025), and is dictated by necessity rather than choice.
With ParaRNN, we aim to overcome the constraint of linearity for SSMs, while unlocking paralleli- zation also for nonlinear RNNs, thus enriching the space of viable options for sequence modeling.
The method acts by re-casting the sequential application of an RNN to an input sequence of length L as the solution of a nonlinear system of equations. Said solution is recovered via Newton's method: the system is linearized, and its solution approximated iteratively. In light of the Markovian nature of classical RNNs, the resulting linearized system obeys a specific block bi-diagonal structure, which leverage to recover the solution efficiently, using custom nigh-performance implementations of parallel reduction operations. We provide detailed description of the method in Sec. 2.
For the sake of favoring the widespread adoption of this method among the research community, we make available the ParaRNN codebase. This takes care of automating the sequence-paralde application of an RNN cell, relying only on the definition of the RNN recurrence formula. The main structure of the framework is outlined in Sec. 4 and App. D.
Finally, we demonstrate the validity of ParaRNN in allowing the training of competitive RNNs for language modeling. To this end, in Sec. 3 we introduce two adaptations of the classical LSTM and GRU cells, tailored for efficient application within the ParaRNN framework. In See. 5, we compare their performance against the Transformer and Mamba architectures, proving their competitiveness both in terms of runtime and of effectiveness in language modeling tasks.
Contributions To summarize, our main contributions can be listed as follows:
1. We adapt and develop efficient parallel reduction algorithms to enable training of nonlinear RNNs at unprecedented scales, expanding prior work on Newton-based RNN parallelization.
2. We demonstrate for the first time that classical nonlinear RNN models can be trained on lan- guage modeling tasks at scales of 7B parameters, and achieve competitive performance with Transformers, thus enriching the space of available architecture choices for LLMs.
3. We enable the exploration of new nonlinear RNN architectures for language modeling at scale, by introducing ParaRNN: a high-performance PyTorch+CUDA library that implements sequence-paralle, training for any nonlinear cell from only the specification of its recurrence step, thereby automating the underlying parallelization complexity.
Previous work There is flourishing line of research around RNNs that can be applied in parallel (Bradbury et al., 2016; Martin & Cundy, 2018; Qin et al., 2023; Gu et al., 2020; 2022a:b: Orvieto et al., 2023), which culminates more recently in Mamba and Mamba2 (Gu Dao, 2023; Dao & Gu, 2024), perhaps the most successful examples of modern SSMs. In a nutshell, SSMs revolve around a simplification of RNNs, that constrains the recurrence step to be purely linear in the state. The associativity of linear operations enables then the use of parallel scan (Hillis & Steele, 1986; Blelloch, 1990) to parallelize the RNN application. In principle, however, this approach to paralleli- zation comes the cost of expressivity: recent theoretical works warn in fact about the limitations of expressive power of linear SSMs (Merrill et al., 2025; Huang et al., 2025; Keiblinger, 2025). With our work, we effectively remove such linearity constraint.
2

---

Published as a conference paper at ICLR 2026
Recent advantages in SSMs have also spawned renewed interest in RNN adaptations. This is most noticeably represented by minGRU/minLSTM (Feng et al., 2024), and xLSTM (Beck et al., 2024). However, the RNNs proposed in Feng et al. (2024) still preserve linear recurrence at their core, although wrapped in additional nonlinearities that however-ww we remark-do not interest its state. On the other hand, xLSTM proposes alternative RNN cells: sLSTM and mLSTM. The first is purely linear, and parallelized similarly to Mamba; the second, being nonlinear, is still applied sequentially, without any attempt to parallelize its application sequence-wise: this comes with constraints on the size of the hidden states considered, otherwise its training- time application would become too costly. Another approach to scaling up nonlinear RNNs is outlined in Chaubard & Kochenderfer (2025), which mitigates the cost of backpropagation through time via gradient approximation, but still relies on efficiently computing sequential forward passes through the RNN.
The parallelization approach of ParaRNN takes inspiration from methods originally developed for parelaaallelizing the solution of Ordinary Differential Equations (ODEs); see Gander (2015) for an overview. The specific idea of combining Newton's iterations and parallel reduction to solve for nonlinear sequential operations in parallel has been first introduced in Danieli et al. (2023), where it was applied ResNets and Diffusion Models, but only hypothesized for language modeling and other sequence models. An application this same approach has been investigated in Lim et al. (2024) and improved in Gonzalez al. (2024), where GRU cell is trained in parallel for time- series modeling. There, the authors consider the use of a quasi-Newton method, where a diagonal approximation of the full RNN Jacobian is employed for computational efficiency, but requires approximating the backward pass as well, which might affect training dynamics. In this work, we show empirically that directly adapting the GRU and LSTM cells definition allows achieve stable and scalable training, and successfully extend the methodology to more complex cells that were not previously tested.
Lastly, our work shares some looser connections with other approaches to speedup via parallelization of sequential operations. In Song et al. (2021); Santilli al. (2023), Jacobi and Gauss-Seidh solvers are applied to accelerate ResNets applications, auto-regressive image generation, and Transformers inference; in Shih et al. (2023), Picard iterations are used to parallelize denoising steps in Diffusion.
PARALLEL APPLICATION OF NONLINEAR RNNS
In this work, we build upon the blueprint of Danieli et al. (2023); Lim et al. (2024), where the sequential application of RNN is re-cast system of nonlinear equations which is then solved via Newton iterations, and adapt it to LLM-scale training. For completeness, in the following we briefly describe the main components of the method used.
Newton's 's Method for Computing Nonlinear RNN Applications Consider the application of an RNN an input sequence [<1]/_1 length L, whose elements are themselves vectors 2; E Rdd, The RNN cell tracks the evolution of hidden state h4 Rdd through its recurrent application:
h = f(1-11,
VI = =1 L,
starting from an all-zero initial state hg = O, with f determining the cell action, or recurrence step. By collating the relationships in (1) we can define a single system of I. equations in the unknowns [hi] This has the form
h f(0,
h2 f(11, x22) = 0
(2)
hr. ---(hh.-i..)
To solve this system, we rely on Newton's method (Ortega & Rheinboldt, 2000). The method
recovers approximate solution of (2) by iteratively solving its linearization, until convergence. More in detail, given the approximate solution the k-th Newton iteration [ -== this is updated
3

---

Published as a conference paper at ICLR 2026
solving
8h 8h2
f(0, h [(hi,022) h
J in
1
y ho
5h
f shh h
so that 1b/+11/= [ + 1/ In (3), Jrlnt denotes the Jacobian of the recurrence step (1) with respect to the hidden state h, evaluated at h," Pseudo-code for the method is given in Alg. la.
The second ingredient for effectively tackling the solution of (2) consists of an efficient solver for the highly-structured, block bi-diagonal linear system in (3). This is described in the following section.
Parallel Reduction All-at-once Solution Block Bi-diagonal Linear Systems As shown in the previous section, thanks to Newton's method, we can re-cast the application of nonlinear RNN! into the iterative solution of a linear system in the form (3), In principle, its solution [öh 1==1 could be explicitly recovered by forward substitution, sequentially unrolling the formula
th = Jzln; oh + (f m)) - hf),, VI - 11.1....,4
(4)
but this would defeat the purpose of parallelization. Instead, notice that (4) itself corresponds to the application of a linear RNN, where the Jacobians J ln; cover the role of the state matrix, while the
residuals Tı = f(h xτ) - h cover that of the input (or foreing term). As such, the solution can be recovered explicitly by leveraging associativity of matrix multiplication when unrolling the recursion (4). We have in fact
th h = Jshng δh + P1 Jsln; (Jrlne. oh + rr-1) +
J Jeln
(5)
There is some redundancy in the terms appearing in the sum in (5) for different l's: this is what allows specialized algorithms such as prefix sum (also known as parallel reduction, or parallel scan
Algorithm la: Newton's 's method
Algorithm 1b: Parallel reduction
Initialization:
Input: Jacobians Jyln, residuals T
Assemble initial guess h
I/ Initialization:
// Newton iterations for k = 1 to Ni. do
A Jshn; 8h, rú;
// Reduction steps:
// Assemble system (3):
for i = 0 to logz L 1 do
Compute Jacobians J J
3
parfor = 2 1o T. do
Compute residuals r/:
öh: <- ōh: A1ōh4-2441
// Solve parallel
A - AIA
h
h k-1
Pera-eleeeeeee rrf); return ōh:;
return h
Nis.
Algorithm 1: Pseudo-code for the parallel application of a nonlinear RNN cell: we use Newton (left) as outer solver for the corresponding nonlinear system (2), while parallel reduction (right) is used as inner solver for the lincarized system (3).
While the method presented here considers the general case where the cell action (1) is nonlinear in the state hy, also naturally encompasses linear RNNs, such as Mamba and other SSMs. In these cases, have
f feo m)) A Bi
and the Jacobians in (3) reduce to the state transition matrices themselves, J hy A:. Consequently, Newton's update reduces the vanilla SSM application, and the target output is recovered in single iteration.
4

---

Published as a conference paper at ICLR 2026
(Hillis & Steele, 1986; Blelloch, 1990)) to compute the whole solution [öh at once, in parallel, in (log, L) steps, rather than sequentially in (L) steps. Pseudo-code for a naive implementation of the algorithm is provided in Alg. Ib, but refer to App. D.2 for the complete description of the more efficient CUDA implementation developed in this work.
Parallel Backward Pass Through RNN Applications Unlike the forward pass, which for general RNNs requires the solution of a nonlinear system, the backward pass through an RNN cell is an inherently linear operation. As such, we do not need to rely on Newton iterations, and the necessary gradients can be recovered directly in a single parallel reduction step. More in detail, given the partial derivatives of the loss function L with respect the cell hidden states [0p, C) computing the full gradients amounts to unrolling the backward recurrence (starting from Vng L On. L)
V hy 1 -=
- Jflh.-. Vh, C On. L,,
V== L L..,,
(6)
Notice this shares the same structure of (5), except it involves transposes of Jacobians, and it unrolls backwards. This notwithstanding, (6) can be solved using the same parallel reduction algorithm Alg. 1b, with minimal modifications.
2.1
LIMITATIONS
While the ParaRNN framework described in this section is general enough to be applied to virtu- ally any Markovian RNN, it presents two major limitations that must be kept into account when deploying it in practical applications.
Newton's Convergence Our method relies the convergence of Newton's iterations recover the result from the RNN cell application. Gonzalez et al. (2024) showed that, for systems like (2), convergence in L Newton steps guaranteed. However, such slow convergence is not practical: it would break any advantage from parallelization, making a vanilla sequential application of the small O(1) number of iterations.
RNN cell more efficient. In practice, we need convergence Fortunately, for the RNN considered in this work, three iterations proved to be sufficient in all cases, but this must be verified for newly defined RNN wanting to leverage ParaRNN. For additional theoretical guarantees in this sense, we refer to the recent work in Gonzalez et al. (2025), while in App. we report more information regarding the Newton solver setup and convergence for our case.
Computational Efficiency While ParaRNN allows to attain sequence- paraeeiiiiaallllizabiittttt it does so the cost of additional overhead computations, which would not be necessary in the sequential case. This is common trade-off in the design of paraliel algorithms. In our case, the brunt of this overhead is associated with the assembly and manipulation of J terms in (5). fact, if we were naively consider RNNs with dense Jacobians, their memory requirement would scale (Ld)) d ) and each pair-wise multiplication (line 5 of Alg. Ib) would incur a computational cost of (dd making it de-facto unfeasible. In practice, then, one should inject some additional structure in the Jacobians definition, so to make these operations more tractable. Notice that similar problem arises when dealing with SSMs: to overcome this, e.g. in S4D (Gu et al., 2022b) the authors impose diagonal structure the SSM state matrix, to render the application of the prefix sum algorithm computationally feasible. This assumption is inherited by the Mamba model (Gu & Dao, 2023), for analogous considerations. Similarly, Gonzalez al. (2024) opt where the RNN Jacobian is approximated with its diagonal, again rendering the resulting matrix multiplications leaner. In our case, we follow an approach similar to Mamba, outlined next.
quasi-Newton method,
3 ADAPTING GRU AND LSTM FOR PARARNN
As mentioned in Sec. 2. the main obstacle to the efficient application of the prefix scan algorithm to (5) lies the computation its product terms Ii= Jsing Rather than considering Jacobians approximations in Gonzalez al. (2024), in our approach choose instead to directly tune the RNN definition to yield Jacobians with a simplified structure. The adapted GRU and LSTM cells described next.
5

---

Published as a conference paper at ICLR 2026
RNN Cell Definitions In defining the main RNNs used in our work, start from two of the most broadly established architectures available in the literature: the Gated Recurrent Unit (GRU) (Cho et al., 2014) and the Short Term Memory (LSTM) (Hochreiter & Schmidhuber, 1997) cells. For the latter, we consider its variant equipped with peephole connections and combining input and forget gates (Greff et al., 2017). Their actions defined as, respectively:
fi og (A 1 h1 By Cyct-1 + by)
žt Og - + B5xi + bz)
21 0g (A h1 B + bb.)
GRU
Ti = dg (A, B, + b,)
LSTM
C f1 -1 (1 f1)
C (Ac(h ① rm) Bax + be)
0 0g (Aoh411 Bo + Coci + bo)
h 21) 3 1111 + O
h: 01 2 (c))
(7a)
(7b)
for - 1..... L. The highlighted equations represent the actual hidden state of the cells (that is, the variables involved in the recurrence step), while the other ones are auxiliary temporary variables needed perform the update itself. Notice in particular that for LSTM the hidden state is given by the collation of two variables: h, = |c h ] R.2d.
Cell Jacobians To apply the ParaRNN machinery described in Sec. 2 to the cells in (7), we need to compute the Jacobians of their recurrence step with respect to the hidden state. These can be recovered explicitly with some algebraic operations, and come in the form, respectively:
JORU = diag(1 - 21) + diag((a h1-1) o (2))
diagdli..uuo.. + diag(zı o ((i) A. + of and
(8a)
Joc (ia(())))))((((-11-__) ,
Jch -
d
Joe
Jch
JISTM =
with
Jhc (diag(on(e) (()
(8b)
JAc
Jnh
diag(or @ 00 (ct)))eec
Jnh diag(an (c) @ (oi)) + CoJen)
+ diag(o: oo Je)
where diag(t) denotes the diagonal matrix with the elements from v as its diagonal, o' is the derivative of o., and for a given temporary variable v. ò denotes its value before the nonlinearity application. Notice that the LSTM Jacobian features a 2 x 2-block structure, directly stemming from the fact that the hidden state split into two variables that each sub-Jacobian J., indicates the partial derivative 0(*): /0())))()
Jacobians Structure Simplification Our goal is to make Jacobians (8) leaner to store and to multiply together. Notice that the overall structure of these Jacobians is ultimately linked to that of the various state and peophole-commection matrices A., C., since all other matrices involved are purely diagonal. An apt simplification consists then in picking these matrices, too, to be diagonal; in other words, we substitute (7):
A. = diag(a.), C. diag(c.),
for some learnable parameters
a., E IR.h.
(9)
This reduces JORU in (8) diagonal matrix, and further simplifies the 2-block structure of JISTM to one with diagonal blocks. In both cases, the Jacobians thus simplified occupy only O(Ld,) memory overall (we verify this also empirically in App. B.3), and each pairwise multiplication can be performed with (dd)) complexity, in perfectly parallclizable fashion over the hidden state components.
While this simplification allows to satisfy the requirements (outlined in Sec. 2.1) to apply parallel scan efficiently, we point out that it has the effect of inhibiting any mixing of the hidden state com- ponents within the RNN cell: indeed, with (9) we effectively reducing single dj -imeensional RNN cell to d,, independent 1-dimensional ones. This has an impact on expressivity, as we ex- plore in Sec. 5.2. We stress however that this choice of structure is not inherent requirement for

---

Published as a conference paper at ICLR 2026
ParaRNN. Alternative Jacobian shapes can be accommodated by the framework, so long they allow for efficient parallel reduction (see Walker et al. (2025) for an excellent overview of feasible structure choices and their implication expressivity, or Yang et al. (2025); Siems et al. (2025) for applications with Householder matrices). We adopt the diagonal constraint here merely to keep work contained, and align with standard practices in modern SSMs: Mamba employs a di- agonal state matrix (further simplified to single scalar parameter in Mamba2), and diagonalized RNN cells appear works ranging from Lei al. (2018) (to boost efficiency in the sequential cell application) to more recent efforts like Farsang et al. (2025) (targeting parallelizability, as we do). This notwithstanding, exploring alternative structures within ParaRNN remains an exciting avenue for future work.
4 EFFICIENT IMPLEMENTATION OF PARALLEL REDUCTION
As part of the main contributions of this work, we make available ParaRNN, fully modular, Py PyTorch-CUDA-bssed software framework allowing to readily parallelize the application of vir- tually any user-defined RNN cell. The user needs only to prescribe the recurrence step defining the cell action (including any parameters necessary to its definition), and the ParaRNN framework takes care of automatically assembling the resulting all-at-once system. Particularly, it leverages autogras rd assemble the Jacobians required in (3), and then proceeds to efficiently solve the system via Newton's 's method and parallel reduction (Alg. 1), parallelizing the RNN cell application.
In ParaRNN we provide three different implementations of the parallel solver, allowing users to trade off ease-of-use against performance:
1. Pure PyTorch A reference implementation of parallel reduction which only relies on native Py- Torch operations, and mainly thought for prototyping, debugging, and exploring new RNN cell definitions. While not directly optimized for performance, this version provides full automatic differentiation support and seamless integration with existing Py PyTorch models. By manually providing the explicit formula for the Jacobians assembly, the user can still choose to boost performance by overriding the default aut.ograd-an. add-baed ones.
2. CUDA-Accelerated Parallel Reduction A performant implementation featuring custom CUDA kernel for the parallel reduction solver, specialized for Jacobians with diagonal or N x N- block-diagonal structure in (8). Jacobian assembly remains in PyTorch, preserving automatic differentiation capabilities while accelerating the computational bottleneck. The CUDA solver employs a hardware-aware hybrid algorithm combining forward substitution and parallel re- duction at distinct levels in the GPU hierarchy (thread, warp, block and device, see NVIDIA (2025)). The goal to maximize register utilization and take full advantage of efficient warp- level directories, while minimizing shared and global memory traffic: we refer to App. D.2 for full implementation details.
3. Fully-fused Our most performant implementation where the whole Newton routine, including iterative Jacobian assembly and parallel solution, is fused within a single CUDA kernel. This eliminates intermediate memory traffic and overhead from kernel function calls, although it re- quires users to provide a CUDA implementation of the RNN cell. Nonetheless, we still preserve modular approach, so that the user only needs to define those operations which directly per- tain to the cell itself (namely, recurrence step and Jacobian assembly: see App. D.3). Our code provides reference implementation for the ParaGRU and ParaLSTM cells described in Sec.
RESULTS
The experiments reported in this section aim at evaluating the competitiveness of ParaRNN in Lan- guage Modeling tasks, particularly when applied to the ParaGRU and PSraL.TTM cells. To this purpose, we focus on two main metrics. First, in Sec. 5.1, we investigate the computational cost (primarily in terms of wall-time) associated with utilizing ParaRNN: our goal is to demonstrate that, by unlocking eqquence-parallelizabll with a runtime comparable to that of Transformer and Mamba architectures; and (ii) achieve infer- ence speed much faster than Transformer, and similar to Mamba. Secondly, in Sec. 5.2 we analyze the overall performance of ParaGRU and STM as language models, and show that they can achieve competitive perplexity and downstream task performance.
ParaRNN allows to (i) train nonlinear-RNWg -artear--NN-baeee models
7

---

Published as a conference paper at ICLR 2026
PARALLEL REDUCTION OPERATIONS
FOLL RNN CELL FORWARD PASS
BickDlog2s2
Pasd
10
100
Timings [ms]
Timings [ms]
10
10
lox
10-2
10
22
24
210
21
288
2.22
2
25
2
21
2
2
Sequence length L
Sequence length
Figure 2: Timing results for the parallel reduction operations Alg. lb (left), and the full parallel RNN application Alg. Ia (right), applied to sequences of varying length. Colors refer to different variants of the solvers: for diagonal Jacobians (blue, used by ParaGRUU and 2x2 block-diagonal ones (green, used by ParaLSTM), Heavier lines refer progressively more efficient implementa- tions: in PyTorch, in CUDA, and (for the RNN forward pass) with the fully-fused CUDA kernel. For reference, equivalent timing results for Mamba are also included (yellow). Actual speedup mea- surements reported in the main text refer to L = 2 (highlighted).
5.1 IMPLEMENTATION PERFORMANCE
Competitive Runtime Here we report profiling results for the main functions involved in the ap- plication of ParaRNN. First, in Fig. 2 (left) time different implementations of the parallel reduction operation Alg. lb. Notably, our CUDA implementation for diagonal Jacobians slightly improves on the equivalent paralles wwtaalee_ssee esccn shipped with Mamba (~ 1.1x speedup L 2"); the one for block-diagonal Jacobians is instead slower (~ 0.84x slowdown at L 2 as ex- pected, being it more memory and computationally intensive. Either case, the brunt of the kernel cost is associated with the memory reading overhead, rather than the computation itself, as indicated by the lines being mostly flat until I. 211, The inherent complexity of the operation is instead more visible from the profiling of the PyTorch implementation: this closely follows logarithmic growth, before collapsing to linear (from L 2 21x), when the GPU is close to capacity and opera- tions must be sequentialized. The profiling curves also mirror the three regimes of parallel scan implementation described in App. D.2: (i) the sequence fits in a single GPU block and is analyzed in parallel within it (L S 20) (ii) the sequence is analyzed in parallel within each block, but quentially over a few blocks (2 S L L 216, where the curve rises to a linear regime); and (iii) the sequence is analyzed in parallel across blocks (L 2-6 where we observe the curve flattening again to log). By contrast, Mamba's implementation collapses to sequential across blocks for L 2 220
Similar observations relative performance hold also when consider the complete parallel application of the RNN cell forward thss-that is, the full Newton routine Alg. instead of the sole parallel reduction operation Alg. lb. Profiling results for this setting are reported Fig. (right). Already our implementation combining Jacobians assembly in PyTorch and parallel reduction in CUDA provides runtimes equivalent to the Mamba SSM application: this is basically comparable to applying ParaGRU (we attain 1.2x speedup for L = 29), while ParaLSTM still lags behind (with a 0.5x slowdown), again following expectations. However,
by fully-fusing the whole Newton routine in CUDA that achieve the largest speedups with respect Mamba: of 2.6x and ~ 1.5x for ParaGRU and ParaLSTM respectively, at L = 29.
Fast Inference For inference, we observe a runtime similar to that of Mamba, shown in Fig. 3: regime, we attain a throughput of 37 338 tokens per second, versus the of Mamba. Notice moreover that, unlike for Transformers, for RNNs the time necessary to autoregressively produce the next token the output does not grow with sequence length.
8

---

Published as a conference paper at ICLR 2026
Overall, the measurements reported in this sec- tion support our claim that ParaRNN represents viable option for accelerating the application of RNNs, with runtime competitive with (and often beating) Mamba. We refer to App. for
MODEL THROUGHPUT
10a
Throughput [Tokeessss]
additional timing results and more details on the setup for the experiments considered here.
10
PunGRU
tuch 37 Brr
5.2 LANGUAGE MODELING PERFORMANCE
PodL5TM
Merthe
28 Mir
Having established runtime parity between Mamba and our RNN cells, we now analyze different axis of ParaGRU and ParaLSTM ef- fectiveness: their performance guage models. Our goal is to show that, with se- quentiality constraints removed and training at scale enabled, even classical RNNs can achieve competitive performance language modeling tasks.
10
2
22
24
26
2
2
Number of new tokens generated L
actual lan-
Figure 3: Tokens throughput at generation for ParaGRU (blue), ParaL.STM (green), Mamba (yellow), and Transformers (purple).
The Importance of Nonlinearities We begin by evaluating ParaGRU and ParaLSTM on synthetic tasks that isolate fundamental capabilities required for language modeling. We focus on retrieval and state-tracking tasks (Arora et al., 2023), where linear SSMs like Mamba are known to strug- gle (Merrill al., 2025) even though minor modifications the eigenvalues distribution of their state matrix can help. at least for specific tasks (Grazzi et al., 2025). The goal of this section is to es- tablish the critical motivation behind our work: nonlinear RNNs provide computational expressivity that linear RNNs fundamentally lack, and hence merit continued development.
To directly assess the expressivity of the recurrent cell in isolation, we use single-layer architec- tures consisting only of an embedding layer, normalization, the RNN cell considered (or Attention, for comparison), and a linear layer-see a detailed description in App. C.2. As shown in Tab. 1, ParaGRU and ParaLSTM can solve perfectly the majority of tasks considered, with near-perfect solutions recovered for Cycle Navigation and Modular Arithmetic tasks. However, both cells (as well as Mamba) struggle the more complex As task. This limitation is directly attributable the diagonal structure adopted in our work (see also Merrill et al. (2025)), as evidenced by the perfect scores achieved by vanilla LSTT and GRU cells with full facobians-all reported in Tab. 1 for reference. Most importantly, though, the nonlinear cells consistently outperform the linear Mamba2 across all tasks considered, highlighting the critical role of nonlincarities in enhancing expressivity.
Matching Transformers Without Attention Moving beyond synthetic tasks, we now evaluate our cells on language modeling to assess their effectiveness. To this end, we train four models: our ParaGRU and ParaLSTM presented in Sec. 3, plus a Transformer with the DCLM architecture (Li et al., 2024) and Mamba2 (Dao & Gu, 2024) as baselines. For each architecture, we consider four different model sizes (of 400M, 1B, 2.9B, and parameters), and follow Chinchilla-optimal (Hoff- mann al., 2022) scaling for their training setup, using SlimPajama (Soboleva et al., 2023) without the Books3 split as training dataset: a complete description of the training procedure, including
Table I: Single-layer accuracy on synthetic tasks. ()* denotes results computed training only the RNN cell (see App. C.2 for details).
Model
Keep-n"
Parity!
MQAR K
k-hop
k-hop
Cycle Nav
Mod Arithm (w/o brackets)
Copy Mcm
Aa
== ||== =1=00 10
k 2
M |V]= 10
IV| = 128
[V] = 2
IV 128
IV= = 5
C 6
|V 10
LSTM
100%
100%
100%
100%
100%
100%
100%
100%
100%
GRU
100%
100%
100%
100%
100%
ParaL.TMM
100%
100%
100%
100%
100%
95%
94%
67%
38%
ParaGRU
100%
100%
100%
100%
100%
97%
90%
63%
40%
Mamhu2
52%
100%
98%
Transformer
100%
53%
100%
82%
78%
100%
57%
100%
28%

---

Published as a conference paper at ICLR 2026
Table Perplexity, parameters count, and evaluation scores on 1m Lm-eval eval-harnes: (Gao al., 2021) tasks, for 7B models. Accuracies in percentages. Shot counts are reported in brackets.
HSSaa (10) (0)
1 OBQA
WinoG (5)
PIQA (0)
MMLU
Model
*params
PPL.
(25)
(3)
(10) (0)
(0)
(0)
Mamba2
6.96B
8.62
40.02
39.59
69.78 69.68
42.20 42.20
65.19 63.77
76.66
26.61
Paral STM
6.76B
9.16
37.46 36.52
62.47 62.85
42.20 38.80
57.70 59.12
75.19
25.31
ParaGRU
6.76B
39.68 36.77
65.85 65.75
42.20
61.40 59.83
76.66
25.29
Transformer
6.89B
9.55
34,30
33.36
62.98 62.20
40.00 37.20
61.48 60.85
74.97
23.12
hyprraarameters configuration and training times is reported in App. C.I. In designing our RNN models, we build upon the DCLM Transformer backbone, replacing attention with our RNN cells, but including also the causal convolution and gated residual layers from Mamba: detailed architec- ture specifications are listed in App. C.I, and a schematic of the RNN block illustrated Fig. 10.
Table 2 presents the downstream task performance of our 7B- 7B-parameters trained models on reference tasks from the Im-eval -hh rness suite (Gao al., 2021) (see Tab. for complete results across scales). Results reveal consistent ordering, in line with the perplexities reported in Fig. I: Mamba2 achieves the strongest performance, followed by our ParaGRU and ParaLSTM models, with the DCLM Transformer trailing behind. To properly contextualize these results, we remark that our goal not to propose a novel, more performant RNN, but to prove that also classical nonlinear RNNs can be trained efficiently at scale through parallelization, achieving performance competitive with modern language models. Ultimately, we are showing that the training inefficiencies that historically sidelined classic RNNs can be overcome, reviving them viable options for language modeling.
CONCLUSION AND FUTURE WORK
In this work, demonstrate that the paraleiiaaaallllizaaioonn barrier for nonlinear RNNs can be effectively overcome even at scale, the linearity constraint no longer precluding additional explo- rations in efficient sequence modeling. Our empirical results show that even classical RNNs, when trained with ParaRNN, can achieve performance competitive with modern without sacrificing training efficiency. This parity represents an important first step: establishes that nonlinear RNNs are again a viable alternative in the modern landscape of sequence modeling.
ceatiecttrrs-ruuciille
In addition, the ParaRNN codebase makes this alternative readily available to practitioners, while delivering performance that matches or exceeds existing parallel implementations. By proving that classical RNNs can compete on equal footing with current architectures once computational barriers removed, aim for this work to further reignite exploration into nonlinear recurrent models.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We would like to extend our thanks to: Hadi Pour Ansari, for all the help during the reviewing pro- cess of this manuscript; Eeshan Gunesh Dhekane and Jagrit Digani, for the extremely valuable feed- back, particularly regarding the CUDA kernels implementation; and Nick Apostoloff and Jerremy Holland, for their support throughout this project. Finally, we are grateful the ICLR Reviewers and Area Chair for their recommendations and contribution in increasing the quality of this work.
REPRODUCIBILITY STATEMENT
To ensure reproducibility, all our models were trained using public datasets and open source libraries. The code reproduce the algorithms described this work is available on GitHub at https: g thub. cwm/ca////////oo sammapple; e/mlm=mm= aal-pararmn /nn/. In addition, this document and its appendices contains all the necessary details to reproduce the experiments. We also include pseudo-code for the Newton and Parallel reduction algorithms in Alg. 1, as well as details and samples of the software package implementation in App. D, and hyperparameter details in App. B and C.I.
ETHICS STATEMENT
Our work primarily focuses on advancing the computational efficiency and scalability recurrent neural networks through novel algorithmic and architectural innovations. As such, it is founda- tional technical contribution that does not directly involve the collection, processing, or deployment of sensitive user data, nor does it present immediate ethical concerns related to bias, privacy, or misuse in its core methodology. Furthermore, in our efforts to promote responsible AI develop- ment, have taken proactive steps regarding data integrity and copyright, such as the removal of copyrighted material like Books3 from SlimPajama. We are committed to responsible research and development, and by open-soureing the ParaRNN codebase, we aim to encourage transparent and collaborative progress in the field of generative modeling.
REFERENCES
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language mod- cls, 2023. URL https: aaaaiiivvv.v www.iii.rgg//....e.ooooo.ooooete htte/abs/2333 04927.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xLSTM: Ex- tended long short-term memory, 2024. URL https: wwg. /400 04517.
Shane Bergsma, Nolan Dey, Gurpreet Gosal, Gavia Gray, Daria Soboleva, and Joel Hestness. Straight to zero: Why linearly decaying the learning rate to zero works best for LLMs, 2025. URL https::////// /tt///axivvv org/ abs 220.225022222 15938.
Guy E. Blelloch. Prefix sums and their applications. Technical Report CMU-CS-90-190, School of Computer Science, Carnegie Mellon University, November 1990.
James Bradbury, Stephen Merity, Caiming Xiong, and Richard Socher. Quasi-recurrent neural net- works, 2016. URL https:
pitetps://rxxiy/ //tp://aai...//../.aa/eae/aaae.a wwtp:bs.////oootooooooo... 611. 01576.
Francois Chaubard and Mykel Kochenderfer. Scaling recurrent neural networks billion param- eters with zero-order optimization, 2025. URL http:: / http:///wwwwwwwwww///sooooess 22055 7852.
Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation, 2014. URL peattss://ariis org/ abs/ 11.406..0770 1078.
Nicola Muca Cirone, Antonio Orvieto, Benjamin Walker, Cristopher Salvi, and Terry Lyons. Theo- retical foundations of deep selective state-space models, 2025. URL https: / w/aariivvoorr.....e..........aaa abs /2402 9047.
11

---

Published as a conference paper at ICLR 2026
Federico Danieli, Miguel Sarabia, Xavier Suau, Pau Rodríguez, and Luca Zappella. DeepPCR: Parallelizing sequential operations in neural networks, 2023. URL https:: htttpp::/aaxii.orrt abs (tss/2099.1331,
Tri Dao and Albert Gu. Transformers SSMs: Generalized models and efficient algorithms through structured state space duality, 2024. URL https://arxx.v.rrg//attt 2405. 21 21060.
Mónika Farsang, Ramin Hasani, Daniela Rus, and Radu Grosu. Scaling up liquid-resistance liquid- capacitance networks for efficient sequence modeling, 2025. URL https: htttp:://arxiv.orgg
abs/2505.2174 217117.
Leo Feng, Frederick Tung, Mohamed Osama Ahmed, Yoshua Bengio, and Hossein Hajimirsadeghi. Were RNNs all we needed?, 2024. URL https: http:///ww......aaeeeoeee abs 2010. 01201.
Martin Gander. 50 Years of Time Parallel Time Integration, volume 9, 69-113. Springer Inter- national Publishing, 2015. ISBN 978-3-319-23320-8. doi: 10. 00.10077978-3331--2331111
Leo Gao, Jonathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot lan- guage model evaluation, September 2021. URL https: dei. org/ /eeooo 5371628.
Xavier Gonzalez, Andrew Warrington, Jimmy T. H. Smith, and Scott W. Linderman. Towards scalable and stable parallelization of nonlinear RNNs, 2024. URL https::/taaaae /ttps///xxxvvv/////// abs/2407 19115.
Xavier Gonzalez, Leo Kozachkov, David M. Zoltowski, Kenneth L. Clarkson, and Scott W. Lin- derman. Predictability enables parallelization of nonlinear state space models, 2025. URL
https://arxiv.ct org/ wttpg//////////e abs/ 16817
Riccardo Grazzi, Julien Siems, Arber Zela, Jörg K. H. Franke, Frank Hutter, and Massimiliano Pontil. Unlocking state-tracking in linear rnns through negative eigenvalues, 2025. URL https:
//arxiv. org abs/ //////11..111 2411..21. 12537.
Klaus Greff, Rupesh K. Srivastava, Jan Koutnik, Bas R. Steunebrink, and Jürgen Schmidhuber. LSTM: A search space odyssey. IEEE Transactions on Neural Networks and Learning Systems, 28(10):2222 2232, October 2017. ISSN 2162-2388. doi: 10.1 10.1109/mnss.2066225528
URL
http: /ttp://dxx doi. org 1t/g/11.00 11199/TNNLS 2016 6. 2582924.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv
preprint arXiv:23/2.00752, 2023.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. HiPPO: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. In The International Conference on Learning Representations (ICLR), 2022a.
Albert Gu, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models, 2022b. URL https: arrivv.ivvv//////////...///aaavv wttp/rrg//abb////// 11893.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on ImageNet classification, 2015. URL https://arxiv..d abs/1502. 01852.
W. Daniel Hillis and Guy L. Steele. Data parallel algorithms. Commun. ACM, 1700:111770 December 1986. ISSN 0001-0782. doi: 10.1145/7902.7903. URL https::t//// //:o////oiii /00 1145/7902.7903.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural Computation, 9:1735- 1780, 11 1997. doi: 10.1 162/neco.. 1997.9.8. 1735.
12

---

Published as a conference paper at ICLR 2026
Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, Tom Hen- nigan, Eric Noland, Katie Millican, George den Driessche, Bogdan Damoc, Aurelia Guy, Simon Osindero, Karen Simonyan, Erich Flsen, Jack W. Rae, Oriol Vinyals, and Laurent Sifre. Training compute-optimal large language models, 2022. URL hnttps://arxiv.... obs//bb/////// 2203.15556
Ningyuan Huang, Miguel Sarabia, Abhinav Moudgil, Pau Rodriguez, Luca Zappella, and Federico Danieli. Understanding imput selectivity in mamba: Impact approximation power, memoriza- tion, and associative recall capacity, 2025. htt.:tttes//aaa...orr//b//////// 118911
11891.
Michael Keiblinger. Recurrence-compler frame-based action models, 2025. URL nttps://
arxiv.org/ /btpgg//bb///10/0 06828.
T. Anderson Keller, Lyle Muller, Terrence Sejnowski, and Max Welling. Traveling waves encode the recent past and enhance sequence learning, 2024. URL 08045.
Endre László, Mike Giles, and Jeremy Appleyard. Manycore algorithms for batch scalar and block tridiagonal solvers. ACM Trans. Math. Softw., 42(4), June 2016. ISSN 0098-3500. doi: 10. 145/ 2830568. URL hwtph:ttp///mm /ttttp:::/di...o /ttp///oii//a org.orgg////ooo 11411115555 444/2830568.
Tao Lei, Yu Zhang, Sida I. Wang, Hui Dai, and Yoav Artzi. Simple recurrent units for highly parallelizable recurrence, 2018. URL https:ttp//// /ttp:ss///rxvvv////ee org.abs.orgaass.. /1099 02755.
Jeffrey Li, Alex Fang, Georgios Smymis, Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guba, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bit- ton, Marianna Nezhurina, Amro Abbas, Cheng- Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer, Kyle Lo, Alaaeldin El-Nouby. Hadi Pouransari, Alexander Toshev, Stephanie Wang, Dirk Groen- eveld, Luca Soldaini, Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. DataComp-LM:: In search of the next generation of training sets for language models. arXiv preprint arXiv:2406. 1194, 2024.
Yi Heng Lim, Qi Zhu, Joshua Selfridge, and Muhammad Firmansyah Kasim. Parallelizing non- linear sequential models over the sequence length, 2024. Uttphttss://exxiv.coo abs/ 2309. 12252.
Eric Martin and Chris Cundy. Parallelizing linear recurrent neural nets over sequence length, 2018. URL https: http///axxv..r////////..//.// abs/ /bss/1799 1709..1..000 04057.
William Merrill, Jackson Petty, and Ashish Sabharwal. The illusion of state in state-space models, 2025. URL
Corporation NVIDIA. CUDA C++ Programming Guide, 2025. URL htLthtss://doo nvidia.. /./.de/... Version 13.0.
com/cuda wwtaacuaaaaac cosiaaa--prrggaamme
J. M. Ortega and W. C. Rheinboldt. Iterative Solution of Nonlinear Equations in Several Variables. Society for Industrial and Applied Mathematies, 2000. doi: 10.1137/1.97808988 hittps://epube. swmppuussssaam.....m.... siam.
117780888799488.... 19468. URL /ttp:////...111111..177999966661
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences, 2023. URL
06349.
Zhen Qin, Songlin Yang, and Yiran Zhong. Hierarchically gated recurrent neural network for se- quence modeling, 2023. https:::: / wrw.obrgss......oooo. ablaabas21211 /2311 04823.
13

---

Published as a conference paper at ICLR 2026
Andrea Santilli, Silvio Severino, Emilian Postolache, Valentino Maiorca, Michele Mancusi, Ric- cardo Marin, and Emanuele Rodola. Accelerating transformer inference for translation via par- allel decoding. In Proceedings of the 6lst Annual Meeting the Association for Computa- tional Linguistics (Volume I: Long Papers). Association for Computational Linguistics, 2023. doi:
URL http://ddd.....oooo/...888/////0
acl acl-long. 689.
Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffusion models, 2023. URL htwps: / /tt./vaxvvvv.////ee org/ abs /2305 6317.
Julien Siems, Timur Carstensen, Arber Zela, Frank Hutter, Massimiliano Pontil, and Riccardo Grazzi. Deltaproduct: Improving tracking in linear mns via houscholder products, 2025.
URL httph:t/////. httph::t///aatte /tte:///axiivii org/ abs/ /200 10297.
Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. SlimPajama: 627B token cleaned and deduplicated version of Red- Pajama. https://cerebras: www.eebrassaaabbbba.o aaap//aa///iimaaaaaaaaaaaaaaaaaa aaaa-a--27bbttoeee
June 2023.
URL https://
huggingface. cooddataeets. wttptt/eeeeeebbrreeeeeeee /mimma jama- 627B.
Yang Song, Chenlin Meng, Renjie Liao, and Stefano Ermon. Accelerating feedforward computa- tion via parallel nonlinear equation solving, 2021. URL htKLittpee//rrriv.orr///aas/
03629.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017. URL https://arxi:: org/abs/ /ww..110.... 62.62.
Benjamin Walker, Lingyi Yang, Nicola Muca Cirone, Cristopher Salvi, and Terry Lyons. Structured linear cdes: Maximally expressive and parallel-in-time sequence models, 2025, URL https: 17761.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- with the delta rule over sequence length, 2025. URL htL http://aawxvv.../bbb.b.bee /bclgg/bb/24// 2406. 06484.
Jingwei Zuo, Maksim Velikanov, Dhia Eddine Rhaiem, Ilyas Chahed, Younes Belkada, Guillaume Kunsch, and Hakim Hacid. Falcon Mamba: The first competitive attention-free 7B language model, 2024. URL https: / //arxivv.vv////////a//eeeeaaaa org/or/g////.ooo.oooooo.oe htsp/b1s/24000 05355.
14

---

Published as a conference paper at ICLR 2026
A
NEWTON CONVERGENCE FOR PARAGRU AND PARALSTM
Summary
In this section, empirically verify that ParaRNN's Newton algorithm is stable and ef- ficient when applied to ParaGRU and ParaLSTM, consistently converging in iterations, regardless of RNN cell used and stage of training considered.
In Sec. 2.1 we hinted at the necessity of the Newton's method Alg. converge quickly for ParaRNN to be appealing as parallelization method for RNN cells upplications. In practice, solving (2) in parallel comes at a computational complexity of (Ns log2 L), Nis is the number of Newton iterations required for convergence: keeping Nus ~ ((1) is then paramount. Ultimately, the convergence behavior of Newton's s method will depend on the actual definition of the cell action in (I), and providing comprehensive theoretical study of the properties of necessary for fast convergence is beyond the scope of this project and a topic for future research-although we refer to Gonzalez al. (2025) for the relevance of Lyapunov stability in this sense. Here, limit ourselves to demonstrating empirically that, when applied to the ParaGRU and ParaLSTM models considered in this work, Newton's 's method consistently converges around the third iteration.
To this end, in Fig. we examine the convergence behavior of Newton's 's method across different stages of model training and for varying sequence lengths. By default, in all forward applications of any RNN cell (including during the training of the models described Sec. 5.2), we consider as an initial guess & for the Newton iterations the quantity
h = f0, 23),
VI 1, L.
(10)
That is, for each input ø,, we apply the cell action considering an all-zero previous hidden state 0. Notice this is computed in perfectly parallel fashion along the sequence length L.
Figure considers two complementary regimes, to show that Newton iterations are stable through- out the training procedure. On the left we show residual norm convergence for freshly-initialize RNN cells (see App. C.1 for the initialization scheme considered), while on the right we consider RNN cells extracted from the last layer of the fully-trained 400M-parameter models described in App. C.I. As evident from the graphs, the method exhibits rapid convergence: the residual drops to machine precision within 3-4 iterations, regardless of RNN cell considered training stage, and convergence rate is only very marginally affected by an increase in the sequence length. For freshly initialized ParaLSTM, the residual at the very first iteration increases, indicating potential marginal
AT
END TRAINING
Residual mose, ||A] J0 and bo
ParCCRUU
10-5
109
PuraLSTM
10-3
3
10
3
10
Newton sterations Nis
Newton iterations Nin
Figure 4: Newton's method convergence behavior for ParaGRU (top) and ParaLSTM (bottom) cells with different input sequence lengths L., for freshly initialized RNN cells (left), and for the RNN cells in the last layer of our trained 400M models (right). Inputs randomly selected from the evaluation datasets, and the residuals reported are the max on batches of size 8.
15

---

Published as a conference paper at ICLR 2026
for improvement in defining initial guess for Newton's method. For ParaGRU, just 2 iterations suffice at initialization, growing to at the end training. These results guided our choice of Nis- for the experiments in Sec. 5.2.
This consistent behavior across both initialization and trained states suggests that the nonlinearity structure of the ParaGRU and ParaLSTM cells lends itself well to solution via Newton's method, but-we we remark- is not in general guaranteed for every choice RNN cell.
16

---

Published as a conference paper at ICLR 2026
B DETAILS ON PROFILING RESULTS
Summary
In this section expand the results in Sec. 5.1 reporting timings for all application modalities of the RNN cells, including sequential and pure PyTorch implementations. The fully-fused forward pass applications of ParaRNN are always faster comparable to Mamba for all the sequence lengths tested, and beat Mamba consistently when the back- ward pass is also taken into account. We also present ablation over the chunk_size hyperparameter, controlling thread-level workload in our CUDA implementations. The resulting optimal values are model- and application mode-specific, with values generally yielding best runtimes.
We additionally confirm the throughput results in Fig. 3 hold across all scales con- sidered, and that the parallel application of the RNN cells requires ( dr)) memory.
Experiments Setup All timings in Sec. 5.1 are collected on NVIDIA A100 GPU. For Fig. 2, profiling is performed by tracking CUDA Events (NVIDIA, 2025, Sec. 6.2.8.8.2).. The profiler first warms-up the GPU with 20 dry runs of the target kernel measure, before proceeding to collecting measurements from 100 runs processing random input sequences of a given length L. In the figures report the minimum? timings for each sequence length considered. To provide fairer comparison against production-ready code, all PyTTorch baseline measurements (including Mamba) performed using torch. comp le () -optimized code, rather than eager execution mode, as this represents the standard deployment configuration for real-world applications. Also for the sake of fairness, for the Mamba baselines in Fig. we consider: on the left, the application of the sole parallel scan routine at the core of Mamba's SSM application; on the right, the application of simplified Mamba module, stripped of gate and convolution components (Gu & Dao, 2023) in other words, we track only those operations necessary for the assembly of the SSM parameters and the application the SSM cell. Finally, full RNN cell application timings consider the default setup with Nir- = 3 Newton iterations.
For the throughput results Fig. 3, we consider the 1B models discussed in Sec. 5.2. profile, simply use Python's 's timeit module, and report the minimum timings over 5 repetitions an NVIDIA A100 GPU. The reduced number of repetitions in this case mainly due to the lengthy generation times required by the Transformer; nonetheless, the reported curves appear smooth, so we consider them sufficient to provide an accurate measure.
B.1 ADDITIONAL TIMING RESULTS
Forward / Backward Pass Varying Application Modality In Fig. we expand on the results reported in Fig. 2. including timings for the sequential application of the RNN cells, as well as from the pure PyTorch implementation. The former results were also used to compute the speedup measurements reported in Fig. 1. Overall, the curves behave as expected: applying the RNN cells sequentially becomes quickly computationally unfeasible due to its linear-growth complexity, and moving to progressively more efficient implementations (from sequential, parallel in PyTorch, to computing the parallel reduction operation in CUDA, to a fully-fused CUDA kernel for the whole Newton routine) lowers the overall runtime for the RNN application. Perhaps surprisingly, the curve for the PyTorch implementation of ParaGRU in Fig. 5 (left) is ragged. We speculate that this is due to the effect of torch. comp le (), which for this cell is able to unlock certain optimizations for specific sequence lengths but not for others. Notice also that the CUDA implementation of ParaGRU surpasses the fully-fused one for L 212 To explain this unexpected result we need to refer to the specifics of the CUDA kernel implementation outlined in App. D.2. Fully-fused kernels exert larger pressure on the GPU registers, in light of the larger number of computations they must
We report the minimum following established benchmarking practices it represents a cleaner metric, and a more faithful indication of the actual performance of a kernel operating under ideal conditions. The average, on the other hand, might be polluted by transient system interference (e.g., synchronization artifacts, GPU thermal throttling events, OS-induced delays, unrelated to algorithmic efficiency.
nigh-peeformance computing:
17

---

Published as a conference paper at ICLR 2026
PARAGRU FORWARD
PARALSTM FORWARD PASS
10a
1yTrch CUDA
PyTieck
CUDA
10
Foxt
Paxed
Timinns [ms)
Timings [ms]
10c
100
10-1
10
2
10-2
22
25
2
211
224
247
22
2
25
2
214
217
Soquence length
Soquence longth L
Figure 5: Timings of different application modalities for the ParaGRU (left) and ParaL.TMM (right) RNN cells, applied sequences of varying length. Heavier lines refer to progressively more efficient applications: sequential, using parallel reduction in PyTorch (Alg. Ia), using parallel reduction in CUDA and system assembly Py PyTorch, and using fully-fused CUDA kernel for the whole Newton routine. Equivalent timing results for Mamba are also included for reference (yellow). Same setup as for Fig. 2 (right).
perform (Jacobian assembly and Newton iterations, on top of parallel reduction). To ensure that GPU block has enough resources carry these computations, then, chunk.s ze hyperparameter which controls the number of equations single thread processes. This in turn triggers the block-sequential solution (hence the collapse linear regime) for shorter sequence lengths than for the CUDA implementation (L > 20 instead of > 2 2),, thus explaining the behavior of the curves in Fig. (left).
forced reduce the
In Fig. we further provide timings for full parallel forward and backward application of our RNN cells, using dummy loss function C(h) = E, h, The setup is analogous to that for Fig. 5, and the observations regarding relative performance of the various implementations remain sub- stantially unchanged. Comparing with Fig. 5, however, we can notice that including the backward pass affects the timings for Mamba much more than that for ParaGRU or ParaLSTM. This in line with expectations: in Mamba, both forward and backward passes involve the application of single parallel reduction operation Alg. On the other hand, for nonlinear RNNs like ParaGRU or
PARAGRU BOCKARDDBACKWARD PASS
PARALSTM FORWARD+ PASS
PyTeech
PyTieck
CUDA
CUDA
10a
Fixci:
Faeed
10a
Timings [ms]
Timings [ms]
10a
10
-0
22
25
2
211
2
217
22
25
28
211
2
217
Soquence length
Soquence length L
Figure 6: Timings of different application modalities for full forward and backward pass of the ParaGRU (left) and Paral (right) RNN cells, applied sequences of varying length. Same setup as for Fig. 5.
18

---

Published as a conference paper at ICLR 2026
DIAGONAL
2 x BLOCK-DIAGOON
Manba
Martha
Timings [ms)
10
10-1
Timings [mN8]
10
10-2
22
20
210
1
2
222
25
220
214
218
222
Sequence length L
Sequence length
Figure 7: Timings of the CUDA implementation of parallel reduction, for diagonal (left) and 2 x block-diagonal (right) Jacobians, varying its chunk_s ze hyperparameter Timings for Mamba implementation of parallel.scan are also reported as reference.
ParaLSTM, the forward pass requires multiple parallel reductions-one per each Newton iteration, that is Nis 3 in our case case-while the backward pass only requires single parallel reduction, as outlined in Sec.
Ablation Over Kernel chunk_size siize Hyperparameter In App. D.2.2 we comment on the vari- ous compile-time hyperparameters one can tune achieve the highest performance out of ParaRNN for their target setup. Likely the most relevant one is chunk size, which defines the amount of work per thread-thi that is, the number of equations each thread must solve individually. Here we report the results of the sweep we conducted with the goal of selecting its optimal value. In par- ticular, Fig. we show timings for the sole parallel reduction operation applied to diagonal (left, used by ParaGRU) and 2 x 2 block-diagonal Jacobians (right, used by ParaLSTMM,, when varying chunk.s i ze. For the diagonal case, we pick a value of 2 as the optimal, as achieves the best performance for L S 2 but 4 would perform marginally better for L > 2 Also for the block- -iiao-daaaonal counterpart, we consider c = 2, as it performs the best in the regime relevant to our training 210 $ I. S 211, although for shorter sentences
might be preferable.
PARAGRU FORWARD PASS
PARALSTM FORWARD PASS
10c
Marla
Marha
109
Timings [ms]
Timings [ms]
10-1
10
1
10-2
2a
25
2
211
2
217
22
25
28
211
2
217
Soqueace length
Soquence length L
Figure 8: Timings of the implementation of the fully-fused kernel for parallel RNN application, for ParaGRU (left) and ParaL.ST' (right), varying its chunk_s size hyperparameter Timings for Mamba SSM application are also reported as reference.
19

---

Published as a conference paper at ICLR 2026
We conduct a similar sweep also for the fully-fused application of ParaGRU and ParaLSTM,, and report its results Fig. 8. For ParaGRU, the main advantage in increasing delaying the triggering of the linear regime, and indeed the choice of c 2 is optimal for the L = 2 2 used in our training; for shorter sequences, however, sticking to e = gives better runtimes. For ParaLSTTM this advantage is instead negligible, and using throughout gives the best runtimes overall.
1 consists in
B.2 ADDITIONAL THROUGHPUT RESULTS
In this section we expand on the results in Fig. 3, by reporting in Tab. how throughput measure- ments vary with the scale of the models considered. The setup follows the same used for Fig. 3.
Table 3: Throughput results at regime (L = 211) for the model types and sizes consid- ered.
Overall, the trend is confirmed: the constant token generation cost of recurrent models allows them to reach much higher throughputs than Transformers, even for reasonably short sequence lengths (we re- port results for T. 2 and RNNs perform marginally better than Mamba at all scales. No- tice moreover that, for the recurrent cells considered, runtime is barely affected by the hidden dimension- ality d,. This is in line with the cell application be- ing perfectly parallelizable across hidden state com- ponents, as granted by their (block-)diagone) )diagonal structure. Indeed, when increasing only the hidden state size (which we do when moving from 400M to IB models, and from 3B to 7B models see also Tab. for an outline of the model parameters considered at the various scales), throughput only marginally impacted, mostly due to the presence of the additional MLP components (whose application runtime instead does scale with dy.) ). Conversely, throughput is mainly affected by in- creasing the model depth (passing from IB to 3B), in line with expectations.
Model
400M
1B
3B
7B
ParaGRU
38.12
38.09
28.58
28.45
ParaLSTM
36.77
36.93
27.06
27.68
Mamba2
29.76
27.61
19.33
18.16
Transformer
6.45 1.98
1.00
0.42
B.3 MEMORY PROFILING
In our work we consider the simplification to the Jacobian structure outlined in Sec. 3, so that by design the cell application requires storage scaling linearly in both sequence length L and hidden state dimension d,. In this section, we verify that this indeed holds in practice. To this end, in Fig. 9 we report how peak GPU memory usage recorded during the application of the ParaGRU and ParaLSTM cells varies when scaling L or da. In both cases, perfect linear scaling is measured, with ParaLSTM requiring proportionally more storage, in light of its more complex Jacobians.
FIXED d, 210
FIXED 210
PenLSTX
POLLSTM
Peak Peak Memory Usage [MiB]
Peak Menury Usage [MiB]
10
10a
du
10a
10P
10-1
10
21
22
23
24
210
23
25
26
27
29
210
Sequence length L.
Hidden dimension da.
Figure 9: Peak GPU memory usage upon parallel application of the ParaGRU and Paral STM cells, varying sequence length L (left) and hidden state dimension dA (right).
20

---

Published as a conference paper at ICLR 2026
C
DETAILS ON LANGUAGE MODELING RESULTS
Summary
In this section we provide more information about the architectures we used, the training parameters, timing protocol, and results regarding the language modeling experiments. In particular, in Tab. we expand upon the results illustrated Tab. 2. providing the down- stream tasks evaluation scores across all model scales considered (400M, IB, 2.9B, and 7B). Performance remains consistent across different parameter counts: with Mamba2, ParaGRU and ParaLSTM often outperforming the Transformer. Importantly, our RNN cells show consistent improvement with scale, and the relative performance gaps between architectures remain stable.
C.I LANGUAGE MODELING
Models Description Schematics of the complete RNN block used for the lan- guage modeling experiments in Sec. 5 are shown in Fig. 10. The module inherits some components from Mamba (a pre- pending short sequence-wise causal con- volution, as well malization layer following the actual RNN cell application), and includes residual connection with a learnable scaling of the input. Still, our block follows closely the Transformer structure: particularly, MLP layers with residual connections are inter- woven with the sequence mixer, which in our case is an RNN cell, rather than Atten- tion. The core RNN cell is either ParaGRU or ParaLSTM from Sec. 3, but all other architectural components remain identical between the two variants. To aid with scaling, we use multi-head implemen- tation of our RNN cells, where the input 27 Rd is separated into Ahadds indepen- dent heads, each of size da = dia/ Theaads The whole block is then repeated raayers times to generate the full model.
Linear
gated RMS nor-
Gated
ParaRNN Cell
Scale
Convolution
Linear
Figure 10: Schematics of RNN block.
The baselines considered in our experiments are the Transformer architecture in DCLM (Li et al., 2024), and the Mamba2 model from HuggingFace. See Tab. for summary of hyperparameters further detailing the architectures used.
Training Setup For experiments consider four architecture types the ParaGRU and Par- aLSTM introduced in this work, plus Mamba2 and Transformer as baselines: see Tab. 4 for archi- tecture -pecifics-each four different scales, of 400M, 1B, 2.9B and 7B parameters. As training dataset, we use SlimPajama (SPJ) (Soboleva et al., 2023), from which we removed the Books3 split (as it is known contain copyrighted material). All models are trained using AdamW with de- tached weight decay and cosine learning rate schedule. incorporating a warm-up up for 10% of the total training iterations and decaying to a final learning rate of (Bergsma et al., 2025). Training conducted on nodes with 8 NVIDIA H100 GPUs each using PyTorch 2.6 with FSDP and native au- tomatic mixed precision, storing weights in bfloat16 6 while maintaining gradients and reduction operations in f f1oat32. Batch sizes and total number of training tokens follow Chinchilla scaling laws (Hoffmann et al., 2022), and we refer to Tab. 5 for the complete training configuration. Hyper- parameter selection is conducted by sweeping over three learning rate and weight decay values for each model and each size, selecting the configuration resulting in lowest perplexity on the SPJ test set (see Tab. 4 for the optimal values selected in runs). Model initialization follows architecture-
21

---

Published as a conference paper at ICLR 2026
Table 4: Architecture specifications and optimal hyperparameters for our experiments. We report structural parameters (number of layers, number of heads and model width) and training hyperpa- rameters (learning rate and weight decay) that achieved lowest perplexity for each model type and scale. Time measurements show average single-step training cost NVIDIA H100 GPUs, while Memory measurements show the per-sample peak memory usage during one optimization step.
Learning
Weight decay
Time mms/step]
Memory [GiB/sample]
Model type
#params
flayers
Thheads
dasodel
ParaGRU
406M
24
4
1024
0.005
1e-3
95.7
6.2
400M
406M
24
4
1024
0.003
le-3
116
6.2
Mamba2
409M
48
32
1024
0.005
115
5.2
Transformer
412M
24
8
1024
0.003
le-3
81.0
5.3
ParaGRU
1.42B
24
4
2048
0.003
lc-4
192
9.8
ParaLSTMM
1.42B
24
4
2048
0.003
le-4
235
Mamba2
1.43B
48
64
2048
0.005
le-4
221
9.1
Transformer
1.44B
24
16
2048
0.002
le-4
157
8.5
ParaGRU
2.90B
32
2
2560
0.002
le-4
326
16.8
2.9
Paral STTM
2.90B
32
2
2560
0.002
le-4
400
16.9
Mamba2
2.83B
64
80
2560
0.003
le-4
393
17.2
Transformer
2.80B
32
32
2560
0.003
le-4
267
14.1
ParaGRU
6.76B
32
4
4096
0.003
le-4
613
38.2
ParaLLTM
6.76B
32
4
4096
0.003
le-4
715
39.3
Mamba2
6.96B
64
128
4096
0.003
le-4
9475
39.8
Transformer
6.89B
32
32
4096
0.003
le-4
520
31.0
specific strategies: for Mamba2, we use the default initialization available in HuggingFace. For Transformers, we use DCLM initialization (Li et al., 2024). Since our ParaGRU and ParaLSTM blocks in Fig. 10 mimic the Transformer architecture, we opt for DCLM initialization for them as well, except for their core RNN cells, whose initialization depends on the parameter considered: for the cell input matrices B, in (7) we use Kaiming Uniform (He et al., 2015) initialization; for the (diagonal) state and peephole matrices a, and c, in (9) we choose Xavier Gaussian, since we found small values improved stability; the bias terms b, are instead initialized to O. We also clip the norm of a, and c, to maximum 0.5 to prevent overflows long sequences: this effectively bounds the Jacobians in (8), similarly to what is done in Mamba, where the state matrix eigenvalues are constrained in (0, 1).
the computation of the hidden state for
To illustrate the stability of our training procedure, Fig. 11 shows training and validation loss curves for our 18-parameter models. Notably, both ParaGRU and ParaLSTM exhibit smooth convergence throughout training especialy when compared to the baselines suggesting that the gradient flow in our recurrent models is well-behaved at scale.
Additional Results Tab. we expand upon the results illustrated in Tab. providing the down- stream tasks evaluation scores across all model scales considered (400M, 1B, 2.9B, and 7B). Over-
Table Scale-specific training configurations on nodes with 8 NVIDIA H100 GPUs each All mod- els use 2048-token sequences and follow (~ 1x) Chinchilla-optimal token budgets (Hoffmann et al., 2022) for batch-size selection. The z-loss regularization coefficient taken from DCLM (Li et al., 2024).
Model Scale
Batch Size per GPU
Effective Batch Size
Total Iterations
Warm-up Iterations
Training Tokens [B]
Nodes
Z-Loss
400M
32
256
le-04
20,000
2,000
10.5
IB
32
256
le-04
55,000
5,500
28.8
2.9B
512
le-04
55,000
5,500
57.7
7B
2,048
5c-06
33,000
3,300
138.4
22

---

Published as a conference paper at ICLR 2026
TRAINING
VALIDATION
RIT Paralo Adamba
ParaCRII ParakST3E
2.13
Kamba
Cross Entropy
Cross Entropy
0
10k
15k
20k
25k
30k
10k
15k
20k
25k
30k
Iteration
Iteration
Figure 11: Cross-entropy loss evolution for the 7B models considered, evaluated the training (left, collected every 150 iterations), and validation (right, collected every 3.3k iterations) datasets.
all, the performance observed is consistent across different parameter counts: while Mamba2 still achieves the lowest perplexity, ParaGRU and ParaLSTM remain competitive, with all three recur- rent architectures often outperforming the Transformer. Importantly, our RNN cells show consistent improvement with scale, and the relative performance gaps between architectures remain stable.
C.2 SYNTHETIC TASKS
To recover the results in Tab. 1, we trained single layer of ParaGRU, ParaLSTM, Mamba2, and Transformer across variety of synthetic tasks that have been used in the literature to evaluate the expressivity of sequential models. The evaluation tasks considered are:
Keep-nth Given input sequence [zi]/1 of values from dictionary V. output its n-th element, Zn.
Parity Given an input sequence of values in {0, 1}, output mod (Ei 21, 2).
Hluttppl-Query Associative Recall (MQAR) Given an input sequence of key-value pairs (k, v) e
K V, followed by noise interwoven with the keys, output for each key its associated value.
Table 6: Final perplexities, and evaluation Im-eval-harnen (Gao et al., 2021), for model types and scales considered.
reference downstream tasks from
Are-C (25)
HSwag (10) (0)
+ OBQA (10) (0)
WinoG (5)
PIQA (0)
MMLU
Model
#params
1. PPL
(3)
(0)
(0)
Transformer
412M
17.00
24.23
23.81
36.11 35.91
26.60 29.40
51.77
$1.69
64.64
25.37
MIOM
PurGGRU
406M
18.56
22.27
21.93
38.69
39.21
25.00
25.80
53.04
52.41
67.30
23.68
Pural.STM
406M
18.40
23.12
24.23
37.43
37.52
26.20
26.20
51.22 49.49
65,94
23.37
Mamba2
409M
16.49
24.32
25.09
40.53
41.00
30.00
34.00
50.99
51.14
66.97
23.41
Transformer
1.44B
12.90
26.87
25.51
47.25
47.14
34.00
33.40
$4.06
54.22
68.49
23.12
PuraGRUU
1.42B
13.20
24.15
23.81
48.71
48.85
24.80
25.80
53.43
53.28
71.27
24.55
STM
1.42B
13.11
24.15
24.32
45.82
46.53
24.40
27.00
51.14 51.07
70.46
24.75
Masubu2
1.43B
12.26
28.41
27.47
50.90
51.21
36.20
36.00
$5.25
$5.17
71.71
24.39
Transformer
2,80B
10.97
30.72
29.52
$5.59
$5.20
36.80
35.40
$7.06
58.41
73,12
23.12
299
PuraGRUU
2.90B
10.92
23.98
23.12
57.58
57.94
26.00
26.60
53.83
56.51
74.48
23.75
ParaLSTMM
2.90B
24.06 34.47
54.22
55.14
53.75 39.43
71.71
23.99
Mambu2
2.83B
10.13
33.70
61.13
61.12
42.20
37.80
59,51
73.78
24.89
Transformer
6.89B
9.55
34,30
33.36
62.98
62.20
40.00
37.20
61.48 60.85
74.97
23.12
ParaGRU
6.76R
9.19
39.68
36.77
65.85
65.75
42.20
40.40
61.40 59.83
76.66
25.29
ParaL.TMM
6.76B
9.166
36.52
62.47
62.85
42.20
38.80
$7.70
59.12
75.19
25.31
Mamba2
6.968
8.62
40.02
39.59
69.78
69.68
42.20
42,20
63.19 63.77
76.66
26.61
23

---

Published as a conference paper at ICLR 2026
k-hops Given input sequence of values from dictionary V, for each element in the sequence "hop back" and report the value following its previous occurrence the sequence. Keep hopping back & times. The special case of 1-hop is also known as induction heads.
Cycle Navigation Given a sequence of steps E {-k,....ke k) on cycle of size C, output the final position mod (Etn (E #),
Modular Arithmetic (w/o brackets) Given sequence of digits in V and operators {+,- forming an arithmetic expression, output the result modulo [V], respecting standard order of operations.
Copy Memory Given sequence [ra)]/ of values from dictionary V, followed by blank tokens and a trigger at æt, output the original sequence.
A5 Given sequence of elements from the alternating group A5 (even permutations), output the cumulative group products at each position.
Models Description For all models we train a single layer. The models under test are structured as: embedding layer with size matching the task vocabulary size, RMSNorm normalization layer, attention / SSM/ RNN block, another RMSNorm normalization layer, and linear layer with size matching the task vocabulary size. For Keep-nth and Parity described below, the SSM and RNN blocks contain only the mixer (in the tables we highlight this with !), that is we remove the convo- lutional layer, normalization layer, and gating mechanism from the block definition (see Fig. 10). We do this to test the expressivity of the mixer themselves, which is the main difference between the models. For the remaining tasks we use instead the full block (as described in App. C.I and Fig. 10) since the performance of both Mamba and the RNNs dropped significantly without these additional components. Moreover, for Keep-n"h we add positional encoding to Mamba and the RNN cells since it is not possible to solve the task without it (the Transformer includes positional encod- ing in all tasks). Across all models, we set a model dimension of 64, split across 4 heads. The convolutional layers (when used) have kernel size of 4.
Training Setup For MQAR, k-hop, Keep-n and Parity, we randomly obtain 10k samples as the training set and 100k samples as the test set (which report on), with sequence length fixed to == 100. We train each architecture for up to 3,000 epochs with a batch size of 16. For the remaining tasks, the training set contains IM samples and the test set 100k samples, with sequence length varying by task: L = 100 for Cycle Navigation, L 50 for Copy Memory, and L = 20 for A5 and Modular Arithmetic. We train these tasks for up to 1,000 epochs with a batch size of 512. For optimization, we employ an Adam optimizer (3, = 0.9, 32 = 0.999), learning rate starting at 10-4, and weight decay 10-6 All tasks are trained on a single NVIDIA A100 GPU, across three different seeds (we report on the best result). The norm of a, and C. in ParaGRU and ParaLSTM is clipped to a maximum of 0.90, except for Parity, where we found that norm clipping was detrimental to the cells' ability to solve the task.
24

---

Published as a conference paper at ICLR 2026
D PARARNN SOFTWARE PACKAGE OVERVIEW
In this section we provide an overview of the ParaRNN software package developed in the course of this research work. In its design we focused on three principles: modularity, ease of use, and extensibility. Consequently, the code separates the definition of the RNN cell and the solver, and drastically reduces the amount of code needed to be implemented by practitioners. Additionally, we provide multiple hooks for extension and generalization of various RNN cells.
D.1
MAIN CLASSES
The core interface of the package revolves around three main classes, respectively responsible for providing abstractions to RNN Cell, the implementation of its action, and the desired method for its application. More in detail:
RNN Cell The BaseRNNCell class represents the abstraction of an actual RNN cell, imple- mented as a torch module: all user-defined cells should inherit from it. This is the only stateful class in the package, and as such it is responsible for initializing and storing the parameters defining the cell itself. These are collected in single SystemP arameters dataclass, to aid interoperability with the other classes in ParaRNN by providing a standardized interface for parameter access. The actual implementation of the key RNN cell functions is kept separate from the class definition itself, and delegated to the static class RNNCe 1 1 Impl. BaseRNNCell is also responsible for coordinat- ing code execution, invoking the correct method from RNNCe wwA.CellApppicatiie depending on the selected RNN cell application modality.
RNN Cell Implementation The base static class RNNCe11 Impl collects the implementation of the functions necessary to describe the action of an RNN cell. Its main responsibility is to imple- ment the receurrence step (1). Starting from this, the class leverages PyToorch autograd fune- tionalities to automatically define the necessary methods to assemble the target all-at-once system (2), and particularly the computation of residuals and Jacobians (3). For additional performance. it is possible to override the latter function, and provide an explicit definition of the Jacobians we do this ourselves in the implementation of the ParaGRU and ParaLSTM cells. This becomes particularly relevant when said Jacobians present specific structure, that can be leveraged for mak- ing the reduction operations involved in the solution of (3) more computationally tractable. To this aim, in ParaRNN we provide two specializations for RNNCe Impl: RNNCe lieeelDDaaga Impl and RNNCe 1 1elllBlcca BBlockDiiagggg Impl, which simplify operations for those cells whose Jacobians are respectively diagonal (like for ParaGRU), composed of N N diagonal blocks (like for ParaL- STM, where N = 2). These specializations implement the corresponding reduction steps (lines 4 and 5 in Alg. 1b), described more in detail in App. D.2. Our software allows to further special- ize RNNCe 11 Impl to encompass additional Jacobians structures (c.g., sparse, or Circulant (Keller et al., 2024)), but the corresponding assembly, storage, and reduction operations should be provided (see also App. D.4).
RNN Cell Application The third main component of ParaRNN is RNNCe MANCellApplicat: This collection of static classes defines the different ways in which the RNN cell can be ap- plied, and effectively implements the forward and backward passes for BaseRNNCel1. The most vanilla version, RNNCell 1eeellleeuutttt wpw.pppppcaatiiiaaallei
simply evaluates the cell output by sequentially applying (1) to the input sequence: this is the classical way PyTorrc RNN cells act, although our work this is mostly used for testing and debugging, or at inference time. RNNCe11P arralllll www.leelppppaaat.ccmmpmppm iin implements the sequence-paralee cell, following the method described in See. 2, while using only PyTorch directives: the class ef- fectively acts from RNNCe 1 Impl to iteratively assemble and solve the linearized system (3) via parallel re- duction. A first optimization is provided in RNNCe leellllllle 1eeecuUDAppp Sippplioen cat ion, where the parallel solve is handled by a custom CUDA kernel, while the system assembly is still performed in Porch. Finally, RNNCel1 1P aarallllle www.eeeeeeeppplce.cc cattpppppiiiiioooom ion bypasses all PyTorch- im- plementations and computes the RNN application via a fully-fused CUDA kernel (implementing Newton's 's iterations, system assembly, and parallel reduction, end-to-end).
application of the RNN
wrapper for our implementation of Newton's 's method, and leverages the methods
25

---

Published as a conference paper at ICLR 2026
D.2 SOLVER IMPLEMENTATIONS
In addition to programming interface for RNN cells, the main contribution of ParaRNN consists in providing ready-to-use efficient implementations of the parallel solver outlined in Sec. 2.
D.2.1 PURE PYTORCH PARALLEL REDUCTION
Our PyTorch reference implementation is a direct adaptation of the pseudocode in Alg. la and lb. Particularly, the parallel reduction loop reduces to
@staticmethod
def =eduuce
jac: torch. Tensor,
(B), 7, -2 shapea will depend on
che: Tensor, # (,),
specific jac structuze
reduction_steep: -yp.collable
porch.Tensor:
num_stepa = math. (math. 1og2 shape
for step range steps)))
jac, the redursion_ateppjac,, ths, step) return ths
The reductiionn peductionnster method defines how to perform pairwise reduction of the equations in (3): namely, it implements lines and of Alg. Ib, adapted to a specific structure of the cell Jacobians. This is the only cell-specific method used in the solver, and by treating it abstraction, allowing for modularity in its utilization. In particular, in ParaRNN we provide implementations of this method for dense, diagonal, and block-diagonal Jacobians, covering both the general and the specific cases analyzed in this paper (c.f. Sec. 3).
The main focus of the development is dedicated to the design and implementation of the CUDA- accelerated solver. We describe this in detail in the following.
D.2.2 CUDA-ACCELERATED PARALLEL REDUCTION
Each level of GPU execution has access progressively slower but larger memory systems (NVIDIA, 2025, Sec. 5.2-5.3). Threads can access their own registers the fastest, warps can shuffle data between threads seamlessly, and blocks can use dedicated shared memory paying small over- head. Only when communication occurs at grid level must we rely on the slow global memory of the device. implementing our kernel for parallel reduction, we minimize overhead by closely mir- roring this hierarchy: the resulting algorithm directly adapts the reduction operations each level. On top of this, rather than blindly maximizing parallelism at every level of the hierarchy, we follow hybrid approach interweaving both sequential and parallel reductions, inspired by similar work on solvers for tri-diagonal systems (László et al., 2016). The action of our algorithm as it traverses the GPU hierarchy is outlined next.
Thread-level At the lowest hierarchy level in our algorithm, we associate each GPU thread with small group of chunk_size adjacent equations in the target system (3). Threads are responsible for reading from global memory to their registers the information (namely Jacobians and residuals) associated with their chunk of equations. At this stage in the algorithm, threads must also manip- ulate said information and start reducing equations so that they all end up depending solely on the value of the very first unknown in their chunk. This is obtained by forward substitution, i.e., by quentially reducing each equation in a chunk using the previous one. Note that, while this operation is sequential, it only requires chunk_st ze> 4 steps, and is nonetheless conducted in parallel across threads. See Fig. 12 for a diagram of the action of the algorithm at this level.
Warp-level The first truly parallel reduction operations start occurring at warp level. Threads within a same warp share information from the last equations in their chunk and apply parallel reduction over them. At the end of this operation, each of these equations depends solely on the very first one in the warp. Explicitly leveraging this level in the hierarchy of the algorithm enables the use of warp-specin information stored in other threads registers, eliminating the need for expensive memory transfers or explicit synchronization. See Fig. 13 for diagram of the action of the algorithm at this level.
directives (such warp shuffle instructions) for threads efficiently access
26

---

Published as a conference paper at ICLR 2026
[Teaa Throal
warp-loo reductioo
Toread reslaction
Stop
Step
Figure 12: Schematics drafting the evolution of the target block bi-diagonal system (3) as it under- goes the first stage of our algorithm. Light-blue squares represent its main diagonal (i.e., identity matrices), while the dark-blue once represent the Jacobians different l's. Threads independently apply forward substitution to chunk_sized =e=== 3 in the figure) equations. At each step, this amounts to using equation I reduce the following one + (i.e., applying lines and in Alg. 1b). Notice this makes all equations in a chunk depend on the value of the very first unknown, as exemplified by the progressive shift towards the left of the sub sub-diagonal block. At the end of this stage, we can effectively extract a simplified system involving only the last equation of each chunk, which is then passed onto the next stage.
Block-level Parallel reduction proceeds then in a similar fashion at block level, where the last equations of each warp are reduced so that they depend only on the very first one in the block. For this stage, we need to resort to the shared memory within the block to share information among warps, access to which must be explicitly synchronized: while this incurs some overhead, it remains significantly faster than accessing the device global memory. Once this stage is completed, we can traverse the hierarchy in reverse. Specifically, we the first equation in the block solve in parallel for the last equations of each warp, then use these to solve in parallel for the last equations of each chunk, and finally use these solve in parallel for each equation within chunk. If a single block is large enough to cover the whole input sequence, then at this stage the solution to (3) is recovered, and the algorithm terminates here, by writing the solution back to global memory. not, the next stage depends on the total sequence length. If it is smaller than certain threshold, then the same block sequentially processes consecutive segments of the sequence, restarting the procedure above. For very long sequences, instead, the next level in the hierarchy is triggered.
Grid-level For extremely long sequences, rather than having a single block covering sequentially the whole sequence length (while still applying parallel reduction within the block), organize the GPU kernel grid so that multiple blocks are fired to cover the sequence in parallel, applying the
To block--eo level reduction
Warp levell reduction
Step
Step
Stop
Figure 13: Schematics drafting the evolution of the simplified block bi-diagonal system stemming from thread-level reduction, it undergoes the second stage of our algorithm. The threads within same warp (8 in the figure, but 32 in practice) apply parallel reduction to further simplify the target system. At each step this amounts using equation to reduce the 1 + 2'-th one. Notice this progressively shifts the sub-diagonal blocks by 2", 2, 2222..... positions, until all equations depend only on the very first unknown in the warp. At the end of this stage, simplified system involving only the last equation of each warp, which is then passed onto the next stage. At the block level, the algorithm proceeds similarly: the last threads of each warp perform parallel reduction, operating on the simplified system.
effectively extract
27

---

Published as a conference paper at ICLR 2026
same procedure listed above. In this case, however, the procedure leaves each block dependent the value its first unknown (whose solution is explicitly available only for the very first block, which covers the start of the sequence). An additional parallel reduction is then necessary solve for the last equations of each block. This requires waiting for the completion of the kernel performing reduction within block, firing second kernel to reduce equations across blocks, and then third one to perform the final substitutions within each block again once their solutions are available. The associated overhead grows accordingly, not only due to the additional kernel launches, but also because in in-between kernels execution the device must be synchronized, and partial results must be written global memory.
The size of the chunk of equations solved by each thread (chunk size), the number of threads allocated to a single block (threads per bl blk),, and the threshold number of sequential block- wise reductions (max_sequent al_steps) are all hyperparameters defined at compile-time, which can tuned to squeeze additional performance from the code by changing the per-thread workload and registers availability. The most optimal combination of parameters is function of the actual GPU device in use, the underlying scalar datatype considered, the target input sequence length, and of the Jacobian structure of the RNN cell. In Fig. we report timing results for sweep on chunk_ss size, which we used to select the most optimal values in this work. The defaults pro- vided in the code good starting point for typical LLM training setup in mind, (i.e., sequences of I. = 29~12 float32 scalars), on a NVIDIA A100 GPU.
D.3 EXAMPLE USE
In the following, provide example implementation of simplified version of the ParaGRU cell used in this work, as a use-case for the ParaRNN codebase.
As first step, need introduce the parameters involved in the cell definition itself:
@dataclass
class wwt.iaggystennaattteeoo aaraametee
(SyatemParameterss
A: Torch.Tensor
combines (As, Ar, Ac)
B: Tenser
combines (Bz, Bry Bc)
b: torch. Tenser
combines (bz, br, bc)
nonlin tonpdatee typ.cal Cal.calllalll
nonlin_reset: typ.Ce Capabellb
nenlin_state: typ. Callable
These will be stored and initialized inside our new cell:
class GRUDiag (BaseRNNCe111
@RUDiagCenfig, s2UCiagSys:en2 wwwpeeeenneeesss....mm.. GRUDiegImp1 1)
invoked by constructor def Lspecific_init (self, config: :RuDiagtonfig))
def innit_ (self, config: GRCDiagConfig)
auper iupet( [nit__cooiigit
System parameters
collated peraneters for computing s,ryc
kwazgs ["device": ewwrrchhepty(().. empty self.atate_dim) .=kwrrga))
device, "étype": seye.dtype)
self.A A Parameter (torch.
self.R nn.Parameter (torch, (torch.empty) " vikwargs!
[3, self.stace_de
self.b (tramh ctttorrccheeey(...... e[3, aif.sstat_ddmm) =*kwrggs))
se________ooniiit (config. nonlin in_update)
self.nonlin_res:: het___eet_ooinerrrriitt (config, nonlin re_eett) solf.nomlin
(config. coonfignncnll. nonlin (n__ttte)
# Parameter inftielisation self.reset.parrmetters
Sproperty
def Systen__arameeee (sel:):
# hanoi ly col leat systemn parameters
return GRUD ammbiagyyteenaaameeeee
A-self.A B-self.B, B, b-self.b,
canlin_uddate=see wen.se...nnnnnnudateeee nonlin_resee siin_resstssel wwa.tssell..no.liio... www.fooonnnnnnnn reset,
28

---

Published as a conference paper at ICLR 2026
The final ingredient to be able to start using the class, is prescribing its action (7a):
class GRUDiagImpl MiliCellDiaggngg [GRUDiegSystent) eramettral)))))
@staticmethor
def
Xt torch, torch.Tenaor,
(R), (1), Din
Tenaor,
(3), (1),
aystem__aate parameterr GRUD Settnniosyyttemaarmmtrrrr
Lorch. Tensor: Torcho
Expb
(orch.einsum.
"...3.vaj=>...2.
pystsenpporrmmters..
aystemn_arrmettrrs.
Ahpüxpb torch. einsum!
'v)....
[sseteen_perammeeeee...e :l, hi
Expbl. :2, :],
2, x torch. anrchh (Ahpexpb, dimm-2)
syaten_paramett
pasten__pt.ameeer wtw.aaammeteeeeaall_eeeeeeee (1)
pysteen_parametet aonlin_statee /tee
aystem_patameters
Bxobl..., 2, :]
This the minimum implementation user must provide to able apply the RNN cell. Just with this, the user can already toggle between sequential and parallel application' by se- lecting model. modeL,.ode-SEEUNNIIN
or model.me modee-AAALLLLL.. Moreover, since the Para- GRU cell specified above presents diagonal structure, the user can directly leverage the effi- cient implementations provided in RNNCelIDian Impi for Jacobians assembly (via autograd) and manipulation, and the corresponding CUDA-acceleraan Still, by manually providing the formula for the Jacobians (8),
solver is also available via selecting
model. PARALLEL. can boost performance by avoiding the reliance autograd:
class GRUDiagImpl (RNNCe11Diaggngg (GRUDiegSystens arameters))
8classmethod def camputt_jjaobbiane
cla,
h: Torch.Tensor,
x: torch.Tenso::
system_
syston__arameter:::
CRUDiogHSSyttomPaaammtttss
>> torch, torch.Tenso:: hml (t)
shift along
Rxph torch. torchheinaumt
"...3,v1
(x, system__araameeers
pre_nl_2, pre_nl_rr_ sorch.anbind
corch.einsur
awaatemparameters.o s), hml
1) Expb :2, s], dim--2
syaten_parametere nonl in (pro_nn_
r pysten_perameters www.ooonnnnnnoe.oo reset (pre_rl__)
pre_nl_e aystem parametars.A
hml
Exph 2,
о aysten_parameter www.eeeerr..ooonnnnoo..m
(pr (preee
grad_ sesten_parameterss ceaamnters.ddriveeee derivative www.nooonnnnnnneeeeeeee _])
grad_ systen_parameters. parameters.. vative wattvv_nonnlieeeeeeeee www_nononnneeeneeeeeeeee reset (pre_nl_r)
g syatem_parameter dew.err.eeeeaaeeeeeeoooooee weweeenooninnseeeeeeeeeeeeeeeeee state (pre_ni_o)
J_,, Je unrch.unbidd
pateen_parameterss torch, z grad_e, grad_e), dimm-2
'Notice that the recurrence_stert method is designed to work regardless of whether the input is batched or not, and whether it is full sequence single element: this allows to reuse the exact same method for the different application modes of the RNN cell.
29

---

Published as a conference paper at ICLR 2026
>, dimm-2
J_ JLe
ir
hm1 J_T1
jac
hml)) 2.
Finally, to be able to use the most optimized application mode model mode=PARALLEL FUSED, suffices for the user to provide equivalent CUDA implementations for the functions above. In particular:
template styenaame scalar_t> // impl ements CSIP polymorphion
class GRJGeliDiegImpl public RiNCellii scelar_t, GRUcellviagipple scaler_t
public:
static void t const scelar_tr... glbvars, ielvara_ts lelVers
11
// Read torch tensors in aystem parameters (glbvars) and atore // thread-specl rie info registers (leiVars)
recurrencoštep
const rha_t sh, const lelVars_ts lelVars, 6hp1
la, Expb] lclVers;
scalar_t z nenlinupd( a[0] expb [0]
scalar_t r inneinReat a["] Rxpb(1] [1]
scalar_t nonlinstater a[2] Bxpb[2]
hpt (1.-21 + return;
device static void computedacchiand const rha_t sh, const ris_t shml, const ieiVars_ts lolVars, Jac_t 4jac
autos la, Exph] lclvars; soalar_t a[0] hml Axpb[0];
scalar_t 1 - =[l]] hml Bxpb[]];
scaler_: a[0] derivattiveNenlinddddt)))
scalar_t jr a[]] derivattivnooniinnee (r);
nonlinUpdiz):
= non]inkes
scalar_t a[2]
expb[2];
scalar_t j0 a[2]
verivativeNon int.iatinnntttttteee to)
hm1 de
o nonlinState(c);
(2. z)
hmL) jz
return;
1:
D.4 EXTENSIONS OTHER JACOBIANS STRUCTURES
For RNNs with Jacobians structures different from (block-) )diagonal, the user has two choices: either fall-back to the implementation for generic dense Jacobians (the base RNNCel Impl), or provide their own specialization. This in particular requires defining two core methods:
/ystrrccturedcel..mpp..
RNNCe Tmp) [SysremPaaanttorra]] typ. Syp.Generrf g (SyereemPrrmmmterre
@clasemethod def compute_ iasobiians
cla,
30

---

Published as a conference paper at ICLR 2026
h: torch.Tensor,
X: sorch.Iensor.
system parameters: SyctenFarametero,,
Lorch, Tensoz:
@classmathed
def sercllel_ -> typ. Syp.Callabie:
portial
saductioo__topp...
The method compute-jacobians not only should take care of the actual assembly of the Ja- cobians, but also implicitly defines how they will be represented. This must be taken into account in the definition of the specific reduction responsible describing how to conduct the reduction operations (see App. D.2.1 and lines and in Alg. 1b) for this specific Jacobian structure. For the largest part, the structure of our CUDA code follows similar pointer-to conter-to-imme to-implemennation design pattern, mirroring that the Py PoTohh code. In principle, then, one provide similar specialization for the Jacobians structure also in CUDA:
function (which also must be implemented),
template typename scalar_t, typename class Mystructurodiall.mg std: :condit conaiitioonl.. ional_ts at [edd:11_voidd vovidd_ Derived >, pustructureddel1Impll scalas_t >, Derived
public RNNCe Impselmlk
Impiements
static polymozphiim
Reduce current eg using
davice static void reduceRqst const jac_te jacother, Jac_ts jac, const ris_tá shsOther, rhs_tä ths
However, to directly employ the current CUDA implementation of parallel reduction, the Jacobians structure should be "lean enough" that a single thread can hold in its registers information about (at least) one full equation of (3). this not the case (e.g., with dense Jacobians), then the provided CUDA implementation of parallel reduction cannot be employed, and an alternative must be implemented.
31

---

Published as a conference paper at ICLR 2026
E USE OF AI WRITING ASSISTANCE
We acknowledge the use large language model (LLM) to assist in refining the language, gram- mar, and clarity of this manuscript. All content was originally drafted by the human authors, and all Al-generated suggestions were critically reviewed, edited, and approved by the authors, who retain full responsibility for the final text.
32


Reviews:
[{"review_id": "aC19HChpRA", "reviewer": "Reviewer_BeBu", "summary": "This paper introduces *FlashRNN*, a framework that reformulates the entire recurrent computation in nonlinear RNNs over a sequence as a *single nonlinear system of equations*, which can then be solved in parallel via Newton iterations and prefix-scan parallel reductions.  To make this practical, the authors constrain the RNN Jacobians' structure (e.g., diagonal or block-diagonal) so that the prefix-scan remains efficient and take care of the channel mixing via a MLP, similarly as done in SSMs. They further provide a CUDA implementation that automatically parallelizes user-defined RNN cells and demonstrate training of 7B-parameter nonlinear RNNs (FlashGRU, FlashLSTM) with competitive perplexity to Transformers and Mamba2 while achieving significant (up to 665×) speedups over naïve sequential RNNs.", "questions": "1. How difficult is it for practitioners to define a new RNN cell compatible with FlashRNN’s structural constraints?  \n\n2. Have you tried experiments with FlashRNN in mixed or low precision? \n\n3. Do you plan and how easy would it be to integrate fused CUDA kernels from [Pöppel et al. (2024)](https://arxiv.org/pdf/2412.07752) in FlashRNN?\n\n4. Would the quasi-Newton approximations as in [Gonzalez et al. (2024)](https://arxiv.org/pdf/2407.19115) work also at scale instead of imposing the diagonal matrix structure?\n\n5. How do GPU memory and throughput scale with hidden size?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "1. Tackling the core sequential bottleneck of nonlinear RNNs is a long-standing challenge, so the motivation of the paper is clear. The provided solution, i.e. recasting recurrence evaluation as a global nonlinear system of equations and applying Newton iterations combined with parallel prefix reductions, is conceptually clean and theoretically grounded from previous work.\n\n2. The authors release **FlashRNN**, a PyTorch + CUDA framework that generalizes to arbitrary RNN cells, lowering adoption barriers and fostering community experimentation.\n\n3. The authors train the efficient versions of GRUs and LSTMs up to 7B parameters with substantial wall-clock speedups. Few works have shown nonlinear RNNs trained at this scale and doing so has potential to reshape how sequence models are trained.\n\n4. The experimental (including runtime) analysis is thorough. The authors provide profiling, ablations, and comparisons to Transformers and SSMs, which add credibility. \n\n5. The authors openly discusses convergence challenges, overheads, and structural constraints in the Jacobian.", "weaknesses": "1. As the authors write, imposing diagonal or block-diagonal Jacobians simplifies parallelization but may severely limit expressivity due to channel mixing. The solution taken is similar as in SSMs, i.e. using downstream MLPs to \"restore\" expressivity. I believe this needs stronger empirical support on more state-tracking tasks where xLSTM performs relatively well.\n\n2. The authors assume that a small, fixed number of Newton iterations (e.g., 3) suffice. However, this is not guaranteed across other recurrent architectures not studied in this paper or different tasks (e.g., time series forecasting, speech, RL, etc.). Where there any scenarios where convergence fails or oscillates? Otherwise, formal guarantees on truncated Newton convergence or stability under limited precision would benefit the paper.\n\n3. Newton iterations can be sensitive to numerical precision.\n\n4. While the authors briefly claim reduced memory scaling relative to BPTT, the paper lacks a quantitative analysis of GPU memory usage across model sizes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3BrgISp2YY", "reviewer": "Reviewer_5Qbd", "summary": "The paper introduces FlashRNN, a framework that parallelizes training-time application of nonlinear RNNs by recasting the unrolled recurrence as a single all-at-once system solved with Newton iterations and a custom parallel reduction (prefix-scan) solver over a block bi-diagonal linearization. Implementations span (i) pure PyTorch, (ii) CUDA-accelerated reduction, and (iii) a fully fused CUDA kernel. \n\nThe authors adapt GRU and LSTM (FlashGRU/FlashLSTM) to make their Jacobians diagonal or 2×2 block-diagonal, enabling efficient reductions. They report large speedups over naïve sequential application versus traditional RNNs and show competitive perplexity and downstream results vs. Mamba and a Transformer baseline at 7B parameters.", "questions": "See the weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The proposed idea casting nonlinear RNN application into a Newton+scan routine with a specialized CUDA solver is well motivated and clearly explained.\n\nThoughtful GPU hierarchy design, e.g. Appendix D2.\n\nThe author promised code release which would encourage the community try the proposed method.", "weaknesses": "It's unclear to me how the the diagonal (GRU) and block-diagonal (LSTM) Jacobians limitation been overcomes. Any study shows those limitation not matters at scale?\n\nI'd suggest add more opensource baseline in Table 2.\n\nIt's unclear numerical stability when further scale up the model.\n\nIt would be interesting to see more ablation on newton iterations, how does the memory footprint change as context length change, and how far the context length can be pushed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "b9I4qqLh1U", "reviewer": "Reviewer_bSyi", "summary": "The paper proposes FlashRNN, which parallelizes nonlinear RNN application by recasting the full sequence recurrence as a single nonlinear system (Eq. (2)) and solving it with Newton iterations whose linear subproblems have a block-bidiagonal structure solvable via parallel scan.", "questions": "N/A", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Novelty:  Prior work typically applies Newton/scan to given RNNs (e.g., Lim; Gonzalez). This paper instead redesigns LSTM/GRU so their Jacobians are diagonal or 2×2 block-diagonal, which makes each Newton step’s linear system amenable to efficient parallel reduction without runtime Jacobian approximations.\n2. Scale of Experiments: This is the most exciting part. The paper successfully trains 7B-parameter FlashGRU/FlashLSTM models and reports their downstream task results. This proves that FlashRNN is not just a toy model. Moreover, the efficient CUDA implementation further highlights the significance of this contribution.", "weaknesses": "N/A", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "xIsbZBCy0R", "reviewer": "Reviewer_M9m5", "summary": "This paper adds to the body of literature on efficient recurrent neural networks (RNNs). Although RNNs have potential for inference efficiency, particularly at long sequence lengths, the training of RNNs has been difficult due to the inherently sequential nature of the forward pass, unlike a Transformer forward pass which can be trivially parallelized.\n\nThis paper offers a new approach to parallelizing the training of RNNs, FlashRNN. The whole forward pass is solved in parallel via a fixed-point iterative method. This allows the efficient training of non-linear RNN cells, albeit the method is only tractable if the RNN cell has a Jacobian with a particular block-diagonal structure. The paper provides efficient implementation of the method in CUDA and scales FlashRNN to the 7B parameter scale.", "questions": "+ Could you demonstrate that the diagonalization in equation (9) does not significantly affect the expressivity of the RNN layer?\n+ Could you elaborate on why it is important to have a nonlinearity in the sequence mixer specifically, given the use of multiple layers with feature-mixing MLPs in all modern architectures? It is not clear to me, given the strong results of Mamba on most tasks (when equipped with MLPs, not just evaluated as a single layer).", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "+ The overall approach is a good contribution, and as far as I am aware this is the first work to allow training of non-linear RNNs via an exact parallelization method with no numerical instability.\n+ The extensive implementation of a custom CUDA kernel which can be applied to a range of RNNs is a great contribution to the open-source community, and allows comparison with other architectures at a large scale.\n+ The extensive computational results are quite impressive, including training a 7B parameter model from scratch in order to compare the performance of FlashRNN at language modelling tasks.\n+ The paper is up-front and clear about the limitations of FlashRNN in section 2.1", "weaknesses": "+ The conceptual contributions of the paper could be argued to be somewhat limited. As far as I can tell (although I'm not an expert in this exact subject), the parallel solution of the sequence model by forward substitution is not novel and is described in Gonzales et al, Lim et al, etc. It is a bit hard to tell which of section 2 is novel and which is from previous work--we know that at least some of section 2 is previous work from the description, but it's not clear if it all is previous work. It seems like the conceptual contribution of this work could be characterized as following Gonzales et al, but simplifying the RNN cell instead of trying to solve the recurrence for general RNNs with a clever method.\n\n+ The paper doesn't empirically justify the diagonalization of the RNN mixer in equation (9). While it's true that diagonalized matrices are often used in SSMs, the motivation for this work is that SSMs are not necessarily expressive enough. Therefore the justification for using a diagonalized matrix is not internally consistent. I think the authors could perhaps add extra ablations on the table 1 experiments where they use a non-diagonal GRU/LSTM (trained sequentially), to demonstrate that there is not much difference between diagonalized and non-diagonalized RNNs. \n\n+ [minor] The paper is missing some earlier work on linear RNNs and parallelizations of linear RNNs, particularly [1,2,3]\n\n\n  [1] Bradbury, James, et al. \"Quasi-recurrent neural networks.\" arXiv preprint arXiv:1611.01576 (2016).\n\n  [2] Martin, Eric, and Chris Cundy. \"Parallelizing linear recurrent neural nets over sequence length.\" ICLR 2018\n\n  [3] Qin, Zhen et al, Hierarchically Gated Recurrent Neural Network for Sequence Modeling, NeurIPS 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)