ParaRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language Models (2510.21450v1)

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 665x over naive 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.

• The paper demonstrates a novel Newton-based parallelization that reformulates sequential RNN computations into efficiently solvable block bi-diagonal systems.
• It introduces ParaGRU and ParaLSTM variants that leverage diagonal or block-diagonal Jacobian constraints to reduce computational overhead while preserving expressivity.
• Empirical results show up to 665× speedups and competitive language modeling performance, positioning nonlinear RNNs as viable alternatives to Transformers.

ParaRNN: Parallel Training of Nonlinear RNNs for LLMs

Motivation and Context

Recurrent Neural Networks (RNNs) have historically been limited in scalability due to their inherently sequential computation, which precludes efficient parallelization along the sequence dimension. This limitation has led to the dominance of architectures such as Transformers and State Space Models (SSMs), which enable parallel sequence processing during training. However, SSMs achieve this by enforcing linearity in the recurrence, which restricts their expressive power and limits their ability to model complex, nonlinear dependencies. The ParaRNN framework addresses this bottleneck by enabling efficient, sequence-parallel training of nonlinear RNNs, thus removing the linearity constraint and expanding the design space for sequence models.

Methodology: Newton-Based Parallelization of Nonlinear RNNs

ParaRNN builds on the insight that the sequential application of an RNN over a sequence can be reformulated as the solution to a nonlinear system of equations, one per sequence position. Specifically, for a sequence of length $L$, the RNN recurrence is cast as a system:

$$\begin{cases} h_1 - f(0, x_1) = 0 \ h_2 - f(h_1, x_2) = 0 \ \vdots \ h_L - f(h_{L-1}, x_L) = 0 \end{cases}$$

This system is solved using Newton's method, which iteratively linearizes the system and solves the resulting block bi-diagonal linear system at each iteration. The key technical contribution is the development of a highly efficient, parallel solver for these structured linear systems, leveraging parallel reduction (scan) algorithms to achieve $O(\log L)$ complexity per Newton iteration.

Figure 2: Schematics drafting the evolution of the target block bi-diagonal system as it undergoes the first stage of the ParaRNN algorithm.

Figure 4: Schematics of the parallel reduction process at the warp level, showing how the system is further simplified for efficient parallel solution.

The backward pass, required for gradient computation, is inherently linear and can be parallelized using the same reduction techniques, further improving training efficiency.

Architectural Adaptations: ParaGRU and ParaLSTM

A practical challenge in applying parallel reduction to nonlinear RNNs is the computational and memory overhead associated with dense Jacobians, which scale as $O(L d_h^2)$ in memory and $O(d_h^3)$ in compute per reduction step. To address this, ParaRNN introduces architectural modifications to classical GRU and LSTM cells, enforcing diagonal or block-diagonal structure on the state and peephole matrices. This reduces the Jacobians to diagonal or $2 \times 2$ block-diagonal form, enabling efficient storage and computation ($O(d_h)$ per step) without significant loss of expressivity, as feature mixing is handled by subsequent MLP layers.

Software Implementation and Engineering

ParaRNN is released as a modular, open-source PyTorch+CUDA library. It provides three main solver implementations:

1. Pure PyTorch: Reference implementation for prototyping and debugging, leveraging autograd for Jacobian computation.
2. CUDA-Accelerated: Custom CUDA kernels for parallel reduction, optimized for diagonal and block-diagonal Jacobians, with system assembly in PyTorch.
3. Fully-Fused CUDA: End-to-end CUDA implementation, fusing Newton iterations, Jacobian assembly, and parallel reduction in a single kernel for maximal performance.

The CUDA implementation is hardware-aware, exploiting the GPU hierarchy (thread, warp, block, grid) to minimize memory traffic and maximize register utilization. The chunk size and other kernel parameters are tunable for optimal performance on different hardware and sequence lengths.

Empirical Results

Training Efficiency

ParaRNN achieves substantial speedups over naive sequential RNN application, with up to $665\times$ acceleration reported. The CUDA implementation for diagonal Jacobians outperforms the equivalent Mamba parallel scan kernel, and the fully-fused implementation achieves up to $2.6\times$ speedup over Mamba for ParaGRU at $L=512$.

Figure 6: Cross-entropy loss evolution for the 7B models considered, showing stable and competitive convergence for ParaRNN-based models.

Inference throughput is also competitive, with ParaRNN models generating tokens at rates similar to or exceeding Mamba, and significantly faster than Transformers, especially for long sequences.

Language Modeling Performance

ParaRNN enables the training of nonlinear RNNs (ParaGRU and ParaLSTM) at scales up to 7B parameters. On standard language modeling benchmarks (e.g., SlimPajama, lm-eval-harness), ParaGRU and ParaLSTM achieve perplexity and downstream task performance comparable to similarly sized Transformers and Mamba2 models. Notably, on synthetic tasks requiring nonlinear computation (e.g., Parity, k-hop), ParaRNN models achieve perfect accuracy, while linear SSMs like Mamba2 fail, highlighting the expressivity advantage of nonlinear recurrence.

Key empirical findings:

• ParaGRU and ParaLSTM match or outperform Transformers on several downstream tasks at scale.
• Mamba2 achieves the lowest perplexity, but ParaRNN models are consistently competitive.
• Nonlinear RNNs solve tasks that are provably intractable for linear SSMs, confirming the necessity of nonlinear recurrence for certain sequence modeling problems.

Limitations and Trade-offs

• Newton Convergence: The efficiency of ParaRNN depends on rapid convergence of Newton's method. Empirically, three iterations suffice for ParaGRU and ParaLSTM, but this must be verified for new RNN cell designs.
• Expressivity vs. Efficiency: The diagonal/block-diagonal Jacobian constraint may limit intra-state mixing within the RNN cell, but this is mitigated by downstream MLPs and is consistent with practices in modern SSMs.
• Overhead: Parallelization introduces additional overhead in Jacobian assembly and manipulation, but this is offset by the substantial gains in sequence-parallel throughput.

Theoretical and Practical Implications

ParaRNN demonstrates that the sequence-parallelization barrier for nonlinear RNNs can be overcome, enabling their use in large-scale language modeling. This reopens the design space for recurrent architectures, allowing exploration of nonlinear sequence models with competitive training and inference efficiency. The approach generalizes to any Markovian RNN with suitable Jacobian structure, and the open-source codebase facilitates further research and practical deployment.

Theoretically, ParaRNN provides a concrete instantiation of time-parallel ODE solvers in deep learning, connecting advances in numerical analysis with neural network training. The empirical results challenge the prevailing assumption that only linear SSMs or attention-based models are viable for large-scale sequence modeling, and suggest that nonlinear recurrence remains a powerful and practical modeling primitive.

Future Directions

• Generalization to Other RNN Cells: Extending ParaRNN to support more complex or structured Jacobians (e.g., sparse, circulant) could further enhance expressivity without sacrificing efficiency.
• Theoretical Analysis of Newton Convergence: Formal characterization of the conditions under which Newton's method converges rapidly for different RNN cell designs.
• Hybrid Architectures: Combining ParaRNN with attention or SSM modules to exploit complementary strengths.
• Applications Beyond Language Modeling: Adapting ParaRNN to other domains requiring efficient sequence modeling, such as time series forecasting or reinforcement learning.

Conclusion

ParaRNN provides a scalable, efficient framework for training nonlinear RNNs in parallel, matching the training efficiency of SSMs and Transformers while restoring the expressive power of nonlinear recurrence. The empirical results establish nonlinear RNNs as competitive alternatives for large-scale language modeling, and the open-source implementation lowers the barrier for further research and application. This work repositions nonlinear RNNs as viable, efficient, and expressive sequence models for modern AI systems.