Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Thanks for the encouragement!

State compression and memory capacity

Absolutely – we believe that the tradeoff between state size and model capacity is central to sequence models, and there are many exciting research directions leading from here.

• The Based paper [1] found that, when plotting memorization capability against effective state size, that various finite-state models including SSMs actually have comparable performance to attention (as a function of state size). Of course, pure softmax attention can get larger effective state sizes.
• Recent SSM follow-ups have found ways to get much larger state sizes (GLA [2], HGRN2 [3], Mamba-2 [4], RWKV-6 [5], xLSTM [6]), which significantly improves their performance on synthetic memorization tasks.
• We think that attempting to theoretically characterize the memory capacity of models in terms of their effective state size (e.g. proving a lower bound) would be a very interesting theoretical question on sequence models.
• Another interesting direction for improving models is whether a model’s state size can grow with the context length. For example, the linear growth (i.e. KV cache) of a pure Transformer might be wasteful, whereas the constant size state of a pure SSM might be too restrictive; can a model’s state grow slowly through the sequence length instead?

[1] Simran Arora et al. “Zoology: Measuring and Improving Recall in Efficient Language Models” ICML 2024.

[2] Songlin Yang et al. “Gated Linear Attention Transformers with Hardware-Efficient Training” ICML 2024.

[3] Zhen Qin et al. “HGRN2: Gated Linear RNNs with State Expansion” https://arxiv.org/abs2404.07904

[4] Tri Dao and Albert Gu. “Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality” ICML 2024.

[5] Bo Peng et al. “Eagle and Finch: RWKV with Matrix-Valued States and Dynamic Recurrence” https://arxiv.org/abs/2404.05892

[6] Maximilian Beck et al. “xLSTM: Extended Long Short-Term Memory” https://arxiv.org/abs/2405.04517

Initial states and length generalization

We are very excited about passing states through the model so that forms of truncated back-propagation can be used to potentially learn on indefinitely long sequences. This functionality and other large improvements will be open-sourced to continue contributing to the community!

We thank the reviewer for the detailed feedback and insightful questions.

1. Ablation Study: The short conv is a simplification of H3’s “shift SSM”. RWKV has “token shift” (equivalent conv with width=2). The short conv does improve model quality (at 350M with 7B tokens, ppl 8.71 with conv and 8.97 without conv). The primary improvement comes from selectivity (Table 6). Extensive ablation experiments are included in Appendix E.3 (in Table 5, 6, 7, 8, 9).

2. Baselines: We compared to Transformer++ in our scaling law study, showing that Mamba can match Transformer++ on perplexity. For table 2 (models up to 3B trained for 300B tokens), we compared to public models trained with similar numbers of tokens. Due to compute budget, we did not train our own Transformers. The fact that an SSM ( linear in seqlen) can match the strongest Transformer (quadratic in seqlen) and outperform other Transformer variants, which has gone through 7 years of intense tuning from the whole research community, is very encouraging to us.

3. Training Efficiency: Mamba has the same number of layers as Transformers (e.g. Mamba 2.7B has 64 S6 layers, Pythia 2.8B has 32 MLP and 32 attn layers, same total params). Mamba’s training cost is comparable to that of Pythia (Pythia 2.8B takes 14240 A100 hours, Mamba-2.7B takes 14370).

4. Relationship to other input selection mechanisms: We view this as analogous to “hard attention” (e.g. in Mnih et al., 2024 "Recurrent models of visual attention"), and “soft attention” (e.g. Bahdanau et al., 2015 “Neural machine translation …”), nowadays just called “attention” (e.g. by Vaswani et al., 2017).

5. Length extrapolation: We will make clear that our claims are with respect to specific empirical results in (1) synthetic tasks (2) audio (3) genomics. Follow-up works have further investigated long-context tasks, such as needle-in-a-haystack (https://github.com/jzhang38/LongMamba) and more realistic NLP tasks such as question-answering (Jamba), where Mamba (or variants) perform well.

6. Settings of the scaling law experiments: Full experiment settings are in Appendix F. For Transformer FLOPs: 6 * nparams * ntokens + attention FLOPs (12 * nlayers * hdim * ctx_len * ntokens). For Mamba: 6 * nparams * ntokens + scan FLOPs (27 * nlayers * hdim * state_dim * ntokens). The Chinchilla (20 tokens per params) protocol was originally tuned for Transformers, and we do not change this (1) due to compute budget (2) to be generous to the Transformer baselines.

We are very glad that the reviewer enjoyed the paper and model. Thanks for the encouragement!

Thanks for the feedback and questions!

We would like to note that we view this work as re-invigorating interest into alternative architectures, and opening new directions in this space. The reviewer raises many interesting questions about other tasks/applications and potential efficiency improvements; while we could not address all of them in this first work, we are optimistic that much more progress and empirical validation will continue to be made.

Long-context tasks

To our knowledge there are not standard and high-quality benchmarks that are widely agreed to measure long context in text. Thus our goal was to focus on more standard LM pretraining objectives, as well as other modalities (e.g. DNA) where long context matters. We note that follow-ups have further investigated long-context tasks, such as needle-in-a-haystack [LongMamba] and question-answering [Jamba], where Mamba (or variants) perform well.

Training efficiency

We note that Mamba is much faster than Transformers in wall-clock time for longer sequences. However, as the reviewer observes, on a FLOP-for-FLOP basis, it is still less hardware-efficient than Transformers because Mamba does not utilize as many matrix multiplications, which GPUs are extremely specialized for.

This is in fact an area of active research with very positive progress. For example, several follow-up works have been able to develop closely related models that leverage matrix multiplications and are much faster than Mamba [GLA, HGRN2, Mamba-2], comparable to Transformers in wall-clock time even for shorter sequences.

Fixed state size.

Indeed, the tension between state size and recall ability is a fundamental tension of sequence models. We expect SSMs and Transformers to have different tradeoffs in inductive bias and capabilities. For example, SSMs naturally have more trouble with rote memorization [1], but have better ability to learn other types of formal languages [2]. In fact, we view SSMs and attention as complementary, and follow-up work has found that hybrid architectures are broadly very strong (e.g. Jamba, Zamba, Griffin).

[1] https://arxiv.org/abs/2402.01032 [2] https://arxiv.org/abs/2405.17394

Static A

The discrete A dynamics $\bar{A}$ (which govern the SSM recurrence in eq. (2a)) depend on both $A$ and $\Delta$. Since the latter is context-aware, so is $\bar{A}$, which is the true recurrence. In early explorations, we found that additionally making $A$ input-dependent did not help performance.