To Infinity and Beyond: Tool-Use Unlocks Length Generalization in State Space Models
Authors: Eran Malach
Omid Saremi
Sinead Williamson
Arwen Bradley
Aryo Lotfi
Emmanuel Abbe
Joshua M. Susskind
Etai Littwin
Venue: ICLR 2026 Oral (2026)
Keywords: State Space Models
Mamba
Length Generalization
LLM
Transformers
Primary area: foundation or frontier models, including LLMs

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.

Published as a conference paper at ICLR 2026
TO INFINITY AND BEYOND: TOOL-USE UNLOCKS LENGTH GENERALIZATION IN STATE SPACE MODELS
Eran Malach
Omid Saremi
Sincad Williamson
Apple
Apple
Apple wi williamoo@@aplee.co
malachSapple.com
osaremi
@apple............. Dsaremi@apple..o com
Arwen Bradley
Aryo Lotfi
Emmanuel Abbe Apple e.abbe@apple. com
Apple
Apple
awwen.bradley@apple.com
alotfigapple.com
Josh Susskind
Etai Littwin Apple
Apple
jsusskind@apple..com
elittwin@apple.co
ABSTRACT
State Space Models (SSMs) have become the leading alternative to Transform- ers for sequence modeling tasks. Their primary advantage is efficiency in long- context and long-form generation, enabled by fixed-size memory and linear seal- ing of computational complexity. We begin this work by showing simple theoret- ical result stating that SSMs cannot accurately solve any "truly long-form" gener- ation problem (in sense 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 boblem-deependnt training data, SSMs can learn to solve any tractable problem and generalize arbitrary problem length/complexity (i.e., achieve length generalization). Following our theoretical finding, we demon- strate that toul-augeented SSMs achieve remarkable length generalization on a variety of arithmetic, reasoning, and coding tasks. These findings highlight SSMs as potential efficient alternative to Transformers interactive tool-based and agentic settings.
INTRODUCTION
Transformers (Vaswani et al., 2017), the main architecture powering large language models, have well-known limitation: due to the attention mechanism, their computational complexity scales quadratically with the sequence length, and their memory scales linearly with lenguf This quadratic dependency becomes major limitation for tasks that require long-contes and long-fore generation. As test-time scaling paradigms that involve the generation long Chain of Thought (CoT) (Wei et al. 2022 have become the leading solution for improving reasoning capabilities (Jacch et al. 2024
Guo al. 2025), the ability to efficiently generate long sequences becomes even more crucial.
To solve this limitation, various works suggested replacing the attention mechanism with other modules where memory and per-token compute are fixed as a function of the sequence length (Choromansii et al. 2020). Examples of such architectures include variants of Linear Transformers Katharopoulos et al. (2200) and State Space Models (Gu et al., (022]) such as Mamba (Gu & Dao
2023 Dao & Gul 2024). DeltaNet (Yang et al. 2024c) and GatedDeltaNet (Yang et al. 2024bb. These architectures achieve performance similar ar to Transformers across wide range of domains Qu al. 2024) at lower inference cost. However, some works have also pointed out significant limitations of these architectures in certain tasks that involve memorization of long sequences and
While naive implementation of attention requires quadratic memory complexity, efficient implementa- tions such Flash Attention (Dao al. 2022) and KV caching enable close linear complexity.

---

Published as a conference paper at ICLR 2026
Single-Turn Tool-Use:
3) Distillation:
cat ppy
Comparing Langth Generalization
For bug malis You basn commundd and cai Maa Impalement
for 3 Tool-Use Strategies
foo.py print((v1))
100-
and v/'vs" /Vvi"> Foo sel
Poss Rate (16)
5/ 4/"v""/"v0'0 P904.
SWE-agent 32B
60-
2) Interactive Usol-Use:
python mxin.py
Let's purr cromisving bie ropooloryy
20-
Traceback (eost recent
File code/166..p"
-1a
0
10
20
25
30
v/"v'"/vvo0"/ FL69.
total
-Pa-T--P-- Aug 25 1068
Number of Functions
python eafn.ay
look contonr understand hoss.
Hambs- Sing e Tum [1]
Pythio Single Tum (3)
Traceback (most recent
-w: Hamba-1.48 interactive (2)
Pyctia-1.4b interactvv 12)
"code/f996.p"".
car wafn.ay
Mambs-1 da Detii lntion (3)
Pyrrie1] de Distlation (5)
We finetune Mamba and Pythia (Transformer) on trajectories collected from different tool- use agents for solving a coding problem. 1) Single-Ture 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 (Yang al., 2025) instructed to solve the bug in the code. Models are trained on codebases of up to 16 files (dashed red line), with context length 8K, 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 and 3), extrapolating beyond the training distribution.
Figure 1:
in-context learning (Jelassi et al. 2024 Park et al. 2024 Akyürek et al. 2024). Possibly due to these limitations, linear-time models are still not widely adopted replacement to Transformers.
The goal of this work is to understand the capabilities and limitations of SSMs, focusing on tasks that require long-form generation. We formally define long-form generation tasks to be problems where the effective number of outputs grows with the complexity of the problem. We focus on such tasks as these are the tasks where SSMs display clear benefit over Transformers in terms of inference efficiency. However, we show that this efficiency comes at a cost of inherent performance degra- dation. Namely, we prove that SSMs fail to solve long-form generation tasks when the complexity of the task increases beyond the capacity of the model, even if the model is allowed to generate CoT of any length. This limitation arises from the bounded memory of the model, which limits the expressive power when generating long sequences. This is in contrast with Transformers which, using CoT, can in principle solve any computationally tractable problem, utilizing their unbounded memory (Meriil& & Sabharwal 2023}. So, to solve long-form generation tasks Transformers and suffer quadratic scaling of compute, or use SSMs and suffer performance degra- dation. Another alternative is to use hybrid models that mix attention and SSM layers and have been recently shown to achieve state-of-the-aat performance at large scale (Blakeman et al.) 2025). However, this ultimately does not eliminate the quadratic dependence on the sequence length.
either use
Following the observation above, we explore another alternative: allowing SSMs to interactively use external tools. LLMs are now increasingly used as agents that interact with external tools for solving tasks such as coding, math or question answering (Luo et al.] 2025) Yehudai et al. 2025). These tools can allow agents to query and read from external resources, and write information that can be used later. Therefore, such tool-use can naturally augment the internal memory of the model, allowing it access to practically unbounded external memory. We introduce new theoretical framework for studying ReAct (Yao et al. 2023) agents, and show that allowing SSMs access to external memory through interactive tool-use makes them much more powerful. We prove that tool-augmented SSMs trained on task-specific trajectories can achieve length generalization generation task. That is, we show that for any such task we can construct training data with tool- use trajectories such that simple training paradigm learns to execute the task with high accuracy, even when evaluated beyond the length of the training data. Importantly, this result only holds for interactive tool-use, and we show that single-turn tool-use SSMs are still limited.
any tractable long-form
2

---

Published as a conference paper at ICLR 2026
pointer1 pointer2
Step
pointerl. read
pot wwffiner...oom read
1
pointer1
pointer2
Addition
Step
6,0
1.0
poinceri. 1eftO
0.8
pointeri. read(
2
pointer2.read0
sumie, carrysu
Accuracy
Mamba
0.6
pointer1
pointer2
Pyttia
Mistral
Step 20:
3
3,0
6,0
1,0
pofnzer2. readO
pointerl
pointer2
10a
10a
10a
Step 21:
3,0
6,0
1,0
Sequence length
Figure 2: Left: Illustration of 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. 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, Mis- tral) trained on trajectories for s 5-digit addition, evaluated up to 1,000-digits (log scale).
Experimentally, show that SSMs trained to interactively use external memory tools achieve length generalization on tasks such as arithmetic, logical reasoning and coding. For example, Mamba model trained solve a simple coding task extrapolates codebases larger than those seen during training when trained on trajectories with interactive tool-use (Figure D Additionally, Mamba model trained execute long-form form multi-digit addition using pointer- based memory can generalize from 5-digit addition to 1,000-digit addition (Figure 2). We observe similar results on multiplication and on logical reasoning task, and more modest extrapolation on solving the Tower of Hanoi task (a task which proved to be difficult for reasoning models, see (hooaaee et al.) (2025) Taken together, our theoretical and experimental results highlight the potential advantage of using SSMs agents with interactive tool access, instead of using them standalone systems.
1.1 RELATED WORK
Chain-of- Thoughh and Scratchpad When solving problems that require reasoning. LLMs are known significantly benefit from generating CoT, detailing the step-by-step process required for solving the target task (Wei et al., 2022; Nye et al. 2021). Indeed, many datasets used for training models on mathematical problems include such CoT in the training data (Toshniwal al., 2024b a Cobbe et al./ (021). Theoretically, CoT is shown to improve both the expressive power of language models (Merill Sabharwal 2023) and their optimization and learnability (Wies et al.) 2022 Malach 2023). Additionally, it was shown that choices of CoT training data that "localize" the computation enable efficient learning and length generalization (Abbe et al. 2024). another work, using CoT that encodes the operation of Turing machine used to improve length generalization on various tasks (Hou et al. 20244] Hua et al. 2025). our work, we follow a similar approach for improving length generalization capabilities of language models. However, focus SSMs instead of Transformers, and study the effect of interactive tool-use for improving learning and generalization.
Emulations and Neural Turing Machines The goal of learning to execute general algorithms with neural networks has been discussed in various works. Abbe Sandon (2023) and Abbe et al. (22T)) show universality learning properties of poly-size neural networks trained by stochastic gra- dient descent. Graves et al. (2014) introduces the Neural Turing Machine (NTM), neural network that can simulate Turing machines and thus execute computable algorithms. NTMs were studied in different settings (Malekmohamad) Faradonbe et al.] 2020), with some improvements such as the Neural GPU (Kaiser & Sutskever| 2015), but were ultimately not widely adopted. Similar works suggested augmenting LSTMs (Hochreiter Schmidhuber 1997) with external stack or external tape (Delétang et al. 2022
Joulin & Mikolov 2015). We similar ideas to study algorithmic
3

---

Published as a conference paper at ICLR 2026
learning and length generalization capabilities of SSMs in the setting of tool-augmente interactive agents.
Length Generalization The problem of length generalization, training models short/simple problems and evaluating them on longer/complex instances, has been studied in many works. These works often focus on training Transformers on arithmetic or algorithmic tasks such sorting, copy- ing or multi-digit addition (Jelassi et al., 20233 Nogueira et al. (02T). Different works suggest various techniques for improving length generalization capabilities of Transformers, including vari- Ch2
ous choices of positional encodings and output format (Zhou et al., [024] Cho et al. [0224] McLeish
et al.
2024| Kazemnejad ad et al. 2023 Ruoss et al., 2023), scratchpads (Nye et al., 2021 Lee et al.
2023
Zhou et al. 2023 Abbe et al.
2024), architecture (Ontanon et al. 2021
Li & McClelland
2023)
mixing different tasks for "task hinting" (Awasthi & Gupta 2023) or using looped Transform-
ers (Fan et al. 2024). Some works aim to give scientific or theoretical explanation to the capability
and limitation of Transformers in extrapolating beyond the context length (Golowich et al. 2025
Zhou et al., 2023 Huang al., 2024 Bhattamishra al., (022). SSMs have been shown to display robust length generalization capabilities in certain cases. Gu & Dao (2023) demonstrate that Mamba achieves significantly better length generalization performance compared to Transformers on some tasks. Other works show that the length extrapolation of SSMs can be significantly improved with modifications to the model (Ben-Kish al., 2024) or the training pipeline (Ruiz & Gu 2025). In this paper. we study the length generalization of SSMs when trained on data with tool-use trajectories. We show that SSMs can achieve perfect length generalization in this setting on various tasks.
THEORY
In this section formally define the notion of long-form generation tasks: tasks that require gen- erating longer output sequences as their complexity increases. Following this, we define a family of functions that generalizes the class of SSMs, and theoretically analyze their limitation and capabili- ties in different tool-use settings.
Definitions and Notation. Fix some set and some distribution P over Z. For some subset S C 2, we denote by P(S) the probability mass of S under P, i.e. P(S)= == Przup [z E 8). For some function f 2 + Z', we denote by fPP the probability distribution of f(z) for P. For some set B, we denote by A(B) the set of probability distributions over B.
Definition 2.1. For some finite set Z and some distribution P over Z, we define the minimum support size of mass 0 a 1 for P to be the size of the smallest set that covers a probability mass: supp, (P) == min S| S C 2, P(S) a).
2.1 LONG-FORM GENERATION
Let E be a dictionary of tokens, and denote by h' the set of strings of tokens. Let X1, X2,... " be a sequence of input spaces, and let V/, . C n be a sequence of output spaces. We assume that the input and output spaces are finite, i.e. |,,, pal < oc for all n. Let D De.... be a sequence of distributions, such that D, is distribution over X,. Finally, let f E' + Σ' be some ground-truth function that satisfies, for all n, that ((x)) C Di. We think of the parameter n as a complexity parameter, and so the distribution Dn generates more complex inputs as n - co. We give the following definition of long-form generation tasks:
Definition 2.2. We say that f. {Da) is a long-form generation task with coverage 10, supp.. (P..)) is monotonically increasing with n and lim (J = cc.
Namely, we require that as the complexity n increases, the effective number of possible outputs (i.e.. outputs that have non-negligible probability mass) increases as well. We note that many natural long-form generation tasks indeed satisfy these conditions, for example: 1) Multi-Digit Addition (or Multúplication): D, is a distribution over strings of the form a + b = (or a x b =), where a, b uniformly random numbers with n-digits. The function if maps the input strings to the solution, e.g. f( + b =-=) ") where
+ b (or e = a b). 2) Sorting: D, is a distribution over n items,
"The condition that the support size is monotonically increasing makes the theoretical results slightly casier introduce, and holds for practically all reasonable long-form generation problems.
4

---

Published as a conference paper at ICLR 2026
maps to the sorted list of items. 3) Code Fixing: D, is distribution over python codes that have bugs that require changing n lines of code, f maps the code to the necessary changes.
2.2 GENERALIZED STATE SPACE MODELS
In this section, we follow similar definitions and notations as in Delassi al. (2024). We define state space to be some finite set S with |S| < Do. A generalized state space model (GSSM) is a (potentially probabilistie) function h Σ" + A(E)) defined by an initial state 8o E 8 and two rules; an update rule : S x Σ + S, and an output rule r S + (E)). Given some input æ E E" the function h generates sequence y E Σ". We define the state and the output of h at time recursively s.t. = (s if s [x] and = (s 1: Vt- lae if > |0z|, and we sample g/: ~ r((]a]+))) We terminate when an end-of-sequence token [EOS] E Σ is observed.
Note that any model that has fixed memory as function of the sequence length satisfies the defini- tion of GSSM. This includes common choices of recurrent models, such as: LSTM (Hochreiter &
Schmidhuber 1997). Linear Transformers (Katharopoulos et al. 2020). H3 (Fu et al. 2022). RetNet
(Sun et al. 2023). Mamba-I and Mamba-2 (Gu & Dao 2023] Dao & Gu. (024) and other variants
Yang et al. 2024b). Additionally, Transformers where all the attention layers have local (sliding window) attention with fixed window size, are also GSSMs. Other computational models that use Transformers to process fixed length sequences and update fixed-size "memory" vectors (Hutchins) et al. 2022) are also GSSMs. Transformers and hybrid-SSM models are not GSSMs, since their memory increases with the sequence length.
CoT, Single-Ture Turn and Interactive Tool-Use. We analyze multiple settings where the model can invoke reasoning and tool-use. We follow the popular ReAct framework (Yao et al., 2023) and let the model generate either thoughts, that capture the internal reasoning of the model, actions that are followed by observations from the environment. The thoughts and actions can be interleaved during the runtime of the model. We specify two types of actions: command actions, that are sent to a tool-oracle O that returns an observation following the execution of the command, and oulput actions that are simply tokens appended an output stream and do not result in an observation. The output stream captures the final response of the model which is then evaluated against the ground- Sruth Thoughts, commands and observations are placed between dedicated open/close tags (e.g- [THINK], [3 THINK). We define more formally the tool-oracle and the interaction protocol in Appendix A We analyze three settings for problem-solving agents. 1) CoT-only: The model is allowed to use only thoughts or outputs, but cannot issue commands or receive external observa- tions"] 2) Single-Turn Tool-Use: The model is allowed to issue a single command, followed by an observation, and then generate the output. The model can use thoughts before and after the tool call, and during the output generation. 3) Interactive Tool-Use: The model is allowed to use as many commands it needs, and freely interleave thoughts, commands and outputs.
2.3 LEARNING ALGORITHMS AND LENGTH GENERALIZATION
Fix some task f, (D, We now define training data distributions for learning the task f. We note that for many downstream tasks, is common to collect training data that contains CoT rea- soning and/or tool-use traces for solving the problem. We therefore allow the training distributions to contain a task cask-specific reasoning and tool-use trajectories. Given some trajectory z Σ', we denote by (out) the value of the output stream after execution of the trajectory. Formally, a training distribution for the task f, { j is a sequence of distributions [Pr) ] s.t. P, is a distribu- tion over X, x E* satisfying that: 1) Dn is the marginal distribution of Pn w.r.t. Xn and 2) For (ar, 2) Pa, with probability 1 it holds that (out) f(x) (i.e. the output stream at the end of generation evaluates to the correct answer).
A learning algorithm A is algorithm that, for some given length n. draws sample of size m from Pl..... P. and returns some hypothesis h : X' + A(E*) that, given an input problem, can gen-
'We focus on agents for solving input-output problems, where the task of the model is to generate some output given the input problem (e.g. question answering, coding, mathematical proofs etc.). This is a different setting from agent that performs actions and collects rewards, as in many Reinforcement Learning problems.
"This setting also includes the case where the model generates the output immediately, without using CoT. We let the algorithm choose freely how to sample from these distributions.
5

---

Published as a conference paper at ICLR 2026
erate reasoning and tool-use trajectory. We denote the output of A in this case by A(P1, Pn).
We say that A is a GSSM learning algorithm it always returns a GSSM. We define the error of w.r.t/ f for some complexity n by errn(h) = Pr [N (x:) # f(x)), with probability over x ~ D,, and the randomness of f. We now define length generalization of an algorithm:
Definition 2.3. We say that A achieves length generalization, for every e, E (0, 1) there ex-
ists some minimal complexity no and sample size m s.t. w.p. Z 1 1- 8 we have that hno A(PI..... Pro)) satisfies erT, (hmg) S for all n Z no-
Namely, we require that the algorithm returns a hypothesis with low-error on problems with arbi- trarily large complexity n, long as it sees "complex enough" inputs sequence in the training data (with complexity larger than no). This requirement may seem relatively strong. as we could expect that the error of the learned model would grow with the complexity of the problem. However, will show theoretically (and to some extent, empirically) that with a carefully constructed training data, achieving such "infinite" length generalization is possible.
2.4
MAIN RESULTS
In this subsection, we state the main theoretical results in the paper. We begin by showing negative result, stating that GSSMs cannot solve long-form generation tasks, if they operate in the CoT- only or single-turn tool-use setting. Following this, we show a positive result, proving that for any computable long-form generation task we can construct training data such that a simple learning algorithm achieves length generalization on the target task in the interactive tool-use setting.
GSSMs cannot Solve Long-Form Generation Tasks without Interaction. We begin by stating the negative result. The proof is relatively simple: since the model has fixed memory, and outputs are a function of the state of the memory, the model cannot generate all outputs as complexity grows.
Theorem 2.1. Let f be a long-form generation task over (Dr) with coverage parameter a e (0, 1). Then, for any CoT-only for all n 2 no the model h has error: erI,,h) 2 1 a.
Single-Turn GSSM h there exists some problem complexity no
The full proof is given in Appendix B An immediate implication of this result is that GSSM learning algorithms cannot achieve length generalization on long-form generation tasks without interaction.
GSSMs with Interactive Tool-Use can Length Generalize on Long-Form Generation Tasks. For some function / E¹ - Σr, we say that is computationally tractable if there exists Turing machine T s.t. for any x E , T begins with æ written on its tape, it halts with f() written on its tape. The following result shows that GSSM learning algorithm can achieve length generalization with interactive tool-use, given proper training data:
Theorem 2.2. There exists memory-tool oracle O and a simple GSSM learning algorithm A s.t. for any computationally tractable long-form generation task f. (h) there exists sequence of training distributions 1 for which A achieves length generalization in the interactive setting, with sample complexity m R&M (MgMM/// (where M is constant that depends the Turing machine for computing the function f).
To show the above result, we define a simple tool that allows read/write access to some external memory, using a pointer that can move left or right between the memory cells. Using this tool, we can simulate the operations of a Turing machine, where use the external memory as the tape of the Turing machine, use thoughts to track the state of the machine and commands to move the head and read/write symbols. Since the transition function of the Turing machine is defined for every pair of state and symbol, to prove that length generalization is achieved we show that, for large enough no. most of these pairs are seen in the training data. We give the complete proof in Appendix [
To conclude, the above results show that interactive tool-use is both necessary and sufficient for GSSMs to achieve length generalization on tractable long-form generation problems.
'The algorithm that we analyze is "simple" in the sense that it learss a function that operates using simple string-mu matching with the training data, similar to e.g. n-gram models. While this is not a "standard" learning algorithm, we believe that similar results can be obtained for more natural algorithms (e.g., gradient descent on simple RNN), the cost of making the analysis much more involved.
6

---

Published as a conference paper at ICLR 2026
Table 1:
Experimental results for synthetic tasks for different models. The notation n + mm(%))
means a model trained on length n achieves accuracy on length m (for the largest m s.t. p 2 5%).
Model
n X 1
n X
Logical Graph
Hano 7
Mamba
10 >K (100%)
10-+1 (100%)
10~>1K (98%)
8->12 (49%
LSTM
10-+500 (100%)
10-+100 (100%)
10~>1K (100%)
8->8 (100%)
GRU
10-+500 (100%)
10-+100 (100%)
10->1K (100 %)
8-8 (100%)
Pythia
10-+20 (79%)
10-+14 (12%)
10-> (5%)
8 (100%)
Mistral
10-+13 (25%)
10-+20 (33%)
10~>500 (9%)
8-8 (100%)
EXPERIMENTS
In this section we evaluate the length generalization capabilities of GSSMs and Transformer-based language models various tasks, including arithmetic,, reasoning and coding. We experiment with different choices of tools that allow read/write memory access, using either a pointer-basee memory access, search tool, or arbitrary bash commands for reading and changing files. We use both tasks where we synthetically generate the ground ttuth trajectory and tool commands, as well as a cod- ing task where we collect the trajectories from SWE coding agent. In our experiments, we largely follow a similar framework for ReAct agents defined in the previous section, where the model can in- terleave thoughts, outputs ('final answer" tokens), and commands that are followed by observations from the environment. In our experiments we use Mamba SSM (Gu & Dao 2023), LSTM (Hochre- iter& & Schmidhuber [197], GRU (Cho et al., 2014). Pythia Transformer (Biderman et al. (0023 and a Transformer with sliding window (local) attention based on the Mistral architecture (Jiang et al., 2023). In all experiments, we see that SSMs/RNNs achieve length generalization performance that is much better compared to Transformers. See Appendix C for experimental details.
3.1 ARITHMETIC TASKS
In the following set of experiments, we augment the model with a pointer-based memory tool that gives the model access to past tokens in the input/output context. In this setting, the model can execute the following commands: 1) initialize a new pointer. 2) move the pointer left right by single token and 3) read the token under given pointer. By default, a new pointer is initialized to the first token position of the input context. The thoughts and outputs are therefore accessible by the pointers (if they reach beyond the length of the input), but commands and observations are discarded (i.c., they are not appended the context memory and cannot be read by the pointers). We give a detailed description of how thoughts, commands and outputs are specified in Appendix D.I The final answer is written in the output stream at the end of the generation.
appended to the context, and
We train the model using the standard next-token prediction objective with saccher-orriigg while masking from the loss the imput question and the observations (the outputs of a read operation, which will be generated by the memory tool). For the training data, we construct synthetic trajectories that simulate the desired algorithm, and train the model to exactly execute the algorithm required for solving the problem using the memory- memory-tool interaction.
Multi-Digit Addition. For this task, we train the model to perform multi-digit addition. We fix some maximal training length n, and for each training example we sample uniformly ni, m2 ~ [I....,n), n.. then sample two numbers 22 where z; is uniformly random nę-digit number. We construct a training example with the trajectory for solving #1 + 22, essentially mimicking the long addition algorithm (see Appendix D.3 for details). For evaluation, we choose n' 2 and evaluate on addition of two n' -diiit numbers. In evaluation, we measure the accuracy of exact-recovery of the trajectory and the final answer (i.e., we measure the probability of generating a solution that exactly matches the desired algorithm). Figure (right) shows the results of this experiment. We observe that Mamba and LSTM trained on 5-digit demonstrations learn to perfectly perform 1,000-digit addition (we did not measure the accuracy beyond this). A Transformer trained in the same setting fails to extrapolate. Additional ablations, such as training with no CoT, no tool-use
Appendix[B<
Due to the sensitivity of the Hanoi experiments to the initialization seed, use 10 seeds for Mamba and Pythia and 3 seeds for other models. Best performing seed is reported. See Appendix D.4 for more details.
7

---

Published as a conference paper at ICLR 2026
and single-turn tool-use, result in little to length generalization, and are discussed in Appendix D.7 We also compare Mamba to two alternative architectures on the addition task: Hybrid-Manne which interleaves Mamba layers with attention layers, and Recurrent Memory Transformer (RMT) (Bulatov et al. 2022), which augments a Transformer with learnable memory tokens and segment- level recurrence. We find that Hybrid-Mamba performs similarly Mamba, while RMT does not achieve meaningful length generalization. See Appendix D.9 for details.
Multi-Digit Multiplication. For this task we use the same pointer-based memory tool described above for learning the multiplication algorithm. In this task, we increase the length of only the first operand, keeping the second operand fixed. Specifically, we fix some maximal training length n, choose n1 [I....,n} n to be the length of the first operand and choose 722 ~ {1. 2} to be the length of the second operand (i.e. we multiply an n1-digit number by I-digit or 2-digit number). We sample 21, 22 where 2; for solving multiplication, for some n' 2 n, and report exact recovery accuracy. We train different SSMs/RNNs and Transformers where first operand has S 10 digits, and evaluate on multiplications of numbers with up to 1,000-digits (Table[T) I.) Here too we see that Mamba models maintain high accuracy when evaluated on numbers that have orders of magnitude more digits than in the training data (also Appendix D.6 for ablations on training steps and maximum number of digits seen during training).
uniformly random Ri-digit number, and construct the trajectory 22 (see details in Appendix D.3). We then evaluate nt' x 1 and n' x 2 digit
Task Mixture. We examine whether co-training a primary task with an auxiliary task that shares related computational structure yields synergistic benefits (Awasthi & Gupta 2023). Our exper- iments indicate that such co-training improves the length generalization of the primary task under limited training budgets. In experiments, the primary task is multiplication (n-digit x 2-digit), co-trained with addition (n + digits) as an auxiliary task. Both tasks share structural similari- ties when expressed as sequences of tool calls. The training distribution for multiplication contains samples up to 20 digits. We compare the accuracy as a function of test length for various training budgets (250, 500 or 800 steps) and various choices of task mixtures (see Appendix D.S.. We ob- serve that under limited budgets (250 steps), introducing auxiliary addition samples yields minor improvements. At intermediate budgets (500 steps), the benefit becomes more pronounced, with certain weights extending generalization to much larger However, with sufficient training (800 steps), all settings converge to strong generalization.
3.2 ALGORITHMIC/REESONNSS TASKS
We next turn to evaluate the tool-use paradigm on tasks that test certain "reasoning" capabilities.
Tower of Hanoi. This task is based on the popular puzzle game, which was also recently used for testing reasoning capabilities of frontier LLMs, showing that they struggle to solve this task as complexity increases (Shojaee et al. 2023). In our setup, we randomly sample (without replacement disks of sizes E {1...., 100). These disks are placed on the first rod (labeled A), ordered from the largest to the smallest, with rods B and C being empty. The input to the model is the list of disks. which captures the initial state of the game. The model then needs to output a sequence of valid moves that result in placing the pile on rod C. We use the same pointer-based memory tool as in the previous experiments, and train the model trajectories with up to n disks, evaluating on larger n' (see Appendix D.3). In this experiment we observe more limited length generalization (Table II. but note that unlike other experiments, here the length of the output increases exponentially with n.
Logical Graph. In this task, we construct a prrected-acyclic computation graph with n nodes. The graph has l input nodes (for some fixed k). and each internal node node computes a Boolean operation (AND/OR) on one or two input variables or their negations. We construct the graph by iteratively adding new internal nodes and randomly choosing their Boolean operation and their connectivity to existing nodes in the graph. We take the last node that is added to be the output node. All nodes are randomly labeled, and the model receives the graph structure and an assignment for the imput variables as a python code (see Figure 4 In this task, instead of using the pointer-base memory tool as in previous tasks, we use search tool: the model can issue command fi find(r), and gets a list of all occurrences of the pattern x in the context. As before, all thoughts and outputs generated by the model are appended to the context and therefore searchable in future iterations.
8

---

Published as a conference paper at ICLR 2026
We fix k = 3 and train the model on trajectories for solving this problem for graphs with < 10 nodes. We evaluate on graphs with n' 2 n nodes, and report the exact-match accuracy in Table We observe that Mamba and recurrent models extrapolate to graphs with = 1K nodes.
3.3 CODING TASK
For the previous tasks, we trained models "from scratch" on synthetic trajectories that invoke tool use for solving arithmetic and algorithmic problems. This allowed us to demonstrate the length generalization capability of SSMs equipped with tool-use in a clean and controlled setting, resulting in perfect recovery of the underlying algorithm in many cases. We now turn study extrapolation of tool-use agents in more realistic coding setting. Importantly, this setting will allow us to go beyond programmatically generated trajectories and collect trajectories from an existing SWE coding agent. This demonstrates that our results and observations can also be applicable in settings where the underlying algorithm/method for solving the task are not known or well-specified.
Our task will be fixing "bug" in given codebase. To construct the codebase we generate python functions, each function saved
separate python file. The functions form a dependency graph, with one root function called main (stored in main. py). Each function declares variables (named vO, vl,...,, v9), gets some variables as inputs and passes some variables to other functions it im- ports. We generate this codebase by randomly generating dependency graph, iteratively adding nodes (functions) to this graph and connecting each node to existing nodes, where each edge rep- resents a parent function importing a child function. Function names are randomly selected from f 9999, except for the last function added to the graph which is called foo. We then ran- domly assign variables and print them and/or pass them from parent functions to child functions. The code always has the following "bug": there is special variable v1C that is declared in main py and is called in foo. . py without properly passing from ma in.. In order fix the code, we need to pass the variable v10 through all the dependency paths from ma in to foo (ideally without changing other functions, though we do not enforce this). See Figure for an illustration.
We start by running a coding agent and collecting its trajectories when attempting to solve this code- fixing task, as we are varying the number of functions n in the codebase (choosing n E [4..... (66)). We use three types of agents for generating trajectories, (illustrated in Figure D: 1) Single- Turn Agent: Hard-coded agent that prints all the files and immediately generates the correct code edits. 2) Interactive Agent: Hard -orddd agent that iteratively runs the code, resolves the issue in up to 3 files, then runs the code again, and keeps going until the code runs without errors. 3) Distillation: An agent based -WE-ageet-IM--3 323 (Yang et al. (025), capable open-source coding model that we couple with mini-SWE-agee gives the model access to the code through bash commands. We instruct the model fix the bug in the code, specifically telling it what the bug is and how it should fix it (pass the variable v1 O from main to foo). See the full prompt and further details in Appendix E. We observe that while this task is relatively simple, the model's 's performance degrades as the complexity (number of functions) in the codebase increases (see statistics in Appendix F.. We therefore filter the trajectories to include only trajectories that correctly fixed the code, and also filter for short trajectories (shorter than the average lenghh for a given size n)).
Yang et al. 2024a) as a simple agent environment which
After collecting around 100K trajectories from each coding agent, finetune two comparable models on these trajectories: Pythia-1.4B (Transformer-basee Mamba-1.4B (Gu & Dao 2023), both pretrained on The Pile (Gao et 2020). We train both models with context length 8,192, on codebases of up to 16 functions (if the trajectory is longer than the context length, we train only on the first 8,192 tokens). We then evaluate both models on codebases of different sizes, letting the models generate beyond the context length | We measure the probability of correctly fixing the code (using the same environment used for collecting the trajectories). As shown in Figure]] [ we observe that for codebases with small number of functions, both Transformer and Mamba models perform well in all settings. Notably, the Transformer-basee model outperforms the Mamba SSM for small codebases in the agent distillation setting, achieving over 90% pass rate. However, for larger codebases, beyond the training distribution (both in terms
model, Biderman et al. (2023)) and
https: / / thub. com/ Swk- agent swe- sww-agent..ttteaeeeeeeeooeeoe
Wee experimented with applying RoPE scaling when using the Transformer beyond the training context length, both in finetuning and evaluation, and observed mixed results. We report the accuracy for the hest choice (with or without RoPE scaling) in each setting.

---

Published as a conference paper at ICLR 2026
of number of functions and trajectory length), we see that the Mamba model maintains much better accuracy as the complexity increases when trained to imitate interactive agents (agents 2 and 3), but fails complex codebases when trained in the single-turn setting (agent 1). This finding aligns with our theoretical results, and also matches the previous synthetic experiments.
3.4 LONG-CONTEXT NATRUAL-LANGUAGE
Long-Context Natural Language
1.0
To further compare the length generalization of SSMs and Transformers in long-context tool-use setting, we train our models on solving a task that involves natural language and long-context reasoning. In particular, we adapt a task from the Oolong benchmark (Bertsch et al. 2025), a very recent bench- mark that proposed a novel set of Long-con contecontext tasks that prove to be hard for frontier models. We choose a representative task from this benchmark, in which the model is presented with a dataset of natural-langue
0.9
Accuracy
0.8
0.7
0.6
Mamba (1.4B)
Pythia (1.4B)
text examples [T]] each with par- ticular date, and needs to answer questions of the form: "In the above data, was a question of category more common, less common, or the same frequency before date Y, as compared to after date Y?". We train the models on hard-coded agent trajectories that use a simple tool call that retrieves the next example the dataset (starting from the first example), along with thought tokens for analyzing each example. We finetune 1.4B Pyhita Transformer and Mamba SSM (similarly to the coding experiment) on problems with up to 5 data examples, and compare the performance of the models when tested on larger num- ber of examples. As shown in Figure 3 here too we observe that while both models have similar performance in-distribution, the Mamba SSM extrapolates better when evaluated on longer datasets.
0.4
10
15
20
25
Number of Examples
Figure 3: Accuracy of Pythia and Mamba trained on in-context datasets with 5 examples, evalu- ated on up to 25 examples.
4 CONCLUSION AND DISCUSSION
We started this work by comparing two families of models for long-form generation: Transformers and SSMs. Transformers inefficient for long-context and long-form generation, their compu- tational complexity scales quadratically with the sequence length. SSMs, on the other hand, offer linear scaling of compute but, as we showed, cannot accurately solve long-form generation tasks (without tools). This demonstrates a clear trade-off between efficiency and accuracy that seemss to be inescapable. Indeed, several works have observed that SSMs are inferior to Transformers in var- ious tasks that require memorization of long sequences (ellassi et al. (2024): Waleffe et al. (2044)
On the positive side, we show that in the agentic/tool-use setting. SSMs can leverage tools to over- come their memory bottleneck, thus offering efficiency. accuracy. and generalization to longer se- quences. In hindsight, SSMs seem to quantities of content, which SSMs can parse efficiently, and also involve multi-turn interactions that can quickly overflow the context of a standard Transformer. However, it seems that there is little work building SSM-based agents, and thus their evaluation is restricted to the "standalone" set- ting, where they are inherently limited. We do not believe this is due to any inability of SSMs to learn tool-use behavior. For example, while Mistral's Haamba-Codestrll-255 naively support function calling. it is able to achieve comparable tool-use performance to several dunction-calling-enable transformer- should encourage the development a tool-based SSMs that operate in various agentic settings, such coding, search or reasoning. This application could potentially unlock the full capabilities of these models, making them competitive with, or even superior to, Transformer-bast agents.
natural fit for tool-use settings: tools often generate large
model does not
models (Appendix F.. We therefore believe this work
Finally, we want emphasize that our work is one of the first analyze the performance of language modeling architectures in a system, rather than as standalone models. To some extent, our analysis shows that certain architectures can be "weaker" when operating standalone, but in fact perform better when incorporated as part of system. Since L.MSs are now rarely used standalone tools, we believe that this aspect of language modeling deserves more attention and focus the field.
We use examples from the Yahoo Answers Topic Classification dataset (Zhang et al. 2015
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Emmanuel Abbe and Colin Sandon. Polynomial-time universality and limitations of deep learn- ing. Communications on Pure and Applied Mathematics, 76(11)33433-349, 2023, doi: htpss:
URL httpp:://oniieellbiittt wttp:/arwwwweeyycoo///ooo/
com/ www.a//oi/ae abs/
10. 10.1002/cpa. 22121
Emmanuel Abbe, Pritish Kamath, Eran Malach, Colin Sandon, and Nathan Srebro. On the power of differentiable learning versus PAC and SQ learning. In Advances in Neural Information Process- ing Systems, volume 34, 2021.
Emmanuel Abbe, Samy Bengio, Aryo Lotfi, Colin Sandon, and Omid Saremi. How far can trans- formers reason? the globality barrier and inductive scratchpad. Advances in Neural Information Processing Systems, 37:27850-27895. 2024.
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Archi- tectures and algorithms. arXiv preprint arXiv:2401. 12973, 2024.
Pranjal Awasthi and Anupam Gupta. Improving length-generrlization in transformers via task hinting, arXiv preprint arXiv:2310.00726, 2023. URL htttp:://aax...oo////aaaaa/
00726
Assaf Ben-Kish, Itamar Zimerman, Shady Ausseissein, Nadav Cohen, Amir Globerson, Lior Wolf, and Raja Giryes. Decimamba: Exploring the length extrapolation potential of mamba. arXiv preprint arXiv:2406. 1....145 14528, 2024.
Amanda Bertsch, Adithya Pratapa, Teruko Mitamura, Graham Neubig, and Matthew R. Gormley. Oolong: Evaluating long context reasoning and aggregation capabilities, 2025. URL https:
// aaxxv.orq/aba. http://abb/2511....gm
2511.02817
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. arXiv preprint 2xXiv:221..12316 2022.
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O' Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, et al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430, PMLR, 2023.
Aaron Blakeman, Aarti Basant, Abhinav Khattar, Adithya Renduchintala, Akhiad Bercovich, Alek- sander Ficek, Alexis Bjorlin, Ali Taghibakhshi, Amala Sanjay Deshmukh, Ameya Sunil Maha- baleshwarkar, et al. Nemotron-h: A family of accurate and efficient hybrid mamba-transformer models. arXiv preprint arXiv:2504.0360 2025.
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35:1 1079-110991, 2022.
Hanseul Cho, Jaeyoung Cha, Pranjal Awasthi, Srinadh Bhojanapalli, Anupam Gupta, and Chulhee Yun. Position coupling: Improving length generalization arithmetic transformers using task structure. Advances in Neural Information Processing Systems, 37:22233-22315, 2024.
Kyunghyun Cho, Bart Van Merriënboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint 14X.vv1409.155: 2014.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009. 14794, 2020.
Karl Cobbe, Vincet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reüchiro Nakano, et al. Training verifiers to solve math word problems. arXiv preprint 1xXv:2210014468. 2021.
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality. arXiv preprint arXiv:2405.2 1060, 2024.

---

Published as a conference paper at ICLR 2026
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in neural information processing systems, 15.16344-16359, 2022.
Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Kevin Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Legg, Joel Veness, et al. Neural networks and the chomsky hierarchy. arXiv preprint rrXiv:2207.02098, 2022.
Ying Fan, Yilun Du, Kannan Ramchandran, and Kangwook Lee. Looped transformers for length generalization. arXiv preprint arXiv:2409. 2024.
Daniel Y Fu, Tri Dao, Khaled K Saab, Armin W Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212. 14052, 2022.
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang. Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
Noah Golowich, Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. The role of sparsity for length generalization in transformers. arXiv preprint arXiv:2502. 16792, 2025.
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv: 1410.540, 2014.
arXiv preprint
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.007/22 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, arXiv preprint arXiv:2111.00300 2021.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. Deepseek-rl: Incentivizing reasoning capability in llms via reinforcement learning. arXiv preprint 1:Xiv:2501.12945 2025.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):
1735-1780, 1997.
Kaiying Hou, David Brandfonbrener, Sham Kakade, Samy Jelassi, and Eran Malach. Universal length generalization with turing programs, arXiv preprint arXiv:4407.03310, 2024.
Zhouqi Hua, Wenwei Zhang, Chengqi Lyu, Yuzhe Gu, Songyang Gao, Kuikun Liu, Dahua Lin, and Kai Chen. The imitation game: Turing machine imitator is length generalizable reasoner. arXiv preprint 23X1v:2507.13332, 2025.
Xinting Huang, Andy Yang, Satwik Bhattamishra, Yash Sarrof, Andreas Krebs, Hattie Zhou, Pree- tum Nakkiran, and Michael Hahn. A formal framework for understanding length generalization in transformers. arXiv preprint arXiv:2410.. 1....22220 140, 2024.
DeLesley Hutchins, Imanol Schlag. Yuhuai Wu, Ethan Dyer, and Behnam Neyshabur. Block- recurrent transformers. Advances in neural information processing systems, 35:33248-33261 2022.
Aaron Jaech. Adam Kalai, Adam Lerer, Adam Richardson, Ahmed El-Kishky, Aiden Low, Alec Helyar, Aleksander Madry, Alex Beutel, Alex Carney, et al. Openai ol system card. arXiv preprint arXiv:2412.16726, 2024.
Samy Jelassi, Stéphane d' Ascoli, Carles Doniih,-Enrrich, Yuhuai Wu, Yuanzhi Li, and François Charton. Length generalization in arithmetic transformers. arXiv preprint arXiv:2306. 15400, 2023.
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after Trans- formers are better than state space models at copying. In International Conference on Machine Learning, pp. 21502-21521. PMLR, 2024.
12

---

Published as a conference paper at ICLR 2026
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023. URL https:
//arxivv.rrg//bb///111 06820
Armand Joulin and Tomas Mikolov. Inferring algorithmic patterns with stack-augmenn recurrent nets. Advances in neural information processing systems, 28, 2015.
Eukasz Kaiser and Ily Sutskever. Neural gpus learn algorithms. arXiv preprint arXiv:1511.0822. 2015.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are mns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165, PMLR, 2020.
Amirhossein Kazemnejad, Inkit Padhi, Karthikeyan Natesan Ramamurthy, Payel Das, and Siva Reddy. The impact of positional encoding on length generalization in transformers. Advances in Neural Information Processing Systems, 36:24892- 2023.
Nayoung Lee. Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos. Teaching arithmetic to small transformers. arXiv preprint arXiv:2307.03381, 2023.
Yuxuan Li and James McClelland. Representations and computations transformers that support generalization on structured tasks. Transactions on Machine Learning Research, 2023.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, al. Jamba: A hybrid transformer- mamba language model. arXiv preprint arXiv:2403. /9887, 2024.
Junyu Luo, Weizhi Zhang, Ye Yuan, Yusheng Zhao, Junwei Yang, Yiyang Gu, Bohan Wu, Binqi Chen, Ziyue Qiao, Qingqing Long, al. Large language model agent: A survey methodology, applications and challenges. arXiv preprint arXiv:2503.21460) 2025.
Eran Malach. rego-segressive next-token predictors are universal learners.
arXiv preprint
arXiv:2309.0mm 06999, 2023.
Soroor Malekmohamadi Faradonbe, Faramarz Safi-Esfahani, and Morteza Karimian-KKlishadrooee A review on neural turing machine (ntm). SN Computer Science, 166):333, 2020.
Sean McLeish, Arpit Bansal, Alex Stein, Neel Jain, John Kirchenbauer, Brian Bartoldson, Bhavya
Kailkhura, Abhinav Bhatele, Jonas Geiping, Avi Schwarzschild, et al. Transformers can do arithmetic with the right embeddings. Advances in Neural Information Processing Systems, 37: 108012-108041, 2024.
William Merrill and Ashish Sabharwal. The expressive power of transformers with chain of thought. arXiv preprint arXiv:2310.07923, 2023.
Rodrigo Nogueira. Zhiying Jiang, and Jimmy Lin. Investigating the limitations of transformers with simple arithmetic tasks. arXiv preprint arXiv:2 13019, 2021.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowyez, Maarten Bosma, David Luan, et al. Show your work: Scratchpads for intermediate computation with language models. 2021.
Santiago Ontanon, Joshua Ainslie, Vaclav Cvicek, and Zachary Fisher. Making transformers solve compositional tasks. arXiv preprint arXiv:2108.04333 2021.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kang- wook Lee, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks. arXiv preprint arXiv:2402.0422 04248, 2024.
Shishir G Patil, Huanzhi Mao, Fanjia Yan, Charlie Cheng-Jie Jie Ji, Vishnu Suresh, lon Stoica, and Joseph E Gonzalez. The berkeley function calling leaderboard (bfcl): From tool use to agentic evaluation of large language models. In International Conference Machine Learning.
13

---

Published as a conference paper at ICLR 2026
Haohao Qu, Liangbo Ning, Rui An, Wenqi Fan, Tyler Derr, Hui Liu, Xin Xu, and Qing Li. A survey of mamba. arXiv preprint arXiv:2408.01129, 2024.
Ricardo Buitrago Ruiz and Albert Gu. Understanding and improving length generalization in recur- rent models. arXiv preprint rrXi:22507.02782, 2025.
Anian Ruoss, Grégoire Delétang, Tim Genewein, Jordi Grau-Moya, Róbert Csordás, Mehdi Ben- nani, Shane Legg, and Joel Veness. Randomized positional encodings boost length generalization of transformers. arXiv preprint arXiv:2305.16845 2023.
Parshin Shojace, Iman Mirzadeh, Keivan Alizadeh, Maxwell Horton, Samy Bengio, and Mehrdad Farajtabar. The illusion of thinking: Understanding the strengths and limitations of reasoning models via the lens of problem complexity. arXiv preprint arXiv:2506.0664 06900, 2025.
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.08622 2023.
Shubham Toshniwal, Wei Du, Ivan Moshkov, Branislav Kisacanin, Alexan Ayrapetyan, and Igor Gitman. Openmathinstruct-2: Accelerating ai for math with massive open-source instruction data. arXiv preprint arXiv:2410.01556/ 2024a.
Shubham Toshniwal, Ivan Moshkov, Sean Narenthiran, Daria Gitman, Fei Jia, and Igor Gitman. Openmathinstruct-:: A 1.8 million math instruction tuning dataset. Advances in Neural Informa- fion Processing Systems, 37:34737-34774, 2024b.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, al. An empirical study of mamba- based language models. arXiv preprint arXiv:2406.00 07887, 2024.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837 2022.
Noam Wies, Yoav Levine, and Amnon Shashua. Sub-task decomposition enables learning in se- quence to sequence tasks. arXiv preprint arXiv:2204.02892, 2022.
John Yang, Carlos Jimenez, Alexander Wettig, Kilian Lieret, Shunyu Yao, Karthik R Narasimhan, and Ofir Press. SWE-agent: Agent-computer interfaces enable automated software engineering. In The Thirty-eighth Annual Conference on Neural Information P Processing Systems, 2024a. URL
https: /
wwwg/bss...........o 2405 15793
John Yang, Kilian Leret, Carlos E Jimenez, Alexander Wettig, Kabir Khandpur, Yanzhe Zhang, Binyuan Hui, Ofir Press, Ludwig Schmidt, and Diyi Yang. Swe-smith: Scaling data for software engineering agents. arXiv preprint arXiv:2504.21799 2025.
Songlin Yang, Jan Kautz, and Ali Hatamizadeh. Gated delta networks: Improving mamba2 with delta rule. arXiv preprint arXiv:2412.06446, 2024b.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. Advances in neural information processing systems, 37:115491-11552. 2024c.
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. In International Conference on Learning Representations (ICLR), 2023.
Asaf Yehudai, Lilach Eden, Alan Li, Guy Uziel, Yilun Zhao, Roy Haim, Arman Cohan, and Michal Schheur. Survey on evaluation of llmed agents. arXiv preprint arXiv:2503. 16416, 2025.
14

---

Published as a conference paper at ICLR 2026
Xiang Zhang. Junbo Zhao, and Yann LeCun. Character-level convolutional networks for text clas- sification. Advances in neural information processing systems, 28, 2015.
Hattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin, Omid Saremi, Josh Susskind, Samy Bengio, and Preetum Nakkiran. What algorithms can transformers learn? study in length generalization. arXiv preprint arXiv:2310. 160.220,,,,,,, 2023.
Yongchao Zhou, Uri Alon, Xinyun Chen, Xuezhi Wang, Rishabh Agarwal, and Denny Zhou. Trans- formers can achieve length generalization but not robustly. arXiv preprint arXiv:2402.09371, 2024.
A MORE DEFINITIONS
Here give more complete and formal definition for models with tool-use. We start by defining a tool-use oracle, which will receive tool-use commands and will return an observation that corre- sponds to the execution of the command. This oracle will be stateful, meaning that its responses can vary depending on the memory of the oracle, which can be updated based on the commands it re- ceives" Let .M be some set which will correspond to the of memories of the oracle. We denote by Mr EM the memory of the oracle after receiving / commands, and let Mn be the initial memory of the oracle. Importantly, we let the initial memory of the oracle depend on the input (e.g,, the input can be stored in the memory of the oracle). For some memory Mt E .M, we define OM, E" + E" to be the mapping from tool calls to observations, given memory Mt.
We augment the dictionary with additional tokens: [TOOL], [TTOL), [OBS], NOBS), [THINK], [[INKK,[[TTINII [[HHIN E E. At any point in the generation, the model can issue a call tool by generating a sequence of the form [TOOL], z, [XTOOL]. for some z E >" which encodes the tool command. The command then appended to the context of the model, with the format: [OBS), M. (z), [OBS). The model can also generate thoughts/reasoning, generated as a sequence [THINK], z, [[HHNK]] All other tokens (tokens that are not col-commands, observations or thinking tokens) considered output tokens, and are appended to the output stream.
passed to the tool oracle, and the resulting observation that is
In the CoT-only setting, the model is only allowed use thinking and output tokens. In the single- turn setting. the model can issue a single tool command, and can start generating the output after receiving the observation from the command (but can think before, after and during the output generation). In the interactive setting, a model can issue a tool call at any point in the generation, possibly interleaved with output tokens and tool commands. When evaluating the output, we ignore all tool commands, observations and thoughts and only consider the output stream at the end of generation.
B PROOFS
B. PROOF OF THEOREM 2.1
Before we introduce the complete proof of the theorem, we prove simple version of the Theorem for the case where the output rule the GSSM is deterministic. This proof simpler than the more general stochastic setting, and captures the key principles for this result.
Proof of Theorem 2.7 for deterministic GSSMs. Let S be the state space of h. By definition, there exists some ng such that for all Z ny it holds that supp.. (D,)) > (S]. Now, fix some n 2 no. and let A C Σ* be all the possible outputs of h. Note that the output is determined by the state of h before the output tokens are generated. Therefore, |A| S S|| < sJ and so we have
[x E A] = Pr D>DD
Pr
[y E A = (A) <a
(DD)
For example, the oracle can receive command for updating the content of file disk, which will affect its memory and hence future requests for reading the contents of the changed file.
15

---

Published as a conference paper at ICLR 2026
Finally, we have
Pr wnDn [f(z) [ A E A]
(())
+ Pr [f(x) ( h()))f(())
Pr A] s (n(A)
<a
and so err, (h) 2 a.
We now show the proof of the theorem for the general (stochastic) setting:
Proof of Theorem 2.7 Let S be the state space of h, and we assume that h operates either in the CoT-only or the single-turn setting. Denote by (xx) the state of the model before generating the first output token. The model h can generate thinking tokens and/or single tool command before generating the first output token, and therefore U(x) is a random variable that depends on æ and the randomness of h Let R(s) be the distribution of outputs (i.e. values of the output stream) generated by the model h it at state s before generating the first output token. Treating h (nou)) (x) as a random variable over outputs, we note that depends only on the state after parsing æ, and therefore ( (x) RUU (x)). Additionally, we denote the conditional distribution over outputs induced by (outt with p(y ar). Again, since the future generation of the model depends only on the state, we have p = p (+)).
Now, by definition, there exists some no such that for all no it holds that supp, (f(),)) > |S]. Fix some n 2 no. Fix some E S. Let be an output with maximal probability under the distribution D.. conditioned on the event that Uxx) = s:
V arg Pr [f(x) y))(x)) s]
w>
Denote by A the set of maximal probability outputs = {ys : S S}. Note that |A| S |S| < (fD,)) and so have
[f(x) E A|= D
Pr
Pr ly E A = ((D))(A) <a g>r((,))
Observe that:
(1))) = =]= Σ
[[] (
Σ
- =
yE
Σ
[ R()==y (x)) =
DE
Σ p(yls) Pr [f(x) y(U(x) s]
MEV.
Σ p(y|s) Pr(f(z) y. (x() s]
vEV.
S Pr[f((=) = y, Uix) = s]
Where in the 4th equality we the fact that the variables independent. Therefore, we have:
Pr
[f(x) = [[(x)]
Σ
Pr[f(=) = h(z))U(()) = s] Pr(U(z) =]
D.
SES
Σ Pr[s(w)= = y. Uxx) s) Pr(U(()) =
<
Σ Prif(z) (U())) = (x) = = = Pr(f(() E A] <a
SES
and so (hr 2 a.
We assume the oracle (e.g., the environment) is deterministic, but we can be easily extend this result to capture stochastic oracle.
16

---

Published as a conference paper at ICLR 2026
Proof of Theorem 2.2 First, let define the oracle O The memory of the oracle at iteration t
holds a sequence tokens m, E E", and additionally some index is E N. At first iteration, we set mg and in 0. The oracle O accepts the following commands:
read: outputs the i,-th token in my If i |mel, output [EOS].
write a: updates the i, token of m, to be a.
move left, move right: adds/subtracts from it.
Next, we describe the training distributions Pn. Since is tractable, there exists some Turing machine T that computes f. By definition, the machine halts for every input, and we can as- sume w.l.o.g. that it halts when the head is position Let Q be the (finite) set of states of T. and let Yo be the initial state. We assume the dictionary Σ contains the following symbols: {0, 1, [STATE] I [STATE]] E L. For each state E Q, we define the encoding of the state enc(g) = [STATE] 2 (SSATE)) where Zn E {0, (l is a binary encoding of the state. Then, for some input æ D., we construct CoT of æ. denoted by F(x), that will capture the "trace" of the machine T:
The sequence F(x) begins with: (aan:///eee////HNNN//OOOOOOOo/
In each step of the Turing machine processing æ, we add to F(x) the sequence:
where q is the current state, and a' the next symbol write when reading o in state q- Additionally, we add [TOOL]move.lerf [000sm/oeeeeff((TOOO [(TOO) if the machine moves the head to the left, and otherwise [TOOL]]ove_rig])) (it.eeeorrrtttttto. [[OOO.].
When the machine reaches a halting state, for every i = 1... [f()) we add:
[TOOL]move
Note that since the machine computes f(x) will be written on its tape when it reaches a valid state. Therefore, it is easy to verify that the memory of the oracle O at step t will hold the state of the tape and the correct position of the head, and that all the tool observations will be correct. Finally, Da and F(x:) together define the distribution Pn, for all n, and it is indeed training distribution for the task (since the non-tool tokens after the [ANS] token correspond to the correct output fx))
Next, we will show that a simple tool-SSM algorithm can achieve length generalization on this task. Let {ax) F(x:)):..- ... (æm. F(x be the set of examples observed by the algorithm. Let A be the set of all pairs of state encodings and symbols that appear together in F(x:)) in some æ,;
À : ((q,0) 3i, [UUNIN/CCK(/II/I////oo/////////O
F(æ,))
Note that for every (q.a) E À there is a single symbol o' and single command d e rriht) and a single state o that follow (q.o) in the trace (corresponding to the operation of the Turing machine). Let R be the function mapping (q, (q to (g', o', d). Note that both À and R can be encoded with fixed (finite) memory. Therefore, we define GSSM h ARR that generates tokens follows: Immediately after the input, generate: [HIXXXXXXCC///IHHXX/OOOOOooo///O
Following each response to read command, generate:
[THINK]enc(q]) HwTG:///M(K/YNNNN//OOOOOOooo. http:///O/DD/tte///ooooooooo/o//
where o/, o'''', d = R(g, a), for the returned by the tool oracle.
17

---

Published as a conference paper at ICLR 2026
When halting state is reached, generate the sequence:
and following the observation [OBS]][]OBS], output a (if a = [EOS] we halt the genera- tion).
If at some point we observe pair a A, output [EOS].
Denote by A(x) S Q x E the set of state-symbol pairs observed by T when processing æ. Its easy to verify that for every s.t. A(x) E A, the GSSM h R will exactly recover F(x). Therefore, the following lemma suffices for proving the theorem:
Lemma B.1. Fix some e, E (0,1). There exists some no s.t. with sample size m = no Q [ ( /5)/e w.p. at least 1 over sampling from PI; PI.... Pno it holds that:
A(x)) CA >1-6
Vn no Pr D
Proof. For every pair of symbols Σ and state q Q denote pala, Pra-D, ()) A(x)] the probability over sampling x ~ Dn that the machine TT reads symbol a while it is in
state q, when processing æ. Let M Q |id Denote:
A. = (q,0) E EQ xQ s.t. suppa(g,)) 2 2c/M]
Now, for every 9. E A.. let no(g, 0) be the minimal n s.t. (=) 2 c/M. Let no = max(no(q, (o(ono(lg))))ne A,
A, Let m no.M (M/() and we will sample m' m/no M Mog(M/6) examples from each of
Fix some (g, E Ac.
Claim: w.p. at least 1 - 8/M we have (1, E Ā.
Proof: Note that nolg, o) S no. and therefore we sample m' examples from Dna(g.o). Let p = "Tee-Proto,,)(44 E A(x)] and by definition p 2 e/M. Therefore, for the m' samples we draw, the probability that do not encounter (g.o) o) in any of the traces is at most (1 p) (1 e/M)m' S exp(-m'e/MM 8/M.
From the above claim, using the union bound, we get that w.p. least Assume this holds, and fix some 2 no- For every (g,o) E Q Σ A. it holds that Præ-D,, (( r) E A(x)] S e/M. From the union bound, the probability over æ D, that there exists some (q) & A. C A s.t. (g,e) E A(x) is at most IQ x Σ A. 4/M s Therefore, the required follows.
we have A Å.
From the above lemma the proof of Theorem 2.2 follows.
C ARCHITECTURE AND TRAINING DETAILS
We train the following architectures for the synthetic experiments:
Mamba-130M
huctps://huqghggaace.oo//
mamba 1130m hf): selective state-space (SSM) language model. We 24-layer. 1536-d intermediate size and 768-d model size configuration to match the Transformer baselines while retaining linear-time sequence modeling.
LSTM: multi-layer recurrent baseline sized roughly comparable capacity (4 layers, hidden size 1536) to probe how classical RNNs fare on our trajectory-style tasks.
GRU: a gated-recurrent baseline (4 layers, hidden size 1536) offering a stronger RNN comparator with fewer parameters per unit than LSTM.
18

---

Published as a conference paper at ICLR 2026
Pythia (GPT-NeoX style)
pyytha-1660m a decoder-only Transformer from the Pythia scaling suite. We adopt 24-layer, 1536-d intermediate size, 768-d model size and 8-head model variant with RoPE, roughly matching Mamba's scale.
co/mi stralai/ Mi fisttal-----v . a modern decoder-only Transformer with sliding-window (512) sparse attention, utilizing RoPE. We use a scaled-down 8-layer, 1536-d intermediate size and 768-d model size.
Mistral-style Transformer
For the synthetic experiment, we perform hyper-parameter search over learning rate, batch size and weight decay. We choose learning rate (0.00005, 0.0001, 0.0003, 0.0005, 0.001, 0.003, 0. 0.005), batch_size E { {128, 256, 512, 1024), tright_decap E (0. 0.01] and fix the number training steps be 2, 000. We run each experiment with seeds, and report the accuracy of the best model. For Tower of Hanoi experiments, due to their sensitivity, we exceptionally used 10 seeds for Mamba and Pythia models and seeds for other architectures. For the code fixing experiment, we finetune a pretrained Mamba- comba-1.4b and Pythia- 1.4b, both trained on The Pile (Gao et al. 22020), with learning rate 0.0001, weight decay 0.01, batch size 512 and 200 training steps. For all experiments, we use single node with 8 H100 GPUs
D SYNTHETIC EXPERIMENTS DETAILS
D.1 MEMORY TOOL DEFINITIONS
As discussed in Section we use either pointer-based or a search-based memory tool to augment the memory of the model. We now describe how the model interacts with the memory tool, and how we differentiate between thoughts, outputs, commands and observations. We generally try to reduce the number of tokens by using dedicated command tokens, and differentiate between output, thoughts and observation tokens based on the context (instead of using open/closing tags).
Pointer-based Memory. The commands for this memory are given special tokens that the model can output, e.g.
or [pointer2.me] .oover...vvv..mmmmm left(). A read command will be im- mediately followed by a single observation token, that is the token read by the pointer its current position. All other tokens (tokens that are not command tokens or observation tokens, which al- ways immediately follow read command)
either thoughts or outputs. We use a single token [ANS] that indicates the final answer of the model, where all tokens before the [ANS] token are considered thoughts and all tokens after the [ANS] token are considered outputs. Both thoughts and outputs are appended to the context memory, and the pointers can move beyond the input context, and start reading thought or outpul tokens that were previously generated by the model. Commands and observations are discarded and are not appended to the external memory (but of course do affect the internal memory and representation the model). The model can freely interleave commands, thoughts and outputs, and therefore the model can interact with the memory while producing the answers.
Search-based Memory. This memory tool allows the model search for given pattern in the context. A search command is sequence of tokens of the form: [COMMAND] f ndd (VALUE)) where is some sequence of tokens to search for. Following the search command, the model will receive a set of observations of the form: [OBSERVATION) 21, Zk. where _ ....... z are all the lines in the memory context that contain the string (similar to the operation of a grep command). As before, all other tokens are either thoughts or outputs, and are appended to the memory and can be searched for future iterations. In this case we take the output be the last line generated by the model.
D.2 LOGICAL REASONING TASK
As described in Section 3
we generate a random logical computation graph with i: = 3 input nodes, two variables their negation.
where each intermediate node is boolean expression The graph is encoded as a python code given to the model as input. Illustration of the graph and the code shown in Figure
19

---

Published as a conference paper at ICLR 2026
V630
v191
v630 True
v872 False
v622 True
Logical Graph
v240
v539
v191
v872
v240 v872 v622
v539 v622 and not v191
v539 ?
f169
code
code + 1352.py 2...
main
F189.py
import
Code Fixing
F352
foo
(643,py
r352.py
def f3521v31: v2 "kivi"
ra16.py
v1 "nectarine"
F643
f816
foo.py
print(vi]
main.py
foolv3, v2, v11
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 these graph structures.
Figure 4:
D.3 TOOL-USE ALGORITHMS
We describe the synthetically generated tool-use trajectories for solving the different tasks presented in Section 3
Multi-Digit Addition. We follow the standard long addition algorithm, summing one-digit from right to left while keeping the "carry" digit. The model uses pointer-baset memory with two point- ers, and performs the following steps:
1. Move each pointer to the least significant digit of each summand, where the first pointer points to the digit of the first summand, and the second pointer to the second summand. To do this, move the first pointer until we read the token and move the second pointer until we read the token =.
2. Read one digit from each summand, compute the sum of the digits and the carry from previous iteration (if it exists), output the new sum and carry as thoughts, and move both pointers to the left. Stop each pointer reaches non-digit token. both pointers reached non-digit tokens, output [ANS] and move the next step.
3. At this step we have the sum written in reverse in memory, along with the carry digits from each iteration. To write the final output reverse, we move the second pointer the right until it reaches the [ANS] token, then start moving to the left, each iteration outputting the sum digit, until the pointer reaches the token.
Multi-Digit Multiplication. We follow the long multiplication algorithm for multiplying n- digit number by k-digit number (for fixed k). We use a pointer-based memory with max(k, 2) pointers. The algorithm exccutes the following steps:
1. Move the first pointer to the least significant digit of the first operand, and the second pointer to the least significant digit of the second operand.
2. Move the first pointer to the left, each time reading a digit and multiplying with the digit of the second pointer. Add the result to the previous carry, and write it together with the new carry. we reach the most significant digit of the first operand, move the first pointer back to the least significant digit, move the second pointer
position the left, output + sign and zeros as required (depending on which digit from the second operand we read). If the second pointer reached the x sign, move the next step.
3. At this step have summation problem with k summands, where the summands written in reverse and also contain carry digits that we should ignore. We move each pointer to the least significant digit of its respective summand, read all digits, compute the
20

---

Published as a conference paper at ICLR 2026
sum and the carry, and move each pointer to the right and skip carry digits. We continue until we reach the most significant digit of all summands, and then output an [ANS] token.
4. Finally, we have the answer written in reverse with carry digits. We move the first pointer to the [ANS] token, then move it one token to the left and output the tokens read by the pointer, skipping carry digits.
Tower of Hanoi. The Tower of Hanoi puzzle can be solved by a simple recursive algorithm. Let be the number of disks in the puzzle. The recursive algorithm involves three steps: (1) recursively moving the top n - disks from rod to rod B: (2) moving the largest disk from rod to rod C:; and (3) recursively moving the - disks from rod B to rod C. Therefore, the puzzle can be solved with 2 1 moves. This algorithm can also be stated iteratively:
At the first step, the smallest disk is moved from rod A to rod or C depending on whether T. is even or odd, respectively.
At the second and other even steps, the only legal move not involving the smallest disk is performed.
At the third and other odd steps, the smallest disk is moved to the rod was not on two turns ago.
Our model uses the iterative algorithm described above to solve the puzzle. The model uses pointer- based memory with one pointer. At each step the algorithm, the model outputs the next move and the subsequent state of the rods. Specifically, the model takes the following steps:
1.
The pointer traverses rod A from its base, reading disks one by one and outputting B and C' alternatively. This step computes the parity of n, which is crucial for the first move. After reaching the end of rod A, the pointer rewinds to the beginning of the current state representation. The model is now ready to predict the next move.
2. At this step, the model outputs the next move, e.g., AC) 13 (meaning disk (5) is moved from rod A to rod C). Next, the model outputs the new state of the puzzle and goes to the step described below. In case that all disks are on rod C, no move is predicted and the output ends at this step.
3. The model uses the pointer to output the new state. The pointer starts by reading the previous state one disk at a time, outputting the new state. Note that the new state differs from the previous state only by the position of single disk. After processing the previous state, the pointer is then advanced past the outputted move (e.g., (5)AC)) to position itself at the beginning of the newly generated state. The model then goes back the move prediction step above.
We also trained our models using the recursive algorithm: however, their length generalization per- formance was weaker. Details of this experiment are presented in Appendix D.4
Logical Reasoning. We use a search-based memory tool to solve the logical reasoning problem detailed in Appendix [D.2] We try to resolve variables' truth value recursively using depth-first- search (DFS). Namely, starting with the output variable, variables in a given expression. If we find a variable with a boolean (True/False) value, we update the expression, replacing the variable's 's name by its value. If we find a child variable that is still not resolved, we search for the variables the child's expression, while also logging the value of the parent variable (that we use for "backtracking"). When we are done resolving variable's 's value, we backtrack to its parent, trying to resolve the parent's value. When we resolved the output nodes value, we finish the generation.
recursively search for the values of
D.4 TOWER OF HANOT EXPERIMENT DETAILS
In contrast other tasks presented in this paper, the solution length for the Tower of Hanoi puzzle scales exponentially with the number of disks. In particular, solving the puzzle for 9 and 12 disks
138 is used after each move as delimiter.
21

---

Recursive Implementation. For puzzle with disks, the model outputs the list of moves for puzzles size 1, 2, ..... n sequentially and uses the move list generated for puzzle of size i- to output the list of moves for puzzle of size using the recursive pattern. The model uses pointer- based memory with two pointers. While outputting the list moves for the puzzle of size i, the first pointer points the vth disk (largest disk moved in the moves of the puzzle of size i). The second pointer used for implementing the recursive pattern and iterating the moves the puzzle of size Published as a conference paper at ICLR 2026
Table Experimental results for the Tower of Hanoi puzzle solved recursively by different models.
Model
Hanoi (recursive)
Mamba
7-38 (100%)
LSTM
7-39 (83%)
GRU
7-28 (100%)
Pythia
7-57 (100%)
Mistral
7-88 (87%)
would require over 42,000 and 385,000 tokens, respectively. A solution is considered correct only if all of its tokens are correctly generated. As a result, even if model has 99% token accuracy, it may not show any length generalization capability. In agreement with this intuition, we found that Hanoi length generalization performance is very sensitive to the random seed. Hence, we used 10 seeds for Mamba and Pythia models and seeds for the rest of the models. In our experiments with Mamba models, we noticed that token accuracy is always high (c.g., 2 99. 99.75% for 12 disks), however, the actual accuracy varies based on the seed. The performance of our 10 seeds (all trained on puzzles with to 8 disks) is shown in Figure We note that Pythia did not show any length generalization, in particular, even the token accuracy did not exceed 93%.
1.0
0.8
Accuracy
0.6
0.4
0.2
Disks
Figure 5: Performance of different seeds when training Mamba model Tower of Hanoi puzzles with up to 8 disks.
As discussed in Appendix D.3 the Tower Hanoi puzzle can also be solved recursively. We also tried the recursive variant of the algorithm, however, its length generalize performance was weaker than the iterative algorithm. Here, we include the implementation and results of the recursive implementation.
i while generating the moves for size i. More precisely, the input gives the list of disks (e.g.. (2) and the model executes the following steps:
1. Both pointers are moved to the smallest (top) disk, and the model outputs the first move, i.e., moving the top disk from rod A to rod C. This solves the problem for single disk. First pointer moves one step back (now pointing to the second smallest disk) and the second pointer advances, pointing the beginning of the first move. The model is now ready to output the moves for solving the puzzle for two disks.
2. At this step, the model copies the last list of moves, swapping rod labels B and C. To achieve the latter, the second pointer traverses the last list of moves and the model reads
22

---

Published as a conference paper at ICLR 2026
Multplication by 1-Digit
Mutiplication by 2-Digt
1.0
kora
0.8
15TH
0.6
OND
3..4
0.4
0.2
0.0
0.0
10
10
10
10
10'h
20
Sequencelength
Sequence length
Logical Graph
Tower Hanol (Recuraive
1.0
1.0
0.8
0.8
0.6
0.4
Maria
0.4
Mario
1STV
0.2
0.2
1y06
7,0
0.0-
Navul
Sequerce length
Figure 6: We train various transformer (Pythia, Mistral) and SSM (Mamba, LSTM, GRU) models on Multi-Digit Multiplication, Logical Graph and Tower of Hanoi tasks, with CoT + pointer-based memory tool. Multi-Digit Multiplication: We train models multiplying number of up to 10- digit by a I-digit number or 2-digit number, using the pointer-based memory tool. Logical Graph: We train models perform a logical graph reasoning problem using search-based memory tool, training on graphs with up to 10 variables. Tower of Hanoi: We train models to solve the Tower of Hanoi (recursive implementation) reasoning problem using cearch-based memory tool, training on problems with up to 7 disks. The first point in each plot is the maximal problem size seen during training (i.e., all other points are -fitt--distriuution extrapolation).
and outputs one token time (performing the swap if needed). This step corresponds to the first step of the recursive algorithm. At the end of copying, the second pointer is rewound to point to the beginning of the list of moves again.
3. Next, the middle move, i.e., moving the largest disk (ith disk while outputting the moves of size i) is performed. This disk is identified by the value of the first pointer and the move is always from rod A to C. This step corresponds to the second step of the recursive algorithm.
4. Similar to step 2, the model copies the list of moves again, swapping B and A. The second pointer is used again for iterating the list of moves and copying. This step corresponds to the third step of the recursive algorithm. After the copying is finished, the second pointer advances and points to the beginning of the newly constructed move list. The first pointer goes one step back so that it points to the next larger disk. This completes the generation for size i, and the process iterates by returning to step for size i+ 1. The generation terminates if there is no disk remaining for the first pointer, indicating that the lists of moves have been generated for all puzzle sizes 1,...,77
The performance of the recursive solution for the Tower of Hanoi puzzle is reported in Table 2
D.5 SSMS AND TRANSFORMER BASELINES
In Figure 6 we report accuracies of our baseline models on the Multi-Digit Multiplication, logical reasoning and Tower of Hanoi tasks. We train each model on the same trajectories, using CoT and tool use. We perform hyperparameter optimization for each model as described in C Our results
We note that we a delimiter (e.g. #) between the list of moves for different number of disks so that they become separahle.
23

---

Published as a conference paper at ICLR 2026
Task 1, max Swepsttppppp. 250
Task 500
Task stepsm800
1.0
train
1.0
1.0
0
train 1on=10
train Jenm20
0.8
0.8
0.8
Accuaa
Acura
train_lenes
train_enn-5
0.2
0.2-
train_len=10
0.2
train Jen==00
train Jenm20
Jeiin_enn200
10
10
0.0
10
10
10
Sequence Length
Sequence Length
Sequence Length
Task
Task stops-500
Task 2.
800
train lan
1.0
train Ien-10 10
0.8
0.8-
0.9
Accurss
0.6
Accuracy
0.6-
train Jen=5
Accuret
train Jen=10
train
train Jen-S
0.2
0.2
0.2
trin Jen=10
train Jen=20
0.0
10'
30'
10
0.0
16'
10' Sequence Longth
10a
0.0
10'
10
10"
Sequence Length
Sequence
Figure 7: Multiplication generalization performance for Mamba across different training config- urations. Each subplot shows accuracy as function of sequence length for specific maximum training steps value. Different colored lines represent different training sequence lengths, with error envelope indicating median absolute discrepancy across 5 runs.
generally point to length generalization advantage for state space models over baseline transformer models.
D.6 ABLATING TRAINING STEPS AND DIGIT LENGTH FOR MULTIPLICATION
In Figure 7 we investigate the impact of different training configurations generalization for multi- digit multipiccation, varying the training budget (250, 500, or 800 steps) and the maximum number of digits seen during training for the first operand (5, 10, 20). For this experiment, the learning rate was set to 0.003 based on validation. Results indicate that increasing the maximum number of digits shown during training for the first operand improved stablity of OOD generalization consistently, with results improving with more training steps. In particular, training with up to 20 digits improves generalization stability perfectly up the maximum digit size tested (1000 digits). However, even training with up to 5 and 10 digits show progressive improvements as the number of training steps increases.
D.7 ADDITIONAL ABLATIONS
We run the following ablations on the multi-digit addition task:
1. No-CoT: we train the model to directly output the final answer, without any CoT or tool- use.
2. No-CoT, reversed answer: we train the model to directly output the final answer in reverse (reverse format was shown to improve length generalization in certain settings, e.g. (Zouu et al. (2024)..
3. No Tool-Use:
The model is trained on similar trajectories as the main experiment, but now needs to predict the output of the memory tool instead of receiving these as observa- tions. Namely, the trajectory used as CoT data.
4. Single-Turn Tool-Use:
train the model with "caleulator", where the model needs to generate a single addition command following the input (i.c., given an input a b needs generate a add(a, b)).
24

---

Published as a conference paper at ICLR 2026
We train the Mamba model in all settings with extensive hyper-parameter tuning 5-digit addi- tion. Experiments 1,2 and 4 results perfect accuracy on 5-digit addition, but little to no length generalization. Experiment 3 results in poor performance even in-distribution.
D.8 TASK MIXTURE
In Figure 8 eeacch panel shows accuracy as a function of test length for various training budgets (250, 500, or 800 steps). The curves correspond different mixing weights, where w 0 denotes the baseline trained only the main task and higher values indicate normalized fraction of auxiliary samples. The error bars indicate variability across random seeds.
Multiplication/ Steps:
Multiplication Steps:
Multipiication Steps: 800
w=0.0
1.0
pasegee 18am
1.0
W=0.33
0.3
W=0.5
0.8
0.8
W=0.6
Accuracy
Accuracy
0.6
0.6
0.6
0.4
0.4
0.4
w=0,33
W=0.33
0.2
0.2
0.2
W=0.5
Ww
w-0.6
W-0.6
103
102
10
10
102
10
10a
10a
10a
Number of Digits
Number Digits
Number of Digits
Figure 8: Multiplication task accuracy under co-training with varying training budgets (see Sec 3.1
The accuracy plots for the main task (multiplication) in the task mixture experiment were presented in section 3.1 For completeness, show the auxiliary task accuracy in Figure 9
Addition / Max Steps: 250
Addition / Max Steps: 500
Addition Max Steps: 800
1.0
1.0
2
1.0
0.8
0.8
0.8
Accurscy
0.6
w=0.33
w=0.0
Accuracy
0.6
w=0.33
w=0.0
Accuracy
0.6
w= 0.33
w=0.0
0.4
w=0.6
0.4
w=0.6
0.4
0.6
0.2
0.0
8-8898
0.0
0.0
00000000000000000000000000000000 0-00--0-------0-
10a
10
10a
10a
10a
10a
10a
Number of Digits
Number of Digits
Number of Digits
Figure 9: Addition task accuracy under co-training with varying training budgets (250, 500, 800 steps). Curves show different mixing weights. See (3.I.
D.9 COMPARISON WITH HYBRRID-MAMER AND RMT
In addition to comparing Mamba against Transformers, we also evaluate two alternative architec- tures on the multi-digit addition task with CoT and pointer-based memory tool: Hybrid-Mambe and the Recurrent Memory Transformer (RMT) (Bulatov et al. 2022).
Hybrid-Mamba. The Hybrid-Mamba model follows the Jamba architecture Lieber et al. (0024), which interleaves Mamba SSM layers with periodic Transformer attention layers. In our configu- ration, we use a 24-layer model with hidden size 768 and intermediate size 2048, where every 4th layer is standard Transformer attention layer (with 12 attention heads and 4 key-value heads) and the remaining 20 layers are Mamba SSM layers (with state size 16, convolution kernel size 4, and expansion factor 2).
25

---

Published as a conference paper at ICLR 2026
Sequence Length vs Functions
Pass Rate Functions
30030
77.5
Sequence Length Itakans)
25000
20000
Hir Length
13000
Avarage Lengen
Max Lenggh
70.0
07.5
2000
10 Number Functions
Number Functions
Figure 10: Pass rate and median sequence length for SWE-agem--LM-328 on the code fixing task.
Recurrent Memory Transformer (RMT). The RMT model (Bulatov et al. 2022) augments a GPT-NeoX Transformer (24 layers, hidden size 768, 8 attention heads, intermediate size 1536) with learnable memory tokens and segment-level recurrence. The input is divided into segments of fixed length, and each segment is processed with the layout: [read_mem content write_mem].. The read-memory tokens are prepended to each segment and attend bidirectionally among themselves; content tokens attend causally to themselves and fully to the read- read-memory tokens; write- write-memory tokens attend fully to both read-memory and content tokens, and causally among themselves. The write-memory hidden states from segment & become the read-memory imputs for segment k + 1, enabling information to flow across segments through fixed-size memory bottleneck. We use 10 memory tokens and experiment with a segment length of 16 and 64.
Experimental Setup. We train all three model types (Mamba, Hybrid-Mambaa and RMT) on the addition task with CoT and pointer-based memory tool, training up 5-digit addition and evaluating on longer sequences. We perform a hyperparameter sweep over learning rate f0.00005, 0.0001, 0.0003, 0.0005, 0.001, 0.003, 0.005). effective batch size {128, 256, 512, 1024], and weight decay E {0, 0.013, with a fixed budget of 2,000 training steps.
Results. We find that the Hybrid-Mamba model performs similarly to Mamba, achieving perfect length generalization the addition task, extrapolating from 5-digit to 1,000 digit addition. This suggests that the periodic attention layers do not hurt the length generalization capability of the SSM backbone. In contrast, while the RMT model learns the task in-distribution (achieving 100% accuracy on 5-digit addition), it does not achieve meaningful length generalization.
E CODE FIXING AGENT SETUP
We use the same system prompt and input prompt as in mini-SWE- agentgent (Yang et al. 2024a) from: https: : / github. com/ Sww-tgentttaaaeeeeeeeeeeeee /enent/mmnni seee-swe-agee -weet We instruct the model to solve the bug in i py, and explain how the bug should be solved. We modify the original prompt of mini-SWE-agen to instruct the model interactively debug the code and generate a fix for up to files at time.
Please solve this issue: Fix the bug in main.py. Make sure to pass variable v10 too foo() and all other relevant functions. Pass v10 to ONLY the relevant functions, do not pass it it is not needed.
You can execute bash commands and edit files to implement the necessary changes.
## Recommended Workflow
This workflows should be done step-by-step that you can iterate your changes and any possible problems.
1. Create script to reproduce the issue and run it
2. Spot 3 files that might be causing the issue
26

---

Published as a conference paper at ICLR 2026
3. Read the content of these files.
3. Edit the source code of these files resolve the issue. Do not edit more than 3 files before running the script again, even if the code is not completely fixed.
4. Verify your fix works by running your script again, if not analyze at most 3 more files that might cause the issue and repeat the debugging process
5. Submit your changes and finish your work by issuing the following command: "echo COMPLETE.TAL TIEE.TASKAAN AND BOOKANDDSUBBBA SUBMITFFINAL FINALLOUTPUT: Do not combine it with any other command. (important)Afeee this command, you cannol continue working on this task. (/important)
We plot the pass rate and generated trajectory length the SWE agent as a function of the number of functions in the code in Figure T0
F TOOL USE CAPABILITIES OF PRETRAINED SSMs
At the time of writing, we were unable to find any publicly-availann SSM models that were fine- tuned for function calling. The closest we could find is Mistral's Hambba-Codestral--B--ddd was fine-tuned on coding tasks. We evaluated this model on the Berkeley Function Calling Leader- board (Patil al.), and found an overall accuracy 16.58% comparable with the reported accura- cies of 16.22% for Falcon3-3B 3-BB- In-trutt and 15.58% for Llama 3. --IItruuuct.
which
27


Reviews:
[{"review_id": "SYaUfQJmHR", "reviewer": "Reviewer_ubi8", "summary": "This paper investigates the length generalization capabilities of State Space Models (SSMs) in long-form generation tasks. The authors prove that SSMs cannot solve such tasks with bounded memory alone, but can achieve length generalization when augmented with interactive tool use. Experiments on arithmetic, reasoning, and coding tasks demonstrate that tool-augmented SSMs can generalize from short training sequences to much longer test sequences, outperforming Transformers in certain settings.", "questions": "1. How do the hand-crafted trajectories in arithmetic experiments satisfy (or approximate) the Turing machine trace assumption in Theorem 2.2?\n\n2. In the coding experiments, you filter for successful trajectories. How does this affect the model's ability to handle errors or suboptimal tool use at test time?\n\n3. Have you tried evaluation on actual software engineering tasks (e.g., SWE-bench lite)? What performance would you expect?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper provides formal definitions (GSSM, long-form generation tasks) and proves a crisp dichotomy: SSMs fail without interaction but succeed with interactive tool access.\n\n2. The arithmetic experiments show remarkable generalization.\n\n3. The experimental coverage is quite comprehensive. The paper evaluates multiple architectures (Mamba, LSTM, GRU, Pythia, Mistral) across diverse tasks (arithmetic, logic, coding), with both synthetic and agent-generated trajectories.\n\n4. Connecting SSM limitations to tool-augmented agents is timely and relevant, given the growing interest in agentic AI systems.\n\n5. The progression from theory to synthetic experiments to realistic coding tasks is logical and easy to follow.", "weaknesses": "1. The theory relies on training with complete Turing machine traces, but the experiments use hand-crafted or distilled trajectories with no analysis of this gap.\n\n2. The key coding experiment uses a synthetic benchmark with a single bug pattern, not a complex real-world benchmark (e.g., SWE-bench).\n\n3. Missing Sample Complexity Analysis: The theory proves learnability (Lemma B.1) but gives no practical bounds on the sample size (m) or task complexity (n_0) required to achieve it.\n\n4. Incomplete baselines: The paper omits comparisons to SOTA hybrid SSM-Attention models (e.g., Jamba), which are a key alternative for efficient long-context modeling.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bUhCvLCaE6", "reviewer": "Reviewer_14Xt", "summary": "This paper argues theoretically and empirically for augmenting SSM architectures with tools that can function as an external memory. In principle, this provides a simple way to overcome the memory limitations of recurrent architectures (as opposed to the more conservative approach of hybridizing with attention layers). The authors first formalize the value of memory tools with some basic theory and then show empirically that SSMs using external memory tools can achieve good performance and length generalization on synthetic tasks. They also show an encouraging proof-of-concept result where (distilled) SSM coding agents with tool use better generalize to large codebases than transformer coding agents.", "questions": "Table 1: clarify this is with tool use. To really make your point that tools are necessary, you could show performance without tools.\n\nwhat does it mean to move a pointer left or right? Increment/decrement it?\n\nIt seems your \"logical graph problem\" is really the circuit evaluation problem, which is known to be P-complete and thus outside the representational power of transformers unless the complexity classes TC0 and P collapse. Furthermore, even with CoT, it would be very surprising if transformers could solve this problem with a small number of steps (e.g., even with log^k n steps, CoT transformers cannot solve this problem unless TC0 and P collapse).\n\nDo you have opinions on whether SFT alone would suffice to encourage tool use with SSM models? Or would RL be helpful for teaching a model to acquire tool use?\n\nThe following statements at the abstract level could be made more clear, since it's not sure what long-form generation or tractable mean at this point:\n\n> We begin this work by showing a simple theoretical result stating that SSMs cannot accurately solve any long-form generation problem, undermining their main competitive advantage.\n\n> SSMs can learn to solve any tractable problem and generalize to arbitrary problem length/complexity\n\nRegarding length generalization with SSMs:\n> SSMs have been shown to display robust length generalization capabilities in certain cases.\n\nIMO you could make a stronger claim. It seems quite robust at this point that SSMs show strong length generalization relative to transformers. e.g., the gated delta net, path attention, and samba papers show zero-shot results consistently outperforming transformers on perplexity evals and RULER.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The overall argument of the paper is well motivated and presented. I appreciate the principled framing even if the theory is a bit surface-level\n2. The question of how to extend SSMs with additional memory is very timely, and the community has not converged on a solution. One approach (e.g., used by Qwen-3-Next) is to hybridize by mixing transformer and SSM layers. But this paper proposes and evaluates a more radical solution without attention, which is valuable.\n3. The empirical results in the paper are quite interesting. In particular, the empirical results showing better length generalization for distilled SSM coding agents with tool use (vs. transformer agents) is a promising proof of concept that the synthetic-task results are practically relevant", "weaknesses": "### Long-Form Task Definition\n\nA task is defined to be long-form generation if the output size (number of possible outputs with probability mass) grows with input size n. I wouldn't necessarily agree with this definition. What if there is only one choice for y given x, but that y has to be large, and it's length scales with y - why don't we consider this long-form generation?\n\n### Theory Reinvents the (Automata-Theoretic) Wheel\n\nThe theoretical results are somewhat superficial. While I appreciate the principled framing, they essentially boil down to an argument that, if a model has fixed memory, tools that serve as interface to additional memory are helpful for memory-intensive tasks.\n\n> Note that any model that has fixed memory as a function of the sequence length satisfies the definition of a GSSM.\n\nYour definition of a generalized state-space model with constant |S| is just reinventing the notion of a probabilistic finite automaton (with a potential distinction based on the long-form task definition thing mentioned above, though I'm not sure it's a meaningful one). See, for example, [Merrill et al., Section 3](https://arxiv.org/abs/2404.08819) for a discussion of the connection between SSMs, state tracking tasks, and automata/regular languages.\n\nIn formal language theory terms, your notion of GSSMs with interactive tool use has a natural interpretation as finite automata augmented with some tool data structure, similar to how pushdown automata are automata augmented with a stack data structure (and this lets pushdown automata solve additional languages that require more memory). Interaction is essentially giving the model more memory through interaction with an external data structure. I think making the connection between GSSMs and automata explicit and discussing this interpretation of tools as data structures would help clarify your argument.\n\nTheorem 2.1 is just the Myhill-Nerode theorem (a probabilistic extension of the basic deterministic one)\n\nTheorem 2.2: it's a bit overly verbose to talk about learning algorithms when this is essentially an expressivity construction. i.e., the learning algorithm and training data you have here is basically just hard-constructing an expressivity construction?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "97llGkc9cK", "reviewer": "Reviewer_5b4b", "summary": "This paper proves that SSMs cannot solve long-form generation tasks due to their fixed-size memory, but that interactive tool-use enables them to achieve length generalization—allowing SSMs to handle arbitrarily long and complex problems by leveraging external memory.\nEmpirically, the authors show that tool-augmented SSMs (e.g., Mamba) outperform Transformers in arithmetic, reasoning, and coding tasks, demonstrating scalability and efficiency advantages when used as interactive agents rather than standalone models.", "questions": "- How this proposed methods could generalize to real-world reasoning or language understanding tasks?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The paper proves that SSMs, despite their efficiency, cannot solve long-form generation tasks under standard or single-turn settings.\n- The paper validates the theory on diverse task categories, including arithmetic, algorithmic reasoning, and real-world coding.", "weaknesses": "- The experiments primarily use small or mid-sized models. These experiments cannot prove if the proposed methods could generalize to larger scales.\n- The paper evaluates on synthetic arithmetic, algorithmic, and coding tasks only. There is no testing on natural-language, multimodal, or real reasoning benchmarks.\n- Although tool-use extends memory, it introduces I/O and latency overhead, which may offset the computational gains of SSMs in real-time systems.\n- Typos, such as \"node node\" in line 398", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "HKttO3CFqg", "reviewer": "Reviewer_Mthj", "summary": "This paper studies **length generalization** of **State Space Models (SSMs)** and recurrent controllers on long-form generation tasks. It formalizes a class of **generalized SSMs (GSSMs)**, that are a finite-state controllers whose memory does **not** grow with sequence length, and a family of **long-form generation** problems where the support size of valid outputs increases with task complexity. Two central theoretical results:\n\n* **Impossibility (Thm 2.1):** Any CoT-only or **single-turn** tool-use GSSM must fail on long-form generation beyond some size with a specific error.\n* **Possibility (Thm 2.2):** With **interactive** tool-use (a simple read/write pointer-tape oracle), there exist training distributions under which a GSSM **length-generalizes** on any *computationally tractable* long-form task, i.e., learns a constant-state controller that operates unbounded external memory.\n\nEmpirically, the paper shows that SSMs/RNNs trained on **interactive, ReAct-style** trajectories strongly extrapolate on:\n\n* synthetic arithmetic (multi-digit addition/multiplication up to 1K digits),\n* reasoning (logical DAG evaluation to 1K nodes; some generalization on Tower of Hanoi),\n* a **coding** bug-fix task with bash/file tools (run→inspect→patch→re-run), where Mamba trained on **interactive/distilled** traces maintains high pass rates well beyond training repository sizes, while **single-turn** tool-use collapses and local-attention Transformers degrade.\n\n**Contributions:**\n\n1. A clean **formalization** of long-form generation and GSSMs.\n2. A **negative result** showing fixed-memory controllers (incl. CoT or single-turn tool calls) cannot solve such tasks at scale.\n3. A **positive result**: interactive tool-use + external memory is sufficient for “infinite” length generalization.\n4. Broad **experimental evidence** across arithmetic, reasoning, and coding, with a realistic agentic setup and trajectory distillation.", "questions": "1. **Scope of Theorem 2.2:** The result asserts existence of training distributions {P_n}. How constructive is the procedure? Can you characterize **sample complexity** (number of transitions/examples) needed to learn a given algorithm’s controller?\n2. **Transformers + external memory:** How do results change when comparing against **Transformers with the same pointer-tape tool** and **interactive** protocol? Are failures due to representation (attention) or optimization/data?\n3. **Tool noise & latency:** Real tools return noisy, delayed, or partially structured outputs. How sensitive is the controller to **observation noise**, **command failures**, or **latency**?\n4. **Generalization beyond seen transitions:** For multiplication/addition, do models truly learn **state transition rules**, or are there subtle failure modes (e.g., rare carry chains)? Any adversarial or **worst-case** tests?\n5. **Safety & security:** In coding, could the learned policy execute unsafe commands? Any recommended **sandboxing** or **guardrails** for deploying SSM agents with system tools?\n6. **When is single-turn enough?** Can you give some task classes where single-turn (or CoT-only) suffices, perhaps when supp_alpha grows slowly or when a single observation is maximally informative?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "* **Originality:** Elegant formalization of long-form generation via support-growth; clean GSSM abstraction; sharp **necessary** (no interaction) vs **sufficient** (interactive tools) results.\n* **Quality:** Both **theory and practice** reinforce each other; experiments span arithmetic, graph reasoning, and a realistic coding benchmark with real tool commands and agent distillation.\n* **Clarity:** Clear separation of CoT-only, single-turn, and interactive regimes; easy-to-follow trajectory visualizations.\n* **Significance:** Offers a compelling recipe for **length-generalizing agents**: small SSM controller + interactive memory/tooling + trajectory imitation. Practical implications for efficient, long-horizon tasks.", "weaknesses": "* **Baselines & fairness:** Transformer comparisons are mainly **local/sliding-window**. Missing head-to-head with **memory-augmented** or **recurrent** Transformer variants (e.g., retrieval-augmented, ring/segment memory, recurrent attention, kNN-LM, RMT) that also enable externalized state.\n* **Policy-execution gap:** Training uses teacher forcing on trajectories; it would help to quantify compounding-error effects when **executing** the learned policy end-to-end under stochastic observations.\n* **Scaling breadth:** Coding experiments use 1.4B-scale models; it would be informative to include **larger SSMs** and show whether gains persist or amplify vs similarly scaled Transformer agents with function-calling tool APIs.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)