Apple 提出 ParaRNN：利用牛顿法实现非线性 RNN 并行训练，70 亿参数模型性能媲美 Transformer
Apple 研究人员提出 ParaRNN 框架，通过结合牛顿迭代法与结构化雅可比矩阵，首次实现了大规模非线性循环神经网络（RNN）的并行训练。该方法使传统 GRU 和 LSTM 模型的训练速度提升 665 倍，并成功训练出首个 70 亿参数的经典 RNN 大语言模型。实验表明，ParaGRU 和 ParaLSTM 在困惑度及下游任务表现上已具备与 Transformer 及 Mamba2 等先进架构竞争的能力。相关代码库已开源以推动高效序列建模研究。
事件概述
Apple Machine Learning Research 团队发布新论文《ParaRNN: Large-Scale Nonlinear RNNs, Trainable in Parallel》，该论文已被 ICLR 2026 接收为口头报告（Oral）。研究提出了一种全新的非线性 RNN 并行训练框架，解决了传统 RNN 因计算顺序性而无法扩展至数十亿参数的问题，使得经典 RNN 架构能够在大模型时代重新获得竞争力。
核心突破与技术原理
1. 解决训练瓶颈
传统 RNN 在推理阶段具有 O(1) 的时间复杂度优势，但在训练阶段必须按时间步顺序展开，无法像 Transformer 或选择性状态空间模型（SSM，如 Mamba）那样进行并行计算。现有的 SSM 通过简化为线性递推关系来实现并行化，但这牺牲了模型的非线性表达能力。
2. 牛顿法并行化策略
ParaRNN 的核心创新在于将牛顿法（Newton's method）应用于 RNN 的并行求解：
- 系统重构：将整个序列的递归关系视为一个包含所有隐藏状态的方程组，而非单步链式结构。
- 迭代线性化：利用牛顿法通过局部导数（雅可比矩阵）将非线性系统近似为线性系统。由于线性操作满足结合律，可应用并行扫描（parallel scan）算法加速求解。
- 快速收敛：对于设计良好的 RNN 模型，仅需约 3 次牛顿迭代即可收敛，从而在保留非线性表达能力的同时，实现接近线性的训练效率。
3. 工程优化：结构化雅可比矩阵
为避免通用 RNN 中雅可比矩阵稠密导致的存储和计算开销（二次方/立方级增长），团队设计了 ParaGRU 和 ParaLSTM 变体：
- 对角化设计：简化细胞定义，使 GRU 的雅可比矩阵呈对角分布，LSTM 的呈块对角分布。
- 自定义 CUDA 内核：针对上述结构化矩阵特性，开发了全融合（fully-fused）CUDA 内核，将牛顿迭代、系统组装和并行归约整合在一个核函数中，显著提升了 GPU 利用率。
关键实验结果
训练效率
- 速度提升：相比传统串行方法，ParaRNN 实现了 665 倍 的训练加速。
- 规模突破：成功训练了首个 70 亿参数（7B） 的经典 RNN 模型。
模型性能
在语言建模任务中，ParaGRU 和 ParaLSTM 在 7B 规模下展现出与 Transformer 及 SOTA SSM（如 Mamba2）相当的性能：
- 困惑度（Perplexity）：在同等参数量下，ParaLSTM (6.76B) 的困惑度为 9.16，与 Mamba2 (6.96B) 的 8.62 处于同一量级。
- 下游任务：在 Arc-C、HSwag、OBQA、WinoG、PiQA 及 MMLU 等多个基准测试中，ParaGRU 和 ParaLSTM 的表现均达到或接近 Transformer 水平。
资源开放
为支持社区探索新的非线性 RNN 模型，Apple 已将 ParaRNN 代码库作为开源框架发布，提供自动化的非线性 RNN 并行训练工具，旨在降低高效序列建模的研究门槛。