💬 NLP

MesaNet: Sequence Modeling by Locally Optimal Test-Time Training

原作者： Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Sarthak Mittal, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, Rif A. Sau

发布于 2026-06-04

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Johannes von Oswald, Nino Scherrer, Seijin Kobayashi, Luca Versari, Songlin Yang, Sarthak Mittal, Maximilian Schlegel, Kaitlin Maile, Yanick Schimpf, Oliver Sieberling, Alexander Meulemans, Rif A. Saurous, Guillaume Lajoie, Charlotte Frenkel, Razvan Pascanu, Blaise Agüera y Arcas, João Sacramento

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

以下是对 MesaNet 论文的解释，通过简单的概念和富有创意的类比进行了拆解。

大局观：当前 AI 的问题

想象一位非常聪明的图书管理员（即 Transformer，目前的 AI 标准）正在读一本书来回答你的问题。

优点： 这位管理员非常擅长记住从书的开头到结尾的所有细节。
缺点： 为了做到这一点，他必须在桌子上堆放越来越多的索引卡片。书越长，卡片堆就越高。最终，这叠卡片会变得巨大无比，占满整个房间，管理员在翻找卡片以寻找答案时会被搞得手忙脚乱。这使得他们在处理长文本时既慢又昂贵。

最近，科学家们尝试构建了一种不同类型的图书管理员（称为 RNNs，如 Mamba 或 DeltaNet）。这些管理员只保留一个固定大小的小笔记本。他们很快，不需要巨大的房间。然而，他们在面对很长的故事时，有时难以记住故事开头的细节，或者在故事变得过于复杂时出错。

解决方案：MesaNet 登场

论文的作者构建了一个新的图书管理员，叫做 MesaNet。他们想要的是“小笔记本管理员”的速度，以及“大卡片堆管理员”的记忆力。

为此，他们引入了一个特殊的工具：Mesa 层。

类比：“即时专家” vs. “慢速学习者”

想象你正在尝试根据刚刚看到的一组数字来解决一个数学题。

旧的 RNNs（慢速学习者）： 每当出现一个新数字时，图书管理员都会做一个微小的猜测，检查自己是否猜错，然后稍微调整一下他们的笔记本。他们是一步步进行的。这种方式很高效，但他们可能无法立即得到“完美”的答案，因为他们只是在不断地“猜测并调整”。
MesaNet（即时专家）： 当一个新数字出现时，MesaNet 管理员不仅仅是猜测。他们会立即进行一次超快速的心算，以确定基于目前为止所看到的每一个数字，如何才能完美地调整他们的笔记本。他们一次性解决整个优化问题，以确保答案是对当前语境的最佳拟合。

它是如何工作的（“测试时训练”）

论文称之为 “测试时训练”（Test-Time Training）。

通常，AI 模型是在工厂里训练一次，然后就直接运行了。它们在实际与你交流时不会学习任何新东西。

MesaNet 的窍门： 每当 MesaNet 读到一个新词时，它都会停顿极短的一瞬间，专门针对那一刻进行“重新训练”。它会问自己：“鉴于直到这一秒为止我所读到的所有内容，存储这段信息的绝对最佳方式是什么？”
代价： 这种“重新训练”比旧方法需要更多的计算能力（就像运行一个快速的数学求解器）。
收益： 因为它每次都求解出“最优”答案，所以它比旧的 RNNs 能更好地理解长篇且复杂的文章。

“块状”创新

这个想法的原始版本（来自之前的论文）由于必须逐个进行计算（就像一页一页地读书），导致训练速度过慢。

新的转折： 作者找到了如何以**“块”（Chunks）**的形式进行这些计算的方法。想象一下，不再是一页一页地读，而是图书管理员抓起一整章，利用强大的计算机芯片一次性解出整章的数学题，然后再继续前进。这使得它能够快速地在海量数据上进行训练。

他们的发现

团队在巨大的数据集（数十亿词汇）上测试了 MesaNet，并将其与 Transformer 及其他快速 RNNs 进行了对比。

在开头表现更好： MesaNet 在理解句子或故事的开头方面表现惊人。在最初的几百个词中，它的表现往往甚至优于那些庞大的 Transformer。
在长上下文方面表现更好： 虽然其他快速 RNNs 会随着故事变长而开始遗忘或产生混乱，但 MesaNet 能更好地稳住阵脚。它能比同类模型更准确地理解长文档。
权衡取舍： MesaNet 并非“免费”的。在阅读过程中，它使用更多的计算资源（FLOPs）来解决那些数学问题。然而，作者发现这种额外的投入带来了更高的准确度，尤其是在需要深度理解长上下文的任务中。
动态速度： 该系统足够聪明，知道何时该努力工作。如果一个句子很简单，它可能会做较少的数学步骤；如果句子很复杂，它就会做更多的步骤。它会根据任务的难度动态分配精力。

总结

MesaNet 是一种新型的 AI 架构，它像是一位在阅读过程中不断重新计算“完美记忆方式”的图书管理员。通过在每一步都求解一个复杂的数学问题以寻找“最佳可能的记忆”，它实现了一种超越其他快速、高效内存模型的理解水平，尤其是在处理长篇且复杂的文本时。它用一点额外的计算能力换取了显著的智能提升。

技术摘要：MESANET：通过局部最优测试时训练进行序列建模

问题陈述

目前的序列建模由使用 Softmax 自注意力机制的因果 Transformer 架构所主导。虽然这些模型非常有效，但在推理过程中，其内存和计算需求相对于序列长度呈线性增长。近期的研究试图通过线性化 Softmax 操作来解决这一问题，从而产生具有常数级内存和计算成本的循环神经网络（RNN）模型（例如 DeltaNet、Mamba、xLSTM）。这些模型通常可以通过一个统一的视角来理解：它们的循环动力学源于一个上下文回归目标，该目标通过在线学习规则（例如梯度下降步骤）进行近似优化。然而，这些现有的方法通常依赖于增量更新（一阶方法），而非在每个时间步求解最优权重，这可能限制了它们在保留过去信息的同时快速适应新证据的能力。此外，之前的最优测试时训练层（如原始的 Mesa 层）受限于在使用遗忘门时的数值不稳定性和缺乏并行性，从而阻碍了可扩展训练。

方法论

作者引入了 MesaNet，这是一种以一种新型的、数值稳定且支持块状并行化的 Mesa 层为核心的语言模型架构。该层旨在于每个时间点执行局部最优测试时训练。

Mesa 层

与通过单次梯度下降步更新状态的标准线性 RNN（如 DeltaNet）不同，Mesa 层将线性映射 $\Phi_t$ 定义为测试时优化问题的解。具体而言，它最小化了一个关于迄今为止见到的所有数据（ $e_{1:t}$ ）的累积正则化平方误差损失：
$\hat{\Phi}^{mesa}_t = \arg \min_{\Phi} \left( \frac{1}{2} \sum_{t'=1}^t \zeta_{tt'} \|v_{t'} - \Phi k_{t'}\|^2 + \frac{1}{2} \text{Tr}(\Phi \Lambda \Phi^\top) \right)$
其中 $\zeta_{tt'}$ 代表累积遗忘因子， $\Lambda$ 是一个对称正定正则化项。

输出通过 $\Delta^{mesa}_t = \hat{\Phi}^{mesa}_t q_t$ 计算。至关重要的是，作者没有使用递归最小二乘法（由于在处理遗忘门时存在数值不稳定问题且无法并行化），而是使用共轭梯度（CG）法求解最优查询响应 $q^*_t = (H_t + \Lambda)^{-1} q_t$ 。

并行化： 作者利用了多时间步下的 CG 方法与门控线性注意力（GLA）之间的等价性。通过维护状态变量 $G_t$ 和 $H_t$ ，并在序列的各个块上并行求解线性系统，该层在训练期间实现了 $O(T)$ 的复杂度，在推理期间实现了 $O(1)$ 的复杂度。
动态计算： 达到停止准则所需的 CG 迭代次数 ( $k$ ) 随序列和 Token 的不同而变化。这使得模型能够动态分配测试时计算量，在推理成本与性能之间进行权衡。

架构 (MesaNet)

MesaNet 采用了标准的仅解码器 Transformer 主干网络（类似于 Llama 2），但将序列混合组件（通常是多头注意力）替换为 Mesa 层。该架构包括：

残差块： 堆叠 $N$ 个包含序列混合层（Mesa）和通道混合层（SwiGLU MLP）的块。
预处理： Keys 和 Queries 在进入 Mesa 层之前经过归一化并与一个小窗口进行卷积。
混合变体： 还探索了一种“Hawk-Mesa”模型，通过交替使用 LRU (Hawk) 层和 Mesa 层，结合短序列和长序列建模的优势。

核心贡献

一种可并行且稳定的 Mesa 层： 本文提出了一种版本的 Mesa 层，克服了原始公式中的顺序瓶颈和数值不稳定问题。它利用具有块状并行化的 CG 求解器，能够在现代硬件（GPU/TPU）上实现高效训练，同时保持最优测试时回归的理论优势。
动态测试时计算分配： 该层显式地调用了一个外部求解器，其迭代次数取决于给定序列下线性系统的难度。这提供了一种原则性的机制，根据输入复杂度来动态分配推理计算量。
广泛的实证评估： 作者在 SlimPajama 数据集上训练了 140M、440M 和 1B 参数规模的 MesaNet。他们通过控制架构、分词器和超参数，提供了与 Mamba2、xLSTM、DeltaNet、Gated DeltaNet 和 Transformer 进行的严格、“苹果对苹果”的对比。

结果

语言建模性能： 在所有规模下，MesaNet 的验证困惑度（PPL）均低于其他 RNN 基线（Mamba2、xLSTM、DeltaNet）。在 1B 参数规模上，MesaNet 达到或略微超过了具有相同架构的 Transformer 基线。
Token 位置分析： 关于相对于 Token 位置的负对数似然（NLL）的详细分析揭示了一个独特的模式：
- 早期 Token： 大多数 RNN（包括 MesaNet）在序列的早期 Token 上表现优于 Transformer。
- 长上下文： 虽然 Transformer 在后期的 Token 上表现出色，但 MesaNet 和 Hawk-Mesa 与其他 RNN 相比，显著地将这种优势延伸到了更远的位置（超过 512 个 Token），尽管在极长深度下最终会落后于 Transformer。
下游基准测试：
- 全局推理： 在需要长程依赖的任务（如 Lambada、HellaSwag、RACE）上，MesaNet 优于所有其他 RNN，尽管仍落后于 Transformer 基线。
- 上下文召回： MesaNet 优于其他 RNN，甚至优于具有 1024 个 Token 窗口的滑动窗口注意力（SWA）模型，但仍处于全注意力 Transformer 之后。
- 少样本学习： MesaNet 在单词乱序和翻译任务中展示了强大的性能，通常在乱序任务上超过 Transformer。
效率： 尽管在推理过程中需要额外的 CG 求解器 FLOPs，但当使用适中的 CG 步数（例如 10-30 步）时，MesaNet 在 H100 GPU 和 TPUv5s 上仍保持了具有竞争力的 Token 吞吐量。

重要性与主张

本文将 MesaNet 定位为神经网络中最优测试时训练的重要一步。作者声称，通过在推理时求解序列优化问题（特别是寻找最佳线性联想记忆权重），模型能够比使用简单梯度更新的模型更有效地适应新证据。

这项工作强调了一种权衡：性能提升（困惑度）和下游推理能力是以增加推理计算量（FLOPs）为代价的。作者认为，这符合增加测试时计算量以提高性能的近期趋势。他们指出，Mesa 层提供了一种灵活的机制，可以在网络内部动态分配这种计算量，从而解决复杂的序列问题。

论文总结道，虽然 MesaNet 目前面临计算开销的问题，但其在长上下文理解和全局推理任务上优于其他 RNN 的能力表明，局部最优测试时训练是高效序列建模的一个有前景的方向，有望弥合 RNN 的效率与 Transformer 的性能之间的差距。作者承认当前实现的计算成本是一个缺点，并建议未来的工作可以涉及热启动优化或混合推理策略以减轻这一问题。