Reproduce “Learning to (Learn at Test Time): RNNs with Expressive Hidden States” on OpenResearch

About this paper

Yu Sun, Xinhao Li, Karan Dalal, Jiarui Xu, Arjun Vikram, Genghan Zhang, Yann Dubois, Xinlei Chen and 4 more

The authors introduce Test-Time Training (TTT) layers that use a learner’s weights as recurrent state. Each token updates the learner through a self-supervised reconstruction task, allowing the state to adapt as a sequence arrives. In language-model experiments up to 32k context, TTT perplexity keeps falling by token position while Mamba plateaus after 16k; experiments cover models up to 1.3B parameters.

Key insights

• Allowing the RNN hidden state to be a continually adapting, self-supervised learning model fundamentally changes how RNNs process and compress information over time, enabling superior long-context utilization.
• Making the self-supervised task itself learnable within an "outer loop" frees researchers from hand-crafting objectives, allowing the model to discover the most effective self-supervision strategy for its end goal.
• The TTT layer framework provides a unified perspective that can theoretically recover existing mechanisms such as linear attention (with specific instantiations) and self-attention (with non-parametric learners).

Summary by alphaXiv.