# M²RNN：矩阵状态非线性 RNN 的容量、表达力与系统代价

**论文**：*M²RNN: Non-Linear RNNs with Matrix-Valued States for Scalable Language Modeling*  
**作者**：Mayank Mishra、Shawn Tan、Ion Stoica、Joseph E. Gonzalez、Tri Dao  
**版本**：arXiv:2603.14360v2，2026 年 5 月 13 日修订；本次获取与研读日期：2026 年 10 月 9 日。  
**原文**：[arXiv v2](https://arxiv.org/abs/2603.14360v2) · [本地 PDF](../source/paper.pdf) · [英文汇报幻灯片](slides.en.pdf)

## 论文总览

M²RNN 将传统非线性循环神经网络的向量状态扩展为矩阵：使用 $k_tv_t^\top$ 写入，用 $\tanh(\mathbf H_{t-1}\mathbf W+k_tv_t^\top)$ 变换历史状态，再通过遗忘门保留或更新记忆。这样既继承现代线性注意力的大状态设计，又保留逐步非线性递推。矩阵乘法的形状还天然适合 tensor core，避免小 batch 向量 RNN 为满足硬件要求而大量补零。

论文最有实践意义的结果来自**少量非线性层与线性循环层、注意力层的组合**。在 7B 总参数、约 1.1B 激活参数的 MoE 上，把 Hybrid Gated DeltaNet 的第一层循环模块换成 M²RNN，WikiText 困惑度从 **13.51 降到 13.07**，LongBench 八项平均分从 **23.1 升到 30.4**。作者报告在所测 16K 上下文训练吞吐中，单层替换的下降在 **6% 以内**。来源：Table 2，PDF p. 16；Table 6，p. 21；Sec. 5.7 / Fig. 10，p. 22。

这三个结果要按配置一起理解：单层方案的 LongBench 增益是 **7.3 分**；**8.1 分**对应五层方案。纯 M²RNN 的真实数据检索仍低于纯 Gated DeltaNet，而且训练明显更慢。因此，证据最支持的用途是给混合架构增加少量非线性状态计算能力。

## 结构地图

以下页码均为从 1 开始的 PDF 页码，本版本与印刷页码一致。

| 原文章节 | 作用 | PDF 页码与关键证据 |
| --- | --- | --- |
| 1 Introduction | 提出状态追踪、容量、硬件利用率三个问题 | pp. 1–3 |
| 2 Background | 比较线性/非线性 RNN，说明序列并行与状态容量限制 | pp. 3–6，Eq. (1)–(9) |
| 3.1 M²RNN Layer | 递推、遗忘门、完整模块及训练稳定性 | pp. 6–9，Eq. (10)–(23)，Figs. 1–2 |
| 3.2 Improved State Tracking | 向量 RNN 模拟定理与 $S_3$ 实验 | p. 9，Theorem 1，Fig. 3 |
| 3.3–3.4 | 扩大状态与改善硬件分块的理由 | p. 10 |
| 4 Distributed Training and Systems Optimizations | 前向、反向与两种张量并行 | pp. 11–14，Algorithms 1–2，Figs. 4–5 |
| 5.1–5.2 Models / Model Training | 架构、训练数据和对照协议 | p. 15 |
| 5.3 Language Modeling | 同构、两路混合、三路混合比较 | pp. 15–17，Tables 1–2 |
| 5.4 In-Context Retrieval | RULER 与真实数据检索 | pp. 17–20，Figs. 6–9，Tables 3–4 |
| 5.5 Long-Context Performance | 摘要、few-shot、代码任务 | pp. 20–21，Tables 5–6 |
| 5.6–5.7 | 状态/门控消融与训练吞吐 | pp. 21–22，Table 7，Fig. 10 |
| 6 Conclusion | 总结混合架构与系统设计 | pp. 22–23 |
| Appendix A / B / C | 参数—状态关系、表达力证明、RMSNormTP 推导 | pp. 29–31，Table 8，Eq. (25)–(33) |

## 一：动机——表达力、记忆容量与训练速度需要同时考虑

### 1.1 “有状态”不自动意味着擅长状态追踪

线性 RNN 中的“线性”指对旧状态的依赖，例如

$$
\mathbf H_t=\mathbf H_{t-1}\mathbf A(x_t)+\mathbf B(x_t).
$$

$\mathbf A$、$\mathbf B$ 可以由非线性输入网络生成；只要旧状态仍以仿射形式参与更新，就能利用相应的扫描、分块算法。非线性 RNN 则在历史状态变换后施加激活函数，递推深度随序列增长，能够进行更丰富的状态演化。

可以把状态追踪想成连续执行“交换第 1、3 个球”“交换第 2、3 个球”，最后回答某个球的位置。一次读取过去 token 和正确累计所有交换，是不同的计算要求。论文以置换群任务建立动机，再研究这种逐步计算是否能改善语言模型。

这里需要保留理论条件：论文背景把某些模型的表达力上界放在 $\mathsf{TC}^0$，把 $S_5$ 词问题放在 $\mathsf{NC}^1$-complete。将两者说成严格分离，需要通常采用但尚未证明的 $\mathsf{TC}^0\ne\mathsf{NC}^1$ 假设，以及相关固定深度、数值精度等模型设定。输入相关且非对角的某些线性递推也能越过所讨论的限制，不能把结论套到所有线性 RNN。原论文 Sec. 2.3.1 本身也讨论了这些扩展。背景补充依据：[Merrill et al., ICML 2024，Sec. 3 / Fig. 2](https://raw.githubusercontent.com/mlresearch/v235/main/assets/merrill24a/merrill24a.pdf)。

### 1.2 非线性 RNN 的旧问题可能有相当部分来自状态太小

向量 RNN 保存 $d$ 个状态量，但稠密状态转移需要 $d^2$ 个参数；直接把 $d$ 放大很昂贵。矩阵状态模型保存 $NKV$ 个量，用外积写入避免直接预测全部元素。M²RNN 进一步只用 $V\times V$ 的矩阵变换每一行，避免为 $KV$ 维展平状态学习一个 $(KV)\times(KV)$ 转移。

这也有计算结构上的取舍：同一个 head 内，$\mathbf W$ 混合 value 维度，**不直接混合不同 key 行**；各行共享该转移矩阵，用各自的 $k_{t,i}$ 调制写入，最后由 query 汇总。跨 head 的信息可经输出投影及后续网络交互。

### 1.3 矩阵状态同时改变硬件映射

向量 RNN 常把 batch 拼成 GEMM 的一个维度。假设每 GPU 的 batch 为 4，但 tensor core 分块需要补到 16，则补零占去该递推乘法的 75% FLOPs。M²RNN 用状态自身的 $K$ 行支撑矩阵乘法，即使 batch 很小也有可用的 tile。它改善了每个顺序步骤的效率，但时间维仍然存在串行依赖。来源：Sec. 2.5.3、3.4，PDF pp. 6、10。

## 二：方法——如何写入、变换、保留和读取矩阵记忆

### 2.1 单个 head 的核心递推

| 符号 | 维度 | 含义 |
| --- | --- | --- |
| $x_t$ | $d\times1$ | 当前层输入；短卷积还使用此前局部输入 |
| $q_t,k_t$ | $K\times1$ | 读取权重与写入方向 |
| $v_t$ | $V\times1$ | 当前写入内容 |
| $\mathbf H_t,\mathbf Z_t$ | $K\times V$ | 保存状态、候选状态 |
| $\mathbf W$ | $V\times V$ | 可学习、与当前输入无关的状态转移矩阵 |
| $f_t$ | 标量，$[0,1]$ | 每个 head 的保留比例 |
| $w_r,g_t,y_t$ | $V\times1$ | value 残差权重、输出门、读出结果 |

从 $\mathbf H_0=0$ 开始，原文 Eq. (10)–(12)，PDF p. 7 为：

$$
\mathbf Z_t=\tanh(\mathbf H_{t-1}\mathbf W+k_tv_t^\top),
$$

$$
\mathbf H_t=f_t\mathbf H_{t-1}+(1-f_t)\mathbf Z_t,
$$

$$
y_t=\mathbf H_t^\top q_t+w_r\odot v_t.
$$

第一步先变换旧记忆，再加入 rank-1 外积输入，并逐元素施加 $\tanh$。第二步在旧状态和候选状态之间插值。第三步用 query 沿 $K$ 维加权读取，再加上当前 value 的可学习直接通路。

外积输入为 rank-1，并不要求累计状态始终为 rank-1；时间累积、矩阵变换和逐元素非线性共同改变状态。另一方面，写入经过 $\tanh$ 且作用在已变换的历史上，不能直接套用线性注意力的“历史外积之和”解释。

**一个有用的公式解读**：$f_t=1$ 时状态保持；$f_t=0$ 时完全采用候选状态，但候选状态中仍有 $\mathbf H_{t-1}\mathbf W$。因此，原文称 $f_t=0$ “完全重置”容易使人误以为历史被清空，实际只是去掉直接保留分支。

由上述递推还能得到一个简单性质：当 $f_t\in[0,1]$ 且初始状态为零时，状态元素始终位于 $[-1,1]$，因为它是旧有界状态与 $\tanh$ 输出的凸组合。这是依据公式得到的分析；状态有界仍不能阻止梯度在时间上放大。

### 2.2 输入相关、与循环状态无关的遗忘门

将第 $n$ 个 head 的输入投影记作标量 $a_{t,n}=(\mathbf W_f x_t)_n$，则原文 Eq. (13)、(17) 可写成

$$
f_{t,n}=(1+e^{a_{t,n}+\beta_n})^{-\alpha_n}
=\exp[-\alpha_n\operatorname{softplus}(a_{t,n}+\beta_n)].
$$

$\alpha_n$ 控制衰减速度，$\beta_n$ 移动门的响应位置；要保持该参数化落在 $[0,1]$，需保持 $\alpha_n\ge0$。当 $\alpha_n=1$ 时，门等于 $\sigma(-a_{t,n}-\beta_n)$。不同 head 的初始化使它们具有不同保留时间尺度。原文给出了 Uniform / LogUniform 初始化形式，但没有在此展开数值范围及正值约束实现。

与 GRU/LSTM 常见门控相比，这里的门不读取本层的 $\mathbf H_{t-1}$，所以整条输入序列的门值可以提前并行计算。门的预计算不会消除候选状态中的时间依赖。

### 2.3 完整模块：局部卷积、循环记忆和输出门

![M²RNN 模块结构](figures/fig-02-block.png)

原图 2，PDF p. 8，按最终 PDF 裁切。右图自下向上阅读；左图显示它在带残差的序列模块和 MLP/MoE 中的位置。

完整计算由原文 Eq. (14)–(23) 给出：

1. $q,k,v$ 分别经过线性投影、核大小为 4 的 depthwise causal Conv1D 和 SiLU。卷积只访问当前及过去位置，可在训练时对整段输入并行计算。
2. $f_t=\psi(\mathbf W_f x_t)$；另一路 $g_t=\operatorname{SiLU}(\mathbf W_gx_t)$ 不经过卷积。
3. 执行矩阵递推与 query 读出，得到 $y_t$。
4. $y_{gt}=\operatorname{RMSNorm}(y_t\odot g_t)$，再通过 $\mathbf W_o$ 返回模型宽度。注意是**先乘输出门，再做 RMSNorm**。
5. 模块外仍有 pre-norm 和残差；序列混合模块之后为 SwiGLU MLP 或 MoE。

作者将 $\mathbf W$ 初始化为单位阵：报告其效果接近正交初始化，优于普通正态初始化。训练还对 BPTT 中每一步的循环状态梯度进行裁剪；Sec. 5.2 另报告全模型 gradient clipping 为 1，两者不应混为同一操作。来源：Sec. 3.1.2，PDF pp. 8–9。

### 2.4 Multi-value 共享如何降低扩容的投影成本

所有 value heads 共享一组 query/key：$N_q=N_k=1$，$N_v=N_f=N_w=N$。每个 head 保留自己的 $v_t,f_t,\mathbf W$ 与状态，单样本单层的循环状态共有 $NKV$ 个元素。

论文使用 $K=64,V=16$，所以单 head 状态有 1,024 个元素，转移矩阵只有 256 个元素。附录 A 按线性投影参数分析，在忽略较小项时：

$$
P_{\rm MH}\approx N(2K+3V)d_m,
\qquad
P_{\rm MV}\approx3NVd_m.
$$

固定相同投影预算 $P$ 后，得到

$$
S_{\rm MH}\approx\frac{PKV}{(2K+3V)d_m},
\qquad
S_{\rm MV}\approx\frac{PK}{3d_m}.
$$

代入 $K=64,V=16$，两种近似表达式给出的容量比为 $(2K+3V)/(3V)=176/48\approx3.67$。这是对 Table 8 的推算，比较的是 multi-value 与 multi-head 的投影分配；与作者所述“M²RNN 的状态约为实验 GDN 的三分之一”是不同的比较。

Table 8 的 $P$ 仅覆盖线性投影并使用近似，不能当成模型总参数。以 multi-value 为例，列出的投影项精确相加为 $(2K+3NV+N)d_m$，还未计入循环转移的 $NV^2$、偏置、卷积等参数。来源：Appendix A / Table 8，PDF p. 29。

### 2.5 表达力证明：在矩阵的一行内嵌入向量 RNN

Theorem 1 与 Appendix B 选择 $f_t=0$、$w_r=0$、$q_t=k_t=e_1$。从零初始状态出发，其余行保持零，第一行满足

$$
\mathbf H_{t,1,:}=\tanh(\mathbf H_{t-1,1,:}\mathbf W+v_t^\top),
\qquad
y_t=\mathbf H_{t,1,:}^\top.
$$

转成列向量 $h_t=\mathbf H_{t,1,:}^\top$，就是

$$
h_t=\tanh(\mathbf W^\top h_{t-1}+v_t).
$$

因此，抽象 M²RNN 递推包含同宽度的普通 tanh RNN。这里的论证针对递推核心，且明确允许门端点 $f_t=0$；实际有限参数的正值门函数通常只能趋近零。严格地把该构造推广到带全部投影、输出归一化的实际网络，还需要处理这些实现约束。来源：Theorem 1，PDF p. 9；Appendix B / Eq. (25)，p. 30。

### 2.6 系统实现：哪些地方并行，哪些地方顺序执行

前向把 batch 与 heads 分给独立的 GPU 工作单元，每个 $(b,n)$ 在片上保留 $\mathbf H$ 和 $\mathbf W$，按时间顺序加载 $q,k,v,f$，完成更新并写出读出向量。矩阵乘法形状为 $(K,V)\times(V,V)$，不依赖 batch 维的补齐。对于固定 $K,V,N$，递推成本随 $T$ 线性增长；核心矩阵乘法工作量为 $O(BTNKV^2)$。

为了避免在正常前向中长期保留每一层全部历史状态，作者在反向开始时重新运行一次递推，并将该层各时间步的状态缓存到 HBM，再反向扫描求梯度。这个策略仍会产生 $O(BTNKV)$ 的临时状态缓存。论文明确指出当前 Triton 反向实现受显存带宽约束，把 CUTLASS 优化留给后续工作。**推理的固定循环状态大小，不等于训练总激活内存与序列长度无关。** 来源：Sec. 4.1–4.2，PDF pp. 11–12。

### 2.7 Algorithm 2 的反向公式应按链式法则修正

原文 PDF p. 12 的 Algorithm 2 存在实质性伪代码问题：重算循环没有显式更新局部 $\mathbf H$；反向把 $r$ 写成当前读出贡献，未累加未来时间步的状态梯度；后续又用乘法覆盖了旧状态梯度的保留分支。它不能按字面直接实现完整 BPTT。以下是依据 Eq. (10)–(12) 的教学推导，针对读出 $\bar y_t=\mathbf H_t^\top q_t$，暂时单独移出 $w_r\odot v_t$ 直接通路。

令 $u_t=\partial\mathcal L/\partial\bar y_t$，$\mathbf C_t$ 为来自未来步骤的状态梯度，则

$$
\mathbf R_t=\mathbf C_t+q_tu_t^\top,
\qquad
\mathbf D_t=((1-f_t)\mathbf R_t)\odot(1-\mathbf Z_t\odot\mathbf Z_t),
$$

$$
\boxed{\mathbf C_{t-1}=f_t\mathbf R_t+\mathbf D_t\mathbf W^\top.}
$$

其余关键梯度为

$$
\frac{\partial\mathcal L}{\partial f_t}
=\langle\mathbf R_t,\mathbf H_{t-1}-\mathbf Z_t\rangle_F,
\qquad
\frac{\partial\mathcal L}{\partial\mathbf W}
=\sum_t\mathbf H_{t-1}^\top\mathbf D_t,
$$

$$
\frac{\partial\mathcal L}{\partial k_t}=\mathbf D_tv_t,
\quad
\frac{\partial\mathcal L}{\partial v_t}=\mathbf D_t^\top k_t,
\quad
\frac{\partial\mathcal L}{\partial q_t}=\mathbf H_tu_t.
$$

恢复 value 残差时，还需给 $v_t$ 梯度加上 $w_r\odot u_t$ 等直接分支贡献。共享 query/key 的梯度须跨 heads 累加，$\mathbf W$ 梯度须跨 batch 和时间累加。这里指出的是论文伪代码与链式法则之间的差异，不能据此认定作者训练内核也犯了同样的错误。

### 2.8 两种张量并行策略

| 策略 | query/key 分配 | 额外通信，相对标准 TP | 主要取舍 |
| --- | --- | --- | --- |
| Topology-aware | 每张 TP GPU 各有一组，$N_q=N_k=N_{TP}$ | 不增加论文所述的额外通信 | 投影规模和归一化组织与 TP 拓扑绑定 |
| Topology-independent | 所有 value 分片共用同一组 query/key | 前向 1 次 AllReduce；反向 3 次 | 保持参数数量不随 TP 大小变化，增加同步 |

第二种策略的额外前向通信同步 RMSNorm 的平方和；反向分别同步 RMSNorm 所需的内积、query 梯度和 key 梯度。若共同投影/卷积的复制权重初始化相同，汇总相同的输入梯度后可保持同步。来源：Sec. 4.3、Figs. 4–5，PDF pp. 12–14。

Appendix C 的关键是每个 token 的标量统计量：

$$
s=\left[\frac1d\sum_{n=1}^{N_{TP}}\sum_{j\in D_n}x_j^2\right]^{-1/2},
\qquad
r=\sum_{n=1}^{N_{TP}}\sum_{j\in D_n}p_jx_j,
$$

$$
\nabla_{x_i}\mathcal L=sp_i-\frac{rs^3x_i}{d},
\quad p_i=w_i\nabla_{y_i}\mathcal L.
$$

这里“通信一个元素”是每个归一化向量一个标量；对整个 batch/序列仍有 $BT$ 个统计量。原文为简洁省略了数值稳定用的 epsilon。来源：Appendix C，PDF pp. 30–31，Eq. (26)–(33)。

## 三：实验——比较什么，以及比较条件如何影响解释

### 3.1 预训练设置

| 项目 | 设置 |
| --- | --- |
| 小模型 | 410M 级 dense，24 层 |
| 大模型 | 7B 总参数、1.1B 激活参数 MoE，40 层；表题简写为 1B active |
| 数据 | Nemotron-CC-v2 的高质量子集，100B tokens |
| 序列与 batch | 长度 4,096；全局 batch 1,024 个样本 |
| 优化 | AdamW，余弦学习率；前 10B tokens warmup 到 $3\times10^{-4}$，末端到 $3\times10^{-5}$ |
| 正则与精度 | weight decay 0.1，norm 参数不衰减；全局梯度裁剪 1；BF16，梯度累积/AllReduce 为 FP32 |
| 训练硬件 | 410M 用 32 张 H100，7B 用 64 张 H100 |
| 对照约束 | 同模型宽度、MLP 宽度、层数及其余超参数；同数据、同顺序；只调整序列混合块 |

来源：Sec. 5 / 5.2，PDF p. 15。RNN、GRU 的精确参数量在消融表分别为 406M、474M，不能把所有架构都说成严格等参数。

### 3.2 三类架构

**同构模型**全部使用同一种循环混合模块：Mamba-2、Gated DeltaNet（下文简称 GDN）、向量 RNN、GRU 或 M²RNN。它们比较循环机制自身的语言与记忆能力。

**两路混合模型**每 8 层有 1 层 attention，其余 7 层为某一种循环模块。24 层模型因此有 3 个 attention 层，40 层模型有 5 个。Transformer++ 是全 attention 对照。

**三路混合模型**在两路混合的基础上替换一部分循环层。后缀 `+M²RNN-1` 表示替换第一层循环模块；410M 的 `-3`、7B 的 `-5` 表示替换每个 attention 层前的循环模块。这些方案都从头训练，不能理解为推理时给已训练模型直接换层。来源：Sec. 5.1，PDF p. 15。

另一个会影响长上下文比较的区别是位置编码：原文明确 Transformer++ 使用 RoPE，线性循环及混合模型使用 NoPE。因此，Transformer++ 在长度外推上的差距同时包含位置编码设计因素；GDN hybrid 内部的替换比较更能直接反映 M²RNN 层的作用。

### 3.3 指标口径

| 实验 | 主要指标与方向 | 比较的能力 |
| --- | --- | --- |
| $S_3$ 置换组合 | Accuracy ↑；训练长度 128，测试到 512 | 已学习状态更新规则的长度泛化 |
| WikiText / LAMBADA | Perplexity（PPL）↓ | 概率语言建模质量；不同数据集 PPL 不能直接比较大小 |
| 十项下游评测 | 0-shot accuracy 的表中平均值 ↑ | 包括 LAMBADA、HellaSwag、PIQA、ARC-E/C 等；ARC 分成两项 |
| RULER | 六类 needle 任务 accuracy ↑；评测到 16,384 | 单关联、多 query、多 key 干扰、多 value 召回 |
| 真实数据检索 | SQuAD、NQ、DROP、TriviaQA、FDA、SWDE 的分数及六项均值 ↑ | 自然文本与结构化信息中的检索能力 |
| LongBench | 所选八项任务分数的平均值 ↑ | 摘要、few-shot、代码补全；不是统一的准确率或完整 LongBench 全任务总分 |
| 训练吞吐 | billion tokens/day ↑ | 7B MoE 在 8× HGX H100、NVLink 上的端到端训练速度 |

真实数据检索表未明确展开各数据集的计分实现，因此以下沿用原文“分数”，不擅自统一称为 EM 或 F1。LongBench 表列出的八项是 GovReport、QMSum、MultiNews、TREC、TriviaQA、SAMSum、LCC、RepoBench-P。

## 四：结果——收益集中在哪里，什么地方仍有退化

### 4.1 状态追踪：$S_3$ 上的长度泛化清晰

![S3 置换状态追踪](figures/fig-03-s3.png)

原图 3，PDF p. 9。横轴为序列长度，纵轴为准确率，虚线 128 为训练长度。M²RNN 与 GRU 到长度 512 仍达到作者报告的至少 99.5%；允许负特征值的 GDN 未学好这个任务；使用两个 Householder 因子的 Gated DeltaProduct 在约 256 内接近满分，继续增长后明显下降。

这说明，在相同训练长度下，理论上能够表达某任务并不足以保证训练后稳健外推。图中的直接实验对象是 **$S_3$**；它并未直接展示 $S_5$、代码执行或实体追踪的大模型测试结果。

### 4.2 语言建模：同构 M²RNN 接近现代循环基线

| 同构模型 | 410M Wiki PPL ↓ | 410M LAMBADA PPL ↓ | 410M 平均准确率 ↑ | 7B Wiki PPL ↓ | 7B 平均准确率 ↑ |
| --- | ---: | ---: | ---: | ---: | ---: |
| Mamba-2 | 22.93 | 38.56 | 54.30 | 13.73 | 65.22 |
| GDN | 22.70 | 34.17 | 53.99 | 13.89 | 64.28 |
| RNN | 33.74 | 317.38 | 48.90 | 17.65 | 60.58 |
| GRU | 25.80 | 63.83 | 53.72 | 14.80 | 64.44 |
| M²RNN | 22.92 | 33.63 | 54.24 | 13.80 | 65.17 |

来源：Tables 1–2，PDF p. 16。M²RNN 已大幅摆脱向量 RNN 的困惑度劣势：410M 的 LAMBADA PPL 比 Mamba-2 低 4.93，比 GDN 低 0.54。另一方面，410M WikiText 仍由 GDN 最好；7B 的 WikiText 与下游平均值由 Mamba-2 略领先。纯非线性矩阵递推达到有竞争力的水平，但没有全面胜出。

### 4.3 两路与三路混合：大模型单层替换的性价比最好

| 7B MoE 模型 | Wiki PPL ↓ | LAMBADA PPL ↓ | 十项平均准确率 ↑ |
| --- | ---: | ---: | ---: |
| Transformer++ | 14.94 | 17.26 | 61.76 |
| Hybrid Mamba-2 | 13.10 | 11.05 | 65.03 |
| Hybrid GDN | 13.51 | 11.06 | 64.64 |
| Hybrid M²RNN | 13.00 | 10.58 | 65.56 |
| Hybrid GDN + M²RNN-1 | 13.07 | 10.33 | 65.80 |
| Hybrid GDN + M²RNN-5 | 12.85 | 10.29 | 65.66 |

来源：Table 2，PDF p. 16。两路 Hybrid M²RNN 相对 Hybrid GDN 的 WikiText 优势为 **0.51 PPL**；单层方案优势为 **0.44 PPL**，平均准确率增加 **1.16 个百分点**。五层进一步降低 PPL，但十项平均准确率略低于单层方案，说明层数的最优选择取决于指标。

410M 结果保留了一个容易被概括遗漏的反例：

| 410M 模型 | Wiki PPL ↓ | LAMBADA PPL ↓ | 十项平均准确率 ↑ |
| --- | ---: | ---: | ---: |
| Hybrid GDN | 21.89 | 29.02 | 54.04 |
| Hybrid M²RNN | 21.53 | 37.31 | 54.15 |
| Hybrid GDN + M²RNN-1 | 21.39 | 33.78 | 53.75 |
| Hybrid GDN + M²RNN-3 | 21.26 | 28.93 | 55.23 |

来源：Table 1，PDF p. 16。单层替换虽然把 Wiki PPL 降低 0.50，但 LAMBADA PPL 变差、平均准确率下降 0.29 个百分点。两路 Hybrid M²RNN 的 Wiki PPL 优势为 0.36，不能把论文中的概括写成每个指标都改善。

Mamba 路线的收益也更有限：410M 的 Wiki PPL 为 21.59 → 21.48（单层）→ 21.39（三层）；7B 为 13.10 → 13.01（单层）→ 13.02（五层）。原文所述“增加层数再改善 0.13 / 0.22 PPL”准确对应 GDN 路线，不应套用到 Mamba 路线。

### 4.4 RULER：改善取决于任务、骨干和替换层数

RULER 的 S-NIAH 只需取回一个关联；MQ-NIAH 需要回答多个 query；MK-NIAH 增加多个 key 作为干扰；MV-NIAH 要返回同一个 key 的多个 value。所有语言模型在 4K 训练，图中到 16K 的部分检验长度外推。

![410M GDN 路线的 RULER 表现](figures/fig-06-niah-410m-gdn.png)

原图 6，PDF p. 18。先看上排三个单 needle 任务，再看下排 MQ、MK、MV；虚线为 4,096，绿色是替换三层，橙色是替换一层。三层方案在 16K 的 S-NIAH-2、MQ、MK、MV 保留了更高的准确率，但不同面板没有同一条始终最优的曲线。S-NIAH-1 中，两路 Hybrid M²RNN 保持满分，三路方案却下降。

![7B GDN 路线的 RULER 表现](figures/fig-08-niah-7b-gdn.png)

原图 8，PDF p. 19。GDN 三路混合在多 query/key/value 任务中呈现较明显收益；部分任务到 16K 仍下降到较低水平。注意两个原文内部差异：图 8 的多层曲线图例写作 `-3`，而 Sec. 5.1 和 7B 表格使用 `-5`；这里保留原图标记，不据图例单独确定层数。Sec. 5.4.1 还概括两路 Hybrid M²RNN 与 Hybrid GDN 在 S-NIAH-1 都达 100%，但 Fig. 6 的 Hybrid GDN 与 Fig. 8 的 Hybrid M²RNN 并未在所有长度保持满分，应以对应曲线为准。

Mamba 路线的图 7、9 则显示单层替换不足以稳定改善长距离检索，增加到多层后才在部分外推长度上更有效。图 9 的 7B 多层图例同样保留了 `-3` 标记。作者将骨干差异解释为 Mamba 状态转移较弱、需要更多非线性层补充；这是作者对观察的解释，实验没有单独隔离该机制。来源：Sec. 5.4.1，PDF pp. 17–19。

### 4.5 真实数据检索：混合方式比纯循环替换更有效

| 模型 | 410M 六项平均分 ↑ | 7B 六项平均分 ↑ |
| --- | ---: | ---: |
| 纯 Mamba-2 | 27.4 | 41.8 |
| 纯 GDN | 30.8 | 42.2 |
| 纯 M²RNN | 26.9 | 37.9 |
| Transformer++ | 34.6 | 45.5 |
| Hybrid Mamba-2 | 44.0 | 55.1 |
| Hybrid GDN | 43.1 | 51.4 |
| Hybrid M²RNN | 46.9 | 55.6 |
| Hybrid GDN + M²RNN-1 | 45.5 | 54.2 |
| Hybrid GDN + M²RNN-3 / -5 | 44.1 | 55.1 |

来源：Tables 3–4，PDF p. 20。纯 M²RNN 的检索分数低于纯 GDN，说明非线性状态计算并没有消除有限记忆瓶颈。作者把部分差距归因于 M²RNN 约小三倍的循环状态；这也是容量与算力之间的选择。

两路 Hybrid M²RNN 分别比 Transformer++ 高 **12.3 / 10.1 分**，比 Hybrid GDN 高 **3.8 / 4.2 分**。在 7B 中，FDA 为 78.0 对 67.1，SWDE 为 79.8 对 71.9，构成较明显的增益；自然问答各项并非都同幅度提高。

单层 GDN 方案相对自身基线增益为 **2.4 / 2.8 分**。Mamba 单层方案仅有 44.0 → 44.3、55.1 → 55.2 的平均分变化；410M 三层到 45.6，7B 五层仍为 55.2。替换的效果明显依赖原有骨干。

### 4.6 LongBench：“约 8 分”来自哪些方案

| 模型 | 410M 八项均分 ↑ | 相对 Hybrid GDN | 7B 八项均分 ↑ | 相对 Hybrid GDN |
| --- | ---: | ---: | ---: | ---: |
| Hybrid GDN | 15.4 | — | 23.1 | — |
| Hybrid M²RNN | 15.4 | 0.0 | 24.6 | +1.5 |
| Hybrid GDN + M²RNN-1 | 19.5 | +4.1 | 30.4 | +7.3 |
| Hybrid GDN + M²RNN-3 / -5 | 23.0 | +7.6 | 31.2 | +8.1 |

来源：Tables 5–6，PDF p. 21；差值按表中已舍入均值计算。“单层接近多层收益”在 7B 上很有说服力，但并不意味着每项能力都一起上升：

| 7B 任务 | Hybrid GDN | +M²RNN-1 | +M²RNN-5 |
| --- | ---: | ---: | ---: |
| GovReport | 10.4 | 17.2 | 12.2 |
| MultiNews | 16.7 | 6.0 | 16.1 |
| TREC | 57.5 | 69.5 | 59.0 |
| TriviaQA | 35.1 | 64.9 | 73.5 |
| SAMSum | 23.6 | 34.1 | 36.0 |
| LCC | 9.7 | 20.1 | 18.0 |
| RepoBench-P | 16.3 | 15.1 | 18.9 |

单层方案在 TriviaQA、TREC、LCC 上收益突出，而 MultiNews 下降 10.7 分。五层改善了其中一些退化，但并非每项都优于单层。

Mamba 路线还有更直接的反例：410M 的均分为 10.3 → 12.2（单层）→ 17.1（三层）；7B 为 **18.8 → 13.1（单层）→ 22.7（五层）**。所以“加一层就足够”主要是本文 GDN 混合配置的结果，无法泛化为任意混合骨干的替换规则。

### 4.7 消融：容量解释获得支持，但没有被完全独立识别

| 模型 | 参数量 | 状态大小，按原表口径 | Wiki PPL ↓ | LAMBADA PPL ↓ |
| --- | ---: | ---: | ---: | ---: |
| 向量 RNN | 406M | 1,360 | 33.74 | 317.38 |
| GRU | 474M | 1,360 | 25.80 | 63.83 |
| M²RNN | 410M | 86,016 | 22.92 | 33.63 |

来源：Table 7，PDF p. 21；分析见 Sec. 5.6，pp. 21–22。M²RNN 的状态元素数量是向量模型的约 **63.25 倍**，参数量仍接近 RNN。相较 RNN，Wiki PPL 降低 **10.82**、LAMBADA PPL 降低 **283.75**；相较参数更多的 GRU，分别降低 **2.88 / 30.20**。

RNN → GRU 在相同状态大小下改善很大，显示门控本身也重要。GRU → M²RNN 同时改变状态形式、读写方式和门控结构；表中没有等状态、等架构条件下“仅去掉 M²RNN 遗忘门”的独立对照。因此，结果支持状态容量是重要瓶颈，但不足以精确分离容量和门控各自贡献，或证明收益主要由单一因素造成。

### 4.8 吞吐：省掉 batch 补零后，顺序递推仍然昂贵

![7B MoE 训练吞吐](figures/fig-10-throughput.png)

原图 10，PDF p. 22。横轴是上下文长度，纵轴是**每天十亿 tokens**，测试硬件为 8 张 HGX H100。Mamba-2 在所画长度上最快；纯 M²RNN 随长度增加吞吐明显降低，即使其总工作量关于长度是线性的。

关键比较是 Hybrid GDN 与 Hybrid GDN + M²RNN-1：作者在 Sec. 5.7 报告 16K 时差距在 6% 内，Sec. 3.4.2 同时提到 4K 和 16K。这是**单层三路混合方案**的训练结果，不能转移到纯 M²RNN、两路 Hybrid M²RNN 或五层方案。Fig. 10 未给出五层方案，也没有报告解码吞吐，因而无法从图中计算这些配置的端到端收益。

## 五：结论——最值得保留的设计启发

M²RNN 的贡献在于把三个原本分开的设计问题连起来：外积和 head 共享使非线性递推拥有足够大的状态；小型稠密 value 转移保留逐步状态变换能力；矩阵 tile 使这种递推更适合 GPU。在 410M 与 7B MoE、100B-token 训练规模下，容量扩展明显提升了传统向量 RNN 的语言建模表现。

更有实践价值的是三路混合结果：对本文 Hybrid GDN，少量 M²RNN 层能带来明显的语言建模与长上下文收益，单层配置保留了接近原基线的训练吞吐。作者把这种互补性解释为 attention 承担检索，线性循环层承担高效记忆，非线性循环层补充状态计算；实验支持模块组合有价值，尚未直接测出各层承担的具体算法。

进一步研究最有区分度的问题，是在相同循环状态与训练 FLOPs 下单独控制非线性、门控和读写结构，以及区分“替换层数”和“替换位置”的贡献。当前单层替换第一层，多层替换 attention 前层，两个因素同时变化。对于实际实现，BPTT 的梯度累计、反向临时状态存储和 TP 通信也会决定这种互补性是否值得其成本。

论文给出的官方资源入口：[训练框架 lm-engine](https://github.com/open-lm-engine/lm-engine)、[模型与内核仓库](https://github.com/open-lm-engine/accelerated-model-architectures)、[M²RNN 模型集合](https://huggingface.co/collections/open-lm-engine/m2rnn)。这些链接来自论文首页。