To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models
Eran Malach, Omid Saremi, Sinead Williamson, Arwen Bradley, Aryo Lotfi, Emmanuel Abbe, Joshua Susskind, Etai Littwin
Abstract
State Space Models (SSMs) have become the leading alternative to Transformers for sequence modeling tasks. Their primary advantage is efficiency in long-context and long-form generation, enabled by fixed-size memory and linear scaling of computational complexity. We begin this work by showing a simple theoretical result stating that SSMs cannot accurately solve any "truly long-form" generation problem (in a sense we formally define), undermining their main competitive advantage. However, we show that this limitation can be mitigated by allowing SSMs interactive access to external tools. In fact, we show that given the right choice of tool access and problem-dependent training data, SSMs can learn to solve any tractable problem and generalize to arbitrary problem length/complexity (i.e., achieve length generalization). Following our theoretical finding, we demonstrate that tool-augmented SSMs achieve remarkable length generalization on a variety of arithmetic, reasoning, and coding tasks. These findings highlight SSMs as a potential efficient alternative to Transformers in interactive tool-based and agentic settings.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext a469d28c-6689-4284-976a-7c34917b7d6cCited by top-tier papers1
Ask how each one uses itBuilds on35
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-AwarenessTri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra et al.NeurIPS 2022 · 5,493 citations
- Efficiently Modeling Long Sequences with Structured State SpacesAlbert Gu, Karan Goel, Christopher RéICLR 2022 · 3,482 citations
- Transformers are RNNs: Fast Autoregressive Transformers with Linear AttentionAngelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, François FleuretICML 2020 · 2,665 citations
- SWE-agent: Agent-Computer Interfaces Enable Automated Software EngineeringJohn Yang, Carlos E. Jimenez, Alexander Wettig, Kilian Lieret et al.NeurIPS 2024 · 2,059 citations
Related papers
- Repeat After Me: Transformers are Better than State Space Models at CopyingSamy Jelassi, David Brandfonbrener, Sham M. Kakade, Eran MalachICML 2024 · 176 citations
- From Generalization Analysis to Optimization Designs for State Space ModelsFusheng Liu, Qianxiao LiICML 2024 · 12 citations
- CodeSSM: Towards State Space Models for Code UnderstandingShweta Verma, Abhinav Anand, Mira MeziniEMNLP 2025
- On the Expressiveness and Length Generalization of Selective State Space Models on Regular LanguagesAleksandar Terzic, Michael Hersche, Giacomo Camposampiero, Thomas Hofmann et al.AAAI 2025 · 8 citations
- Block-State TransformersJonathan Pilault, Mahan Fathi, Orhan Firat, Chris Pal et al.NeurIPS 2023 · 33 citations
