Mamba: Linear-Time Sequence Modeling with Selective State Spaces

We are glad that the reviewer appreciates Mamba’s relationship to gating mechanisms and finds the empirical results strong.

Regarding the connection to gating: Theorem 1 establishes the relationship for a special case of $A$ and $B$. More broadly, the $\Delta$ parameter of SSMs plays the role of generalized gating, and we view it as an elegant way to automatically induce gating (instead of the conventional heuristic explanation of gates).

In general, $\Delta$ controls the balance between how much to focus or ignore the current input $x_t$. It is analogous to the role of the gate $g_t$ in Theorem 1, mechanically, a large $\Delta$ resets the state $h$ and focuses on the current input $x$, while a small $\Delta$ persists the state and ignores the current input. SSMs (equation (1)-(2)) can be interpreted as a continuous system discretized by a timestep $\Delta$, and in this context the intuition is that large $\Delta \to \infty$ represents the system focusing on the current input for longer (thus ``selecting'' it and forgetting its current state) while a small $\Delta \to 0$ represents a transient input that is ignored.

We plan to include this discussion, as well as an extended discussion on the intuition and role of the other parameters of selective SSMs ($A$, $B$, and $C$), in an extended preprint that will be released together with the open source release of the model.

We are glad the reviewer found the ideas in the paper simple, compelling, and clear.

Citations and related work

We appreciate the reviewer’s suggestions for related work, which are all very related and appropriate. We have included all of them (and more linear attention variants) such as RFA, cosFormer, Performer, TransNormer, RAFA, and TNN in the extended related work.

Adding extrapolation experiments to the language model would be interesting.

Our original submission includes length extrapolation results for the Induction Heads task, which is hypothesized to be deeply related to generalization abilities of language models [Anth]. In Figure 3 we show that Mamba can extrapolate from sequences of length 256 to length 1000000, while no other model (including Attention with variants of relative positional embeddings) extrapolate beyond length 512.

In the shared response, we include new results on measuring the length extrapolation abilities of Mamba compared to a standard Transformer.

The ablation analysis in Table 6 should be more comprehensive, with a total of $2^3$ possible combinations. I suggest that the authors include the remaining two combinations.

We chose to exclude some partial combinations for the sake of clarity. This table includes the full ablations.

What's your setting of Scaling Law? Why is your ratio of token number and model size is the same as Chicilla's paper? I suppose the FLOPs of Transformers and SSMs would differ. Suppose the FLOPs of Transformers and SSMs would differ given the same amounts of total parameters, is this important to the final performance(accuracy)?

The FLOPs of Transformers and SSMs do not differ significantly at shorter sequence lengths (e.g. length 2048), and we believe that the standard Chinchilla scaling laws should apply. As shown in the difference between Figure 4 (Left) and 4 (Right), the FLOPs of Transformers vs. SSMs do differ at longer sequence lengths. It would be very interesting for future work to explore whether SSMs have different scaling laws especially with differing context lengths.

How did you parameterize the first convolutional layer in the Mamba-Block.

We use a standard local (causal) convolution of width 4. It is a depthwise convolution, i.e. operates independently per channel. This follows similar architecture design in H3, Hyena, and RWKV.

Providing more detailed implementation, such as offering core code, is very helpful.

As mentioned in the shared response, we are planning to release all code and pretrained models for the community.

We appreciate the reviewer’s close reading of the paper and that the reviewer found our results compelling. We respond to the reviewer’s questions below.

The model still has a quadratic memory requirement during training like Transformers.

We would like to clarify that Mamba’s memory requirement scales linearly in sequence length; the memory bottleneck is simply the size of the activation tensors. This is true of most deep sequence models, including Transformers (when using an optimized implementation such as FlashAttention). Mamba is in fact very memory efficient; we have included additional comparisons in Appendix F.5 (see response 4/4 to Reviewer du8a for numbers).

Have you evaluated scaling behavior beyond 1.4B parameters? How does it compare to Transformers at 10B scales?

We have managed to scale Mamba to the next standard model size of 2.7-2.8B parameters. Results are shown in the general response; it continues to strongly outperform open-source Transformers of the same size, and match Transformers of the next larger size (7B parameters).

The input selection mechanism introduces additional hyper parameters. How sensitive are the results to hyperparameters like the projection rank?

This is an excellent question; many new models come with new hyperparameters that must be tuned. Mamba has a few main axes of variation:

• The projection rank $R$, as the reviewer noted
• The state size $N$; larger states can perform better but slow down the model
• Prior works on SSMs have had many variants of the initialization.

Ablations for all of these have been included in the original submission, Tables 7, 8, and 9. We see that increasing $R$ and $N$ both offer small improvements for a modest increase in parameters; however, they do not offer a significant performance difference. In fact, all our experiments did no tuning on any of these hyperparameters (fixing $R$ to be a constant fraction $1/16$ of the model dimension; fixing $N=16$; and fixing the initialization). We strongly believe that Mamba can be a promising drop-in replacement for attention in general foundation models.

One of the main contributions that the authors claim is long sequence modeling. The authors should compare with more baselines on LRA (Long Range Arena), which is essentially the standard benchmark for long sequence understanding.

As elaborated in the first part of our response, in this paper we have deliberately focused on more impactful settings than small benchmarks such as LRA. We offer several additional points of justification:

• Our emphasis on the importance of pretraining is underscored by recent works such as [Amos], which indicate that Transformers are competitive with state-of-the-art SSM models on LRA when allowed pretraining. This supports our claim that small-scale train-from-scratch benchmarks such as LRA may not be the most appropriate for measuring the quality of new foundation model architectures at scale, where thus far only Transformers have been effective.
• LRA is not actually long-range compared to our selection of tasks. Our shortest-range task is standard Language Modeling (Figure 4), for which we have sequences of length 8192, which is already longer than (or on the order of magnitude of all LRA tasks). On the other hand, we have multiple settings (induction heads, genomics, and audio) involving sequences of length >1000000.

[Amos] “Never Train from Scratch: Fair Comparison of Long-Sequence Models Requires Data-Driven Priors”. https://arxiv.org/abs/2310.02980

Memory benchmarking is missing. Even though Section 4.5 is titled “speed and memory benchmark”, only speed comparisons are presented.

The memory usage is simply the size of the activation tensors, as with most deep sequence models (including FlashAttention). Mamba is in fact very memory efficient; we have included additional measurements of the training memory requirements of 125M models on 1 A100 80GB GPU. Each batch consists of sequences of length 2048. We compare to the most memory-efficient Transformer implementation we are aware of (with kernel fusion from torch.compile and with FlashAttention-2).

Mamba's memory requirement is comparable to a similar-sized Transformer with an extremely optimized implementation, and we expect further improvement in Mamba's memory footprint in the future. We have added this Table to Appendix F.5.

Also, the authors should provide more detailed setups of Figure 8 left, e.g., model layers, model sizes, details of the convolution, etc.

These details have been reported in the original submission. The main body directs to Appendix F.5, which includes these details:

• These measure only the core operation (scan, convolution, attention), not the full block (e.g. do not include the QKV projections of attention). There is no concept of layers
• The model dimension is $1024$, with sequence lengths from $2^9=512$ to $2^{19}=524288$.
• The convolution is an FFT-based convolution using PyTorch primitives.

Could the authors provide some intuitions why FlashAttention is the slowest when the sequence length is very large (Figure 8 left)?

Attention (including FlashAttention) is a quadratic-time operator, and therefore is significantly slower than subquadratic operators (e.g. convolutions and linear recurrences) such as Mamba when the sequence length is long.

Many attention-based Transformer models show length generalization ability, i.e., models can be trained on a shorter sequence length and tested on a longer sequence length. Some examples include relative positional encoding (T5) and Alibi [6]. Because SSMs are in general sequential, does Mamba have this length generalization ability?

Our original submission includes length extrapolation results for the Induction Heads task, which is hypothesized to be deeply related to generalization abilities of language models [Anth]. In Figure 3 we show that Mamba can extrapolate from sequences of length 256 to length 1000000, while no other model (including Attention with variants of relative positional embeddings) extrapolate beyond length 512.

In the shared response, we include new results on measuring the length extrapolation abilities of Mamba compared to a standard Transformer.

[Anth] “​​In-context Learning and Induction Heads.” https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html

The authors need to compare with stronger baselines. The authors acknowledge that H3 was used as a motivation for the model architecture. However, they did not compare with H3 in the experiments. From Table 4 in [7], ppl of H3 is 8.8 (125M), 7.1 (355M), and 6.0 (1.3B) on the Pile dataset, which are considerably better than Mamba. The authors need to show comparisons with H3.

H3 is one of our main baselines, and is featured prominently in almost all results in the paper:

• The synthetic tasks, Selective Copying (Table 1) and Induction Heads (Figure 3).
• For Chinchilla scaling laws (Figure 4), we used a much stronger version of H3 that we call H3++ (details in Appendix F.2.1).
• For genomics, we compare against Hyena (Figure 5), which is the successor of H3.
• We have since downloaded the pretrained H3 models from [7] and ran the suite of evaluations (see shared response).

The reviewer mentions that the reported perplexity numbers from [7] are better than ours. However, as mentioned in the caption of our Table 2, perplexity scores are only comparable when trained with the same tokenizer. The H3 paper used the GPT2 tokenizer and is not comparable to ours which uses the now-standard NeoX tokenizer. Instead, the downstream evaluations are directly comparable, where Mamba performs significantly better than H3.

For the pre-trained models, the authors only show results on zero-shot inference. This setting is quite limited and the results cannot support the effectiveness of Mamba well. I suggest the authors run more long-sequence experiments such as document summarization, where the input sequence is naturally long (e.g., the average sequence length of the arXiv dataset is greater than 8k).

We agree that the reviewer’s suggestion is a great idea, and in fact are planning to pursue such long-context applications in follow-up work. Unfortunately, these evaluations require an extensive fine-tuning pipeline in a non-autoregressive setting which fall beyond the scope of this paper; however, we strongly believe that Mamba will be effective in such settings.

The motivation of Mamba is to address the drawbacks of recurrent models while improving the efficiency of attention-based models. There are many works following the same direction: S4-diagonal [1], SGConv [2], MEGA [3], SPADE [4], and many efficient Transformer models (e.g., [5]). All of these models achieve near linear complexity, and the authors need to compare Mamba with these works in terms of both model performance and efficiency. For model performance, some simple experiments such as language modeling on Wikitext-103 should suffice.

As stated in the shared response, we deliberately focused on the perplexity of larger-scale pretraining rather than small-scale benchmarks. Nevertheless, Mamba strongly outperforms all suggested models and many more on WikiText-103, as one might expect from our general results on language.

First, we compare Mamba in the exact same setting as the Hyena paper [Poli, Table 4.3]. Beyond their reported numbers, we tune our own strong Transformer baseline. We then swap out the model for Mamba, which improves over our Transformer by 1.7 ppl and the original baseline Transformer by 2.3 ppl.

Therefore all of these subquadratic architectures do not improve over a comparable Transformer baseline. Moreover, our own Transformer baseline (18.0 ppl) above is stronger than that of prior work, and our Mamba model (without tuning) is even significantly better (16.3 ppl), all in the same setting.

Next, we show results from several of the reviewer’s suggested baselines [2,3,4] as reported by their original papers. Because these results were obtained in potentially different codebases and training setups, we also report each paper’s own Transformer baseline as well as the improvement.

The improvements of these methods are all small (up to -0.3 ppl) compared to our results above (-2.3 ppl). Additionally, all three of these methods are not linear-time: they are hybrid models explicitly using quadratic attention. Without attention, these methods are S4-based LTI models that do not perform well on language, which is the main thesis of our paper.

Finally, we analyze the reviewer’s last suggested baseline FLASH [5]. While [5] does not report WikiText-103, we found that the recent preprint [Qin] performs controlled comparisons of a wide range of models including FLASH. Their results [Qin, Table 1] are reproduced here.

Consistent with the rest of these results and the main thesis of our paper, every other existing subquadratic architecture does not improve over the basic Transformer.

[Poli] “Hyena Hierarchy: Towards Larger Convolutional Language Models.” https://arxiv.org/abs/2302.10866

[Qin] “Hierarchically Gated Recurrent Neural Network for Sequence Modeling.” https://arxiv.org/abs/2311.04823

We greatly appreciate the reviewer’s feedback and suggestions. We would like to first address the reviewer’s main high-level concerns, and then respond to each detailed point individually.

The reviewer’s main concerns are about our baselines and benchmarks, with suggestions to compare against a broader range of attention-free or hybrid models on small-scale benchmarks such as WikiText-103 and Long Range Arena. We would like to reiterate the main point of our new model as stated in the Introduction of the original submission:

We propose a new class of selective state space models, that improves on prior work on several axes to achieve the modeling power of Transformers while scaling linearly in sequence length.

We would like to note that almost all of the reviewer’s suggested baselines are either:

• Already included in our baselines: For example, S4D, DSS [1], and SGConv [2] are special cases of linear time-invariant (LTI) models that are nearly identical and perform very similarly to S4 and Hyena, two of our main baselines. Similarly, H3 [7] as well as Hyena are some of our main comparisons for language modeling (Figure 4).
• Not Transformer quality, or not linear-time: A central thesis of our paper is that prior linear-time models (which are almost all LTI) do not match the performance of Transformers on information-dense data such as language. Indeed, models such as SGConv [2], MEGA [3] , and SPADE [4] explicitly use hybrid models incorporating quadratic attention to achieve their results on WikiText-103 (which are still worse than ours; see new results below). If the attention is removed, all of these models become LTI state-space models like the above and would perform similarly to S4 (significantly worse than Mamba/S6).

In the detailed response below, we expand on comparisons on WikiText-103 for all the above suggested models and many more. Nevertheless, we would like to emphasize that our original submission has taken care to compare against the most modern and well-known subquadratic models that have been evaluated on large-scale pretraining such as Hyena, RWKV, and RetNet.

Additionally, the Introduction states:

We empirically validate Mamba’s potential as a general sequence foundation model backbone, in both pretraining quality and domain-specific task performance, on several types of modalities and settings.

Foundation modeling – e.g. demonstrating scaling abilities on large-scale pretraining data – has become a fundamental paradigm in modern machine learning that is arguably more important than individual benchmarks. While small-scale benchmarks such as WikiText-103 (around 100M tokens) and Long Range Arena (LRA) have been popular for demonstrating the potential of new models such as SSMs, our work deliberately moves beyond these toward showing that SSMs – with our new contributions – can actually be effective at scale (e.g. on the Pile dataset - 300B tokens - 3000x the size of WikiText-103).

Length extrapolation

We have included an attached figure that evaluates length extrapolation of pretrained 3B parameter language models: https://ibb.co/XVT0hGJ

This figure was generated by taking the validation set of the Pile dataset, feeding in each example with no padding or concatenation, and measuring the loss per token. The mean loss (log perplexity) for each position is plotted. Note for example that the perplexity for the first token is high because it has no context, and improves up to the training context length (2048) for both Mamba and the baseline Transformer (Pythia). Interestingly, Mamba’s perplexity improves significantly past its training context up to around length 3000.

We would like to emphasize that length extrapolation was not a direct motivation of our model and we view these capabilities as a bonus feature. We note that

1. The baseline model (Pythia) here was not trained with length extrapolation in mind, and there are perhaps other Transformer variants that are more generalizable (e.g. T5 or Alibi relative positional encodings).
2. We are not aware of any open-source 3B models trained on the Pile using relative positional encodings, and thus could not create this comparison.
3. Mamba, like Pythia, was not trained with length extrapolation in mind, and are thus comparable. Just as Transformers have many techniques (e.g. different positional embeddings) to improve their capabilities on axes such as length generalization, in future work it may be interesting to derive SSM-specific techniques for similar capabilities.

We thank reviewers du8a and 5ZBk for the interesting suggestion.

Code and model release

We are planning to open-source all model code as well as the weights of our pretrained models of all sizes (125M to 2.8B parameters). We hope that Mamba presents a valuable contribution to the community.