To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models

Abstract: State Space Models (SSMs) have become the leading alternative to Transformers for sequence modeling. 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.

• The paper demonstrates that interactive tool-use enables SSMs to generalize arbitrarily long sequences by overcoming fixed-size memory constraints.
• It provides a rigorous theoretical framework and extensive empirical validation across coding, arithmetic, and reasoning tasks.
• The study highlights that multi-turn interactions with external memory tools are essential to achieve Turing-complete performance in long-form generation.

Tool-Use Enables Length Generalization in State Space Models

Introduction

This paper rigorously investigates the limitations and capabilities of State Space Models (SSMs) for long-form sequence modeling, particularly in the context of interactive tool-use. SSMs, such as Mamba, have emerged as efficient alternatives to Transformers due to their fixed-size memory and linear computational scaling. However, the authors provide a formal negative result: SSMs are fundamentally unable to solve "truly long-form" generation tasks when operating in standalone or single-turn tool-use settings, regardless of chain-of-thought (CoT) reasoning. This limitation arises from the bounded memory inherent to SSMs, which restricts their expressive power as task complexity and output length grow.

The central contribution is a theoretical and empirical demonstration that interactive tool-use—where the model can issue multiple commands to external memory or computational tools—enables SSMs to generalize to arbitrary sequence lengths and problem complexities. The paper formalizes this claim, proves that interactive tool-use is both necessary and sufficient for SSMs to achieve length generalization on any tractable long-form generation task, and validates these findings across arithmetic, reasoning, and coding domains.

Theoretical Framework

The authors define long-form generation tasks as those where the effective number of possible outputs grows with problem complexity. Examples include multi-digit arithmetic, sorting, and code fixing. They introduce the Generalized State Space Model (GSSM) abstraction, encompassing SSMs, RNNs, and Transformers with local attention, and analyze their behavior under three regimes: CoT-only, single-turn tool-use, and interactive tool-use.

A key theorem establishes that GSSMs cannot solve long-form generation tasks in the CoT-only or single-turn tool-use settings. Specifically, for any coverage parameter $\alpha$, there exists a problem complexity $n_0$ such that the error rate of any GSSM is at least $1-\alpha$ for $n \geq n_0$. This result is proved by showing that the fixed-size state space cannot encode the exponentially growing output space required for long-form tasks.

Conversely, the authors prove that with interactive tool-use—specifically, access to an external memory tool with read/write and pointer operations—GSSMs can simulate a Turing machine and thus achieve length generalization. They construct training distributions that cover the necessary state-symbol transitions and show that a simple learning algorithm can generalize to arbitrary problem lengths, provided sufficient coverage in the training data.

Experimental Validation

The empirical section evaluates SSMs (Mamba, LSTM, GRU) and Transformer baselines (Pythia, Mistral) on synthetic and realistic tasks, using pointer-based and search-based memory tools.

Coding Task: Tool-Use Trajectories

The authors finetune Mamba and Pythia on trajectories collected from three types of tool-use agents for a code fixing task: single-turn, interactive, and agent distillation. Models are trained on codebases of up to 16 files and evaluated on larger codebases.

Figure 1: Mamba and Pythia finetuned on tool-use agent trajectories for code fixing; Mamba generalizes to larger codebases when trained on interactive agent data, outperforming Pythia in out-of-distribution settings.

Results show that while Pythia outperforms Mamba on small codebases and single-turn tool-use, Mamba exhibits superior extrapolation to larger codebases when trained to imitate interactive agents. This supports the theoretical claim that interactive tool-use is essential for length generalization in SSMs.

Arithmetic Tasks: Multi-Digit Addition and Multiplication

Models are trained on synthetic trajectories for multi-digit addition and multiplication using pointer-based memory tools. For addition, Mamba and LSTM trained on 5-digit examples generalize perfectly to 1,000-digit addition, whereas Transformers fail to extrapolate.

Figure 2: Left: Interactive tool-use agent trajectory for multi-digit addition with pointer-based memory. Right: Accuracy of SSMs and Transformers on addition tasks, showing SSMs generalize to 1,000 digits while Transformers do not.

Similar results are observed for multiplication, with SSMs maintaining high accuracy on numbers orders of magnitude longer than those seen during training.

Reasoning Tasks: Logical Graphs and Tower of Hanoi

The logical graph task uses a search-based memory tool, and SSMs generalize from graphs with 10 nodes (training) to 1,000 nodes (testing) with near-perfect accuracy. For Tower of Hanoi, which has exponentially growing output length, SSMs show more limited but still significant length generalization, with performance sensitive to initialization seed.

Figure 3: Illustration of logical graph reasoning and code fixing tasks, highlighting the graph-based structure and synthetic problem generation.

Figure 4: Example logical reasoning graph and its encoding.

Figure 5: Performance variability across seeds for Mamba on Tower of Hanoi, indicating sensitivity in length generalization for tasks with exponential output growth.

Aggregate Results

Figure 6: Length generalization performance of SSMs and Transformers on multiplication, logical graph, and Tower of Hanoi tasks; SSMs consistently outperform Transformers in out-of-distribution extrapolation.

Figure 7: Multiplication generalization for Mamba across training configurations, showing stable accuracy as sequence length increases with sufficient training coverage.

Ablations and Task Mixtures

Ablation studies confirm that CoT-only, reversed answer, and single-turn tool-use settings do not yield length generalization. Only interactive tool-use with sufficient coverage in training enables SSMs to generalize. Task mixture experiments show that co-training with structurally related auxiliary tasks (e.g., addition for multiplication) can improve generalization under limited training budgets.

Realistic Coding Agent Evaluation

The code fixing task is further evaluated using trajectories from SWE-agent-LM-32B, demonstrating that Mamba maintains higher pass rates and longer sequence handling than Transformer baselines as codebase complexity increases.

Figure 8: Pass rate and median sequence length for SWE-agent-LM-32B on the code fixing task, illustrating degradation with increasing codebase size.

Implications and Future Directions

The findings have significant implications for the design of efficient, scalable language modeling systems. SSMs, when equipped with interactive tool-use, can overcome their memory bottleneck and achieve both efficiency and accuracy in long-form generation tasks. This positions SSMs as strong candidates for agentic and tool-based applications, such as coding assistants, mathematical reasoning agents, and complex search systems.

The analysis also highlights the importance of evaluating model architectures in system-level, agentic contexts rather than as standalone predictors. As LLMs are increasingly deployed as agents interacting with external tools, the synergy between model architecture and tool-use becomes critical.

Future research should focus on developing pretrained SSMs with robust tool-use capabilities, exploring hybrid architectures that combine SSM efficiency with Transformer expressivity, and formalizing training paradigms that maximize coverage of state-symbol transitions for length generalization.

Conclusion

This paper establishes that SSMs are fundamentally limited in standalone and single-turn tool-use settings for long-form generation tasks due to bounded memory. However, interactive tool-use enables SSMs to generalize to arbitrary sequence lengths and complexities, matching or exceeding Transformer performance in agentic settings. Theoretical proofs and extensive experiments across arithmetic, reasoning, and coding domains validate these claims. The work encourages a shift toward system-level evaluation and development of tool-augmented SSM agents, with broad implications for efficient, scalable AI systems.