ParaRNN实现非线性RNN并行训练助力大型语言模型

循环神经网络（Recurrent Neural Networks, RNNs）作为序列建模的奠基模型，在深度学习领域具有重要意义。然而，RNN固有的序列特性限制了其并行计算能力，成为规模化应用的关键障碍。这一问题促使高度可并行的架构（如Transformer）以及最新的状态空间模型（State Space Models, SSMs）占据主导地位。虽然SSMs可以通过结构化的线性递归实现高效并行化，但它们的线性约束限制了模型的表达能力，使其难以捕捉复杂的、非线性的序列依赖性。

为解决这一痛点，我们推出了ParaRNN，一个突破非线性RNN序列依赖瓶颈的全新框架。通过将非线性递归关系描述为一个单一的方程组，并结合牛顿迭代法（Newton’s Iterations）和自定义的并行化简操作，ParaRNN实现了非线性RNN的前所未有的训练速度提升。我们的实验表明，与传统的逐步序列运算相比，ParaRNN在并行训练中达到了高达665倍的加速性能。

ParaRNN基于现有研究，创新性地将序列中非线性的递归关系转化为一个系统化的方程组。这一方法得益于牛顿迭代的快速收敛特性，同时凭借专门设计的并行操作，成功打破了传统RNN结构中无法并行化的限制。

通过ParaRNN框架，我们首次成功对非线性RNN模型进行了超大规模参数训练（7B参数）。在实验中，采用ParaRNN优化的LSTM和GRU架构不仅能够与同等规模的Transformer和Mamba2模型在困惑度（Perplexity）上达到可比的表现，还展示出了其更强的非线性建模能力。

为了推动高效序列建模领域的进一步发展，我们已将ParaRNN的代码库以开源框架形式向公众发布。这一框架提供了自动并行化非线性RNN的能力，为研究人员和从业者探索大规模非线性RNN模型提供了全新工具，有望加速相关研究和应用落地。