End-to-End Test-Time Training voor Lange Context
End-to-End Test-Time Training for Long Context
December 29, 2025
Auteurs: 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
cs.AI
Samenvatting
Wij formuleren taalmodellering met lange context als een continu leerprobleem in plaats van een kwestie van architectuurontwerp. Binnen deze formulering gebruiken we enkel een standaardarchitectuur – een Transformer met sliding-window attention. Ons model blijft echter leren tijdens de testfase via next-token prediction op de gegeven context, waarbij het de gelezen context comprimeert in zijn gewichten. Daarnaast verbeteren we de initialisatie van het model om leren tijdens de testfase mogelijk te maken via meta-leren tijdens de trainingsfase. Onze methode, een vorm van Test-Time Training (TTT), is End-to-End (E2E) zowel tijdens de testfase (via next-token prediction) als de trainingsfase (via meta-leren), in tegenstelling tot eerdere vormen. Wij voeren uitgebreide experimenten uit, met een focus op schaaleigenschappen. Voor 3B-modellen getraind met 164B tokens schaalt onze methode (TTT-E2E) met de contextlengte op dezelfde manier als een Transformer met volledige aandacht, terwijl andere methoden, zoals Mamba 2 en Gated DeltaNet, dat niet doen. Vergelijkbaar met RNN's heeft TTT-E2E echter een constante inferentielatentie, ongeacht de contextlengte, waardoor het 2,7 keer sneller is dan volledige aandacht voor een context van 128K. Onze code is openbaar beschikbaar.
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.