ai

[Paper] 端到端 Test-Time Training 用于长上下文

发布: (2025年12月30日 GMT+8 02:30)
8 min read
原文: arXiv

Source: arXiv - 2512.23675v1

概览

本文将长上下文语言建模重新定义为持续学习问题，而不是依赖不断增大的注意力机制。通过让具有滑动窗口注意力的标准 Transformer 在推理时持续 学习——即时预测下一个标记，它将庞大的上下文直接压缩到模型权重中。训练期间的元学习步骤为模型的测试时适应做准备，产生一种 端到端测试时训练 (TTT‑E2E) 方法，其规模与全注意力 Transformer 相当，同时保持推理延迟不变。

关键贡献

  • Continual‑learning formulation for long‑context LM: 将输入上下文视为模型持续更新的流。
  • Test‑time training loop：在当前上下文上执行下一个 token 预测，实质上将上下文写入模型参数。
  • Meta‑learning pre‑training：优化模型的初始权重，以便在推理期间快速适应。
  • Empirical scaling study：对最高 3 B 参数、在 164 B token 上训练的模型进行的实证扩展研究，表明 TTT‑E2E 在保持推理成本不变的情况下匹配全注意力的扩展效果。
  • Speed advantage：在 128 K token 窗口下比全注意力 Transformer 快 2.7 倍，延迟特性与 RNN‑style 模型相匹配。
  • Open‑source release：发布代码和训练配方，促进可复现性和社区扩展。

方法论

  1. Base Architecture – 一个仅配备 滑动窗口注意力（例如，4 K 令牌窗口）的普通 Transformer。这使得每一步的计算量保持受限，无论整体上下文长度多长。

  2. Meta‑learning 预训练 – 在标准语言模型训练之前，作者进行一次 元学习 阶段（类似于 MAML）。目标是找到一种初始化，使其能够在新的文本块上通过少量梯度步骤快速微调。

  3. Test‑time 训练循环 – 推理时，模型按顺序处理输入。对于每个新令牌，它会：

    • 执行前向传播以预测下一个令牌（标准语言模型损失）。
    • 对该损失进行一次梯度更新，就地更新自身权重。
    • 将注意力窗口向前滑动，丢弃最旧的令牌。

    实际上，模型在前进的过程中将长上下文“写入”其参数中，因此后续预测能够受益于完整的历史信息，而无需直接对其进行注意。

  4. 端到端 (E2E) 训练 – 在测试时使用的相同基于梯度的更新也被纳入训练目标，确保模型在生成文本的同时学会自我改进。

Results & Findings

模型上下文长度每标记延迟扩展趋势（性能 vs. 长度）
Full‑attention Transformer (baseline)up to 128 K↑ linearly with lengthPerformance improves with length, but latency explodes
Mamba‑2 / Gated‑DeltaNetup to 128 K~constantPerformance plateaus early, unable to exploit very long context
TTT‑E2E (this work)up to 128 Kconstant (RNN‑like)Matches full‑attention scaling – perplexity keeps dropping as context grows
  • 对于一个在 164 B 令牌上训练的 3 B 参数模型，TTT‑E2E 在将上下文从 8 K 增加到 128 K 令牌时，实现了与全注意力 Transformer 相同的困惑度下降。
  • 推理速度：在 128 K 上下文下，TTT‑E2E 比全注意力基线 快 2.7 倍，且质量相当。
  • 消融实验证实，元学习初始化和测试时的梯度更新两者缺一不可；去除任意一项都会削弱扩展行为。

Practical Implications

  • 成本效益高的长上下文语言模型 – 开发者可以部署一个中等规模的 Transformer（例如 30 B 参数），仍然能够利用 100 K‑token 上下文的优势，而不会出现全注意力导致的内存爆炸。
  • 实时应用 – 聊天机器人、代码助手或文档分析工具在需要摄取大量文本时，能够保持低延迟，使其适用于交互式场景。
  • 边缘部署 – 由于每步计算保持有界，该方法适用于内存受限的硬件（例如 16 GB 显存的 GPU，甚至是专用推理芯片）。
  • 持续学习流水线 – 测试时的训练循环可以扩展为即时适应特定领域词汇或用户数据，为无需完整微调的个性化语言模型打开可能。
  • 兼容性 – 不需要异构架构；现有的 Transformer 代码库可以通过加入元学习和测试时更新钩子进行改造。

限制与未来工作

  • 梯度开销 – 虽然延迟保持不变，但每个 token 仍然需要一次反向传播，这在缺乏高效混合精度自动求导流水线的 GPU 上可能更为沉重。
  • 在线更新的稳定性 – 该方法依赖于推理时精心安排的学习率调度；学习率调得不当会导致漂移或在噪声输入上性能下降。
  • 优化器状态的内存 – 在测试时训练期间存储每个参数的优化器动量（例如 Adam）会增加一定的内存占用。
  • 超过 3 B 的扩展 – 本文聚焦于最多 3 B 参数的模型；尚不清楚该方法在 10 B 以上模型上表现如何，因为优化器状态大小可能成为瓶颈。
  • 作者提出的未来方向 包括：探索更轻量的更新规则（如 SGD 或低秩适配器），结合检索增强机制进一步提升长程推理能力，以及将该框架应用于多模态序列（音频/视频），在这些场景中上下文长度更为关键。

作者

  • Arnuv Tandon
  • Karan Dalal
  • Xinhao Li
  • Daniel Koceja
  • Marcel Rød
  • Sam Buchanan
  • Xiaolong Wang
  • Jure Leskovec
  • Sanmi Koyejo
  • Tatsunori Hashimoto
  • Carlos Guestrin
  • Jed McCaleb
  • Yejin Choi
  • Yu Sun

论文信息

  • arXiv ID: 2512.23675v1
  • 类别: cs.LG
  • 发布时间: 2025年12月29日
  • PDF: 下载 PDF
Back to Blog

相关文章

阅读更多 »