Rethinking Expressivity and Efficiency in Test-Time Training
2026-08-21 • Machine Learning
Machine Learning
AI summaryⓘ
The authors present E²-TTT, a new method that improves how language models update their understanding during reading by combining detailed updates for each word with faster, chunk-based processing. They developed a way to keep important timing details in updates while allowing training to happen in parallel chunks, which previous methods missed. Their approach matches the speed of fast methods but better handles very long texts, especially shown by keeping high accuracy when reading much longer than trained. The authors tested this with large models and found it performs well in language tasks and memory retrieval.
Test-Time TrainingLong-Context ProcessingFast WeightsChunk-wise ApproximationLanguage ModelingIn-Context LearningLength ExtrapolationParallel TrainingMomentum StatesAttention Mechanisms
Authors
Zeyun Zhong, Joya Chen, Manuel Martin, Frederik Diederichs, Juergen Gall, Juergen Beyerer
Abstract
Test-Time Training (TTT) enables long-context processing via continuous weight updates during inference, but current methods struggle to balance the expressivity of per-token update dynamics with the hardware efficiency of chunk-wise approximations. We propose E$^2$-TTT (Expressive and Efficient TTT) to bridge this gap. Under the standard approximation of taking gradients at the chunk-start weights, we derive a closed-form state transition that exactly reproduces the chunk-end fast-weight and momentum states of the per-token recurrence. This enables fully parallelized chunk-level training while preserving the temporal structure of the update rule that prior chunk-wise methods discard. We validate E$^2$-TTT by training models up to 1.3B parameters from scratch. It performs on par with previous TTT and hybrid attention baselines in language modeling while outperforming them on in-context retrieval. Its advantage is most pronounced in length extrapolation: on the standard ``Needle in a Haystack'' passkey test, it retains over 90% accuracy at $8\times$ the training context length. Meanwhile, E$^2$-TTT can match the training throughput of efficient chunk-wise methods, demonstrating that it effectively reconciles expressivity with efficiency. The code is available at https://github.com/zeyun-zhong/E2-TTT.