Work

Long-context Test-Time Training

Studying whether pretrained full-attention Transformers can warm-start TTT-E2E models to reduce long-context training cost while preserving quality.

Georgia Tech collaboration Active
Authors

Luis M. Montoya, Vijay K. Madisetti

Affiliation

Georgia Tech

Topics
long context test-time training

Key results

Books32K validation loss vs naive SWA conversion at 125M
≈40% lower
than scratch TTT-E2E when a full-attention seed exists
≈7.4× cheaper

Test-time training (TTT-E2E) is a promising architecture for long context, but training a long-context TTT-E2E model from scratch is expensive. This project, a Georgia Tech collaboration with Vijay K. Madisetti, asks: if you already have a pre-trained full-attention Transformer, can you warm-start a TTT-E2E model from it and keep most of the quality at a fraction of the cost?

Comparison

The experiments compare four long-context paths under matched budgets. The first extends the full-attention seed directly to 32K context as a continuation control. The second naively converts the seed to sliding window attention before extension, a mechanism-swap control. The third is the proposed path, bridge the seed through an 8K TTT-E2E stage, then extend to 32K. And the fourth, pre-trains TTT-E2E from scratch as the most close in-family reference.

Results

At 125M parameters, warm-started TTT-E2E reduces Books32K validation loss by about 40% relative to the naive sliding-window conversion, and costs about 7.4x less than the full scratch TTT-E2E path when a full-attention seed already exists. Scratch pre-raining remains the best final-quality path at both 125M and 760M at same compute budget, but the warm-start tax, the second checkpoint versus the third checkpoint gap, roughly halves from 125M to 760M, which is the encouraging direction, as this tells us that the shortcut gets relatively better as models grow.

Measurement

Every reported scalar comes from checkpoint restore plus float32 evaluation on a fixed Books32K surface. The public repository contains the JAX runtime, the Hydra experiment registry for all four paths, and the scripts that regenerate the paper's tables and figures.

Status

The manuscript is in progress, while deciding the venue to submit it.