我们在新闻中不断看到拥有更大上下文窗口的大型语言模型(LLM),它们承诺能够同时查看完整的对话历史、大量书籍或多个代码库。然而,这些模型仍然会重复相同的错误。我们依然需要将之前的上下文复制并粘贴回对话中,才能让LLM“理解”。一个聪明的同事会学习这些模式,不断适应并持续进步。为什么LLM却难以做到这一点?
在这篇博文中,我们探讨了大语言模型的内存与人类记忆之间的关键差异。随后,我们介绍了测试时训练的一种端到端方法(TTT-E2E),这是我们的最新研究成果,其中大语言模型通过下一词预测,将所读取的上下文信息压缩到其模型权重中。
图 1 突出显示了我们的主要结果,该结果根据损失(左)和延迟(右)衡量上下文长度的扩展情况。注意力完全集中的 Transformer 在损失方面表现优异,但在延迟方面表现欠佳。递归神经网络(RNN)(例如 Mamba 2 和门控 DeltaNet)在延迟方面能够实现良好扩展,但会带来损失。TTT-E2E 是唯一在两方面均能实现良好扩展的方法。<!–
左侧面板: TTT-E2E 可将性能最差的线(灰色)提升至 128K 上下文长度下的最优线(浅绿色)。损失(Δ)(y 值)的计算方式为(报告方法的损失)-(全注意力 Transformer 的损失),因此全注意力损失本身(深绿色)对应 y = 0 处的水平直线。尽管其他方法在较长序列上会产生更显著的损失*,但 TTT-E2E 在长上下文场景中仍保持与全注意力机制相当的优势。
右侧面板: 与 RNN 类似,无论上下文长度如何,TTT-E2E 均保持恒定的推理延迟,在 NVIDIA H100 上,相比 128K 上下文的全注意力机制速度快 2.7 倍,在 200 万上下文时快 35 倍。所有模型均具有 3B 参数,并经过 164B tokens 的训练。
在损失和延迟方面,按上下文长度缩放是长上下文与大语言模型研究中的一个基础性问题。TTT-E2E 是首个在此问题上展现出积极进展的方法,而所有其他方法则表现出本质上不同的趋势。此外,我们在严格且广泛的实验中观察到,TTT-E2E 的扩展趋势未出现任何衰退迹象。这些结果表明,研究社区或许有望在 2026 年达成对长上下文问题的基本解决。
LLM 内存与人类内存有何不同?
尽管人类对确切细节的回忆并不完美,但在改善生活体验形式的“上下文”方面却表现得相当出色。例如,回想你的第一次机器学习讲座,你可能已记不清讲师说的第一个词,但当时建立的直觉或许仍有助于你理解这篇博文,即使那已是几年前的事。
另一方面,具有自注意力机制的 Transformer 在长上下文中效率较低,部分原因在于其设计初衷是为了实现近乎无损的信息召回。自注意力的基本形式称为全注意力,它通过缓存并比较每个 token 的键和值,来保留每个 token 的完整记忆。因此,每个细节都会被充分关注,但每个 token 的计算成本会随上下文长度呈线性增长。处理第 1000 万个 token 所需的时间,是处理第 10 个 token 所需时间的大约一百万倍。
为了在不燃烧行星的情况下处理长背景,现代架构通常将全局注意力与近似方法相结合,例如滑动窗口注意力、Mamba 和门控 DeltaNet 层。这些近似方法对每个 token 都具有固定的成本,但在较长语境下的表现相比完全注意力仍有明显下降。具体而言,这些近似方法会丢失对预测未来有帮助的重要信息,如图 1 所示。
我们的方法是:将上下文压缩为权重
我们如何设计一种以 token 为成本的方法,使其仍能保留长语境中关键的预测性直观信息?
关键机制是压缩。例如,人类会将大量经验压缩至大脑中,从而在保留重要信息的同时忽略许多细节。对于语言模型而言,我们已知通过 next-token 预测进行训练,也会将海量数据压缩到其参数权重中。那么,若在测试阶段继续基于上下文的 next-token 预测来训练语言模型,会发生什么情况?
在添加另一个缺失部分后,我们发现这种简单的测试时训练(TTT)形式非常有效。在训练过程中,我们通过元学习(而非标准预训练)为 TTT 准备模型初始化。这一改进使我们的方法具备了两种端到端(E2E)的优化路径:内循环直接优化网络末端的 next-token 预测损失,外循环则直接优化经过 TTT 后的最终损失。与此前在长上下文 TTT(例如 Titans)上的工作相比,这种方法在优化目标上更具一致性与直接性。
RAG 将发挥什么作用?
TTT 类似于更新人脑,而 RAG 等基于检索的方法则类似于记事本或日历中的记录与查找功能。这种记事本将持续作为大脑的实用补充,尤其是在细节至关重要的情况下,例如购买大量日用品时。然而,人类的生产力主要取决于其大脑的运作能力,而非所使用的记事本。同理,AI 智能体的生产力在很大程度上取决于它将大量上下文压缩为预测性与直觉性信息的能力。
限制
在训练过程中,TTT-E2E 的元学习阶段需要梯度支持。由于 FlashAttention 的标准 API 不支持梯度计算,我们当前实现的元学习效率相比针对短上下文(8K)的标准预训练慢了 3.4 倍。为突破这一限制,可采取以下两种方案:一是开发支持梯度计算的自定义注意力核函数,二是从未经 TTT 预训练的标准 Transformer 模型初始化 TTT-E2E。我们诚邀社区成员共同参与这一探索!
总结
如需更深入地了解方法、结果及实现细节,请参阅完整论文 端到端 Test-Time Training for Long Context。所有实验均可通过我们公开库中的代码与数据集复现。