← 全部论文2506.05233v2 / r1
版本与下载

MesaNet:把上下文回归求解到收敛的序列模型

论文:MesaNet: Sequence Modeling by Locally Optimal Test-Time Training
作者:Johannes von Oswald、Nino Scherrer(共同第一作者),Seijin Kobayashi、Luca Versari、Songlin Yang 等,共 17 位作者。机构为 Google、Google DeepMind 和 MIT CSAIL。
版本:arXiv:2506.05233v2,2026-06-03 修订,ICLR 2026;首次提交于 2025-06-05。研读与获取日期:2026-10-20。
来源:版本固定的论文页面 · 本地原文 PDF · 英文汇报 PDF · 可编辑 TeX。下文页码均为从 1 开始的 PDF 页码。

论文总览

MesaNet 研究的是一个具体问题:既然现代循环序列层可以理解为在上下文中训练一个小型关联记忆,为什么每个 token 只更新一步,而不直接求出当前上下文的最优线性映射?

它沿用 Transformer 的残差块和门控 MLP,用 Mesa 层替换序列混合模块。每个注意力头维护两个固定大小的矩阵:键值交叉统计量 $G_t$ 和键的二阶统计量 $H_t$。读取记忆时先解线性方程 $(H_t+\Lambda)q_t^*=q_t$,再输出 $G_tq_t^*$。这样得到的是带遗忘、带正则的累计最小二乘解。本文的主要新增贡献,是用共轭梯度(conjugate gradient, CG)和分块并行实现,把早期串行 Mesa 层变为能训练至约 1B 参数的模型,同时保留固定正则和输入相关遗忘。

结果有三条主线。第一,在统一骨干、1B 参数、50B 训练 token 的比较中,MesaNet 的六项平均 PPL 为 13.79,与 Transformer 相同,优于 Gated DeltaNet 的 13.87。第二,位置分解揭示循环模型的优势更多来自序列前部;上下文召回均分仍为 41.79% 对 52.27%,存在明显差距。第三,动态 CG 停止标准可把平均迭代数从 30 减到约 9,保持接近原模型的表现,说明“求解精度”可以成为推理时计算预算的一项控制量。出处:Sec. 4–5,PDF pp. 7–10,Tables 2–3、Figs. 3–5。

标题中的 locally optimal 指每个层、每个头所定义的局部回归目标。这个二次目标本身严格凸,具有唯一的全局最优解;它不意味着整个语言模型达到全局最优,也不意味着有限次 CG 在数值上精确等于闭式解。

结构地图

原文章节论证作用PDF 页码与主要证据
1 Introduction从快权重学习引出累计回归pp. 1–2,Fig. 1
2 A Parallelizable Mesa Layer目标、充分统计量、CG 与分块并行pp. 2–5,Eqs. (1)–(10)
3 Train and Inference Time of the Mesa Layer额外计算与实测速度pp. 5–6,Fig. 2
4 MesaNet in a Language World统一骨干与语言实验pp. 6–9,Tables 1–3
4.1 Language Modeling (Within and Beyond Train Sequence Length)平均 PPL、位置差异、长度外推pp. 7–8,Figs. 3–4
4.2 Language Benchmarks全局/局部推理、召回、少样本学习pp. 8–9,Table 3
5 Test-Time Compute Analysis固定与动态 CG 步数pp. 9–10,Fig. 5
6 Discussion计算开销、学习与记忆的取舍pp. 10–11
A–B Related Work / Derivation of previous test-time training rulesGLA、DeltaNet、LongHorn、Atlas 的目标比较pp. 20–23,Table 4
C–D Rank-One Update CG / Chunkwise Parallel FormCG 伪代码、前向与反向计算pp. 24–28,Algorithm 1
E–G Full Description / Experimental Details模块细节、复杂度、训练与精度pp. 28–35,Tables 5–8、Fig. 6
H–J Original RLS Mesa / State Tracking / Further Discussion旧算法问题、奇偶校验、后续方向pp. 34–38,Fig. 8
K MesaNet Trained in Synthetic EnvironmentsMAD 与 RegBenchp. 39,Table 9、Fig. 9
L Extended Results in Language Environment跨规模数据、召回细分、翻译与 NIAHpp. 39–51,Tables 10–19
M–O CG Steps during Training / Evaluation Methodology / Internal Analysis训练消融、指标口径、头的分工pp. 51–53,Figs. 13–14

一:动机——固定容量记忆如何更好地利用上下文

1. 从保存全部 KV 到学习一个映射

标准因果 softmax attention 在推理时保留过去 token 的键和值,因此单个新 token 的注意力计算与 KV 存储随历史长度增长。线性注意力把历史压缩进矩阵,读取时只需做矩阵向量乘法:

$$ o_t=\Phi_t q_t,\qquad \Phi_t=\gamma_t\Phi_{t-1}+\beta_t v_tk_t^\top. $$

这可以看成快权重(fast weights)学习:$k_t$ 是输入,$v_t$ 是要建立关联的目标,$\Phi_t$ 是在当前上下文内形成的映射,$q_t$ 用来查询这份记忆。遗忘门 $\gamma_t$ 控制旧信息的保留,输入门 $\beta_t$ 控制新信息的写入。网络的投影参数跨训练样本学习;快权重则随当前序列变化。出处:Sec. 2,PDF p. 3,Eqs. (1)–(3)。

这种加法式记忆的问题是,相关的键会互相干扰;把更多外积加进去,并不自动让当前映射更符合所有已见关联。

2. DeltaNet 修正最新误差,Mesa 优化累计误差

DeltaNet 对当前样本的平方误差做一次梯度更新:

$$ \ell_t(\Phi)=\frac12\|v_t-\Phi k_t\|^2,\qquad \Phi_t=\Phi_{t-1}+\beta_t(v_t-\Phi_{t-1}k_t)k_t^\top. $$

它根据“新值与旧预测之间的差”写入,比简单叠加更有针对性。但更新某个方向也可能破坏其他相关键上的预测。Mesa 则把所有保留的历史关联放进同一个加权回归目标,求出它们在当前表示空间中的最佳折中。LongHorn 虽然也使用一个解析最小化步骤,其目标仍围绕当前样本与上一步参数的距离;“每步求最优”是否涉及全部历史,是这里的重要区别。出处:Appendix B,PDF pp. 21–23。

本文并非首次提出 Mesa 目标。早期 Mesa 层使用递归最小二乘(recursive least squares, RLS),通过 Sherman–Morrison 更新逆矩阵。加入遗忘后,逆状态会出现 $1/\gamma_t$ 放大,初始正则也会随历史衰减;时间上的串行依赖又妨碍大规模并行训练。本文重构的是如何稳定、并行且可调精度地计算这个目标的解。出处:Sec. 2,PDF p. 4;Appendix H,PDF pp. 34–36。

二:方法——统计量累积与查询求解

1. 整体数据流

MesaNet 原始架构图

原图 Fig. 1,PDF p. 2。左侧是残差块,中间是 SwiGLU MLP,右侧是 Mesa 序列层。阅读顺序为:输入归一化 → 生成 $k,q,v$ 与门 → Mesa rule → 每头归一化与输出投影 → 残差连接。

具体地,Mesa 层对归一化后的表示做投影,使用宽度为 4 的短卷积和 SiLU,并对键、查询做 L2 归一化。每头生成标量输入门与遗忘门;输出经过每头 RMSNorm 后投影回残差流。Hawk-Mesa 变体交替使用 Hawk 的向量状态层与 Mesa 层,以结合不同序列位置上的建模优势。出处:Appendix E,PDF pp. 28–30。

2. 目标、符号与唯一解

下文省略头索引,向量均为列向量。

符号维度含义
$k_t,q_t$$\mathbb R^{n_a}$当前键与查询
$v_t,o_t$$\mathbb R^{n_v}$当前值与序列层输出
$\Phi$$\mathbb R^{n_v\times n_a}$要在上下文中学习的线性映射
$G_t$$\mathbb R^{n_v\times n_a}$累计键值交叉统计量
$H_t$$\mathbb R^{n_a\times n_a}$累计键的二阶统计量
$\beta_t,\gamma_t$标量,范围 $[0,1]$输入权重、遗忘因子
$\Lambda$正定对角矩阵,$n_a\times n_a$训练得到、时间上固定的正则

遗忘权重为

$$ \zeta_{ti}=\mathbf1_{i\le t}\prod_{s=i+1}^{t}\gamma_s. $$

为显式对应原文 Eq. (8) 的输入门,下面把 Eq. (4) 写成等价的带 $\beta_i$ 形式:

$$ \mathcal L_t(\Phi) =\frac12\sum_{i=1}^{t}\zeta_{ti}\beta_i \|v_i-\Phi k_i\|^2 +\frac12\operatorname{Tr}(\Phi\Lambda\Phi^\top). $$

原文在部分推导中省略 $\beta_i$;Appendix D 说明可把它吸收进键和值,即同时替换为 $\sqrt{\beta_i}k_i$、$\sqrt{\beta_i}v_i$。不能只缩放值,因为键的二阶统计量也需要同一权重。

对 $\Phi$ 求导并令其为零:

$$ \nabla_\Phi\mathcal L_t =\Phi\!\left(\sum_i\zeta_{ti}\beta_i k_i k_i^\top+\Lambda\right) -\sum_i\zeta_{ti}\beta_i v_i k_i^\top=0, $$

所以

$$ \widehat\Phi_t=G_t(H_t+\Lambda)^{-1}, \qquad o_t=\widehat\Phi_tq_t. $$

$H_t$ 是半正定矩阵,$\Lambda\succ0$,故 $H_t+\Lambda$ 正定,回归目标有唯一解。这里的“最优”是当前投影产生的键值对、当前遗忘权重、当前正则下的平方误差最优。键值投影是否适合语言预测,仍由外层的 next-token 训练决定。出处:Sec. 2,PDF pp. 3–4,Eqs. (4)–(8);Appendix D,PDF pp. 25–26。

3. 为什么不需要保存全部历史

累计目标所需的统计量可以递推:

$$ \begin{aligned} G_t&=\gamma_tG_{t-1}+\beta_tv_tk_t^\top,\\ H_t&=\gamma_tH_{t-1}+\beta_tk_tk_t^\top,\qquad G_0=H_0=0. \end{aligned} $$

$G_t$ 表示键和值怎样共同出现,$H_t$ 表示键空间中哪些方向反复出现、哪些方向相关。查询变换 $(H_t+\Lambda)^{-1}q_t$ 会利用这些二阶统计量调整读取方向,减少重复与相关键造成的干扰。

正则矩阵在读取时单独加入,因而不会随着 $\gamma_t$ 衰减。这一点区别于把 $\Lambda$ 仅作为 RLS 初始状态的做法。由于目标是二次的,保存 $G,H$ 足以恢复其关于 $\Phi$ 的梯度与最优解;与 $\Phi$ 无关的常数项不影响求解。

教学例子:两个相关键。 设 $k_1=(1,0)^\top$、$k_2=(1,1)^\top/\sqrt2$,目标值分别为 $1,0$,无遗忘、输入权重均为 1。GLA 得到的映射为 $(1,0)$,对第二个键读出 $1/\sqrt2$。一步 Delta 更新后映射变成 $(1/2,-1/2)$:第二个关联正确了,第一个关联的读出却变为 $1/2$。Mesa 在 $\Lambda=\lambda I,\lambda\to0^+$ 的极限得到 $(1,-1)$,能同时满足两个关联。这个例子说明二阶校正的作用;一般的有正则或超容量情形仍是最小二乘折中。

4. 只求一个查询方向,无需显式求逆

最终只需要 $\widehat\Phi_tq_t$,因此可以先求

$$ A_tq_t^*=q_t,\qquad A_t=H_t+\Lambda, $$

再计算 $o_t=G_tq_t^*$。这样无需显式形成逆矩阵,也无需为每个值维度求一遍完整快权重。

CG 迭代的昂贵操作是矩阵向量乘法 $A_tp$;其余主要是内积、向量加减与缩放。若当前近似解为 $x_j$,残差为 $r_j=q_t-A_tx_j$,停止条件是

$$ \|r_j\|_2\le\epsilon\|r_0\|_2, $$

同时设置最大迭代数。注意这是相对初始残差的准则,不是直接比较语言模型损失,也不是固定的绝对误差。

主实验固定 30 步,动态推理实验使用容差提前停止。条件数越高的线性系统通常需要更多迭代;不同头、不同 token 因而具有不同成本。数值解的质量也受有限精度影响,所以理论闭式最优解和实际固定步数输出应区分。出处:Appendix C,PDF pp. 24–25;Sec. 5,PDF pp. 9–10。

原文关于零步退化存在一个实现细节:Sec. 3 用 $q_t^*=q_t$ 描述退化为 GLA;但 Appendix C 又给出对角缩放初值,且文字使用 $H_t$,Algorithm 1 使用 $H_{t-1}$。因此严格的 GLA 等价应写成“令查询变换为恒等”,不能仅凭“CG 迭代次数为零”推断所有初始化实现都完全等价。

5. 分块并行来自同一个 GLA 计算结构

若只按递推式逐 token 实现,训练仍会受到串行计算限制。设块边界为 $c$,块长为 $C$。在给定块首状态 $G_c$ 后,把 Eq. (8) 完整展开:

$$ G_t=\zeta_{tc}G_c+ \sum_{i=c+1}^{t}\zeta_{ti}\beta_i v_i k_i^\top, \qquad c<t\le c+C. $$

已知各查询 $q_t^*$ 时,整块输出都可由“块首状态贡献 + 块内因果矩阵乘法”得到。吸收输入门后,令 $\bar k_i=\sqrt{\beta_i}k_i$、$\bar v_i=\sqrt{\beta_i}v_i$,并定义

$$ D_c=\operatorname{diag}(\zeta_{c+1,c},\ldots,\zeta_{c+C,c}),\qquad (Z_c)_{ij}=\mathbf1_{i\le j}\zeta_{c+j,c+i}, $$

则

$$ O_c=G_cQ_c^*D_c+ \bar V_c\bigl[Z_c\odot(\bar K_c^\top Q_c^*)\bigr]. $$

更重要的是,CG 所需的 $H_tp_t$ 也可写为

$$ H_tp_t =\zeta_{tc}H_cp_t+ \sum_{i=c+1}^{t}\zeta_{ti}\bar k_i(\bar k_i^\top p_t). $$

它正是把上述值向量换成键向量后的 GLA 结构。因此,可以在每轮 CG 中对一个块内的多个时间点并行执行矩阵乘法,只保存块边界的完整矩阵状态,而不为每个时间点物化 $H_t$。CG 的迭代轮次仍有前后依赖;并行的是同一轮内不同 token 的计算。 出处:Sec. 2,PDF pp. 4–5;Appendix D,PDF pp. 25–28。

上面的跨块衰减因子是按 Eq. (8) 补全的展开。原文 Eqs. (39)–(40) 将跨块项简写为 $G_cq_t^*$、$G_cQ_c^*$;若 $G_c$ 指块边界处未经后续衰减的状态,就需要这里的 $\zeta_{tc}$ 或 $D_c$。复现时应以递推定义核对这一点。

6. 外层训练如何穿过线性求解

令 $\mathcal J$ 为语言模型训练损失,$g_t=\partial\mathcal J/\partial o_t$。由于 $A_t$ 对称,反向需要再解

$$ A_tu_t=G_t^\top g_t. $$

于是

$$ \frac{\partial\mathcal J}{\partial q_t}=u_t,\qquad \frac{\partial\mathcal J}{\partial G_t}=g_tq_t^{*\top},\qquad \frac{\partial\mathcal J}{\partial\Lambda_{ii}}=-u_{t,i}q^*_{t,i}. $$

这里用 $g,u$ 重写,避免与原文前面表示 token embedding 的 $e_t$ 混淆。对键、值和遗忘门的梯度再通过统计量累积反向传播,也能组织成分块计算。

作者使用的是以线性系统已被充分求解为前提的解映射梯度,并非逐步反传截断 CG 过程。训练时如果大幅减少迭代,前向误差和这种梯度近似都会变得重要。Appendix J 提出随机 CG 预算与对迭代展开求导作为后续方向,本文没有完成这一方案。出处:Appendix D,PDF pp. 26–28;Appendix J,PDF p. 37。

7. 稳定性与复杂度的实际条件

实验把对角正则初始化为 1,并参数化为 $0.25+\operatorname{softplus}(\cdot)$。此外,重复 token 会使某些键方向反复累积,恶化条件数;作者采用与输入强度相关的遗忘修正:

$$ \gamma_t\leftarrow \gamma_t\bigl[1-(1-0.9975)\beta_t^2\bigr]. $$

这个公式是随 $\beta_t$ 改变的衰减修正,不能简单理解成所有情况下 $\gamma_t$ 都被硬截断到 0.9975。Pallas CG kernel 内部使用 FP32 矩阵乘法,其他激活主要用 BF16;求解完成后再转换输出。上述条件是论文稳定运行的组成部分。出处:Appendix G.5,PDF pp. 33–34。

若每头 $n_a=n_v=d$、CG 轮数为 $K$:

序列层每条序列、每头的持久状态单个新 token 的主要计算
全局 MHA$2td$$O(td)$
GLA / Delta 类$d^2$$O(d^2)$
Mesa$2d^2$$O((K+1)d^2)$

表中只比较序列混合核心,未计共同的投影和 MLP。固定维度、固定迭代上限时,Mesa 的单步状态与计算对历史长度为 $O(1)$;“常数”仍可能很大。训练计算随 $T$ 线性增长的说法也以固定 $C,d,K$ 为前提。原文提到额外状态在其模型总内存中不到 1%,这是包含参数的特定口径;单头循环状态在 $n_a=n_v$ 时实际翻倍,多序列并发时也会逐条增长。出处:Sec. 2–3,PDF pp. 4–6;Table 5,PDF p. 29。

三:实验——比较什么,怎样控制条件

1. 语言模型训练协议

模型在 SlimPajama 上训练 15B 或 50B token,训练上下文均为 2048,使用词表大小 50,257 的 GPT-2 tokenizer。小、中、大三档的架构配置为:

规模残差块数头数嵌入维度每头键维度
约 140M146768128
约 440M2881024128
约 940M28121536128

结果表分别使用 145M、400M、1B 等近似标签,下面沿用表中的标签,不把它们视为额外的模型档位。共同骨干采用独立序列层与 SwiGLU MLP、RMSNorm、共享输入输出嵌入;softmax attention 使用 RoPE。训练采用 batch size 256、AdamW、weight decay 0.1、2000 步 warmup、余弦学习率衰减、梯度范数裁剪到 1。Mesa 主实验使用 30 步 CG。出处:Sec. 4,PDF pp. 6–7;Appendix G,PDF pp. 31–34。

作者固定数据及其顺序,并为主要模型分别扫描学习率;SWA 作为参照使用固定学习率 0.001。由于所有方法共享同一骨干,Mamba2、GLA 等也采用受控的参数化,部分设定不同于各方法原始完整架构。因此,这些实验主要回答同一骨干中哪种序列混合规则更好。它们匹配参数和训练数据,未匹配 FLOPs,Mesa 的额外求解预算须结合速度实验理解。

短文档被用 BOS 分隔后拼接到训练长度;验证数据也采用同样方式。这种数据构造可能偏重短序列,作者在 Appendix G.1 明确指出了这一点。

2. 基线各自回答的问题

基线比较角色
Transformer / MHA保存未压缩 KV 的全局注意力参照
SWA-4、16、64、256、1024通过限制每层注意力窗口,检验分数对上下文范围的敏感性
GLA、Mamba2、xLSTM加法式/门控线性记忆与归一化方案
DeltaNet、Gated DeltaNet基于预测误差的一阶快权重更新
Hawk向量状态循环层,检验很强的局部建模能走多远
Hawk-Mesa交替采用局部与较长范围建模层的循环混合模型

SWA 的窗口是每层直接访问范围。层叠可以扩大有效感受野,故 SWA-4 得分不错不能被解释为模型整体仅利用最后 4 个 token。作者的分组是根据窗口变化观察到的性能敏感性作出的实验诊断。

3. 指标与评估口径

PPL 是所有被评估 token 的平均负对数似然取指数:

$$ \operatorname{PPL} =\exp\left[-\frac{1}{N_{\rm tok}}\sum_{j,i} \log P(s_{j,i}\mid s_{j,<i})\right]. $$

它越低越好。Table 2 的 AVG 是六个数据集 PPL 的算术平均,不是先合并六个数据集的 token 再算 PPL。作者还计算截至位置 $k$ 的累计平均 NLL,以及仅在位置 $k$ 的 NLL;后者能更直接显示性能交叉点:

$$ \Delta\operatorname{NLL}_k =\operatorname{NLL}^{\rm model}_k -\operatorname{NLL}^{\rm MHA}_k. $$

负值表示该模型在该位置优于 Transformer。对变长序列,应只在含有第 $k$ 个 token 的序列上平均。出处:Appendix L.1,PDF pp. 39–41。

多选任务对候选答案的条件 NLL 按 token 数归一化。上下文召回输入截断到 2000 token,贪心生成至换行或最多 48 token,再检查目标字符串是否出现在生成中,忽略大小写。单词变换报告 10 组随机 few-shot 前缀上的均值与标准差;翻译报告 BLEU-sb。NIAH 的标准差则来自键值类型组合,不能混作重复训练的误差条。出处:Appendix L.2–L.3、N,PDF pp. 47–52。

四:结果——从平均分到真实的上下文利用

1. 训练长度内:平均 PPL 很强,但收益并不均匀

下表摘录 1B 参数、50B 训练 token、最大评估长度 2048 的结果。SLIM 为 SlimPajama 验证集,AVG 包含 SLIM、Lambada、WikiText-2、PG19、GovReport、Qasper 六项。

模型SLIM PPL ↓WikiText-2 PPL ↓Qasper PPL ↓六项 AVG ↓
Mamba211.3912.2314.0214.58
DeltaNet11.0111.7313.4414.05
Gated DeltaNet10.8911.5813.2813.87
MesaNet10.8311.4913.1313.79
Hawk-Mesa10.7811.5313.2013.75
SWA-102411.0011.7813.1113.97
Transformer10.8611.4212.8613.79

出处:Table 2,PDF p. 7;完整规模结果见 Table 10,PDF p. 42。

Mesa 相对 Gated DeltaNet 的六项平均 PPL 下降 0.08,约 0.58%;它与 Transformer 的平均数相同,却并非每个数据集都相同。Hawk-Mesa 的 13.75 最低,表明组合不同循环层可进一步改善这个聚合指标。

按 token 位置比较的 NLL 差值

原图 Fig. 3,PDF p. 7。横轴是 token 位置,右面板使用对数横轴;纵轴为相对 MHA 的 NLL 差值,零线是 Transformer。多数循环模型在序列前部更强,随后变为正值。Mesa 与 Hawk-Mesa 把负值区域延长到约 512 token 之后,但较晚位置仍落后。

这解释了平均 PPL 和上下文能力为何可能给出不同印象:序列前部的收益足以抵消后部损失。Hawk 在很早位置突出,却更早失去优势;将它与 Mesa 交替组合,符合 Hawk-Mesa 的设计动机。Appendix L 在多个规模和训练预算上重复观察到类似趋势,而小模型上的 Mesa 并非所有指标都最佳。

2. 长度外推:Mesa 稳定,局部注意力仍是强参照

长度外推的累计 PPL

原图 Fig. 4,PDF p. 8。横轴为序列长度,纵轴是 WikiText、GovReport、Qasper 三项的平均 mean-so-far PPL。必须注意它是累计指标,且图中只含三个数据集。下面的附录表则包含 PG19,共四项,不能混用两个平均数。

模型4K 四项平均 PPL ↓32K 四项平均 PPL ↓
GLA11.9412.96
xLSTM11.9813.05
Gated DeltaNet11.7512.44
MesaNet11.6712.63
Hawk-Mesa11.6615.09
SWA-102411.9812.75
Transformer60.50588.41

配置仍为 1B、50B token;四项为 WikiText-2、PG19、GovReport、Qasper。出处:Tables 11–12,PDF p. 43。

Mesa 能从 2K 训练长度外推到 32K 而不发生严重损失爆炸,但在 32K 的平均值上 Gated DeltaNet 更好,SWA-1024 也很接近。这里的 Transformer 是在 2K 上训练的本文基线,其外推失败不能推广到经过长上下文训练或专门位置外推处理的 Transformer。

Hawk-Mesa 在 2K 的最佳平均 PPL 也没有自动转化成最佳长度外推。能够稳定处理长输入、能利用远距离信息回答问题、平均语言预测损失低,是需要分别衡量的能力。

3. 推理基准拆分:窗口敏感性比一个总均分更有解释力

作者用小窗口 SWA 分组:若分数随窗口从 4 增至 64 上升,且短窗口模型不超过全窗口 Transformer,则归入“全局推理”;否则归入“局部推理”。划分基于约 400M 和 1B 模型、15B/50B token 的实验。

全局组包含 Lambada、HellaSwag、RACE-M/H;局部组包含 PIQA、WinoGrande、ARC-E/C、SIQA、BoolQ、OpenBookQA、StoryCloze。这是该实验范围内的操作性分组,不是对数据集能力要求的永久分类。出处:Appendix L.2.1,PDF pp. 44–46,Tables 13–14。

模型全局推理均分 ↑局部推理均分 ↑上下文召回均分 ↑
Hawk41.1751.5725.04
GLA44.3451.4439.64
DeltaNet43.8651.5840.46
Gated DeltaNet44.8450.7639.54
MesaNet45.0350.4941.79
Hawk-Mesa44.6251.5139.99
SWA-431.2049.6211.87
SWA-102445.6851.3540.84
Transformer45.5451.6252.27

单位为百分比,配置为 1B、50B token。出处:Table 3(b),PDF p. 8。

Mesa 在全局组超过 Gated DeltaNet 0.19 个百分点,在这组循环基线中最高,但低于 Transformer 和 SWA-1024。局部组则没有 Mesa 优势,Hawk 的较高分尤其说明“常识总均分”容易混合不同能力。这些分数差异应按实际大小理解,不能把 0.19 个百分点写成显著的大幅提升。

4. 上下文召回:改善明显,压缩瓶颈依然存在

六个召回任务为 SWDE、SQuAD、FDA、TQA、NQ、DROP。Mesa 的平均分为 41.79%,相比 Gated DeltaNet 提高 2.25 个百分点,相比 SWA-1024 提高 0.95 个百分点,但仍比 Transformer 低 10.48 个百分点。

差距在不同任务上很不均匀。1B/50B 时,FDA 为 Mesa 32.67%、Transformer 70.96%;SQuAD 则是 46.25% 对 46.92%。作者检查发现 SQuAD 待召回信息经常位于最后 1K token,而 FDA 的信息更多在约 2K 长输入的开头;这与窗口对照中的性能变化相吻合。出处:Appendix L.2.2,PDF pp. 47–48,Tables 15–16。

因此,Mesa 的优化器确实改善固定状态的利用,但“累计回归求到最优”无法恢复已经在压缩表示中丢失的任意历史细节。1B/15B 这一配置下 Gated DeltaNet 的召回均分还高于 Mesa,说明主表优势也应保留训练预算条件。

5. 少样本学习:单词变换和翻译呈现不同排序

模型100-shot 单词变换均值,准确率% ↑50-shot 翻译均值,BLEU-sb ↑
Gated DeltaNet8.908.53
MesaNet10.108.17
Hawk-Mesa8.618.81
SWA-10246.6614.17
Transformer6.9813.61

配置为 1B/50B。单词变换均分来自两个有有效信号的任务,翻译均分来自法语、德语、罗马尼亚语到英语三项。出处:Table 3(b),PDF p. 8;Tables 17–18,PDF pp. 49–50。

在循环移位单词任务上,Mesa 的 100-shot 准确率为 14.0±0.5%,Transformer 为 9.2±0.6%;在中间字母重排任务上,二者为 6.2±0.4% 和 4.7±0.2%。标准差来自 10 次随机示例前缀选择。虽然 Mesa 相对优势清楚,绝对正确率仍低,结论应落在“这些小模型更善于从示例推断这种变换”上。

翻译中,Mesa 的 8.17 BLEU-sb 低于 Gated DeltaNet 的 8.53,也明显低于 Transformer 的 13.61;Hawk-Mesa 才是这里最好的循环模型。原文部分文字笼统称 Mesa 在循环翻译模型中领先,本配置的表格支持更细的结论。少样本学习并非所有任务都会因最优线性回归而同幅改善。

6. NIAH:干扰文本的类型会改变结论

论文使用 RULER 的单针任务,在 noise 与 essay 两类背景中测试 words、numbers、UUIDs 的 9 种键值组合,每种组合 1000 个样本。

模型Noise 2K ↑Noise 4K ↑Essay 2K ↑Essay 4K ↑
Gated DeltaNet98.3±4.196.3±8.152.1±33.711.0±9.4
MesaNet99.5±0.595.1±3.966.8±28.917.9±9.0
Hawk-Mesa97.6±3.565.3±21.690.9±10.555.5±28.5
Transformer99.7±0.30.0±0.098.2±2.50.0±0.0

准确率%,均值与标准差跨 9 种键值类型组合,1B/50B。出处:Appendix L.3,PDF pp. 50–51,Table 19。

Mesa 在 Noise 2K 几乎全对,4K 仍保持 95.1%,但换成 Essay 4K 仅为 17.9%。因此,高 NIAH 得分必须连同背景、键值类型与长度一起读。Transformer 的 4K 零分同样反映本文 2K 基线的外推失败;它不抵消在训练长度内的强召回表现。

7. 测试时计算:把求解预算分配给需要它的头

固定与动态 CG 的损失和步数

原图 Fig. 5 右侧,PDF p. 10。左面板表示相对 30 步基线的累计 NLL 增量;中间放大较小差异;右面板给出动态停止时的平均 CG 步数。暖色是固定轮数,蓝色是不同残差容差。

统一减少 CG 步数时,较后位置的损失增加更明显;动态容差造成的损失变化更均匀。作者报告 $\epsilon=10^{-4}$ 时,平均步数约为 9,表现接近固定 30 步。按步数计算约减少 70%,这并不是整个模型端到端耗时减少 70%。

CG 步数与下游得分

原图 Fig. 5 左侧,PDF p. 10。横轴是准确率,三个面板依次为全局推理、局部推理和上下文召回。固定 30 步的三项为 45.03、50.49、41.79;$\epsilon=10^{-4}$ 时为 45.03、50.40、41.96。较充分的求解恢复模型性能,但分数并不随每一次额外迭代严格单调。

为什么动态分配有用?在约 400M/50B 的内部分析中,许多头的条件数低且稳定,少数头遗忘很弱、条件数随位置升高,需要更多 CG 迭代。求解难度因此不只是序列长度的函数,也与当前数据和头有关。出处:Sec. 5;Appendix O,PDF pp. 52–53,Fig. 14。这里内部分析使用的容差 $10^{-3}$ 与上述主消融的 $10^{-4}$ 不应混淆。

训练预算消融另用约 400M/15B 模型:从 30 步减少到 10–25 步时,损失变化相对小;降到 5 步后损失明显上升。5 步训练的模型在推理时增加到 30 步虽有改善,也不能完全消除训练差距。出处:Appendix M,PDF p. 51,Fig. 13。

8. 实际速度:FLOPs 增加没有按比例转化为时间增加

H100 训练吞吐量

原图 Fig. 2 右面板,PDF p. 5。横轴同时改变序列长度与 batch size,每组总 token 数约为 32K;纵轴为千 token/秒。比较的是共享骨干的 1B 模型与 FlashAttention-2、Triton 循环层实现。Mesa 使用 15 或 30 步 CG,速度具有竞争力,但 GLA 更快,Mesa 也并非所有配置都最快。

H100 自回归生成吞吐量

原图 Fig. 6 右上面板,PDF p. 35。固定 batch size 为 8,横轴增加生成长度,纵轴为 token/秒。循环模型的吞吐随长度较平稳;Transformer 在较长长度下下降。Mesa 不同固定 CG 步数的曲线与 Gated DeltaNet 较接近,GLA 仍更快。图中 H100 与 TPUv5 的相对曲线不完全相同。

作者把额外计算没有等比例拖慢解码解释为注意力核心受内存带宽限制:更多算术有时可以在相近数据搬运下完成。但动态 CG 的平均步数与这组固定步数吞吐测试是两项不同实验,不能直接拼成一个实测动态加速倍数。出处:Sec. 3,PDF pp. 5–6;Appendix G.7,PDF pp. 34–35。

9. 合成任务与状态跟踪补充了机制图景

MAD 包含关联召回、模糊召回、训练数据记忆、选择性复制和压缩。小模型在每项任务上分别训练并扫描超参数,报告不同难度子任务的平均准确率。Mesa 的总体平均为 76.1%,Transformer 为 75.8%,xLSTM 为 73.2%,Gated DeltaNet 为 71.0%。Mesa 的模糊召回为 58.5%,但训练数据记忆低于 Hawk,压缩低于 Hawk-Mesa;它并非每项最佳。出处:Appendix F、K,PDF pp. 30–31、39,Table 9。

RegBench 的测试准确率

原图 Fig. 9,PDF p. 39。横轴为训练自动机数量,单位为千;纵轴为测试准确率。模型从上下文中推断概率有限自动机产生的伪语言,在未见过的自动机上评估。Mesa 曲线接近 Transformer,并领先其他循环方法,支持其在上下文规则归纳方面的优势。

奇偶校验则暴露另一个限制。Appendix I 让模型在长度 40 的位串上训练,并测试至 256。允许 $\gamma_t\in(-1,1)$、把正则初始化为 $50I$ 并下界限制到 49 后,模型达到很强的外推,而只允许正门的对照在更长序列上接近随机猜测。大正则使读取近似 $G_t\Lambda^{-1}q_t$,接近 GLA 的功能;负门又破坏了非负加权回归的直接解释。作者把这视为尚待解决的问题:能否同时保留非退化的 Mesa 回归收益和状态跟踪能力? 出处:PDF pp. 36–37,Fig. 8。

五:结论——贡献、取舍与可延伸的问题

本文最有价值的机制贡献,是把“在上下文内学习”推进为可并行求解的累计回归层:$G,H$ 压缩了所需历史,CG 只计算当前查询需要的方向,前向与反向的核心操作又能复用 GLA 的分块结构。它还展示了一个与生成更多推理 token 不同的计算分配方式:在每层内部根据线性系统难度决定求解轮数。

实验方面,Mesa 在统一骨干中改善了循环模型的语言建模、部分上下文学习与召回性能;同时,位置 NLL 和 SWA 对照表明,低平均 PPL 容易掩盖较晚 token 与全局信息利用上的不足。这套诊断方法本身也是论文的重要贡献。

这条路线的主要取舍有三点:

  • 更好的求解不等于更大的记忆容量。 $G,H$ 的维度固定,最优回归只能最优地使用当前表示;它无法像未压缩 KV 那样为未来任意查询保留所有细节。
  • 额外计算需要用端到端预算判断。 Mesa 匹配的是参数和训练 token,求解次数、精度、状态大小以及具体硬件共同决定成本。约 9 步的动态结果很有吸引力,但应与固定预算吞吐实验分别解读。
  • 当前优势有任务与配置边界。 1B/50B 的全局推理优势较小,召回仍落后 Transformer,翻译中也不是最好的循环模型;Hawk-Mesa 的 PPL 优势又不总能保留到长输入。

作者提出用相邻时间点的解热启动、在近乎不遗忘的头中探索训练/推理不同求解器、学习求解器,以及随机预算下对 CG 展开训练。它们都是后续方向。由本文结果进一步得到的一个研究问题是:在相同推理延迟或总 FLOPs下,应把预算用于更精确地解现有回归、扩大状态、增加局部注意力,还是训练更适合低步数求解的表示?这比单独增加 CG 轮数更能检验 Mesa 的实际价值。