StateLoop improves looped transformers with persistent token memory
StateLoop: Persistent State for Looped Transformers
Computation and Language
Summary
Language models can be made smaller by reusing the same part over and over, but this can make it hard for them to remember details across steps. The authors introduce StateLoop, which keeps an extra memory for each word that updates as the model processes information multiple times. This persistent memory helps the model keep important details longer and improves its ability to understand language and perform tasks it hasn’t seen before. The benefits grow as the model repeats its computation more times.
What this means in practice
- •For machine learning engineers: Develop more efficient language models that maintain important context across repeated computation steps for better performance on diverse tasks.
- •For natural language processing teams: Integrate persistent state mechanisms to improve zero-shot task accuracy in resource-constrained language models.
Authors
Feichen Song, Boyi Zeng, He Li, Shixiang Song, Yiqin Hao, Zitong Wang, Yizhou Yang, Yifan Liu, Ziwei He, Zhouhan Lin
Abstract
Looped Transformers offer a parameter-efficient approach to language modeling by repeatedly applying a shared backbone. Vanilla looped Transformers use each backbone output directly as the input to the next iteration. These outputs may contain information useful several iterations later, but retaining it across the loop depends entirely on the single hidden state. We introduce StateLoop, which maintains one additional persistent state per token to retain information across iterations. At each iteration, a learned gate controls how much of the previous state is retained and how much information from the backbone output is incorporated. The updated state is projected back and added to the backbone output for subsequent iterations and final prediction. This design reduces Pile perplexity and improves average accuracy across eight zero-shot tasks over Vanilla Looped. Its perplexity advantage over Vanilla Looped grows with trained loop depth, consistent with persistent state becoming increasingly useful as recurrent computation deepens.