Paper: Mamba: Linear-Time Sequence Modeling with Selective State Spaces (ICLR-2024)

Reviewer comment (Theory):
I am curious if we could build a better understanding of the selection mechanism that you propose.

Author response:
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.

Reviewer comment (Theory):
In Theorem 1 you link that mechanism to gating in RNNs as a special case. Is it possible to understand better the generalization through some discussion / qualitative examples?

Author response:
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.

Reviewer comment (Novelty):
The paper lacks references to some relevant works, such as [1], [2], [3], [4] which discuss some Linear Attention methods, and [5], which is also a LongConv method. These references are completely absent in the paper. It is suggested that the authors consider adding these citations to provide a more comprehensive review of related work.

Author response:
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.

Reviewer comment (Experiments):
Adding extrapolation experiments to the language model would be interesting.

Author response:
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.

Reviewer comment (Experiments):
The ablation analysis in Table 6 should be more comprehensive, with a total of $2^3$ possible combinations. The authors are advised to include the remaining two combinations.

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

| Selective $\Delta$ | Selective $B$ | Selective $C$ | Perplexity |
|--------------------|---------------|---------------|------------|
|                    |               |               | 10.93      |
|                    | +             |               | 10.15      |
|                    |               | +             | 9.98       |
| +                  |               |               | 9.81       |
|                    | +             | +             | 9.49       |
| +                  | +             |               | 9.21       |
| +                  |               | +             | 9.07       |
| +                  | +             | +             | 8.71       |

Reviewer comment (Experiments):
What's your setting of Scaling Law? Why is your ratio of token number and model size the same as Chicilla's paper? It is supposed that the FLOPs of Transformers and SSMs would differ. Given the same amounts of total parameters, is this difference in FLOPs important to the final performance (accuracy)?

Author response:
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.

Reviewer comment (Reproducibility):
How did you parameterize the first convolutional layer in the Mamba-Block?

Author response:
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.

Reviewer comment (Reproducibility):
Providing more detailed implementation, such as offering core code, would be very helpful.

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

Reviewer comment (Theory):
The model still has a quadratic memory requirement during training like Transformers.

Author response:
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).

Reviewer comment (Experiments):
Have you evaluated scaling behavior beyond 1.4B parameters? How does it compare to Transformers at 10B scales?

Author response:
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).

Reviewer comment (Experiments):
The input selection mechanism introduces additional hyper parameters. How sensitive are the results to hyperparameters like the projection rank?

Author response:
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.

Reviewer comment (Experiments):
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.

Author response:
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**.


| Model (125M)   | WikiText-103 ppl | $\Delta$ to Transf. |
|----------------|------------------|-----|
| Transformer    | 18.6             | 0.0 |
| Hybrid H3      | 18.5             | -0.1 |
| Performer      | 26.8             | 8.2 |
| Reformer       | 25.6             | 7.0 |
| AFT-conv       | 28.2             | 9.6 |
| Linear Attn.   | 25.6             | 7.0 |
| Hyena          | 18.5             | -0.1 |
| Transf. (ours) | 18.0             | -0.6 |
| **Mamba** (ours) | **16.3**         | **-2.3** |

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.

|Model             |WikiText-103 ppl  | Transformer ppl | $\Delta$ |
|------------------|------------------|------|-------|
|SGConv + Attn [2] | 18.70            | 18.3 | 0.40  |
|MEGA [3]          | 18.07            | 18.3 | -0.23 |
|SPADE [4]         | 18.5             | 18.8 | -0.3  |

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.

| Model (44M)      | WikiText-103 ppl | $\Delta$ to Transf. |
|---------------   |---------         |---------------------|
| **Transformer**  | 24.78            | 0.0                 |
| FLASH            | 26.70            | 1.92                |
| Linear Attn.     | 28.05            | 3.27                |
| Performer        | 63.16            | 38.38               |
| cosFormer        | 27.06            | 2.28                |
| gMLP             | 29.13            | 4.35                |
| Syn(D)           | 32.43            | 7.65                |
| S4               | 39.66            | 14.88               |
| DSS              | 41.07            | 16.29               |
| GSS              | 30.74            | 5.96                |
| RWKV             | 25.07            | 0.29                |
| LRU              | 31.12            | 6.34                |
| TNN              | 24.67            | -0.11               |
| HGRN             | 24.82            | 0.04                |

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

Reviewer comment (Experiments):
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?

Author response:
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

Reviewer comment (Experiments):
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).

Author response:
**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.

Reviewer comment (Experiments):
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.

Author response:
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

Reviewer comment (Experiments):
Memory benchmarking is missing. Even though Section 4.5 is titled “speed and memory benchmark”, only speed comparisons are presented. Also, the authors should provide more detailed setups of Figure 8 left, e.g., model layers, model sizes, details of the convolution, etc. Could the authors provide some intuitions why FlashAttention is the slowest when the sequence length is very large (Figure 8 left)?

Author response:
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).

| Batch size | Transformer (w/ FlashAttention-2) | Mamba  |
|------------|-----------------------------------|--------|
|          1 | 4.6GB                             | 4.8GB  |
|          2 | 5.2GB                             | 5.8GB  |
|          4 | 6.9GB                             | 7.3GB  |
|          8 | 11.5GB                            | 12.3GB |
|         16 | 20.7GB                            | 23.1GB |
|         32 | 34.5GB                            | 38.2GB |

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.

**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.

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.