M$^2$RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling

Transformer 模型虽具备高度并行性，但其计算能力被严格限制在 TC⁰ 复杂度类内，因而无法胜任实体追踪、代码执行等已被严格证明需要更强表达能力的任务。受此局限性启发，我们重新审视了非线性循环神经网络（RNN）在语言建模中的应用，并提出了一种新型架构——矩阵到矩阵 RNN（M²RNN）：该模型采用矩阵形式的隐状态，并配备表达力强的非线性状态转移机制。我们发现，非线性 RNN 在语言建模任务上的性能瓶颈主要在于其隐状态维度不足；同时，我们也展示了其隐状态维度扩展机制如何高效利用张量核心（tensor cores）。实验表明，M²RNN 能够在训练阶段未曾见过的序列长度上，实现完美的隐状态追踪泛化能力。这一优势同样可迁移到大规模语言建模任务中。在将循环层与注意力机制交替堆叠的混合架构中，混合式 M²RNN（Hybrid M²RNN）在 70 亿参数稀疏专家模型（7B MoE）上，相较同等规模的门控 DeltaNet（Gated DeltaNet）混合模型，困惑度（perplexity）降低了 0.4–0.5 点，且其循环层所用隐状态尺寸仅为后者的三分之一。尤为值得注意的是，在现有混合架构中仅将单个循环层替换为 M²RNN，即可获得与完整 Hybrid M²RNN 相当的精度提升，且对训练吞吐量的影响微乎其微。此外，仅含一个 M²RNN 层的混合门控 DeltaNet 模型，在长上下文泛化能力方面亦表现卓越，在 LongBench 基准测试中，其性能比当前最优的混合线性注意力架构高出最多达 8 分。综上所述，这些结果充分证实：非线性 RNN 层是一种极具潜力的基础构件，可助力构建高效、可扩展的语言模型。

Transformer架构虽高度并行，但其计算能力被严格限制在TC⁰复杂度类内，无法理论上保证完成需要更高计算表达力的任务（如精确实体跟踪、符号推理或可编程状态维护）。论文旨在验证：非线性RNN——通过增强隐藏状态表示能力和状态转移非线性——能否突破该理论瓶颈，在保持训练可扩展性的同时实现强泛化（尤其长程状态跟踪与超长上下文建模），从而为语言模型提供兼具表达力与效率的替代性循环构建模块。

提出Matrix-to-Matrix RNN（M²RNN）：采用矩阵值隐藏状态（而非标量或向量），结合可学习的、高阶非线性状态更新（如双线性门控+矩阵函数组合），显著提升单层状态容量与动态建模能力；关键创新在于‘状态尺寸扩展机制’——将隐状态维度增长解耦于参数量增长，并天然适配GPU张量核（Tensor Core）的矩阵乘加速，实现高吞吐下的大状态规模。这不同于传统RNN的向量状态或线性RNN（如DeltaNet）的受限表达，也区别于Transformer依赖长序列注意力带来的O(n²)开销。

实验表明M²RNN在合成任务（如多跳指代跟踪）上实现完美零样本长序列泛化（测试长度远超训练最大长度）；在7B MoE大模型中，Hybrid M²RNN以1/3的循环层状态尺寸超越Gated DeltaNet基线0.4–0.5 perplexity；仅替换1个循环层即带来显著增益且几乎不降吞吐；在LongBench长文本基准上，含单M²RNN层的Hybrid DeltaNet比最优线性注意力模型（如FlashAttention-2+LR-Attention混合）高8分；论文未明确提及开源代码，但所有实验均基于标准LM基准（WikiText-103, C4, LongBench）和公开MoE架构；值得深挖方向包括：M²RNN与程序合成/形式验证的接口、状态矩阵的可解释性分解、以及其在多模态状态建模中的迁移。

Gated DeltaNet (ICML 2023); Linear Transformers (ICLR 2021); RWKV: Reinventing RNNs for the Transformer Era (NeurIPS 2023); State Space Models (S4, Mamba, ICML 2022/2023); Attention is All You Need (NIPS 2017); Neural Turing Machines (NIPS 2014); Structured State Space Sequence Models (SSM) (TMLR 2023)