Research map / Architectures and efficiency / Attention and transformers
Understanding and Improving Length Generalization in Recurrent Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes the unexplored states hypothesis for why recurrent models (Mamba, Mamba-2, GLA, RWKV-v6) fail to length generalize, and shows that cheap post-training interventions that expose models to attainable initial states (State Passing, fitted noise, TBTT) enable generalization from 2k to 128k with ~500 steps.
Problem
Recurrent models such as state space models and linear attention often degrade sharply on sequences longer than their training context, and it is unclear why and how to fix this cheaply across architectures.
Approach
Length generalization fails when training covers only a subset of the state distributions attainable on long sequences; training interventions that initialize the recurrent state from attainable states (e.g., the final state of another sequence, or Gaussian noise fitted to final states) expose the model to those states and enable generalization without changing the architecture.
Key results
- State Passing post-training for 500 steps (~0.1% of pre-training budget) enables length generalization from 2k training context to 128k validation sequences on recurrent models.
- The 780m Mamba-2 model finetuned with fitted noise achieves perfect passkey retrieval accuracy on sequences of length 256k.
- On synthetic copying (45m models, trained on lengths 50-100, evaluated on 300), State Passing raises character accuracy from 0.27 to 0.47.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Increase training context length | Mamba-2 and GLA · Pile validation | position-wise perplexity | Helped |
| Train for more tokens | Mamba-2 · Pile validation | position-wise perplexity | Hurt |
| Extend the learning-rate warm-up period from 10% to 25% | Mamba-2 45m · Pile validation, context 1024 | position-wise perplexity | Hurt |
| Initialize the training state with random Gaussian noise | Mamba-2 370m · Pile validation | position-wise perplexity | Helped |
| Initialize the training state with random Gaussian noise | Mamba-2 780m and 1.3b · Pile validation | position-wise perplexity | No clear effect |
| Initialize the training state with Gaussian noise fitted to final-state means and variances | Mamba-2 370m and 780m · Pile validation | position-wise perplexity | Helped |
| Initialize the training state with Gaussian noise fitted to final-state means and variances | Mamba-2 1.3b · Pile validation | position-wise perplexity | No clear effect |
| Initialize the state with the final state of a different sequence (State Passing) | Mamba-2 · Pile validation | position-wise perplexity | Helped |
| Initialize each chunk with the previous chunk’s final state (TBTT) | Mamba-2 · Pile validation | position-wise perplexity | Helped |
| Use State Passing | Mamba-1 and GLA · Pile validation | position-wise perplexity | Helped |
| Use State Passing | Mamba-2 · BABILong | task performance | Helped |
| Fine-tune with fitted noise rather than standard fine-tuning | Mamba-2 370m and 780m · Passkey retrieval | accuracy | Helped |
Surprising findings
- Training on long sequences is not necessary: TBTT achieves performance very similar to State Passing, suggesting exploring attainable states is the key.
- Random Gaussian initial state helps only the 370m Mamba-2 model and fails for 780m and 1.3b, while fitted noise fails for 1.3b.
Limitations
- Fitted noise fails to correct the 1.3b Mamba-2 model, and random noise fails for larger models, so the simple variants do not generalize across all sizes.
- Focus is mostly on text modeling perplexity and a few long-context tasks; other distributions are left for future work.
- Does not cover non-linear-recurrence architectures such as xLSTM, which are left out of scope.
Experimental setup
Datasets: The Pile BABILong BABI passkey retrieval task synthetic copying task
Benchmarks: The Pile validation BABILong passkey retrieval synthetic copying
Metrics: position-wise perplexity Effective Remembrance (total variation) accuracy character accuracy
Baselines: Mamba-1 official checkpoint Mamba-2 official checkpoints GLA pre-trained from scratch RWKV-v6 Random Noise initialization Fitted Noise initialization TBTT
Builds on: Mamba Mamba-2 Gated Linear Attention Truncated Backpropagation Through Time passkey retrieval
Closest papers by meaning
- On the Importance of Gating: Memorization vs. In-Context Learning in State Space Models (NeurIPS 2026 · similarity 0.90)
- The Role of Sparsity for Length Generalization in LLMs (ICML 2025 · similarity 0.90)
- Provable Length Generalization in Sequence Prediction via Spectral Filtering (ICML 2025 · similarity 0.89)
- Long-Short Alignment for Effective Long-Context Modeling in LLMs (ICML 2025 · similarity 0.89)
- Beyond Single-Task: Robust Multi-Task Length Generalization for LLMs (NeurIPS 2025 · similarity 0.88)
- How Do Position Encodings Affect Length Generalization? Case Studies On In-Context Function Learning (AAAI 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.