title: Stuffed Mamba: State Collapse and State Capacity of RNN-Based
    Long-Context Modeling
summary: One essential advantage of recurrent neural networks (RNNs) over
  transformer-based language models is their linear computational complexity
  concerning the sequence length, which makes them much faster in handling long
  sequences during inference. However, most publicly available RNNs (e.g., Mamba
  and RWKV) are trained on sequences with less than 10K tokens, and their
  effectiveness in longer contexts remains largely unsatisfying so far. In this
  paper, we study the cause of the inability to process long context for RNNs and
  suggest critical mitigations. We examine two practical concerns when applying
  state-of-the-art RNNs to long contexts: (1) the inability to extrapolate to
  inputs longer than the training length and (2) the upper bound of memory
  capacity. Addressing the first concern, we first investigate *state collapse*
  (SC), a phenomenon that causes severe performance degradation on sequence
  lengths not encountered during training. With controlled experiments, we
  attribute this to overfitting due to the recurrent state being
  overparameterized for the training length. For the second concern, we train a
  series of Mamba-2 models on long documents to empirically estimate the
  recurrent state capacity in language modeling and passkey retrieval. Then,
  three SC mitigation methods are proposed to improve Mamba-2's length
  generalizability, allowing the model to process more than 1M tokens without SC.
  We also find that the recurrent state capacity in passkey retrieval scales
  exponentially to the state size, and we empirically train a Mamba-2 370M with
  near-perfect passkey retrieval accuracy on 256K context length. This suggests a
  promising future for RNN-based long-context modeling.
paper_authors._id: 670801c33f448c140030ee45
paper_authors.name: Yingfa Chen
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2024-10-11T07:43:57.226Z
paper_authors._id: 6144e4667f2544bb450787b2
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/6144e4667f2544bb450787b2/5wSDDqJbI4TGtBLP9IvjY.png
paper_authors.fullname: Yingfa Chen
paper_authors.type: user
paper_authors.user: chen-yingfa
paper_authors._id: 670801c33f448c140030ee46
paper_authors.name: Xinrong Zhang
paper_authors._id: 670801c33f448c140030ee47
paper_authors.name: Shengding Hu
paper_authors._id: 670801c33f448c140030ee48
paper_authors.name: Xu Han
paper_authors._id: 670801c33f448c140030ee49
paper_authors.name: Zhiyuan Liu
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-06-02T07:57:10.247Z
paper_authors._id: 6472f1f7485a7c8e1bcf507b
paper_authors.avatarUrl: /avatars/79dca137583875191ea883b193b630b7.svg
paper_authors.fullname: Zhiyuan Liu
paper_authors.type: user
paper_authors.user: zibuyu9
paper_authors._id: 670801c33f448c140030ee4a
paper_authors.name: Maosong Sun
paper_title: Stuffed Mamba: State Collapse and State Capacity of RNN-Based
    Long-Context Modeling
paper_summary: One essential advantage of recurrent neural networks (RNNs) over
  transformer-based language models is their linear computational complexity
  concerning the sequence length, which makes them much faster in handling long
  sequences during inference. However, most publicly available RNNs (e.g., Mamba
  and RWKV) are trained on sequences with less than 10K tokens, and their
  effectiveness in longer contexts remains largely unsatisfying so far. In this
  paper, we study the cause of the inability to process long context for RNNs and
  suggest critical mitigations. We examine two practical concerns when applying
  state-of-the-art RNNs to long contexts: (1) the inability to extrapolate to
  inputs longer than the training length and (2) the upper bound of memory
  capacity. Addressing the first concern, we first investigate *state collapse*
  (SC), a phenomenon that causes severe performance degradation on sequence
  lengths not encountered during training. With controlled experiments, we
  attribute this to overfitting due to the recurrent state being
  overparameterized for the training length. For the second concern, we train a
  series of Mamba-2 models on long documents to empirically estimate the
  recurrent state capacity in language modeling and passkey retrieval. Then,
  three SC mitigation methods are proposed to improve Mamba-2's length
  generalizability, allowing the model to process more than 1M tokens without SC.
  We also find that the recurrent state capacity in passkey retrieval scales
  exponentially to the state size, and we empirically train a Mamba-2 370M with
  near-perfect passkey retrieval accuracy on 256K context length. This suggests a
  promising future for RNN-based long-context modeling.
paper_ai_summary: Recurrent neural networks (RNNs) are studied for handling long sequences, addressing state collapse and memory capacity issues to improve their performance beyond the training sequence length.
paper_ai_keywords: recurrent neural networks RNNs transformer-based language models state collapse long context modeling passkey retrieval model capacity generalizability