End-to-End Test-Time Training for Long Context

LLM Long Context ML ML Continual Learning
我们将长上下文语言建模问题转化为持续学习问题,而非架构设计问题。基于这一思路,我们仅采用标准架构——即带有滑动窗口注意力的Transformer模型。然而,我们的模型能够在测试阶段通过在给定上下文上进行下一词预测,持续地进行学习,并将所读取的上下文信息压缩到其模型参数中。此外,我们在训练阶段利用元学习来优化模型在测试时学习的初始化状态。总体而言,我们的方法属于一种测试时训练(Test-Time Training, TTT)形式,在测试阶段(通过下一词预测)和训练阶段(通过元学习)均实现了端到端(End-to-End, E2E),这与以往的方法不同。我们开展了大量实验,重点关注方法随规模变化的特性。特别是,对于使用1640亿个token训练的30亿参数模型,我们的方法(TTT-E2E)在上下文长度扩展方面的表现与采用全注意力机制的Transformer相当,而其他方法如Mamba 2和Gated DeltaNet则无法实现这种扩展性。然而,与循环神经网络(RNN)类似,TTT-E2E的推理延迟在不同上下文长度下保持恒定,因此在处理128K长度上下文时,其速度比全注意力机制快2.7倍。我们的代码已公开发布。
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.
许愿