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

TL;DR

This paper introduces a test-time training method for long-context language modeling using a standard Transformer architecture, enabling the model to adapt during inference and scale efficiently with context length.

Contribution

It presents a novel end-to-end test-time training approach with meta-learning that allows Transformers to adapt to long contexts without changing architecture.

Findings

Scales with context length similarly to full attention for 3B models.

Achieves 2.7x faster inference than full attention at 128K context length.

Maintains constant inference latency regardless of context size.

Abstract

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.