PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2512.23675v2cs.LG

TL;DR

Long-context Transformers incur growing costs because full attention scans all previous tokens. TTT-E2E instead continues next-token training at test time with sliding-window attention, improving performance while compressing context into updated weights.

  • Problem

    Full self-attention readily recalls every detail but its per-token cost grows linearly with context length, making long-context processing prohibitive.

  • Method

    TTT-E2E combines sliding-window attention with end-to-end next-token optimization during test-time training and meta-learned initialization.

  • Results

    TTT-E2E improves test loss by 0.018 even when added to full attention, while architectural modifications without TTT remain nearly indistinguishable from full attention.

  • Takeaways & Limitations

    The results support treating test-time weight updates as a complementary long-term memory alongside sliding-window short-term memory.

  • Takeaways & Limitations

    Training latency remains a significant limitation: TTT-E2E is 3.4× slower than full attention at 8K context, despite being 1.2× faster at 128K.

Abstract

from arXiv · show

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.

1 Introduction

Long-context language modeling must balance effective use of longer context with computational efficiency. TTT addresses this by compressing context into updated weights and preparing those updates through end-to-end meta-learning.

  • 1 Introduction: Full attention processes every prior key and value for each new token, making per-token cost grow linearly with context length.It preserves details effectively but becomes prohibitive for long contexts.
  • 1 Introduction: RNNs offer constant cost per token but become less effective on longer contexts, while sliding-window and hybrid methods remain below full attention in long-context language modeling.
  • 1 Introduction: TTT continues next-token training on the given context so the model compresses important information into its weights instead of recalling every detail.
  • 1 Introduction: Meta-learning prepares the initialization for test-time training by optimizing loss after TTT rather than only the model’s out-of-box loss.
  • 1 Introduction: TTT-E2E is end-to-end at both loops: next-token prediction in the inner loop and final post-TTT loss optimization in the outer loop.

2 Method

At test time, language modeling separates context processing from next-token decoding. TTT changes the computational profile from full attention’s quadratic prefill and linear decode to linear prefill and constant-cost decode.

  • 2 Method: Next-token prediction first conditions on the available tokens during prefill, then decodes a distribution over the next token.
  • 2 Method: The test loss is cross entropy between the decoded distribution and the next token generated by nature.
  • 2 Method: In the toy setup, TTT-E2E with b = 1 converts the no-attention baseline into a curve that performs almost as well as full attention.
  • 2 Method: TTT has O(T) prefill and O(1) decode, compared with full attention’s O(T^2) prefill and O(T) decode.

2.1 TTT via Next-Token Prediction

TTT turns next-token prediction on the observed context into a sequence of test-time weight updates. The updated weights then use the accumulated context to predict the next token.

  • 2.1 TTT via Next-Token Prediction: The attention-free baseline is effectively a bigram because it has no memory of tokens beyond the immediate input.
  • 2.1 TTT via Next-Token Prediction: TTT trains the baseline on context tokens by predicting each next token and comparing it with the observed target.
  • 2.1 TTT via Next-Token Prediction: At test time, TTT updates the weights sequentially with gradient descent for each context position.
  • 2.1 TTT via Next-Token Prediction: After processing the context, the model predicts the next token using the final updated weights.

2.2 Learning to (Learn at Test Time)

TTT-E2E trains the initial weights for the model’s behavior after test-time updates, aligning training with the eventual test-time objective. This alignment enables TTT-E2E to approach full-attention performance in the toy comparison, while the basic method still faces efficiency and stability issues.

  • 2.2 Learning to (Learn at Test Time): TTT-E2E optimizes the average post-update test loss over training sequences, matching the training objective to test-time behavior.
  • 2.2 Learning to (Learn at Test Time): Training a static model’s loss without accounting for test-time weight updates creates a mismatch and offers little guarantee of low post-update test loss.
  • 2.2 Learning to (Learn at Test Time): TTT-E2E performs almost as well as full attention in the toy experiment, whereas TTT-naive is only slightly better than the toy baseline.
  • 2.2 Learning to (Learn at Test Time): Computing the end-to-end training gradient requires gradients of gradients, which modern automatic-differentiation frameworks can compute.
  • 2.2 Learning to (Learn at Test Time): The basic TTT-E2E procedure has efficiency and stability problems because its sequential inner-loop updates cannot be parallelized and depend on single tokens.

2.3 Mini-Batch TTT and Sliding Window

The method replaces online TTT updates with mini-batch updates for better parallelism and stability, then adds sliding-window attention to preserve within-batch context. Implementation choices restrict which layers are updated and protect pretrained knowledge while retaining efficient long-context processing.

  • Mini-Batch TTT: Mini-batch TTT improves parallelism and stability relative to online gradient descent, but predictions within each batch miss earlier batch context.The method outputs after T/b updates, and b = 1 recovers the online formulation.
  • Sliding Window: Sliding-window attention restores context within each mini-batch, with k ≥ b required so attention can retain those tokens before TTT updates the weights.For T = 128K, the main results use k = 8K and b = 1K.
  • Implementation Details: TTT updates only MLP layers because updating embeddings, normalization, or attention layers causes instability in the outer loop.Embedding, normalization, and attention layers remain frozen during test-time training.
  • Implementation Details: Updating only the last 1/4 of blocks trades context-scaling capacity for lower gradient-backpropagation cost.The paper frames updated-layer count as a compute-versus-storage trade-off.
  • Implementation Details: A second static MLP in updated blocks serves as safe storage for pretrained knowledge while total parameter count remains matched to baselines.The hidden dimensions are reduced throughout the network to preserve parameter-count fairness.
  • Decoding: The final architecture extends naturally to multiple-token decoding by taking a gradient step only after decoded tokens fill a TTT mini-batch.The method otherwise requires no special decoding procedure under the stated divisibility assumption.

2.4 Alternative Derivation

The alternative derivation connects TTT-E2E to prior Key-Value Binding by replacing layer-wise reconstruction losses with next-token prediction. Removing auxiliary parameters and restricting updates yields a larger effective state with lower inference latency while preserving the end-to-end objective.

  • Starting Point: Key-Value Binding: TTT-KVB stores key-value associations implicitly in a model trained at test time, whereas TTT-E2E uses the network’s final next-token prediction loss.TTT-KVB’s mechanism underlies several related TTT and linear-attention variants.
  • Key Step: E2E at Test Time: Replacing KVB reconstruction loss with next-token prediction substantially improves language-modeling performance, essentially reaching the final method’s level.The intermediate TTT-E2E all layers MH method makes test-time training directly optimize token-level loss.
  • Key Step: E2E at Test Time: The derivation removes layer-wise reconstruction losses and their θK, θV parameters, making the intermediate method end-to-end at test time.The test-time training loss becomes exactly the token-level test loss ℓt.
  • Connection to RNNs: TTT-E2E updates only one RNN-like layer in a single backward pass, unlike TTT-KVB’s independently updated layer units.Both methods can be interpreted through forward and backward passes, but TTT-KVB stops gradients within each block.
  • Final Step: Larger State with Less Compute: Removing multi-head LoRA updates and updating only the last 1/4 of blocks increases hidden-state capacity while reducing inference computation.The final method has a 5× larger hidden state and half the inference latency of TTT-E2E all layers MH for the 760M model.
  • Final Step: Larger State with Less Compute: For the 760M model, TTT-E2E uses an 88M hidden state and 0.0086 seconds per 1K prefill tokens, versus 18M and 0.017 seconds for all-layers MH.These comparisons are reported for H100 prefill inference.
  • Final Step: Larger State with Less Compute: The paper predicts that smaller state leads to worse context scaling, motivating its later ablation of the number of updated layers.The difference is difficult to see at 8K context length, so scaling is evaluated separately.

3 Main Results

TTT-E2E improves long-context language modeling through test-time learning, with gains that persist beyond sliding-window architecture changes and scale similarly to full attention under larger compute budgets. Its advantage is strongest for non-recall evaluations, while full attention remains substantially better at exact retrieval.

  • 3.2 Ablations on Hyper-Parameters: A 0.018 test-loss improvement over full attention shows that TTT-E2E provides an orthogonal gain beyond compensating for sliding-window limitations.This result holds when context length and other factors are fixed.
  • 3.2.1 Number of Layers Updated: Updating the last 1/4 of layers preserves full-attention-like context scaling, whereas updating only 1 or 3 layers does not; updating 12 layers performs roughly like 6.The final method therefore updates the last quarter of layers across model sizes.
  • 3.3 Scaling with Training Compute: TTT-E2E follows a similar scaling trend to full attention at medium-to-large compute, with regime boundaries near 760M parameters and 48B training tokens.Its advantage over full attention decreases in the small-compute regime.
  • 3.4 Loss Breakdown by Token Index: TTT-E2E is the only evaluated method with lower loss than full attention throughout 32K and 128K contexts, with most of its aggregate advantage coming from earlier tokens.Its loss gap is small near the end of the context window but does not reverse at 128K.
  • 3.5 Needle in a Haystack: Full attention dramatically outperforms TTT-E2E on Needle-in-a-Haystack retrieval, indicating that compression sacrifices some nearly lossless recall of seemingly irrelevant details.TTT-E2E instead performs better on the limited Qwen-loss evaluation during decoding.
  • 3.7 Computational Efficiency: Training latency remains a significant limitation of the current implementation despite the method’s favorable long-context inference behavior.The supplied results identify training efficiency as an ongoing limitation rather than a resolved advantage.

4 Related Work

The paper situates Test-Time Training within continual learning, emphasizing instance-specific adaptation, self-supervised learning, and compression of sequential context into model weights. It distinguishes TTT-E2E from prior approaches through end-to-end next-token optimization, meta-learning, and minimal architectural changes.

  • Continual Learning and Test-Time Training: Test-Time Training formulates learning around each individual test instance, extending continual learning beyond models updated from gradually changing distributions.The paper contrasts this formulation with conventional continual learning, which typically samples training and test data from a changing distribution.
  • Forms of Test-Time Training: Prior TTT methods use nearest-neighbor fine-tuning, auxiliary self-supervision, or generated data for reasoning and visual-motor tasks, often increasing effective capacity or improving shifted-distribution generalization.These approaches differ in whether they use retrieved neighbors, unlabeled-instance self-supervision, or generated task-specific data.
  • Continual Learning and Test-Time Training: Unlike conventional systems that reset after independent predictions, sequential TTT can retain prior experience, compressing earlier context into weights for correlated data streams.The paper uses videos and robotics as examples where no-reset adaptation can improve performance.
  • Long-Context Test-Time Training: TTT-E2E avoids memorizing key-value associations and derives long-context adaptation as continual learning with minimal architectural changes.This positions the method against TTT-KVB-style approaches that explicitly maintain key-value associations.
  • Fast Weights: The closest methodological precedent adds trainable fast-weight MLPs updated by next-token loss, but lacks efficiency gains and places fast weights only at the model end.The paper reports that interleaving fast weights with attention layers is critical for preserving gains over larger baselines.

5 Conclusion

TTT-E2E is presented as a general method for long-context language modeling that combines test-time-updated weights with sliding-window short-term memory.

  • TTT-E2E applies Test-Time Training to a Transformer with sliding-window attention for long-context language modeling.
  • The method treats test-time-updated weights as long-term memory and the sliding window as short-term memory.
  • The paper argues that stronger short-term memory could further improve this combined memory hierarchy.

A Recipe for the Toy Example

The toy-example experiments compare full attention with an attention-free Transformer baseline under controlled training settings and learning-rate selection.

  • The toy example compares Transformer with full attention against an attention-free Transformer with fewer parameters.
  • All TTT methods use the attention-free architecture, and both architectures contain two Transformer blocks.
  • The methods are trained on DCLM at context length 128 using a 125K-token outer-loop batch and held-out DCLM evaluation.
  • The selected learning rate is 3e−3 for full attention and 5e−3 for the attention-free Transformer and all TTT variants.

B Basic Recipe

The basic recipe uses a standard Transformer setup with QK normalization and specified pre-training, extension-fine-tuning, positional-encoding, and context-scaling choices.

  • The recipe uses a standard Transformer with QK normalization and the Llama 3 tokenizer, following GPT-3 configurations with Mamba 2-inspired training changes.
  • Extension fine-tuning uses 5% of the pre-training tokens, doubles batch size, and selects a 4e−4 peak learning rate across the listed model sizes and context lengths.
  • RoPE uses θ = 500K at 8K pre-training context, while extension fine-tuning uses θ values from 1M at 16K through 10M at 128K.
  • Longer context improves loss for full attention and hybrid SWA-plus-full models, but hurts SWA, Mamba 2, Gated DeltaNet, and TTT-KVB after 32K.

C Improvements to the Baselines

The baseline improvements update kernel implementations and apply QK normalization to improve efficiency, training stability, and baseline comparability.

  • The authors improve baselines by upgrading their attention layers to the latest FlashAttention 3 kernel.
  • QK normalization stabilizes TTT-E2E training and improves Transformer baselines, so it is also applied to other sliding-window baselines.
  • For 760M Gated DeltaNet, QK normalization reduces pre-training loss from 2.814 to 2.809 at 8K and extension-fine-tuning loss from 2.691 to 2.683 at 32K.

D Additional Details for Decoding Evaluation

Decoding evaluation uses standard sampling settings and repetition-penalty decoding to reduce repetitive generations.

  • Sampling uses temperature 1 and top-p 0.95, following prior work.
  • A repetition penalty of 1.1 is applied to make generated text more reasonable.The authors use the HuggingFace generation API and report manual inspection of the full-attention baseline.
Loading 2512.23675v2…