Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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, Josh Susskind, Etai Littwin

TL;DR

This work analyzes the gap between memory-bounded state space models (SSMs) and Transformers for long-form sequence generation. It proves that fixed-memory GSSMs cannot solve arbitrarily long tasks without interacting with external tools, but that interactive tool-use can enable perfect length generalization on any computable long-form task given appropriate training data. The authors provide a formal framework for CoT and tool-use settings, and validate the theory with experiments across arithmetic, reasoning, and coding tasks, showing strong extrapolation when SSMs are augmented with tools. The findings suggest that, in interactive, tool-using environments, SSMs can be an efficient and scalable alternative to Transformers for long-context, agentic applications such as coding, search, and complex reasoning.

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.

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

TL;DR

This work analyzes the gap between memory-bounded state space models (SSMs) and Transformers for long-form sequence generation. It proves that fixed-memory GSSMs cannot solve arbitrarily long tasks without interacting with external tools, but that interactive tool-use can enable perfect length generalization on any computable long-form task given appropriate training data. The authors provide a formal framework for CoT and tool-use settings, and validate the theory with experiments across arithmetic, reasoning, and coding tasks, showing strong extrapolation when SSMs are augmented with tools. The findings suggest that, in interactive, tool-using environments, SSMs can be an efficient and scalable alternative to Transformers for long-context, agentic applications such as coding, search, and complex reasoning.

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.
Paper Structure (47 sections, 3 theorems, 9 equations, 10 figures, 2 tables)

This paper contains 47 sections, 3 theorems, 9 equations, 10 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Chain-of-Thought and Scratchpad
  4. Emulations and Neural Turing Machines
  5. Length Generalization
  6. Theory
  7. Definitions and Notation.
  8. Long-Form Generation
  9. Generalized State Space Models
  10. CoT, Single-Turn and Interactive Tool-Use.
  11. Learning Algorithms and Length Generalization
  12. Main Results
  13. GSSMs cannot Solve Long-Form Generation Tasks without Interaction.
  14. GSSMs with Interactive Tool-Use can Length Generalize on Long-Form Generation Tasks.
  15. Experiments
  16. ...and 32 more sections

Key Result

Theorem 2.1

Let $f$ be a long-form generation task over $\{{\mathcal{D}}_n\}_{n=1}^\infty$ with coverage parameter $\alpha \in (0,1)$. Then, for any CoT-only or Single-Turn GSSM $h$ there exists some problem complexity $n_0$ s.t. for all $n\ge n_0$ the model $h$ has error: $\mathrm{err}_{n}(h) \ge 1-\alpha$.

Figures (10)

  • Figure 1: We finetune Mamba and Pythia (Transformer) on trajectories collected from different tool-use agents for solving a coding problem. 1) Single-Turn Tool-Use: Hard-coded agent that prints all the files in the repository and then outputs all the required changes. 2) Interactive Tool-Use: Hard-coded agent that iteratively runs the code, changes a few files, runs the code again etc. until all problems are resolved. 3) Distillation: SWE-agent Language Model yang2025swe instructed to solve the bug in the code. Models are trained on codebases of up to 16 files (dashed red line), with context length 8,192, and evaluated on larger codebases with longer context. While Pythia outperforms Mamba on smaller codebases and single-turn tool use, Mamba displays favorable performance on large codebases when trained to imitate interactive agents (agents 2 and 3), extrapolating beyond the training distribution.
  • Figure 2: Left: Illustration of an interactive tool-use agent trajectory with pointer-based memory tool for solving multi-digit addition. The agent can generate thoughts (blue), outputs (purple) or commands (orange), and receive observations (green) from the memory tool. At each step, we show the state of the memory context on the top row, and below it show the sequence of generated tokens. Right: Accuracy of recurrent/SSM models (Mamba, LSTM, GRU) and Transformers (Pythia, Mistral) trained on trajectories for $\le$ 5-digit addition, evaluated on up to 1,000-digits (log scale).
  • Figure 3: Illustration of Logical Graph reasoning task and code fixing task. We generate random graphs that define logical function structure (top) or code dependencies (bottom), and synthetically generate problems according to these graph structures.
  • Figure 4: Example of a logical reasoning graph and its encoding.
  • Figure 5: Performance of different seeds when training a Mamba model on Tower of Hanoi puzzles with up to 8 disks.
  • ...and 5 more figures

Theorems & Definitions (10)

  • Definition 2.1
  • Definition 2.2
  • Definition 2.3
  • Theorem 2.1
  • Theorem 2.2
  • proof : Proof of Theorem \ref{['thm:lower_bound']} for deterministic GSSMs.
  • proof : Proof of Theorem \ref{['thm:lower_bound']}
  • proof : Proof of Theorem \ref{['thm:length_generalization']}
  • Lemma B.1
  • proof