arXiv cs.AI· Mayank Mishra, Shawn Tan, Ion Stoica, Joseph E. Gonzalez, Tri Dao·· 11 小时前AI 评分38
M²RNN：采用矩阵值状态的非线性 RNN，面向可扩展语言建模
M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling
AI 导读
研究者提出 M²RNN，一种隐藏状态为矩阵值、状态转移为非线性表达的 RNN 架构，并指出非线性 RNN 的语言建模性能受状态大小限制，其状态扩展机制可高效利用 tensor core。
正文 · AI 翻译
摘要：Transformer 具有高度并行性，但局限于 TC$^0$ 复杂度类内的计算，无法完成实体跟踪和代码执行等可证明需要更强表达能力的任务。受此限制的启发，我们重新审视用于语言建模的非线性循环神经网络（RNN），并提出矩阵到矩阵 RNN（M$^2$RNN）：一种具有矩阵值隐藏状态和强表达力非线性状态转移的架构。我们证明，非线性 RNN 的语言建模性能受其状态大小限制，并展示状态大小扩展机制如何实现对张量核心的高效利用。在实验中，M$^2$RNN 在训练时未见过的序列长度上实现了完美的状态跟踪泛化。这些优势同样可推广到大规模语言建模。在将循环层与注意力层交替的混合设置中，Hybrid M$^2$RNN 在 7B MoE 模型上比同等的 Gated DeltaNet 混合架构困惑度低 $0.4$-$0.5$ 个点，同时循环层使用的状态大小小 $3\times$。值得注意的是，在现有混合架构中即使仅将一个循环层替换为 M$^2$RNN，也能获得与 Hybrid M$^2$RNN 相当的准确率提升，且对训练吞吐量的影响极小。此外，加入单个 M$^2$RNN 层的 Hybrid Gated DeltaNet 模型也实现了更优的长上下文泛化能力，在 LongBench 上比最先进的混合线性注意力架构高出多达 $8$ 个点。这些结果共同表明，非线性 RNN 层是构建高效且可扩展语言模型的有力组件。
| 学科： | 机器学习（cs.LG）；人工智能（cs.AI） |
| 引用为： | arXiv:2603.14360 [cs.LG] |
| （或本版本的 arXiv:2603.14360v3 [cs.LG]） | |
| https://doi.org/10.48550/arXiv.2603.14360 由 DataCite 颁发的 arXiv DOI |
提交历史
来自：Mayank Mishra [查看邮箱]
[v1]
2026年3月15日周日 12:53:09 UTC（921 KB）
[v2]
2026年5月13日周三 23:29:01 UTC（944 KB）
[v3]
2026年10月16日周五 00:31:17 UTC（794 KB）
来源：arXiv cs.AI · arxiv.org