Shared memory improves test-time training for AI models
Universal Test-Time Training
Machine LearningComputation and LanguageComputer Vision and Pattern Recognition
Summary
Test-Time Training helps AI models adapt as they work, by updating a memory of recent information. Previous methods kept separate memories for each layer in the model, but the authors show that sharing one memory across all layers works better. This shared memory lets different parts of the model learn from each other's updates, improving results in language tasks and image synthesis. Their approach achieves higher accuracy and better image quality than methods with separate memories.
What this means in practice
- •For natural language engineers: Improve language models to better adapt on the fly during inference by sharing memory across layers for higher accuracy without extra compute cost.
- •For computer vision developers: Enhance novel view synthesis systems with shared memory layers to increase image quality at fixed computational resources.
Authors
Zefan Cai, Qinzhe Hu, Ziqiao Ma, Hao Tan, Junjie Hu
Abstract
Recent Test-Time Training (TTT) architectures compress context into fast weights that are updated online and queried as memory. Existing TTT designs keep this memory private to each layer: it recurs only over time, and depth merely indexes L separate memories. We argue that memory ownership need not be tied to depth, and introduce Universal Test-Time Training (uTTT), in which all layers read and write one shared memory while retaining layer-specific backbone parameters. The shared memory thus recurs over two dimensions, time and depth, with chunks and layers as their units: a write by a deep layer in one chunk can be read by a shallow layer in the next. We instantiate this idea as uTTT-MoE and uTTT-Dense. uTTT-MoE routes each token head to a few experts in a pool shared by all layers; uTTT-Dense applies the whole shared memory at every layer without routing. In language modeling, uTTT-MoE reaches 15.5 and 27.9 RULER accuracy at 124M and 760M, 2.6 and 2.1 points above its layer-private counterpart at equal state and active compute, the highest among tested bounded-state models, with per-token loss matching or beating full attention. In novel view synthesis, sharing at fixed per-layer compute gains 0.92 dB in view-23 object PSNR in routed models and 0.76 dB in dense models.