Elamur Research⌕
☰
home / papers / 2026 / StateX: Enhancing RNN Recall via Post-training State Expansion

StateX: Enhancing RNN Recall via Post-training State Expansion

Paper · 2026 · Xingyu Shen, Yingfa Chen, Zhen Leng Thai et al.

Summary

Recurrent language models such as linear attention and state-space models process long inputs at constant cost per token. However, they compress the entire context into a fixed-size state, which limits how well they can retrieve earlier details. Larger states improve recall, but pre-training with them is expensive, so released RNNs keep their states below 2% of total model size.

StateX addresses this by enlarging the state of an already pre-trained model and then continuing training on a modest data budget. For GLA, the method relies on the observation that using fewer heads yields a larger state. For Mamba2, it adds extra channels to the Q and K projections. In both cases the parameter increase is zero or negligible. Only a handful of key layers are modified, and their parameters are reinitialized rather than inherited.

Applied to public 1.3B GLA and Mamba2 checkpoints with 10B tokens of post-training, StateX outperforms ordinary continual training on recall-intensive, in-context learning and needle-in-a-haystack tasks. Common-sense reasoning performance stays intact. This offers an inexpensive way to narrow the recall gap between RNNs and Transformers.

Key points

  • Base models: fla-hub/gla-1.3B-100B and mamba2-1.3b-hf, each post-trained on 10B tokens.
  • Relative gains over plain continual training:
    • Recall-intensive tasks: +3.36% (GLA) and +1.1% (Mamba2).
    • In-context learning: +7.2% (GLA) and +1.0% (Mamba2).
    • Average NIAH accuracy up to 64K context: 26.0% → 42.2% (GLA) and 33.2% → 39.2% (Mamba2).
  • Ablations:
    • Reinitializing the modified parameters beats inheriting them.
    • Expanding 4 layers works best, out of 24 in GLA and 48 in Mamba2. Expanding too many layers causes a sharp drop in performance.
  • Efficiency:
    • Training throughput is close to the vanilla models and about 2× that of MoM.
    • Compared with MoM, prefilling is 1.7–2.5× faster and decoding is 81–147× faster (measured on A800 GPUs).
  • In a from-scratch GLA-340M study (20B SlimPajama tokens), a single head gave the best results: recall of 26.0 versus 15.4 with 16 heads, plus lower training loss.

Our summary of the paper, in our own words — not its text. The work belongs to its authors; read it at the links above.

 Institutions

🎓Tsinghua UniversityBeijing · China333 works · 87,019 citations

 Fields, topics and concepts

📚 Shelves 5The field of science this work belongs to (OpenAlex fields).Computer Science · 16,667 worksArtificial Intelligence · 16,194 worksArts and Humanities · 16,039 worksLanguage and Linguistics · 15,891 worksComputer Vision and Pattern Recognition · 2,869 works

Cite

Cite as Xingyu Shen, Yingfa Chen, Zhen Leng Thai, Xu Han, Zhiyuan Liu and Maosong Sun. 2026. StateX: Enhancing RNN Recall via Post-training State Expansion. Findings of the Association for Computational Linguistics: ACL 2026.

The paper itself, not this page. BibTeX from the library's record of the work. DOI 10.18653/v1/2026.findings-acl.1073

@inproceedings{shen2026statex,
  title = {{StateX: Enhancing RNN Recall via Post-training State Expansion}},
  author = {Xingyu Shen, Yingfa Chen, Zhen Leng Thai, Xu Han, Zhiyuan Liu and Maosong Sun},
  booktitle = {Findings of the Association for Computational Linguistics: ACL 2026},
  year = {2026},
  doi = {10.18653/v1/2026.findings-acl.1073},
  url = {https://doi.org/10.18653/v1/2026.findings-acl.1073}
}

References

37 references, 12 of them with a card of ours. ↗ leads to the work, «find» searches for one we have no id for yet.

  1. Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, James Zou, Atri Rudra, and Christopher Ré. 2024a. Simple linear attention language models balance the recall-throughput tradeoff. In Proceedings of the 41st International Conference on Machine Learning, pages 1763–1840.
  2. Simran Arora, Aman Timalsina, Aaryan Singhal, Benjamin Spector, Sabri Eyuboglu, Xinyi Zhao, Ashish Rao, Atri Rudra, and Christopher Ré. 2024b. Just read twice: closing the recall gap for recurrent language models. Preprint, arXiv:2407.05483.
  3. Aviv Bick, Eric Xing, and Albert Gu. 2025. Understanding the skill gap in recurrent language models: The role of the gather-and-aggregate mechanism. Preprint, arXiv:2504.18574. — [arXiv:2504.18574](https://arxiv.org/abs/2504.18574) find
  4. Yingfa Chen, Xinrong Zhang, Shengding Hu, Xu Han, Zhiyuan Liu, and Maosong Sun. 2025. Stuffed mamba: Oversized states lead to the inability to forget. Preprint, arXiv:2410.07145. — [arXiv:2410.07145](https://arxiv.org/abs/2410.07145) find
  5. Kyunghyun Cho, Bart van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. 2014. On the properties of neural machine translation: Encoder-decoder approaches. Preprint, arXiv:1409.1259.
  6. Tri Dao and Albert Gu. 2024. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning, pages 10041–10071. PMLR. find
  7. Jusen Du, Weigao Sun, Disen Lan, Jiaxi Hu, and Yu Cheng. 2025. Mom: Linear sequence modeling with mixture-of-memories. Preprint, arXiv:2502.13685.
  8. Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Alain Le Noac’h, Haonan Li, Kyle McDonell, Niklas Muennighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lintang Sutawika, and 5 others.… find
  9. Tianyu Gao, Alexander Wettig, Howard Yen, and Danqi Chen. 2025. How to train long-context language models (effectively). In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 7376– 7399, Vienna, Austria. Association for Computationa…
  10. Albert Gu and Tri Dao. 2024. Mamba: Linear-time sequence modeling with selective state spaces. Preprint, arXiv:2312.00752. find
  11. Sepp Hochreiter and Jürgen Schmidhuber. 1997. Long short-term memory. Neural Computation, 9(8):1735– 1780. find
  12. Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang, and Boris Ginsburg. 2024. Ruler: What’s the real context size of your long-context language models? Preprint, arXiv:2404.06654.
  13. Weizhe Hua, Zihang Dai, Hanxiao Liu, and Quoc V. Le. 2022. Transformer quality in linear time. Preprint, arXiv:2202.10447. find
  14. Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach. 2024a. Repeat after me: Transformers are better than state space models at copying. Preprint, arXiv:2402.01032. find
  15. Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. 2024b. Repeat after me: Transformers are better than state space models at copying. In International Conference on Machine Learning, pages 21502–21521. PMLR. find
  16. Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. 2020. Transformers are rnns: Fast autoregressive transformers with linear attention. Preprint, arXiv:2006.16236. find
  17. Team Kimi, Yu Zhang, Zongyu Lin, Xingcheng Yao, Jiaxi Hu, Fanqing Meng, Chengyin Liu, Xin Men, Songlin Yang, Zhiyuan Li, Wentao Li, Enzhe Lu, Weizhou Liu, Yanru Chen, Weixin Xu, Longhui Yu, Yejie Wang, Yu Fan, Longguang Zhong, and 41 others. 2025. Kimi linear: An expressive, efficient attention arch…
  18. Kai Liu, Jianfei Gao, and Kai Chen. 2025. Scaling up the state size of RNN LLMs for long-context scenarios. In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 11516–11529, Vienna, Austria. Association for Computational Linguistic… find
  19. Xingtai Lv, Youbang Sun, Kaiyan Zhang, Shang Qu, Xuekai Zhu, Yuchen Fan, Yi Wu, Ermo Hua, Xinwei Long, Ning Ding, and Bowen Zhou. 2025. Technologies on effectiveness and efficiency: A survey of state spaces models. Preprint, arXiv:2503.11224. — [arXiv:2503.11224](https://arxiv.org/abs/2503.11224) find
  20. Sewon Min, Xinxi Lyu, Ari Holtzman, Mikel Artetxe, Mike Lewis, Hannaneh Hajishirzi, and Luke Zettlemoyer. 2022. Rethinking the role of demonstrations: What makes in-context learning work? In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 11048–11064, Ab…
  21. MiniMax, Aonian Li, Bangwei Gong, Bo Yang, Boji Shan, Chang Liu, Cheng Zhu, Chunhao Zhang, Congchao Guo, Da Chen, Dong Li, Enwei Jiao, Gengxin Li, Guojun Zhang, Haohai Sun, Houze Dong, Jiadai Zhu, Jiaqi Zhuang, Jiayuan Song, and 71 others. 2025. Minimax-01: Scaling foundation models with lightning a…
  22. Yuqi Pan, Yongqi An, Zheng Li, Yuhong Chou, Ruijie Zhu, Xiaohui Wang, Mingxuan Wang, Jinqiao Wang, and Guoqi Li. 2025. Scaling linear attention with sparse state expansion. Preprint, arXiv:2507.16577. — [arXiv:2507.16577](https://arxiv.org/abs/2507.16577) find
  23. Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, and 15 others. 2023. Rwkv:… find
  24. Bo Peng, Daniel Goldstein, Quentin Anthony, Alon Albalak, Eric Alcaide, Stella Biderman, Eugene Cheah, Xingjian Du, Teddy Ferdinan, Haowen Hou, Przemysław Kazienko, Kranthi Kiran GV, Jan Kocoń, Bartłomiej Koptyra, Satyapriya Krishna, Ronald Mc-Clelland Jr., Jiaju Lin, Niklas Muennighoff, Fares Obei… find
  25. Bo Peng, Ruichong Zhang, Daniel Goldstein, Eric Alcaide, Xingjian Du, Haowen Hou, Jiaju Lin, Jiaxing Liu, Janna Lu, William Merrill, Guangyu Song, Kaifeng Tan, Saiteja Utpala, Nathan Wilce, Johan S. Wind, Tianyi Wu, Daniel Wuttke, and Christian Zhou-Zheng. 2025. Rwkv-7 "goose" with expressive dynami…
  26. Zhen Qin, Songlin Yang, Weixuan Sun, Xuyang Shen, Dong Li, Weigao Sun, and Yiran Zhong. 2024. Hgrn2: Gated linear rnns with state expansion. Preprint, arXiv:2404.07904. — [arXiv:2404.07904](https://arxiv.org/abs/2404.07904) find
  27. Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. 2021. Linear transformers are secretly fast weight programmers. Preprint, arXiv:2102.11174. find
  28. Daria Soboleva, Faisal Al-Khateeb, Robert Myers, Jacob R Steeves, Joel Hestness, and Nolan Dey. 2023. SlimPajama: A 627B token cleaned and deduplicated version of RedPajama. https://cerebras.ai/bl og/slimpajama-a-627b-token-cleaned-and-d eduplicated-version-of-redpajama. ↗
  29. Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. 2023. Retentive network: A successor to transformer for large language models. Preprint, arXiv:2307.08621.
  30. Falcon-LLM Team. 2024. The falcon 3 family of open models. find
  31. Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2023. Attention is all you need. Preprint, arXiv:1706.03762.
  32. Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi, and Bryan Catanzaro. 2024. An empirical study of mamba-based language mo… find
  33. Ke Alexander Wang, Jiaxin Shi, and Emily B. Fox. 2025. Test-time regression: a unifying framework for designing sequence models with associative memory. Preprint, arXiv:2501.12352. — [arXiv:2501.12352](https://arxiv.org/abs/2501.12352) find
  34. Songlin Yang, Jan Kautz, and Ali Hatamizadeh. 2025. Gated delta networks: Improving mamba2 with delta rule. Preprint, arXiv:2412.06464. find
  35. Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. 2024. Gated linear attention transformers with hardware-efficient training. In International Conference on Machine Learning, pages 56501–56523. find
  36. Tianyuan Zhang, Sai Bi, Yicong Hong, Kai Zhang, Fujun Luan, Songlin Yang, Kalyan Sunkavalli, William T. Freeman, and Hao Tan. 2025. Test-time training done right. Preprint, arXiv:2505.23884. find
  37. Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Khai Hao, Xu Han, Zhen Leng Thai, Shuo Wang, Zhiyuan Liu, and Maosong Sun. 2024. ∞bench: Extending long context evaluation beyond 100k tokens. Preprint, arXiv:2402.13718. — [arXiv:2402.13718](https://arxiv.org/abs/2402.13718) find

Updated

↑ back to top

This page in Markdown → https://md.research.elamur.ai/free/en/papers/shen-2026-statex-enhancing-rnn-recall-post/