Genghan Zhang

Authors

Yu Sun^†,Xinhao Li^†,Karan Dalal^†,Jiarui Xu,Arjun Vikram,Genghan Zhang,Yann Dubois,Xinlei Chen,Xiaolong Wang,Sanmi Koyejo,Tatsunori Hashimoto,andCarlos Guestrin
† co-first author

Abstract

Self-attention performs well in long context but has quadratic complexity. Existing RNN layers have linear complexity, but their performance in long context is limited by the expressive power of their hidden states. We present a practical framework for instantiating sequence modeling layers with linear complexity and expressive hidden states. The key idea is to make the hidden state a machine learning model itself, and the update rule a step of self-supervised learning. Since the hidden state is updated by training even on test sequences, our layers are called Test-Time Training (TTT) layers. We consider two instantiations: TTT-Linear and TTT-MLP, whose hidden state is a linear model and a two-layer MLP respectively. We evaluate our instantiations at the scale of 125M to 1.3B parameters, comparing with a strong Transformer and Mamba, a modern RNN. Similar to Transformer, TTT-Linear and TTT-MLP can keep reducing perplexity by conditioning on more tokens, while Mamba cannot after 16k context. TTT-MLP still faces challenges in memory I/O, but shows larger potential in long context, pointing to a promising direction for future research.

Article

BibTeX


       @InProceedings{yu2025learning, 
 title = 	 {Learning to ({L}earn at Test Time): {RNN}s with Expressive Hidden States}, 
 author =       {Sun, Yu and Li, Xinhao and Dalal, Karan and Xu, Jiarui and Vikram, Arjun and Zhang, Genghan and Dubois, Yann and Chen, Xinlei and Wang, Xiaolong and Koyejo, Sanmi and Hashimoto, Tatsunori and Guestrin, Carlos}, 
 booktitle = 	 {Proceedings of the 42nd International Conference on Machine Learning}, 
 pages = 	 {57503--57522}, 
 year = 	 {2025}, 
 editor = 	 {Singh, Aarti and Fazel, Maryam and Hsu, Daniel and Lacoste-Julien, Simon and Berkenkamp, Felix and Maharaj, Tegan and Wagstaff, Kiri and Zhu, Jerry}, 
 volume = 	 {267}, 
 series = 	 {Proceedings of Machine Learning Research}, 
 month = 	 {13--19 Jul}, 
 publisher =    {PMLR}, 
 pdf = 	 {https://raw.githubusercontent.com/mlresearch/v267/main/assets/sun25h/sun25h.pdf}, 
 url = 	 {https://proceedings.mlr.press/v267/sun25h.html} 
 }