LongSSM: On the Length Extension of State-space Models in Language Modelling
Open accessAbstract
In this paper, we investigate the length-extension of state-space models (SSMs) in language modeling. Length extension involves training models on short sequences and testing them on longer ones. We show that state-space models trained with zero hidden states initialization have difficulty doing length extension. We explain this difficulty by pointing out the length extension is equivalent to polynomial extrapolation. Based on the theory, we propose a simple yet effective method - changing the hidden states initialization scheme - to improve the length extension. Moreover, our method shows that using long training sequence length is beneficial but not necessary to length extension. Changing the hidden state initialization enables the efficient training of long-memory model with a smaller training context length.
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Integrity signals
Registry
Cites
Statistics
Claims the paper makes
The claims and statistics of this paper haven’t been checked yet.
Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.
Claims other papers make about this one
Each is checked against this paper’s own text: does it say what the citing paper says it does?
A very long context cannot be properly compressed in a small vector-valued or matrix-valued state.
“On the other hand, a very long context cannot be properly compressed in a small vector-valued or matrix-valued states (S. Wang 2024).”
From Titans: Learning to Memorize at Test Time (2025) · cites this paper as S. Wang 2024
The Introduction (no page given) says Mamba tends to have limited length extension beyond its training length, which supports a concern about retaining long-context information. But the paper does not say that a very long context cannot be properly compressed into a small vector- or matrix-valued state; its Conclusion (no page given) reports effective extension to 32K tokens with a different hidden-state initialization scheme. (LLM, conf=0.84)
Experimental — accuracy not yet measured
Self-check
Percentages, sums and reported differences re-computed from the paper’s own numbers.
The self-check needs the paper’s full text.
Cites
No references could be resolved yet.
Research works printed but not linked: 22
- Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D. Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, and Amanda Askell. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
- Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. HiPPO: Recurrent Memory with Optimal Polynomial Projections. In Advances in Neural Information Processing Systems, volume 33, pages 1474-1487.
- Albert Gu, Karan Goel, and Christopher Re. Efficiently Modeling Long Se- quences with Structured State Spaces. In International Conference on Learning Representations, October 2021.
- Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models. arXiv preprint arXiv:2307.08621, 2023. · arXiv
- Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is All you Need. In Advances in Neural Information Processing Systems, volume 30.
- Ofir Press, Noah A. Smith, and Mike Lewis. Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation, April 2022.
- Eric Martin and Chris Cundy. Parallelizing Linear Recurrent Neural Nets Over Sequence Length. In International Conference on Learning Representations, February 2018.
- Ann Yuan, Andy Coenen, Emily Reif, and Daphne Ippolito. Wordcraft: Story Writing With Large Language Models. In 27th International Conference on Intelligent User Interfaces, IUI '22, pages 841-852, New York, NY, USA, March 2022. Association for Computing Machinery. ISBN 978-1-4503-9144-3. doi: 10.1145/3490099.3511105. · DOI · Paper page
- Owen Marschall, Kyunghyun Cho, and Cristina Savin. A unified framework of online learning algorithms for training recurrent neural networks. The Journal of Machine Learning Research, 21(1):5320-5353, 2020.
- Shida Wang and Beichen Xue. State-space models with layer-wise nonlinearity are universal approximators with exponential decaying memory. In Thirty-Seventh Conference on Neural Information Processing Systems, November 2023.
- Zhong Li, Jiequn Han, Weinan E, and Qianxiao Li. On the Curse of Memory in Recurrent Neural Networks: Approximation and Optimization Analysis. In International Conference on Learning Representations, October 2020.
- Haotian Jiang, Qianxiao Li, Zhong Li, and Shida Wang. A Brief Survey on the Approximation Theory for Sequence Modelling. Journal of Machine Learning, 2 (1):1-30, June 2023. ISSN 2790-203X, 2790-2048. doi: 10.4208/jml.221221. · DOI
- Jimmy T. H. Smith, Andrew Warrington, and Scott Linderman. Simplified State Space Layers for Sequence Modeling. In International Conference on Learning Representations, February 2023.
- David E. Rumelhart, Geoffrey E. Hinton, and Ronald J. Williams. Learning representations by back-propagating errors. Nature, 323(6088):533-536, October 1986. ISSN 1476-4687. doi: 10.1038/323533a0. · DOI · Paper page
- Sepp Hochreiter and Jürgen Schmidhuber. Long Short-term Memory. Neural computation, 9:1735-80, December 1997. doi: 10.1162/neco.1997.9.8.1735. · DOI · Paper page
- Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning Phrase Rep- resentations using RNN Encoder-Decoder for Statistical Machine Translation. Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing, pages 1724-1734, September 2014.
- Y. Bengio, P. Simard, and P. Frasconi. Learning long-term dependencies with gradient descent is difficult. IEEE Transactions on Neural Networks, 5(2):157-166, March 1994. ISSN 1941-0093. doi: 10.1109/72.279181. · DOI · Paper page
- Sepp Hochreiter. The Vanishing Gradient Problem During Learning Recurrent Neu- ral Nets and Problem Solutions. International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems, 06(02):107-116, April 1998. ISSN 0218-4885, 1793-6411. doi: 10.1142/S0218488598000094. · DOI · Paper page
- Shida Wang, Zhong Li, and Qianxiao Li. Inverse Approximation Theory for Nonlinear Recurrent Neural Networks. In The Twelfth International Conference on Learning Representations, October 2023.
- Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. ISSN 1533-7928.
- Laurent Demanet and Alex Townsend. Stable Extrapolation of Analytic Functions. Foundations of Computational Mathematics, 19(2):297-331, April 2019. ISSN 1615-3375, 1615-3383. doi: 10.1007/s10208-018-9384-1. · DOI
- Thomas M. Cover and Joy A. Thomas. Elements of Information Theory (Wiley Series in Telecommunications and Signal Processing). Wiley-Interscience, USA, June 2006. ISBN 978-0-471-24195-9.
Cited by
Sorted by citations
Titans: Learning to Memorize at Test Time
Behrouz, A., Zhong, P., & Mirrokni, V. (2025) · arXiv · doi:10.48550/arxiv.2501.00663 · Cited by 24 · see citing works
“…and efficiency (linear vs. quadratic complexity), whose advantages is appeared 39th Conference on Neural Information Processing Systems (NeurIPS 2025). for very long context; On the other hand, a very long context cannot be properly compressed in a small vector-valued or matrix-valued states [10].”
- Authors
- Wang, S.
- Affiliations
- National University of Singapore (SG)
- Corresponding author
- Wang, Shida
- Venue
- arXiv
- Publisher
- arXiv
- Published
- 2024
- ISSN
- 2331-8422 (ISSN-L)
- arXiv
- 2406.02080
- OpenAlex
- W4399415802
- Cited by
- 0
- References
- 0
- Type
- posted content
- Language
- en
- Open access
- green
- arxiv.org
- Locations
- arXiv (Cornell University) (submittedVersion)
- arXiv (Cornell University)
- Topics
- Physical Sciences › Computer Science › Artificial Intelligence › Topic Modeling
- Physical Sciences › Computer Science › Artificial Intelligence › Natural Language Processing Techniques
- Physical Sciences › Computer Science › Artificial Intelligence › Speech Recognition and Synthesis
- Keywords
- state-space model, language modeling, polynomial extrapolation, long memory models
About these checks
Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.
Statistics and arithmetic are recomputed, not judged.
Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.
What counts as a claim ›How we measure confidence ›How we check this ›
Are you an author, or do you see a problem with this page? Request a review
