ParaRNN: A Revolutionary Framework for Parallel Training of Nonlinear RNNs

Background: The Parallelization Dilemma of RNNs

Recurrent neural networks (RNNs) were foundational to sequence modeling due to their ability to memorize historical information. However, their inherent sequential dependency prevents parallel computation, becoming a major bottleneck for scaling. This led Transformers (BERT, GPT, etc.) to dominate large language models (LLMs) thanks to highly parallelizable self-attention — though attention suffers quadratic compute and memory costs on long sequences.

Structured state space models (SSMs) such as Mamba and Mamba2 emerged as strong competitors. By modeling sequence evolution as a linear time-invariant system, SSMs achieve RNN-like state memory while supporting efficient parallel training via parallel scan. However, this efficiency comes from imposing a linear constraint on the recurrence, limiting expressiveness: linear recurrences cannot capture the nonlinear gating and activation mechanisms that give classic RNNs (LSTM, GRU) their power.

This raises the central question ParaRNN addresses: Can we achieve parallel training while preserving the expressive power of nonlinear recurrence?

Core Idea: Turning Recurrence into a System of Equations

ParaRNN (Parallelizable Nonlinear RNN) reformulates the step-by-step nonlinear recurrence of an RNN as a large coupled system of nonlinear equations — one equation per timestep's hidden state. Instead of solving these equations sequentially in time, ParaRNN solves them simultaneously using numerical methods:

1. Unfold the sequence recursion into a system of nonlinear equations. 2. Solve the system with Newton's iteration, which iteratively approaches the solution by computing residuals and solving a linearized system. 3. Accelerate each iteration with parallel reduction on GPUs/TPUs.

Technical Implementation

Newton's Iteration

• Each iteration computes the Jacobian matrix and solves a linear system to update the solution.
• For RNN hidden-state systems, the Jacobian is block lower-triangular (each timestep depends only on its input and the previous state), enabling efficient parallel computation of all blocks and fast forward/backward substitution.
• To manage cost at scale, ParaRNN uses approximate Newton strategies, e.g., recomputing the Jacobian only every few steps or using low-rank approximations.

Parallel Reduction

• Aggregating residuals/gradients across thousands of timesteps is a natural fit for parallel reduction, reducing O(N) serial accumulation to O(log N) parallel steps.
• ParaRNN employs a two-level reduction strategy: each thread block performs a local reduction via shared memory, writes results to global memory, and a final block merges them — minimizing memory conflicts and maximizing parallelism.

Experimental Results

Applying ParaRNN to LSTM and GRU produced ParaLSTM and ParaGRU models, evaluated on large-scale language modeling:

• Training speed: up to 665x faster than conventional serial RNN training under the same hardware and model size.
• Model scale: successful training of RNNs with 7 billion parameters — previously unthinkable for RNNs.
• Performance: 7B ParaRNN models achieve perplexity comparable to same-scale Transformers and Mamba2, and are even slightly better on some tasks.
• Long sequences: with linear complexity and native recurrence, ParaRNN handles long sequences without chunking or attention; memory overhead from storing hidden states for Newton iteration is mitigated via chunking and gradient checkpointing.

Comparison: ParaRNN vs. Transformer vs. Mamba2

| Dimension | ParaRNN | Transformer | Mamba2 | |---|---|---|---| | Training parallelism | Yes (Newton + parallel reduction) | Yes (self-attention) | Yes (linear recurrence + scan) | | Expressiveness | Nonlinear gating (LSTM/GRU) | Global attention, highly flexible | Linear recurrence, limited nonlinearity | | Long-sequence handling | Linear complexity (like Mamba2) | Quadratic cost | Linear complexity, strong | | Scale demonstrated | 7B parameters | Up to hundreds of billions | Billions to tens of billions |

In principle, ParaRNN's expressiveness sits between Transformer and Mamba2: less globally flexible than attention, but better at nonlinear patterns than linear SSMs. On long-sequence modeling, ParaRNN and Mamba2 belong to the same tier, clearly outperforming Transformers.

Conclusion: The Revival of Nonlinear RNNs in the LLM Era

ParaRNN frees nonlinear RNN training from its serial constraints, showing that classic architectures like LSTM and GRU can match Transformer and Mamba2 at the 7B scale in language modeling. It demonstrates that algorithmic innovation can convert the weaknesses of classical models into strengths, removing the need to choose between expressive power and training efficiency. As the framework is open-sourced and refined, it opens the door to richer RNN variants (attention mechanisms, richer gating, cross-layer connections) and a more diverse sequence modeling landscape where Transformers, SSMs, and RNNs each play to their strengths.