长期以来,循环神经网络 (RNN) 一直处于 Transformer 的阴影之下。Transformer 主导了现代序列建模,因为它们能够高效地利用长上下文——拥有越多的历史信息,预测未来词元就越容易。其缺点在于二次复杂度,使得它们在处理长序列时运行缓慢且占用大量内存。相比之下,RNN 具有线性复杂度,但一直难以充分利用更多的上下文。

正如 OpenAI 在 2020 年提出的著名缩放法则论文所展示的那样,经典的 RNN (如 LSTM) 无法像 Transformer 那样扩展或利用长上下文数据。但现代 RNN 架构,例如 Mamba,正在缩小这一差距。情况似乎是如此——直到研究者发现,当序列长度变得非常长时,即使最好的 RNN 也会出现性能平台期。

相反,Transformer 却能持续改进。

图表显示,TTT 层和 Transformer 随着上下文的增加持续改进,而 Mamba 的性能则停滞不前。

与 Mamba 相比,TTT-Linear 和 TTT-MLP 在更长上下文下困惑度持续降低,而 Mamba 在超过 16k token 后停止改进。

这形成了一个奇怪的悖论: RNN 在其性能最差的区间,恰恰是其计算效率最高的区间。问题出在 RNN 压缩信息的方式上。


固定大小隐藏状态的问题

从宏观层面看,每个序列模型都会处理一串输入词元 \(x_1, x_2, ..., x_t\),并输出一串词元 \(z_1, z_2, ..., z_t\)。为了实现这一点,它需要维持一个隐藏状态 \(s_t\),用于总结迄今为止看到的所有信息。

一张示意图,展示序列模型的通用结构: 每个时间步都有输入、隐藏状态和输出。

序列模型逐个处理词元,并更新一个捕捉先前上下文的隐藏状态。不同的架构在状态的表示和更新方式上各不相同。

不同的架构对隐藏状态的处理方式不同:

  • RNN (如 LSTM 或 Mamba) :
    隐藏状态 \(s_t\) 是一个固定大小的向量。每个词元的计算量保持恒定,这很高效——但将成千上万个词元压缩到一个微小的向量中不可避免地会丢失信息。
  • 自注意力 (Transformer) :
    隐藏的“状态”实际上是键值缓存 (KV Cache) ,它随着序列长度增长。每个过去的词元都可被后续计算访问,具有高表达能力,但代价是二次复杂度。

挑战很明显: 如何在保留 RNN 的效率的同时,获得 Transformer 能够长时记忆细节的能力?要在不丢失意义的前提下,将成千上万个甚至上百万个词元压缩成固定大小的表示,又该如何做到?

论文给出的关键见解极其简单而优雅:** 将隐藏状态视为一个机器学习模型本身。**


核心理念: 会学习的隐藏状态

设想隐藏状态不再是一个静态向量,而是一个小模型 \(f\) 的权重。这个模型会随着序列的处理不断演化——在实时进行学习。

TTT 层示意图: 隐藏状态是一组模型权重 W,更新规则是一次梯度下降步骤。

隐藏状态 \(W_t\) 被视为内部模型 \(f\) 的参数。每个输入词元都会触发一次更新,就像在测试时执行一次训练迭代。

测试时训练 (Test-Time Training, TTT) 层将这一理念具体化为一种机制:

  1. 隐藏状态 = 模型权重:
    在时间 \(t\),隐藏状态是 \(W_t\),即模型 \(f\) 的权重。
  2. 输出规则 (预测) :
    每个输出词元由当前模型生成: \[ z_t = f(x_t; W_t) \]
  3. 更新规则 (学习步骤) :
    隐藏状态通过对一个自监督损失 \(\ell\) 执行一次梯度下降进行更新: \[ W_t = W_{t-1} - \eta \nabla \ell(W_{t-1}; x_t) \] 其中,\(\eta\) 为学习率。

每个序列输入都会触发一个在线学习过程。即使在推理阶段——传统网络仅“应用”已学模式——TTT 层也会从测试序列本身进行学习

一种简单的自监督损失是重构任务:

\[ \ell(W; x_t) = \|f(\tilde{x}_t; W) - x_t\|^2 \]

其中 \(\tilde{x}_t\) 是输入的受扰版本。模型需学习各维度间的关系以准确重构原始词元。

随着序列的推进,重构误差会稳步下降——这表明隐藏状态模型确实在学习并不断优化其对数据的理解。

图表显示,自监督 TTT 损失随着序列推进而下降,表明内部模型在学习。

TTT 损失 \(\ell(W_t; x_t)\) 沿序列下降,显示隐藏模型在测试时推理过程中持续学习。


让其在实践中运行

理论概念令人振奋,但要在实际中表现出色还需要若干关键创新。论文提出了几项重要改进,使 TTT 不再只是思想实验,而是成为可扩展架构。

1. 学习自监督任务本身

作者没有手动设计自监督目标,而是让任务本身变得可学习。外部模型会学习内部模型应学习的特征。

他们引入了三个可学习投影: \(\theta_K\)、\(\theta_V\) 和 \(\theta_Q\)——与注意力机制中的键 (Key) 、值 (Value) 、查询 (Query) 相对应。

  • 训练视图: \(\theta_K x_t\) — 输入给 \(f\) 的“损坏”版本。
  • 标签视图: \(\theta_V x_t\) — 用于训练 \(f\) 的目标输出。
  • 测试视图: \(\theta_Q x_t\) — 用于预测的输入。

损失函数变为:

\[ \ell(W; x_t) = \|f(\theta_K x_t; W) - \theta_V x_t\|^2, \]

预测则为:

\[ z_t = f(\theta_Q x_t; W_t). \]

通过在主模型训练期间优化这些视图,系统实际上学会了如何学习——选择哪些自监督信号最有助于最终的语言建模。


2. 提升 TTT 效率: 小批量与对偶形式

朴素的 TTT 逐步更新每个隐藏状态,容易成为瓶颈。研究者提出两种互补技术来加速这一过程。

测试时的小批量梯度下降

不同于每个词元单步更新,梯度会在一小连续词元上计算,且均相对于该批次开始时的状态。这种并行化显著提高了速度,而准确度仅略有下降。

TTT 小批量计算图示,展示梯度如何并行计算。

小批量更新方案: 来自多个词元的梯度可并行计算,实现高效批次更新。

如图 7 所示,批次越大吞吐量越高,但困惑度略有提升。作者发现 b = 16 是最佳权衡。

图表展示 TTT 小批量大小、困惑度与计算时间的权衡。

小批次与性能权衡: 较小批次准确度更高,较大批次计算更快。作者选择 \(b=16\)。

对偶形式

第二项突破是权重更新的对偶形式。与直接计算巨型梯度张量不同,对偶形式将计算重写为矩阵乘法——非常适合现代 GPU 和 TPU。

这种代数重构在保持正确性的同时,将性能提升超过 ,从而实现大规模训练和高效推理。


序列建模的统一视角

TTT 框架的最大魅力在于它如何统一 RNN 与 Transformer 的思想。

TTT = 线性注意力

若选择:

  1. 线性内部模型 \(f(x) = W x\),
  2. 批量梯度下降 (每次更新可访问完整序列) ,

则 TTT 方程会简化为线性注意力形式,一种高效的自注意力变体:

\[ z_t = \sum_{s=1}^t (\theta_V x_s)(\theta_K x_s)^T (\theta_Q x_t) \]

这种等价性提供了一个基准,作者在此基础上逐步加入改进——小批量更新、残差连接、可学习初始化等。

表格展示消融研究结果,从线性注意力开始,逐步加入 TTT 特性以改善困惑度。

消融研究: 从线性注意力出发,加入如小批量更新等 TTT 特性显著改善困惑度。

TTT = 自注意力

若将参数化内部模型 \(f\) 替换为非参数估计器 (如 Nadaraya-Watson 核回归) ,则 TTT 等价于标准自注意力

此时隐藏状态是已接收的所有词元 \(x_1, ..., x_t\) 的列表,更新规则只需追加新词元。这样就能将注意力解释为 TTT 的一个特例: 内部学习器保存了所有训练样本,而非将其压缩为权重。

维恩图展示: TTT 层可同时表示 RNN (参数化学习器) 与自注意力 (非参数化学习器) 。

TTT 层统一了 RNN (固定大小隐藏状态的参数化学习器) 与自注意力 (状态不断扩展的非参数化学习器) 。

这种框架不仅连接模型,还在效率与表达性之间建立了概念桥梁,表明它们是同一学习光谱的两端。


实验: 验证 TTT

作者提出了两个实用版本:

  • TTT-Linear: 内部模型 \(f\) 为线性形式。
  • TTT-MLP: 内部模型 \(f\) 为两层 MLP,使用 GELU 激活。

他们将这些模型与强大的 Transformer 和 Mamba 基线进行比较,模型规模从 1.25 亿到 13 亿参数,并在短上下文和长上下文中进行评估。

短上下文 (2k–8k 词元,The Pile 数据集)

在 2k 上下文下,TTT-Linear 表现与 Mamba 和 Transformer 相当;在 8k 上下文下略有超越,尤以包含时序卷积的 Mamba 主干架构为佳。

在 The Pile 数据集上,2k 和 8k 上下文下的缩放曲线图,显示 TTT 层性能竞争力强。

在短上下文 (2k、8k 词元) 下,TTT 层表现与 Mamba 和 Transformer 相当,并随上下文增长稳步提升。

长上下文 (最高 32k 词元,Books3 数据集)

在长上下文场景中,TTT 层表现出明显优势。TTT-Linear 与 TTT-MLP 均随序列延长持续改进,而 Mamba 在 16k 词元后停滞。

在 Books 数据集上,2k 和 32k 上下文下的缩放曲线图,显示 TTT 层在长上下文中优于 Mamba。

在 32k 词元下,TTT 层超越 Mamba。Transformer 也继续提升,但代价是更高的计算量 (FLOP) 。

在按词元位置评估的困惑度曲线上,TTT 模型的表现与 Transformer 类似——读得越多,预测越好——而 Mamba 在长序列中期即出现平台期。


效率: 保持 RNN 的速度

尽管存在内部学习循环,TTT 层仍保持高效,其每词元延迟随上下文增长保持恒定,而 Transformer 成本则呈线性增长。

延迟比较图,显示 TTT 层在上下文增加时的每词元耗时保持恒定,不同于 Transformer。

在 NVIDIA A100 GPU 上的延迟比较: TTT 层像 RNN 一样,无论上下文长度如何,每词元处理时间保持恒定。

事实上,借助优化的对偶形式,TTT-Linear 的速度比 Mamba 更快。TTT-MLP 由于模型更大稍慢,但复杂度仍保持线性。


重新思考网络设计

《学会 (在测试时学习) 》这篇论文不仅仅是又一个 RNN 变体——它代表着一种新的范式。通过让隐藏状态成为主动学习者,作者将序列建模重新定义为一个嵌套学习问题:

  • 内循环在测试序列上学习,并实时更新权重。
  • 外循环学习内循环应该如何学习——包括自监督任务、投影和超参数。

这一新视角引出了多个令人兴奋的方向:

  • 探索更丰富的内部学习器 (如视觉中的 CNN、机器人中的策略网络) 。
  • 使用多层嵌套——让内部学习器自身也执行 TTT。
  • 通过时间并行优化系统,实现百万词元级的上下文。

本质上,使隐藏状态成为学习者,解锁了一种融合模型: 既具 RNN 的高效率,又具 Transformer 的强表达性能。


核心要点

将动态学习融入隐藏状态本身,测试时训练层 (TTT) 重新定义了网络的记忆与适应方式。TTT-Linear 与 TTT-MLP 表明,通过测试时的学习,我们可以得到模型——

  • 时间复杂度线性扩展,
  • 随上下文增长持续提升,
  • 并在同一框架下统一 RNN 与注意力机制。

或许,“高效”的 RNN 与“高表达”的 Transformer 之间的界限并非不可逾越。
凭借恰当的架构,RNN 能像 Transformer 一样学习——甚至可能超越它们。