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
6 min
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.
Alex: What parts of the model get these nudges, and why not everything?
Sam: They only update specific sections called MLP layers. Those are like simple neural networks that process patterns without the complexity of attention. Updating everything else, like attention or normalization parts, causes instability, so those stay fixed. They limit nudges to the last quarter of the model's blocks. In those blocks, they add a second, unchanging MLP to hold pre-trained knowledge, preventing forgetting.
Alex: How does end-to-end training make the whole thing work as well as full attention?
Sam: Standard training ignores these test-time nudges, so the model starts from a poor spot. End-to-end fixes that with meta-learning. During training, they simulate the full sequence of nudges and optimize the starting weights so the final predictions match what you'd want after rehearsal. It's like practicing bike rides end-to-end, not just straight lines. The paper shows TTT-E2E nearly matches full attention's accuracy on long contexts.
Alex: Does that hold for even longer texts?
Sam: They tested different numbers of layers to update, focusing on the last ones. Updating too few means less space to store compressed info—like trying to summarize a book in one sticky note. But with the last quarter, it scales just like full attention, holding steady up to 128,000 words.
Alex: What about errors across a long text?
Sam: They looked at errors for each word position in texts of 32,000 and 128,000 words. TTT-E2E had lower errors than full attention at every spot. The average advantage came mostly from the early parts. That's because full attention has to guess well for every possible future word, spreading its skill thin. TTT-E2E focuses just on the current chunk, since later nudges handle what's ahead—like prepping for today's homework without worrying about next week's test.
Alex: What about tasks needing perfect recall of random details buried deep?
Sam: They tested that with needle-in-a-haystack—hide a key fact anywhere in a long passage, like a single phone number in a novel. Full attention keeps every detail handy. TTT-E2E drops to much lower success at longest lengths, confirming the trade-off: compression loses some fine details for speed.
Alex: Does this hold when the model generates its own long continuation?
Sam: Yes. They had it continue 8,000 words after an 8,000-word prompt from books. TTT-E2E's generated text earned lower errors overall than full attention's. It even self-trains on its own words once batches fill—like practicing what it just wrote to stay sharp.
Alex: What about costs?
Sam: Training is slower—about 3.4 times at short contexts, due to extra steps computing nudges. It evens out at longer lengths. Most training happens short, so it's a notable drawback needing fixes.
Alex: So it compresses long texts into its weights through step-by-step nudges, matching full attention's skill without the slowdown. But recall for buried details suffers, and training takes longer.
Sam: That's right. The paper positions this as continual learning, where the model adapts to each unique input like a person building personalized notes from a lecture, without forgetting core knowledge. It enables processing book-length documents in real time at steady speed, while nearing Transformer's accuracy. Thanks for listening to ResearchPod.