StateLoop: Persistent State for Looped Transformers
AuthorsFeichen Song, Boyi Zeng, He Li, Shixiang Song, Yiqin Hao, Zitong Wang, Yizhou Yang, Yifan Liu, Ziwei He, Zhouhan Lin
AffiliationsShanghai Jiao Tong University · Sun Yat-sen University · Shanghai Innovation Institute
StateLoop gives repeatedly applied Transformers a persistent memory to carry useful information across computation loops.
Key results
At 1.4B parameters, StateLoop improves on Vanilla Looped.
Percentage-point gain over Vanilla Looped across eight tasks.
Percentage-point gain over Vanilla Looped in zero-shot evaluation.
Percentage-point gain over Vanilla Looped in zero-shot evaluation.
StateLoop advantage over Vanilla Looped at 7 iterations, up from 0.14 at 3.
StateLoop uses less decoding latency than Vanilla Looped with 5 iterations.
What the paper found
StateLoop adds one persistent state per token to a looped Transformer, preserving information across repeated applications of a shared backbone. A learned selective state-space update gates retention of the previous state against new backbone output; the resulting state is projected back into the direct recurrent stream and used in the final prediction. Optional causal-attention and Mamba-2 mixers let states exchange information across tokens. In models trained on the Pile with a LLaMA-style backbone and GPT-NeoX tokenizer, the 1.4B-parameter StateLoop reduces Pile perplexity by 0.18 and raises zero-shot accuracy averaged across eight tasks by 0.97 percentage points over Vanilla Looped. Causal-attention and Mamba-2 mixers raise that accuracy gain to 1.43 and 1.44 percentage points. The advantage also grows with loop depth: at 410M scale, the Pile perplexity gap expands from 0.14 at 3 iterations to 0.30 at 7. At decoding time, StateLoop’s state branch adds just 1.42% latency, while using 18.86% less latency than a Vanilla Looped model with one extra iteration. Ablations show that carrying state across iterations matters: resetting it worsens perplexity from 9.04 to 9.25.
Original 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.
Read the original paperMore in Transformers
Browse all 51 papers →Hierarchical Residuals: A Minimalist Decoupling of Residual Width from Compute Width
Ziteng Wang, Jianqiao Lu, Xinyu Ma, Jintao Zhang, Jun Zhu, Yiyuan Ma, Jianfei Chen
HiRes gives Transformers a wider memory path between blocks without making each layer more computationally expensive.
How to Loop Your Transformer: Loop Scaling, Specialization, and Parallel Inference
Chanakya Ekbote, Artidoro Pagnoni, Devendra Singh Sachan, Shruti Bhosale, Mike Lewis, Maryam Fazel-Zarandi, Shashank Rajput, Sujan Kumar Gonugondla
The paper explores how to make transformers deeper through weight reuse while balancing model quality, training cost, and inference speed.
A Conditional No-Hair Theorem for Transformers: How Language Models Forget Their Prompt
Shuyang Xiang
The paper finds that a transformer’s direct memory of its prompt fades as it generates text, while the text itself can carry that memory forward.