ParaRNN: A Revolutionary Framework Unlocking Parallel Training of Nonlinear RNNs

Introduction: The Parallelization Dilemma of RNNs

Recurrent neural networks (RNNs) were foundational for sequence modeling due to their ability to retain historical information. However, their inherent sequential dependency prevents parallel computation, making scaling difficult. In recent years, Transformer architectures (BERT, GPT, etc.) have dominated large language models (LLMs) thanks to highly parallelizable training via self-attention. Yet attention's quadratic compute and memory cost on long sequences has driven interest in alternatives.

Structured state-space models (SSMs) such as Mamba and Mamba2 have emerged as strong competitors. SSMs use linear time-invariant systems to model sequence evolution, achieving RNN-like state memory while supporting efficient parallel training via techniques like parallel scan. However, to enable parallelism, SSMs impose a linearity constraint on the recurrence, limiting expressiveness—complex nonlinear dependencies like gating mechanisms (key strengths of LSTM/GRU) cannot be captured directly.

This raises the central question: Can we keep the expressive power of nonlinear recurrence while achieving parallel training? ParaRNN answers this long-standing challenge.

ParaRNN's Core Idea: Turning Recursion into a System of Equations

ParaRNN (Parallelizable Nonlinear RNN) reformulates the sequentially executed nonlinear recurrence as one large system of equations and solves it in parallel with numerical methods.

• Each time step's hidden state update becomes one equation involving the current input, the previous hidden state, and the parameters.
• Stacking equations across all time steps yields a large nonlinear equation system.
• Instead of solving it step by step, ParaRNN solves the whole system simultaneously.

To solve this system in parallel, ParaRNN uses Newton's iteration, and accelerates each iteration with parallel reduction on GPU/TPU hardware. In short: unfold the recurrence into an equation system, solve it via Newton iteration, and speed each iteration with parallel reduction.

Technical Implementation

Newton's Iteration

Newton's method iteratively approaches the solution, computing the Jacobian matrix and solving a linearized system at each step. For RNN hidden-state systems, the Jacobian is block lower-triangular (each time step depends only on its input and the prior state), enabling efficient forward/backward substitution with parallel computation of all blocks. To control the cost of computing and storing Jacobians at scale, ParaRNN adopts approximate Newton strategies, e.g., recomputing the Jacobian periodically or using low-rank approximations.

Parallel Reduction

Aggregations over all time steps (residuals, gradients) are natural fits for parallel reduction, cutting O(N) summations to O(log N) parallel steps. ParaRNN uses a two-level reduction scheme on GPUs: each thread block performs local reduction in shared memory, writes results to global memory, and a final block performs the global reduction—minimizing memory conflicts while maximizing parallelism.

Experimental Results

Researchers built ParaLSTM and ParaGRU by applying ParaRNN to classic architectures and trained them on large-scale language modeling:

• Training speed: up to 665x faster than traditional sequential RNN training on the same hardware and model size.
• Model scale: successfully trained nonlinear RNNs with 7 billion parameters—a scale previously out of reach for RNNs.
• Performance: 7B ParaRNN models achieve perplexity comparable to similarly sized Transformer and Mamba2 models, and even slight advantages on some tasks.
• Long sequences: ParaRNN retains the linear-complexity, recurrent structure of RNNs, avoiding the quadratic cost of attention; chunking and gradient checkpointing mitigate the memory cost of storing full-sequence states during Newton iteration.

Comparison with Transformer and Mamba2

• Training parallelism: All three support parallel training. ParaRNN adds iteration overhead from solving nonlinear systems, but optimized implementations keep it manageable, delivering comparable or faster effective training speed.
• Expressiveness: Transformer's attention captures arbitrary pairwise dependencies; Mamba2's linear recurrence limits nonlinear modeling; ParaRNN preserves LSTM/GRU-style nonlinear gating, placing its theoretical expressiveness between the two.
• Long-sequence handling: Transformer faces quadratic cost; Mamba2 and ParaRNN both scale linearly, clearly outperforming Transformer on long sequences.
• Scale: Transformer reaches hundreds of billions of parameters today; Mamba2 operates in the billions-to-tens-of-billions range; ParaRNN has demonstrated 7B-parameter models competitive with same-scale peers, with room to scale further.

Conclusion: The Revival of Nonlinear RNNs in the LLM Era

ParaRNN liberates nonlinear RNN training from its sequential bottleneck, matching the parallel training capability of Transformers and SSMs—no longer forcing a trade-off between expressiveness and training efficiency. Classic architectures like LSTM and GRU regain viability at LLM scale, and the framework opens paths to richer RNN variants (attention hybrids, richer gating, cross-layer connections) without sacrificing efficiency.

ParaRNN also demonstrates the value of algorithmic innovation: clever numerical techniques (Newton iteration plus parallel reduction) can turn a classic model's weakness into strength. With the framework's open-sourcing, a more diverse sequence-modeling landscape—where Transformers, SSMs, and RNNs each play to their strengths—becomes a realistic prospect.