让模型在推理时持续学习，长文本处理更高效
End-to-End Test-Time Training for Long Context
- 用测试时训练让Transformer在推理中动态更新权重
- 30亿参数模型处理128K上下文，速度比全注意力快2.7倍
- 适合需要长上下文、低延迟的实时应用
我们将长上下文语言建模视为持续学习问题，而非依赖架构设计。采用标准Transformer结构配合滑动窗口注意力，模型通过在测试时对给定上下文进行下一步词预测实现持续学习，将读取的上下文信息压缩进自身权重。同时，通过训练时的元学习优化初始化，以提升测试时学习能力。整体方法为端到端测试时训练（TTT-E2E），与以往方法不同。实验聚焦可扩展性：对于使用1640亿个标记训练的30亿参数模型，该方法在上下文长度增加时性能衰减方式与全注意力Transformer一致；而Mamba 2和Gated DeltaNet则不具备此特性。类似RNN，TTT-E2E推理延迟恒定，128K上下文下速度比全注意力快2.7倍。代码已公开。
原文摘要 · Abstract (English)
We formulate long-context language modeling as a problem in continual learning rather than architecture design. Under this formulation, we only use a standard architecture -- a Transformer with sliding-window attention. However, our model continues learning at test time via next-token prediction on the given context, compressing the context it reads into its weights. In addition, we improve the model's initialization for learning at test time via meta-learning at training time. Overall, our method, a form of Test-Time Training (TTT), is End-to-End (E2E) both at test time (via next-token prediction) and training time (via meta-learning), in contrast to previous forms. We conduct extensive experiments with a focus on scaling properties. In particular, for 3B models trained with 164B tokens, our method (TTT-E2E) scales with context length in the same way as Transformer with full attention, while others, such as Mamba 2 and Gated DeltaNet, do not. However, similar to RNNs, TTT-E2E has constant inference latency regardless of context length, making it 2.7 times faster than full attention for 128K context. Our code is publicly available.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。