End-to-End Test-Time Training for Long Context
Authors: name: Arnuv Tandon, name: Karan Dalal, name: Xinhao Li, name: Daniel Koceja, name: Marcel Rød, name: Sam Buchanan, name: Xiaolong Wang, name: Jure Leskovec, name: Sanmi Koyejo, name: Tatsunori Hashimoto, name: Carlos Guestrin, name: Jed McCaleb, name: Yejin Choi, name: Yu Sun

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.