Demystify Mamba in Vision: A Linear Attention Perspective

We would first like to express our appreciation for your time and insightful comments. Please find our response to your concerns in the following:

1. The analysis of single-head attention.

Thanks for your insightful comment.

• The concept "head" in our paper is the same as its original definition in multi-head attention [1], i.e. the number of groups of $Q, K$, rather than the number of dynamic matrices. Since there exists only one set of $C,B$ in S6 (see Eq. 11 of our paper), it resembles single-head linear attention.
• A head can possibly generate multiple dynamic matrices. In multi-head attention, one head can only generate a single dynamic matrix $QK^\top$. However, a head of S6 can produce multiple equivalent dynamic matrices. This can be attributed to its special input gate and forget gate design. Specifically, as shown in Eq. 11 and Eq. 12 of our paper, there is only one set of $C\in\mathbb{R}^{N\times d}, B\in\mathbb{R}^{d\times N}$ in S6. Therefore, without considering the input gate $\Delta_i$ and forget gate $\hat{A}_i$, the model can produce only one dynamic attention matrix $CB\in\mathbb{R}^{N\times N}$. However, the input gate $\Delta_i$ and forget gate $\hat{A}_i$ make things different. Specifically, $\Delta_i\in\mathbb{R}^{1\times c}$ and $\hat{A}_i\in\mathbb{R}^{d\times c}$ are not shared across channels $c$, which enables S6 to produce $c$ equivalent dynamic matrices with the single set of $C, B$. Intuitively, S6 generates one attention matrix $CB\in\mathbb{R}^{N\times N}$, yet each channel filters tokens differently based on its input gate weight and introduces unique local biases with its forget gate weight, leading to varied equivalent dynamic matrices.
• In conclusion, S6 uses the single-head design, but its input gate and forget gate lead to multiple equivalent dynamic matrices.
• Additionally, we can introduce multi-head design to S6 by initially generating multiple sets of $C^h\in\mathbb{R}^{N\times d}, B^h\in\mathbb{R}^{d\times N},h=1,\cdots,\rm{num\\_heads}$, akin to multi-head attention.

[1] Attention is all you need. In NeurIPS, 2017.

2. The comparison on ImageNet-1K dataset.

Firstly, MESA was not utilized in ablation experiments or downstream tasks. These results already fully support the two core findings of this paper: (1) The forget gate and block design are the key designs of Mamba. (2) With the merits of these two key designs, MLLA can outperform various vision Mamba models.

Secondly, we clarify MESA as follows:

• MESA is an overfitting prevention strategy rather than an optimizer.
• Without MESA, MLLA model can also achieve comparable results. For example, after removing MESA and increasing drop path rate—a commonly used strategy to prevent overfitting—to 0.3, MLLA-T yields 83.3 accuracy on ImageNet, which still significantly surpasses various SOTA vision Mamba model, e.g. LocalVMamba-T: 82.7, VMamba-T: 82.5, etc.
• In our early experiments, we found MESA slightly beneficial for MLLA model, but it did not benefit other models as well and even hindered their performance.

3. The motivation for using Swin as the baseline.

• Window-based attention is not used. The shifted window attention in Swin Transformer is replaced with global linear attention to create our baseline model. We only employ the macro structure of Swin, i.e. width, depth and etc., as a clean architecture.
• The baseline model shares the identical configuration as Swin-T, with widths = [96, 192, 384, 768] and depths = [2, 2, 6, 2].

4. The setting of d_state in SSM block design.

The concept of d_state is akin to the head_dim in attention. Therefore, following Swin Transformer, we set it as 32 for all the models in our study.

5. The experiments on downstream tasks.

• MLLA's effectiveness on downstream tasks is fully validated by comprehensive COCO object detection and instance segmentation experiments. Hence, we conduct only one experiment on semantic segmentation task to save computation resources.

• As requested, we provide additional results of semantic segmentation task below.

  Backbone	#Params	FLOPs	mIoU SS	mIoU MS
  VMamba-T	62M	949G	47.9	48.8
  MLLA-T	55M	932G	48.1	49.0

  Due to time and resource constraints, we couldn't provide the MLLA-S results here.

• To the best of our knowledge, the result for Mask R-CNN 3x + base model is not provided by any vision Mamba models we compared with. Hence, we also omit this experiment to save computation resources.

• Full comparison with VMamba-T on object detection is provided in the PDF in general response. MLLA-T achieves comparable results to VMamba-T in the 3x setting but slightly underperforms VMamba-T in 1x training. This can be attributed to MLLA-T having 6M fewer parameters than VMamba-T, possibly requiring longer training to converge. MLLA-S/B significantly outperform VMamba-S/B in both 1x and 3x settings.

6. The comparison in Table 6.

• As detailed in our paper, the block structure is also a key design of Mamba and our MLLA. Hence, we provide a direct comparison with other linear attention methods in Table 6. to validate the effectiveness of our design.

• Here we offer additional comparison under exactly same macro architecture, the Swin-T structure.

  Method	#Params	FLOPs	Acc.
  Hydra Attention	29M	4.5G	80.7
  Efficient Attention	29M	4.5G	81.0
  FLatten Transformer	29M	4.5G	82.1
  MLLA	30M	4.8G	82.6

  MLLA shows obviously better result. It is noteworthy that we are unable to provide results for SOFT under this setting since it requires more than 1 hour to train one epoch.

7. Typo.

Thanks. It will be corrected.

We would first like to express our appreciation for your time and insightful comments. Please find our response to your concerns in the following:

1. Reason for using Swin Transformer.

We offer clarification on employing Swin Transformer:

• Swin Transformer is a widely used architecture in vision tasks. We use it as a clean and fair structure to validate the effectiveness of each special design of Mamba.
• Notably, the core operation of Swin Transformer, shifted window attention, is replaced with global linear attention to from our baseline model. Therefore, we only employ the macro structure of Swin—such as width and depth—without employing its window-based attention mechanism.

2. Replacing the forget gate with positional encoding.

• Thanks for the insightful question. Yes, the forget gate decays previous hidden state, thus providing local bias and position information.
• Local bias and decay do not hinder performance. Instead, they are beneficial for the model. For example, RoPE provides long-term decay and local bias, as demonstrated in its paper [1]. As shown in Table 2 of our study, integrating RoPE into the baseline linear attention model leads to an obvious performance improvement from 77.6 to 80.0. This indicates that local bias and decay can enhance model effectiveness.
• The reason why positional encodings outperform forget gate in Table 2 is that they can enjoy a global receptive field. As analyzed in our paper, the forget gate has to employ recurrent calculation, which restricts the receptive field of each token to the preceding sequence, resulting in a causal mode. As verified by a concurrent work, MambaOut [2], such a causal mode can lead to notable performance drop in vision models (see its Fig. 3). Therefore, while the local bias and decay of the forget gate can benefit the model, its causal mode hinders performance. In contrast, positional encodings enable the model to benefit from both local bias and a global receptive field at the same time, thus yielding better results than the forget gate.

[1] Roformer: Enhanced transformer with rotary position embedding. Neurocomputing.

[2] MambaOut: Do We Really Need Mamba for Vision? arXiv:2405.07992.

Dear Reviewer dwbt, thank you for your insightful review and for engaging with our work. We would like to know if there are any additional questions or concerns. We are eager to engage in further discussion and provide clarification to fully address them.

We would first like to express our appreciation for your time and insightful comments. Please find our response to your concerns in the following:

1. Mamba for image synthesis.

• Thanks for pointing out these works. We will give more credits to them and include detailed discussions in the revised version.
• The exploration in this paper is orthogonal to previous explorations on Mamba, and our work can further benefit previous works in two ways:
  • a. Previous studies on Mamba can substitute their SSM with our MLLA block, which is proved to be more effective and efficient.
  • b. The analyses in our work can help other studies design better Mamba models. Our work reveals some suboptimal designs of Mamba, such as the absence of attention normalization, an important component. Hence, previous works on Mamba, e.g. Zigma, can add attention normalization to their model to achieve better results. This enhancement is compatible with other explorations, such as layer-wise scan paths with 8 directions.
• This paper focuses on analyzing. Currently we mainly conduct classification, object detection and semantic segmentation experiments, which already validate our analyses and findings. In the future, we will apply MLLA to image synthesis task to develop efficient models.

2. Discussion with the mentioned paper.

• Thank you for pointing this out. The mentioned paper was made public on June 10, 2024, after the NeurIPS submission deadline of May 22, 2024.
• The mentioned work proposes a unified framework for efficient autoregressive sequence transformations, demonstrating various linear attention models within this framework. Despite some similarities with our work, there are fundamental distinctions:
  • a. Our work not only reveals the close relationship between Mamba and linear attention, but more importantly, provides analyses and experiments to verify the effectiveness of each design. In contrast, the referenced work mainly develops a framework to include different linear attention models.
  • b. Our analyses ultimately leads to the development of a novel linear attention model, MLLA. In contrast, the mentioned work focuses on the development of efficient training algorithms for existing models.
• We will discuss with this paper in the revised manuscript.

3. Extension to other linear attention-based models.

• The exploration in this work can be extended to other linear attention-based models.

• Our work reveals that Mamba can be viewed as a special variant of vanilla linear attention. Other linear attention-based models, e.g. RWKV, xLSTM, can also be viewed as variants of vanilla linear attention. Therefore, the exploration and insights in this work can naturally extend to these models.

• We take xLSTM as an example. According to Eq. 4 of our paper, linear attention can be written as:

  $$S_i=S_{i-1}+K_i^\top V_i,\quad Z_i=Z_{i-1}+K_i^\top, \quad y_i=Q_i S_i / Q_i Z_i$$

  Using the same notations, xLSTM can be formulated as:

  $$S_i=F_i S_{i-1}+I_i K_i^\top V_i,\quad Z_i=F_i Z_{i-1}+I_i K_i^\top, \quad y_i=Q_i S_i / \text{max}\\{1, Q_i Z_i\\},$$

  where $I_i$ is input gate and $F_i$ is forget gate. Additionally, xLSTM employs multi-head design according to its paper. Therefore, xLSTM resembles linear attention with additional input gate $I_i$, forget gate $F_i$ and modified normalization $\text{max}\{1, Q_i Z_i\}$. As a result, many explorations of this work can be extended to xLSTM. For instance, we can make bold speculations like:

  • a. The forget gate $F_i$ is likely to play a crucial role in xLSTM.
  • b. When applying xLSTM to vision models, the forget gate could possibly be replaced with positional encodings. Just like what we did in Table 2.
  • c. The input gate $I_i$ may not offer significant benefits in vision tasks.
  • d. Enhanced block designs from Mamba could potentially benefit xLSTM.
  • e. The modified normalization $\text{max}\\{1, Q_i Z_i\\}$ might provide greater stability compared to vanilla linear attention, warranting further exploration.

• In conclusion, the analyses and insights of our work can extend to other linear attention-based models, helping them improve interpretability, further enhance models, etc.

We would first like to express our appreciation for your time and insightful comments. Please find our response to your concerns in the following:

1. The name of MLLA model.

Thanks for your comment.

Firstly, we name the final model Mamba-Like Linear Attention because it incorporates two key designs from Mamba: the forget gate and block design. Given that there is no SSM in the final model, maybe Mamba-Inspired Linear Attention (MILA) is a better name?

Secondly, we are very happy to rename the model if a more suitable name is suggested.

2. Previous "Demystify" papers.

Thanks for the valuable suggestion. We offer discussions with these works.

[1] studies the connection between local attention and dynamic depth-wise convolution, validating that dynamic depth-wise convolution can perform on-par with or slightly better than local window attention. It highlights the effectiveness of larger kernel sizes and input-dependent weights for vision tasks. [2] replaces the attention operations in Swin model families with simple linear mapping layers and shows that the macro architecture may be more responsible for high model performance. [3] proposes a unified macro architecture to identify the real gains of popular convolution and attention operators. It demonstrates that different token mixers achieves varying performance under the same macro architecture, with Halo attention and deformable convolution yield optimal results.

Based on these papers and our findings regarding the forget gate and block design, we boldly make the following inferences:

• Both macro and micro structures are important. [1, 2] emphasize the importance of macro structure, while [3] shows that superior token mixers can enhance performance under the same macro architecture. In our study, the forget gate and block design serve as micro and macro structures, respectively, both playing significant roles.
• Input-dependent feature, large receptive field and local bias are beneficial for effective micro token mixer designs. [1, 3] and our paper suggest these properties are effective. This further explains the effectiveness of MLLA, which incorporates all these features while maintaining linear complexity.

We will cite these related works and provide detailed discussion in the revised version.

[1] On the Connection between Local Attention and Dynamic Depth-wise Convolution, ICLR 2022.

[2] What Makes for Hierarchical Vision Transformer? DOI: 10.1109/TPAMI. 2023.3282019.

[3] Demystify Transformers & Convolutions in Modern Image Deep Networks, arXiv:2211.05781.

3. Training cost.

32 GPUs are used to perform this research. Each tiny-scale model, including ablation experiments, requires around 12 hours for training on ImageNet-1K with 32 GPUs.

4. The position of Fig. 1.

Thanks for your question. Fig. 1 is placed on page 2 as it serves as the main figure of our paper. Considering it is cited on page 4, we will adjust its position.

5. The results of Mamba + Swin.

Due to time and computation resource constrains, we are unable to provide the results of Mamba + Swin architecture. However, we find that the initial version of VMamba yields pertinent findings by applying Mamba to Swin architecture in a cross-scan approach. The results are provided below.

Dear Reviewer QFdu, thank you for your insightful review and for engaging with our work. We would like to know if there are any additional questions or concerns. We are eager to engage in further discussion and provide clarification to fully address them.