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

TL;DR

The paper reframes long-context language modeling as a continual-learning problem and introduces End-to-End Test-Time Training (TTT-E2E) using a sliding-window Transformer. It trains a meta-learned initialization so that test-time gradient updates compress the given context into model weights, with inner loops targeting next-token prediction and outer loops optimizing the initialization. Empirical results on 3B models trained with 164B tokens show that TTT-E2E scales with context length comparably to full attention while delivering constant inference latency and substantial speedups for very long contexts; ablations reveal crucial design choices in which layers to update and how to structure mini-batches. Needle-in-a-Haystack evaluations reveal limits on recall tasks, while decoding long sequences demonstrates practical gains in efficiency and maintaining reasonable text quality. The work provides a scalable, end-to-end framework for long-context modeling with open-source code, and situates this approach within a broader landscape of continual learning and fast-weight meta-learning literature.

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.

Figures (9)

• Figure 1: Scaling with context length, in terms of test loss (left) and latency (right). Left: Our method (TTT-E2E) turns the worst line (green) into the best (blue) at 128K context length. Loss $\Delta$ ($\downarrow$), the $y$-value, is computed as (loss of the reported method) $-$ (loss of Transformer with full attention), so loss $\Delta$ of full attention itself (orange) is the flat line at $y=0$. While other methods produce worse loss $\Delta$ in longer context, TTT-E2E maintains the same advantage over full attention. All models have 3B parameters and are trained with 164B tokens. Right: Similar to SWA and the RNN baselines, TTT-E2E has constant inference latency regardless of context length, making it $2.7\times$ faster than full attention for 128K context on an H100.
• Figure 2: Toy example. Left: Given $x_1$ and $x_2$ as context, we want to predict the unknown $x_3$. Our toy baseline, a Transformer without self-attention (using only the upward arrows), is effectively a bigram since it has no memory of $x_1$. TTT (using all the arrows) first tries to predict $x_2$ from $x_1$ as an exercise: It computes the loss $\ell_2$ between $x_2$ and the prediction $\hat{p}_2$, then takes a gradient step on $\ell_2$. Now information of $x_1$ is stored in the updated MLPs (blue). Right: Token-level test loss $\ell_t$ for various methods in our toy example, as discussed in Subsection \ref{['subsec:meta']}, except for TTT-E2E $b=16$ discussed in Subsection \ref{['subsec:main']}. In particular, TTT-E2E $b=1$ turns the green line (our toy baseline) into the blue line, which performs almost as well as orange (using full attention).
• Figure 3: Computation graphs following the setup in Figure \ref{['fig:toy']}: Given $x_1$ and $x_2$ as context, we want to predict the unknown $x_3$. Left: Our main method with the sliding-window attention layers and the implementation details discussed in Subsection \ref{['subsec:main']}. For ease of notation, our illustration uses online gradient descent ($b=1$). The lowest downward arrow is disconnected to the MLP below, since gradients pass through the last $L/4$ blocks but not further down. Right: The first step of our alternative derivation in Subsection \ref{['subsec:alternative']}: a simplified version of TTT-KVB in prior work zhang2025testsun2024learning.
• Figure 4: Ablations on three hyper-parameters: sliding window size $k$, mini-batch size $b$, and the number of layers updated during TTT; see details in Subsection \ref{['subsec:results_ablations']}. Given the trends in these ablations, we set $k=8$K, $b=1$K, and we update 1/4 the total number of layers. Loss $\Delta$ ($\downarrow$), the $y$-value in the rightmost panel, is the same as in Figure \ref{['fig:teaser']}. It is computed as (loss of the reported method) $-$ (loss of Transformer with full attention), so loss $\Delta$ of full attention itself (orange) is the flat line at $y=0$. GDN stands for Gated DeltaNet yang2023gated.
• Figure 5: Scaling with training compute in two axes: model size (left) and number of training tokens (right); see details in Subsection \ref{['subsec:results_train']}. Overall, TTT-E2E exhibits a similar trend to full attention under a large training budget (right of the dotted line). We report results both on DCLM at 8K context length after pre-training (a, c) and on Books at 32K after fine-tuning with the same context length (b, d). Loss $\Delta$ ($\downarrow$), the $y$-value, is the same as in Figure \ref{['fig:teaser']} and \ref{['fig:ablations']}. The legend in the leftmost panel is shared across all panels.