ParaRNN: Unlocking Parallel Training of Nonlinear RNNs for Large Language Models

文章核心总结与翻译

一、主要内容

本文针对传统循环神经网络（RNNs）因固有关联性导致的并行训练障碍，以及状态空间模型（SSMs）为并行化牺牲非线性表达能力的问题，提出了ParaRNN框架——一种实现非线性RNN并行训练的解决方案。

1. 核心挑战：

   • 传统RNN（如LSTM、GRU）需沿序列逐次展开，无法并行训练，难以规模化；
   • Transformer虽支持并行训练，但推理效率较低；
   • SSMs（如Mamba）通过线性递归实现并行化，却受限于线性约束，表达能力不足。

2. 技术方案：

   • 将非线性RNN的序列递归关系转化为单一非线性方程组，通过牛顿迭代法线性化该系统；
   • 利用方程组的块双对角结构，结合自定义并行归约算法，实现O(log₂L)时间复杂度的并行求解；
   • 适配GRU和LSTM架构，通过简化Jacobian矩阵结构（对角线/块对角线）降低计算开销，推出ParaGRU和ParaLSTM模型。

3. 实验结果：

   • 训练速度较传统串行RNN提升最高665倍，可高效训练70亿参数规模模型；
   • 推理吞吐量达35-37 tokens/s，优于Mamba（27 tokens/s），且推理时间不随序列