Demystify Mamba in Vision: A Linear Attention Perspective

Summary

This paper presents a thorough analysis of the key factors contributing to the success of the S6 module in Mamba model, and introduces a new linear attention vision network, MLLA, inspired by the S6’s design. Extensive experimental results show outstanding performance, validating the effectiveness of the proposed model.

Strengths

The motivation and analysis underlying this work are fascinating and perceptive, offering a profound examination of the mathematical relationships between Mamba and Linear Attention. A thorough investigation of the performance contributions of each component from S6 is conducted, with corresponding refinements made to linear attention, ultimately leading to the development of MLLA. The proposed MLLA serves as a generic vision backbone network, which outperforms recent Mamba-based models, demonstrating its superior capabilities

Weaknesses

I appreciate the paper's insightful analyses, but several key concerns regarding methodology and experiments need to be addressed. If the authors fully resolve these issues, I would be willing to reconsider my evaluation. **Possible inaccurate analysis of single-head attention** The analysis of single-head attention appears to be inaccurate and unclear. Generally, multi-head attention in Transformer generates H dynamic matrices during attention computation, where H represents the number of heads. However, the paper seems to interpret S6 as a single-head module without providing sufficient evidence. The authors should explain this point from the perspective of the number of attention matrices generated. Moreover, according to Reference 1, S6 can actually produce multiple dynamic matrices instead of just one. **Unfair comparison on ImageNet-1K dataset** The paper uses MESA as its optimizer, improving accuracy compared to AdamW. In contrast, the methods in Table 3 do not use MESA and may also be prone to overfitting. For fairness, authors should ensure that training strategies are consistent with competitors. Moreover, the authors should not hide this apparent difference in the appendix. **Questionable motivation for using Swin as the baseline** The motivation for using Swin as a baseline is questionable. One important aspect of linear attention is its ability to efficiently establish global dependencies, so it is unclear why window-based linear attention is used. Additionally, the specific configuration of the baseline model, such as channels and depth, is not provided. **Unclear setting of d_state in SSM block design** In the block design, it is unclear what value d_state of SSM is set to. **Insufficient experiments on downstream tasks** (1) Semantic segmentation should include tiny, small, and base models. (2) Mask R-CNN 3x should include the base model. (3) The performance comparison of MLLA-T and VMamba-T on object detection is lacking. **Unfair comparison in Table 6** The comparison in Table 6 is unfair. The authors aim to compare the performance of different linear attention methods, but different block settings have a significant impact. The authors should use the same architecture, replacing only token mixer module. **Typo** In line 294, 'SMM' should be 'SSM'.

Questions

Please see the main weaknesses.

Rating

7

Confidence

5

Soundness

3

Presentation

2

Contribution

3

Limitations

Please see the main weaknesses.

Summary

The paper takes a deep dive into the workings of Mamba in vision related tasks and compares it to Linear Attention. They conclude that Mamba is a special form of Linear Attention and describe the role that certain parts in the architecture and give them a name that better represent their actual use. They come up with input and forget gate that control the flow of information while shortcuts are closely related to the residual connections. The authors take those learnings and apply them to linear attention and show that those concepts do improve the linear attention mechanism up to the point where it beats mamba.

Strengths

The authors draw clear connections between Linear Attention and Mamba which is showcased in terms of raw similarity (eq. 11 and 12). Furthermore, they underline this by showing the input gate values on two example images which helps general understanding. They demonstrate the usefulness of the explained modules in Mamba apply them to Linear Attention and show that they do improve performance. In total a sound paper. The overall gist can be followed very nicely.

Weaknesses

No apparent weaknesses

Questions

- Why did you choose to use Swin Transformer? - The forget gate creates local bias and position information but also decays the previous hidden state if I understood correctly. When replaced by RoPE or similar the performance increases as shown in Table2. Does this mean that the local bias and the decay can hinder performance?

Rating

7

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

The authors mention the limitations of their work, especially that is it no exhaustive.

Summary

This paper explores the similarities and differences between the Mamba and linear attention Transformer models. It redefines Mamba as a variant of linear attention Transformer with six key distinctions. The paper also studies each design's impact, identifying the forget gate and block design as key to Mamba's success. Based on these findings, the paper introduces a Mamba-like linear Attention (MLLA) model, which outperforms other vision Mamba models on various tasks while maintaining fast computation and inference speed.

Strengths

* The paper is well-written and well-motivated from the perspective of linear attention. * The findings are interesting to me.

Weaknesses

* Several works using Mamba for image synthesis should be discussed [1][2][3]. In Line 68, I wonder whether the exploration in this paper is orthogonal to previous explorations on Mamba. For example, Zigma considers layerwise scanpaths with 8 directions. Can the proposed method still yield some improvements on that result? * Another consideration is that the authors should discuss something with this paper: https://arxiv.org/abs/2406.06484. Check Table 4; it provides a much more general framework that includes most of the linear attention models. * My third concern is whether the exploration in this work can extend to other linear attention-based models, such as RWKV and xLSTM.

Questions

I have some concerns (see weaknesses) about this paper. I am inclined to raise my score if my concerns are fully resolved. **Reference:** [1],Diffusion Models Without Attention.CVPR24. [2],ZIGMA: A DiT-style Zigzag Mamba Diffusion Model,ECCV24. [3],Scalable Diffusion Models with State Space Backbone,Arxiv.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

Yes, they are addressed and discussed.

Summary

This paper primarily discusses the similarities and differences between the Mamba model and the Linear Attention Transformer, and conducts an in-depth analysis of the key factors contributing to Mamba's success in visual tasks. The paper elaborates on six major design differences of Mamba compared to the Linear Attention Transformer: input gate, forget gate, shortcut connection, attention-free normalization, single-head attention, and modified block design. For each of these designs, the author meticulously analyzes their advantages and disadvantages, and evaluates their applicability in visual tasks. Through experimental research, the paper emphasizes that the forget gate and block design are the core contributors to Mamba's success. Based on these findings, the author proposes a model called Mamba-Like Linear Attention (MLLA), which integrates the advantages of the forget gate and block design into the linear attention framework. Experimental results show that MLLA outperforms various visual Mamba models in image classification and high-resolution dense prediction tasks, while maintaining parallel computing capabilities and fast inference speed.

Strengths

1. The paper provides a new perspective for understanding Mamba's success by thoroughly analyzing the similarities and differences between the Mamba model and the Linear Attention Transformer. By rephrasing formulas, the paper considers Mamba as a variant of the Linear Attention Transformer and clearly points out Mamba's six design features, including input gate, forget gate, shortcut connection, attention-free normalization, single-head attention, and modified block design. 2. Each design aspect is meticulously analyzed, and its advantages and disadvantages in visual tasks are evaluated. Experimental results verify that the forget gate and block design significantly contribute to Mamba's performance improvement, while other designs either have minimal marginal contributions or may potentially harm model performance. 3. The Mamba-Like Linear Attention (MLLA) model is proposed, combining the advantages of the forget gate and block design in linear attention. The MLLA model outperforms existing Mamba models in image classification and high-resolution dense prediction tasks while maintaining computational parallelization and high-speed inference capabilities.

Weaknesses

There are some minor issues: 1. In the final MLLA model, there is no SSM but a forget gate. The forget gate was originally proposed in RNN and LSTM. Thus calling it a Mamba-like model seems to be improper. 2. Previous "Demystify" papers are not cited in the paper, e.g., [a, b, c]. The final architecture MLLA follows the macro structure of Swin and [a,b,c] all show that the macro structure is very important but the local blocks are not so important. The paper should build a connection with the previous "Demystify" papers and tell us what are the most effective parts? [a] On the Connection between Local Attention and Dynamic Depth-wise Convolution, ICLR 2022 [b] What Makes for Hierarchical Vision Transformer? DOI: 10.1109/TPAMI.2023.3282019 [c] Demystify Transformers & Convolutions in Modern Image Deep Networks, arXiv:2211.05781

Questions

1. How many GPUs are used to perform this research, including the training cost and ablation experiments cost? 2. Why Fig. 1 is placed on page 2 but cited on page 4? 3. What are the results of Mamba + Swin?

Rating

7

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

Nope.