Demystify Mamba in Vision: A Linear Attention Perspective

Authors: dongchen  han
ziyi  wang
zhuofan  xia
yizeng  han
yifan  pu
chunjiang  ge
jun  song
shiji  song
bo  zheng
gao  huang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 07:29:46

Abstract: Mamba is an effective state space model with linear computation complexity. It has recently shown impressive efficiency in dealing with high-resolution inputs across various vision tasks. In this paper, we reveal that the powerful Mamba model shares surprising similarities with linear attention Transformer, which typically underperform conventional Transformer in practice. By exploring the similarities and disparities between the effective Mamba and subpar linear attention Transformer, we provide comprehensive analyses to demystify the key factors behind Mamba’s success. Specifically, we reformulate the selective state space model and linear attention within a unified formulation, rephrasing Mamba as a variant of linear attention Transformer with six major distinctions: input gate, forget gate, shortcut, no attention normalization, single-head, and modified block design. For each design, we meticulously analyze its pros and cons, and empirically evaluate its impact on model performance in vision tasks. Interestingly, the results highlight the forget gate and block design as the core contributors to Mamba’s success, while the other four designs are less crucial. Based on these findings, we propose a Mamba-
Inspired Linear Attention (MILA) model by incorporating the merits of these two key designs into linear attention. The resulting model outperforms various vision Mamba models in both image classification and high-resolution dense prediction tasks, while enjoying parallelizable computation and fast inference speed. Code is available at https://github.com/LeapLabTHU/MLLA.

Main paper: 
 introduction: 
 Recently, state space models, exemplified by Mamba, have rapidly gained wide research interest. In contrast to the quadratic complexity of prevailing Transformer models, the state-space-based Mamba offers effective sequence modeling with linear complexity. This crucial property allows Mamba to handle extremely long sequences with manageable computational costs, making it a promising architecture for both natural language processing [14, 29] and visual recognition [57, 31] . However, Mamba is not the first model to achieve global modeling with linear complexity. Linear attention [26] , an early work, was proposed as an computationally efficient alternative to the widely adopted Softmax attention [42] , namely dot-product attention. Specifically, linear attention replaces the non-linear Softmax function in attention operation with linear normalization. This enables a change in computation order from (QK ⊤ )V to Q(K ⊤ V ), thus reducing computation complexity from O(N 2 ) to O(N ). Despite its efficiency, previous works [4, 39, 15, 16] proved that linear attention suffers from insufficient expressive power, making it impractical for real applications. Surprisingly, we find a very close relationship between the formulas of high-performance Mamba 11)) and single head linear attention (eq. ( 12)).It can be seen that selective SSM resembles single-head linear attention with additional input gate ∆ i , forget gate A i and shortcut D ⊙ x i , while omitting normalization Q i Z i .and subpar linear attention Transformer. Therefore, a compelling research question emerges: What factors contribute to Mamba's success and its significant superiority to linear attention Transformer?In this paper, we offer both theoretical and empirical analyses to unveil Mamba through the lens of linear attention Transformer. Specifically, we rewrite the formulas of selective state space model and linear attention within a unified formulation, depicting Mamba as a variation of linear attention Transformer with six distinctions: input gate, forget gate, shortcut, no attention normalization, single-head, and modified block design. To demystify what factors lead to Mamba's effectiveness, empirical studies on vision tasks are conducted to assess the impact of each special design. The results demonstrate that the forget gate and block design tend to be the two core contributors to Mamba's superiority. While the block design can be easily adopted, the forget gate necessitates recurrent computation, which may not be well-suited for non-auto-regressive vision models. Therefore, we delve into the essence of the forget gate and verify that it can be replaced by suitable positional encoding in vision tasks. Based on our findings, we introduce the two core contributors or their alternatives to linear attention Transformer, presenting our Mamba-Inspired Linear Attention (MILA) model.Experimental results demonstrate that MILA achieves superior results to various vision Mamba models in both image classification and high-resolution dense prediction tasks, validating that linear attention can surpass Mamba with the merits of two core designs.Our main contributions and takeaways are as follows:• We reveal Mamba's close relationship to linear attention Transformer: Mamba and linear attention Transformer can be formulated within a unified framework, with Mamba exhibiting six distinct designs compared to the conventional linear attention paradigm: input gate, forget gate, shortcut, no attention normalization, single-head and modified block design. • We provide detailed analyses of each special design and empirically validate that the forget gate and block design largely lead to Mamba's superiority. Additionally, we demonstrate that the recurrent calculation of the forget gate might not be ideal for vision models. Instead, proper positional encoding can function as the forget gate in vision tasks, while preserving parallelizable computation and fast inference speed. • We develop a series of linear attention vision Transformer models named MILA, which inherit the core merits of Mamba and tend to be more suitable for vision tasks than the original Mamba model. 
related works: 
 Vision Transformer and attention. Originating from natural language processing, Transformer and attention have been highly successful in vision tasks [32, 3, 36, 37] , demonstrating superiority to the conventional CNN models [18, 7, 12] . However, the quadratic complexity of widely adopted Softmax attention [42] poses challenges in handling high-resolution images. Numerous works have attempted to reduce the computational cost by introducing local attention windows [32, 9, 20] or sparsity [43, 45, 56] . Linear attention [26] , another approach, inherently offers linear complexity O(N ) and is capable of modeling long sequences. Despite its efficiency, previous works [4, 39, 47, 17] have shown that linear attention always fails to deliver satisfactory results, limiting its applicability.Mamba [14] is a recently proposed state space model that achieves effective sequence modeling with linear complexity. Motivated by its potential for modeling high-resolution images, many researchers try to apply Mamba to vision tasks [14, 31, 27, 25, 38, 49, 48, 23] . For instance, VMamba [31] introduces a cross-scan module to enable 1D selective scanning in 2D image space. LocalMamba [25] utilizes local windows to enhance local modeling capability. EfficientVMamba [38] designs an atrous-based selective scan approach to enhance efficiency. In addition, MambaOut [51] analyzes whether Mamba is needed for vision, and explainability methods [1] have also been proposed.In contrast to incorporating Mamba into vision, this paper reveals the surprising similarities between the formulas of inferior linear attention Transformer and powerful Mamba model. This interesting finding gives us the opportunity to demystify the key factors behind Mamba's success. 
preliminaries: 
 This section revisits the formulations of attention and selective state space model. To facilitate comparison in Sec. 4, we employ identical notations for the dimensions of certain variables in both linear attention and selective state space model, and make some modifications to the formula formats.Let x ∈ R N ×C denote a sequence of N features with dimension C. Single head Softmax attention [42] , also known as dot-product attention, can be written as:EQUATIONwhereW Q , W K ∈ R C×d , W V ∈ R C×C denote projection matrices, Q, K ∈ R N ×d , V ∈ R N ×Crepresent query/key/value matrices, andQ i , K i ∈ R 1×d , V i ∈ R 1×Care individual query/key/value tokens. Softmax attention computes the similarities between each query-key pair, leading to O(N 2 ) complexity. Therefore, it incurs unbearable computational cost in long-sequence modeling scenarios.Linear attention [26] , another attention paradigm, is proposed to effectively address this problem by reducing the computation complexity to O(N ). Specifically, linear attention replaces the non-linear Softmax function with linear normalization, and adopts an additional kernel function ϕ in Q and K:EQUATIONThis enables the rearrangement of the computation order from (QK ⊤ )V to Q(K ⊤ V ) based on the associative property of matrix multiplication, thus reducing computation complexity to O(N ).Equation ( 2) defines linear attention with a global receptive field, where each query aggregates information from all keys and values. In practice, linear attention can also be implemented in autoregressive models, restricting the receptive field of the i-th token to proceeding tokens, i.e., token j, j ≤ i. This causal linear attention is formulated as follows:EQUATIONThis results in a recurrent linear attention form:EQUATIONState space model (SSM). The classical state space model is a continuous system that maps the input x(t) ∈ R to output y(t) ∈ R through a hidden state h(t) ∈ R d×1 , which can be written as follows:EQUATIONDiscrete SSM. To be applied to deep neural networks, SSM is first transformed into its discrete version through zero-order hold discretization. Specifically, the continuous parameters A, B are transformed into their discretized counterparts A, B using a timescale parameter ∆ ∈ R:EQUATIONTherefore, discrete SSM rewrite eq. ( 5) as:EQUATION)(a) Selective SSM Model 8)) and its equivalent form (eq. ( 9)).Selective State Space Model. Mamba [14] improves SSM with selection, presenting the selective state space model. The parameters B, C, ∆ is set as the function of x i , thus becoming input-dependent parameters B i , C i , ∆ i . As a result, the discretized parametersA i = exp(∆ i A), B i = ∆ i B i are also input-dependent.The selective state space model can be written as:EQUATIONFor the convenience of subsequent derivation, we make three modifications to eq. ( 8):• Mamba practically sets A, A i as diagonal matrices. Therefore,A i h i-1 = A i ⊙ h i-1, where A i = diag(A i ) ∈ R d×1 denotes the matrix composed of diagonal elements of A i .• GivenB i = ∆ i B i and ∆ i ∈ R, we have B i x i = ∆ i B i x i = B i (∆ i x i ) = B i (∆ i ⊙ x i ).• Dx i = D ⊙ x i , where ⊙ denotes the Hadamard product, i.e., element-wise multiplication.Consequently, we rewrite eq. ( 8) as:EQUATIONThe selective state space model formulated in eq. ( 9) can only deal with scalar input x i ∈ R. To operate over an input sequence x ∈ R N ×C , x i ∈ R 1×C , Mamba applies eq. ( 9) independently to each channel, leading to the following formulations:EQUATIONwhere B i , C i , ∆ i are derived from the input. Specifically, Mamba employsB = (xW B ) ⊤ , C = xW C , ∆ = Softplus(xW 1 W 2 ) to produce the parameters B ∈ R d×N , C ∈ R N ×d , ∆ ∈ R N ×C , where W B , W C ∈ R C×d , W 1 ∈ R C×C0 , W 2 ∈ R C0×Care projection matrices. Notably, eq. ( 10) is exactly the selective SSM employed in Mamba, we only make modifications to formula formats. 
connecting mamba and linear attention transformer: 
 In this section, we reveal the similarities and disparities between Mamba and linear attention Transformer from two perspectives: core operation and macro architecture.As detailed in Sec. 3, for an input sequence of N tokens x ∈ R N ×C , the formulations of selective state space model and linear attention are provided by eq. ( 10) and eq. ( 4), respectively. Many underlying similarities exist between the formulas of these two operations. To facilitate comprehension, we rewrite eq. ( 10) and eq. ( 4) with a unified formulation as follows:EQUATIONEQUATIONAs illustrated in Fig. 1 , a close relationship between eq. ( 11) and eq. ( 12) is evident. Specifically,h i ∼ S i ∈ R d×C , B i ∼ K ⊤ i ∈ R d×1 , x i ∼ V i ∈ R 1×d , and C i ∼ Q i ∈ R 1×d. Therefore, selective SSM can be viewed as a special variation of linear attention, indicating a very close connection between these two mechanisms. Furthermore, four major differences can be observed:1. In eq. ( 11), the input x i is augmented by Hadamard product with ∆ i . Since ∆ = Softplus(xW 1 W 2 ), all elements of ∆ i are positive. Therefore, we view ∆ i as an input gate, controlling whether to let the input x i into the hidden state.2. There is an additional A i in eq. (11) . Mamba sets A as a diagonal matrix with negative diagonal elements, thus ensuring all elements ofA i = diag(A i ) = exp(diag(A)∆ i )to fall between 0 and 1. Hence, we interpret A i as a forget gate, which decides the degree of attenuation for the previous hidden state h i-1 .3. A learnable shortcut from the input x i to the output y i is employed in eq. ( 11), i.e. D ⊙ x i .4. As depicted in eq. ( 12), linear attention divides the output by Q i Z i to maintain that the attention weights sum up to 1, while eq. ( 11) does not have such normalization.In addition to these four differences, it is also important to note that eq. ( 12) represents single-head linear attention as there is only one group of Q, K. This indicates that the selective state space model is akin to single-head linear attention and does not incorporate a multi-head design.In a word, the similarities and disparities between selective SSM and linear attention can be summarized as: selective state space model resembles linear attention with additional input gate, forget gate and shortcut, while omitting normalization and multi-head design.Input gate. As discussed before, ∆ i actually functions as an input gate for x i , determining its access to the hidden state. The values of this input gate are predicted from the current input x i as∆ i = Softplus(x i W 1 W 2 ). Therefore, by learning the weight of W 1 , W 2 , the model can discern the "utility" of x i , generating large ∆ i values for "useful" x i and small ones for "less useful" x i . For example, in vision tasks, tokens representing foreground objects may yield larger input gate values, while background tokens may yield smaller ones.Forget gate. A i acts as a forget gate in selective state space model, offering two essential properties: local bias and positional information. Firstly, all elements of A i strictly range from 0 to 1, indicating that the model consistently decays the previous hidden state h i-1 upon the arrival of the current token x i . This results in a strong local bias. Secondly, A i provides positional information for the model. It ensures that the model is sensitive to the order of input sequences. Without this forget gate, rearranging the order of the preceding sequence will not affect subsequent outputs. For instance, in recurrent linear attention, if we change the order of x 1 and x 2 , the outputs y i , i ≥ 3 will not change.Hence, the forget gate A i plays an important role in selective SSM.Despite its effectiveness, incorporating the forget gate also poses significant challenges. Firstly, it forces the model to adopt the recurrent formulation during both training and inference. Previous state space models typically use global convolution for efficient parallelizable training, which is incompatible with selective SSM due to the input-dependency of A i . As a remedy, Mamba [14] proposes a hardware-aware algorithm to speed up computation by performing parallel scan in recurrent mode (see the abstract of [14] ). Although effective, such recurrent calculation unavoidably reduces model throughput and is still slower than parallel linear attention (eq. ( 2)). Secondly, the forget gate inherently functions in causal mode, which may not be very suitable for non-autoregressive vision models. Using the forget gate A i in vision tasks requires transforming the image into a 1D sequence and conducting recurrent computation, which limits the receptive field of each image token to its preceding sequence and incurs extra latency. Therefore, we believe that the forget gate is ideally suited for modeling causal data, which naturally needs auto-regressive training and recurrent calculation. However, it may not be as suitable for non-causal data like images. We further speculate that a suitable positional encoding can substitute for the forget gate, since certain positional encodings, such as LePE [9] and RoPE [41] , can also provide local bias and positional information.Shortcut. Selective SSM employs a learnable shortcut D ⊙ x i , making it resemble a residual block [22] . This shortcut may aid in optimizing the model and stabilizing training.Normalization. The output in linear attention is divided by Q i Z i to ensure the attention weights sum up to 1. We believe this normalization is crucial for stabilizing training and improving model capacity.Let's consider an input αx i , α > 0. It is transformed into αQ i , αK i , αV i through projections. If there is no normalization on attention weights, as α increases, αQ i exhibits larger similarities with all keys αQ i K ⊤ j , ∀j. This indicates that longer tokens will have larger attention scores with every token, leading to longer output. Additionally, as α grows bigger, αK i yields bigger similarities with all queries αQ j K ⊤ i , ∀j. This implies that all queries will focus more on longer tokens. As a result, longer tokens tend to dominate the whole feature map, while shorter tokens may fail to represent their corresponding semantics. This may result in training instability and could possibly lower model's expressiveness. Normalizing the attention weights can significantly alleviate this issue.Multi-head. Linear attention commonly utilizes multi-head design [42] for better outcome. Multihead attention employs multiple groups of Q, K to produce attention matrices and allows the model to simultaneously attend to information from various representation subspaces at different positions, thus enhancing its expressive power.Modern linear attention Transformer models commonly adopt the block design depicted in Fig. 3(a) , which are comprised of a linear attention sub-block and a MLP sub-block. In contrast, Mamba modifies the block design by combining two basic designs, H3 [13] and Gated Attention [24] , resulting in the architecture illustrated in Fig. 3(b) . The improved Mamba block integrates multiple operations such as selective SSM, depth-wise convolution, linear mapping, activation function, gating mechanism, etc., and tends to be more effective than the conventional Transformer block design.Mamba can be seen as a variant of linear attention Transformer with specialized linear attention and modified block design. The special linear attention variation, i.e. selective state space model, has five major distinctions from the common linear attention paradigm, detailed in Sec. 4.2. And the differences in block designs are analyzed in Sec. 4.3. In summary, Sec. 4.2 and Sec. 4.3 reveal the intimate relationship between Mamba and linear attention Transformer, highlighting a total of six differences: five in core operation and one in macro design.  
empirical study: 
 Mamba [14] is seen as a powerful alternative to Transformer [42] , while linear attention models generally being considered inferior [39, 15] . In Sec. 4, we illustrated the surprisingly close relationship between Mamba and linear attention Transformer and pointed out six major distinctions. In this section, we conduct experiments to assess the impact of each distinction, shedding some light on the core contributors behind Mamba's success.We employ the widely used Swin Transformer [32] architecture to verify the effects of the six differences. Firstly, we substitute the Softmax attention in Swin Transformer with linear attention to create our baseline model. Subsequently, we separately introduce each distinction to the baseline model to assess its impact. Based on the results, we further validate whether linear attention can achieve superior results with the merits of Mamba. Specifically, we integrate the useful designs into linear attention Transformer to create our Mamba-Inspired Linear Attention (MILA) model, and assess its effectiveness by comparing it with various vision Mamba designs across multiple tasks, including ImageNet-1K classification [8] , COCO object detection [30] , and ADE20K semantic segmentation [55] . Detailed model architectures and training setups are shown in the Appendix. As shown in Tab. 1, we separately apply each distinction to the baseline linear attention model and assess their performances on ImageNet-1K.Input Gate. Introducing the input gate results in a modest accuracy improvement of 0.2, indicating that it is slightly helpful for the model. Visualizations in Fig. 4 (a) aid in understanding the impact of the input gate. It can be seen that the model tends to generate higher input gate values for more informative regions like foreground objects, while suppressing less useful tokens. However, the model struggles to generate highly effective input gates, since the input gate values ∆ i = Softplus(x i W 1 W 2 ) are predicted solely from the current input token x i without considering the overall semantics of the image. For example, in one image, the dog may be the area of interest, whereas in another, it might simply be part of the background. Without leveraging information from the entire image, assigning large input gate values to the dog in one image while blocking it in another is impractical. Moreover, employing input gate results in a 7% decrease in model throughput.Forget Gate. Employing the forget gate in linear attention leads to an obvious performance improvement from 77.6 to 78.4. However, such accuracy gain comes at a cost: the model throughput drops severely from 1152 to 743. This is because the forget gate has to employ recurrent calculation, which is slower than the parallelizable matrix multiplication in linear attention. It's worth noting that we already utilize the hardware-aware algorithm proposed in Mamba to speed up the recurrent computation. Thus, we believe the forget gate might not be very suitable for modeling non-causal data like [42] , LePE [9] , CPE [5] and RoPE [41] can both help the model yield better results than the forget gate, while maintaining high throughput. We attribute the improved outcomes to a broader receptive field. Specifically, when using the forget gate, we have to adopt the recurrent linear attention format which restricts the receptive field of each token to the preceding sequence. In contrast, without the forget gate, it is natural to utilize parallel linear attention to achieve a global receptive field.Shortcut. As illustrated in Tab. 1, the usage of learnable shortcut in linear attention provides a 0.2 accuracy gain, while decreasing the throughput from 1152 to 1066. Normalization. Without normalization, the model suffers from severe performance degradation from 77.6 to 72.4. This can be attributed to the issue of long tokens dominating, as discussed in Sec. 4.2. To confirm this, we compute the standard deviation of token lengths (l2 norm) in each layer using both the baseline model and the model without attention normalization. As depicted in Fig. 5 , without normalization, the standard deviation of token lengths tends to be much larger than the baseline, particularly in the last two layers. This supports our analysis that without normalization, the difference in token length becomes significant, with some long tokens dominating the model while others struggling to convey their semantics.Multi-head. Modern Transformers typically adopt the multi-head design [42] to enhance their expressive power. As shown in Tab. 1, removing this design reduces computational cost and accelerates the model but significantly diminishes performance. We consider this trade-off unwarranted.Block Design. We employ two ways to assess the effects of Mamba's block design: 1. Substituting the entire Transformer block with Mamba's block architecture. 2. Replacing the attention sub-block with Mamba's block design, while preserving the MLP sub-block. In both settings, the selective SMM in Mamba's block is substituted with linear attention. To maintain similar FLOPs, we employ Mamba expansion factors [14] E = 2.0 and E = 1.0 for the two settings, respectively. The results are presented in Tab. 1 as "Block Design all" and "Block Design sub". Both replacement approaches result in performance improvements, demonstrating the efficacy of Mamba's macro design. Substituting the attention sub-block yields better result, which creates our MILA block shown in Fig. 3(c ). Notably, we omit the V projection before linear attention calculation, as a similar input projection already exists. The module complexity of a MILA block is expressed as:EQUATIONwhich is slightly larger than the complexity of a Transformer block (Fig. 3a ), 4N C 2 +2N Cd+8N C 2 .Based on our findings, we integrate the forget gate and block design into linear attention, introducing our MILA model. Notably, we practically use LePE, CPE, and RoPE to replace the forget gate's local bias, input-dependent positional information, and global positional information, respectively.As shown in Tab. 3, our MILA models consistently outperform various vision Mamba models across all model sizes, owing to the integration of useful designs from both Mamba and linear attention. These results also validate that with the merits of Mamba's two key designs, the inferior linear attention Transformer can surpass high-performance Mamba. Notably, we empirically observe that MILA exhibits greater scalability compared to vision Mamba models, as MILA-B achieves an accuracy of 85.3, surpassing other models by a significant margin. Additionally, MILA also outperforms various CNN and vision Transformer designs. For instance, MILA exhibits better performance than MambaOut [51] , a recent work that removes the selective SSM in Mamba and employs a gated convolution architecture. Inference time. We offer real speed measurements in Fig. 6 . Substituting the forget gate with positional encoding, our MILA models benefit from parallelizable computation, resulting in significantly faster inference speeds compared to vision Mamba models. For instance, our model achieves 4.5x faster inference speed than Mamba2D [27] , while maintaining better accuracy. Compared to the highly optimized VMamba model [31] , our model also delivers a 1.5x speedup accompanied by a 0.5 accuracy gain. These substantial improvements in model speed further support our analysis that the parallelizable MILA is more suitable than Mamba for modeling non-causal data such as images. 13)) and parallelizable computation, making it ideally suitable for high-resolution image modeling scenarios. Notably, MILA outperforms MambaOut [51] by a significant margin, which aligns with the findings in MambaOut [51] .ADE-20K semantic segmentation. We report the results on ADE-20K dataset in Tab. 5, where "SS" and "MS" denote single-scale and multi-scale testing, respectively. Similar to the object detection task, MILA yields better results in semantic segmentation, further verifying our analyses and the effectiveness of MILA model. Transformer 107M 496G 46.9 --42.3 --PVTv2-B5 [44] Transformer 102M 557G 47.4 68.6 51.9 42.5 65.7 46.0 FocalNet-B [50] Transformer 111M 507G 49.0 70.9 53.9 43.5 67.9 46.7 CSWin-B [9] Transformer 97M 526G 48.7 70.4 53.9 43.9 67. 8  
conclusion: 
 This paper reveals the surprisingly close relationship between the powerful Mamba and subpar linear attention Transformer, shedding some light on Mamba's superiority and success. We rephrase Mamba as a variant of linear attention Transformer and identify its six major special designs: input gate, forget gate, shortcut, no attention normalization, single-head and modified block design. Empirical validation shows that the forget gate and block design largely enhance performance, while the other distinctions offer marginal contributions or impair model performance. Based on our findings, we propose our Mamba-Like Linear Attention (MILA) model by incorporating the merits of these two key designs into linear attention. MILA surpasses various vision Mamba models across multiple tasks, while maintaining parallel computation and high inference speed. Specifically, we begin with the formulas and rephrase Mamba as a variant of linear attention Transformer with six major distinctions: input gate, forget gate, shortcut, no attention normalization, single-head and modified block design. Moreover, we meticulously analyze the pros and cons of each design and prove that the forget gate and block design are the core contributors to Mamba's success. Based on our findings, we propose our Mamba-Inspired Linear Attention (MILA) model, which surpasses various vision Mamba models across multiple tasks, while maintaining parallel computation and high inference speed. However, there may be other small differences between the implementation details of Mamba and linear attention Transformer, and this paper is not exhaustive. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks as it is a foundational research focusing on general model structure.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We use all assets properly according to their licenses, and give credits to the creators in appendix A.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: The manuscript has been reviewed by several reviewers, all of whom, after rebuttal, gave positive ratings and agreed that the manuscript meets the bar of NeurIPS.

Essentially, the reviewers were happy to find that the study is well motivated and the design is well conducted. There are a few flaws but the reviewers are convinced after the rebuttal.

The AC agrees with the consensus from the reviewers and therefore commends the manuscript to be accepted. Congratulations!

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: 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?
limitations: limitations: Nope.
paper_summary: paper_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.
strength_weakness: strength_weakness: * 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: 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.
limitations: limitations: Yes, they are addressed and discussed.
paper_summary: paper_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.
strength_weakness: strength_weakness: No apparent weaknesses
questions: 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?
limitations: limitations: The authors mention the limitations of their work, especially that is it no exhaustive.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Please see the main weaknesses.
limitations: limitations: Please see the main weaknesses.

Comments: comment: Thank you for providing additional results. The paper now offers both insightful methods and solid experimental results, which have contributed to my improved score.
Dear Reviewer XBcf, we sincerely appreciate your rigorous research attitude and valuable comments. Here, we provide some additional results we obtained these days to fully resolve your concern.

---

**1. Results without MESA.**

- ***MLLA models trained without MESA.***

  |         Model         | #Params | FLOPs |   Acc.   |
  | :-------------------: | :-----: | :---: | :------: |
  |         Vim-S         |   26M   | 5.1G  |   80.3   |
  |       VMamba-T        |   31M   | 4.9G  |   82.5   |
  |     LocalVMamba-T     |   26M   | 5.7G  |   82.7   |
  | **MLLA-T (w/o MESA)** |   25M   | 4.2G  | **83.3** |
  | **MLLA-T (w/ MESA)**  |   25M   | 4.2G  | **83.5** |
  |                       |         |       |          |
  |       VMamba-S        |   50M   | 8.7G  |   83.6   |
  |     LocalVMamba-S     |   50M   | 11.4G |   83.7   |
  | **MLLA-S (w/o MESA)** |   43M   | 7.3G  | **84.2** |
  | **MLLA-S (w/ MESA)**  |   43M   | 7.3G  | **84.3** |
  |                       |         |       |          |
  |       Mamba2D-B       |   94M   |   -   |   82.3   |
  |       VMamba-B        |   89M   | 15.4G |   83.9   |
  | **MLLA-B (w/o MESA)** |   96M   | 16.2G | **85.0** |
  | **MLLA-B (w/ MESA)**  |   96M   | 16.2G | **85.3** |

  From these results, we observe:

  - The overfitting prevention strategy MESA only provides ***a modest accuracy gain of 0.1-0.3.***
  - Without MESA, MLLA models still ***significantly surpass various SOTA vision Mamba models.***

- ***Downstream tasks using the backbone trained without MESA.***

  |        Method        | Pre-trained Backbone | #Params | FLOPs | $AP^b$ | $AP^m$ |
  | :------------------: | :------------------: | :-----: | :---: | :----: | :----: |
  | MLLA-T Mask R-CNN 1x |       w/o MESA       |   44M   | 255G  |  46.8  |  42.0  |
  | MLLA-T Mask R-CNN 1x |       w/ MESA        |   44M   | 255G  |  46.8  |  42.1  |
  | MLLA-T Mask R-CNN 3x |       w/o MESA       |   44M   | 255G  |  48.7  |  43.7  |
  | MLLA-T Mask R-CNN 3x |       w/ MESA        |   44M   | 255G  |  48.8  |  43.8  |

  The results show that ***the inclusion of MESA during pre-training has a negligible effect on downstream task performance.***

- ***We fully appreciate your rigorous research attitude.*** We will include these results and detailed discussions in the revised version.

---

**2. Conclusions of the discussion on MESA.**

Based on the additional results and our previous discussions, we believe the following conclusions can be made:

- ***MESA serves as a strategy to prevent overfitting.***
- ***Without MESA, MLLA models can also achieve comparable results in both classification and downstream tasks.*** This is validated by the additional results.
- ***MESA does not affect the three main contributions and findings of this work.*** Our work has three main contributions: revealing the relationship between Mamba and linear attention, analyzing and validating the key difference, and presenting MLLA. The first two contributions do not involve MESA. While the last contribution of our work, MLLA, originally used MESA in its training, additional results verify that MLLA models still significantly surpass various SOTA vision Mamba models without it. This further supports our findings. Therefore, MESA does not influence the three core contributions of our work.
- ***We employ MESA to provide the community with more robust models.*** Just like in the early stages of studies on visual Transformer, currently vision Mamba research does not have a well-established and universally accepted training protocol. We additionally employ the overfitting prevention strategy MESA to fully demonstrate MLLA's potential. Our goal is to benefit the community and follow-up works. We will include detailed discussions in the revised version.
I appreciate your clarification. The use of MESA or not does not impact the paper's main contribution. However, I believe it's unfair to directly compare other methods with this work, as it employs additional training techniques. Given the approaching rebuttal deadline, conducting further experiments on small and base models may be difficult. Nevertheless, based on the authors' provided results of MLLA-T (without MESA), I find it acceptable. As a result, I have increased my score. I hope the authors will include comprehensive experimental results in the revised paper.
Thanks to the authors for their clarifications. I will keep my score unchanged.
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.
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.
Thank you again for your time and valuable comments.
Thank you for your response. After reading the response and other reviewers' comments, I believe that my concerns are mostly resolved and I will increase my rating from 4 to 6.
We would like to express our appreciation for your comments. However, we believe there are some misunderstandings and we offer clarifications below.

---

**1. MESA does not affect the three main contributions and findings of this work.**

We would like to emphasize that the three main contributions of our paper remain unaffected by the use of MESA:

- We reveal Mamba’s close relationship to linear attention Transformer.
- We provide detailed analyses of each special design and validate that the forget gate and block design largely lead to Mamba’s superiority.
- We present MLLA, a novel linear attention model that outperforms vision Mamba models.

The first finding is thoroughly analyzed in Section 4 of our paper.

The second one is validated by the results in Table 1 and Table 2, ***where MESA is not employed.***

The last contribution of our work, MLLA, uses MESA in its training. However, as we already mentioned in our previous response, without MESA, MLLA-T can also achieve 83.3 accuracy and still significantly surpasses various vision Mamba models. Notably, the MLLA model is built to validate our last finding,  i.e. whether linear attention can match or surpass Mamba in vision, rather than to compete against SOTA vision Transformers.

***In conclusion, MESA does not influence the core contributions of our work, but rather serves as an additional strategy to help MLLA performs optimally.***

We further clarify our intent and the reason for using MESA in the following.

---

**2. The reason for using MESA.**

We would first like to clarify that MESA is just a strategy to prevent overfitting, and it cannot boost performance like token labeling. Token labeling benefits from ***a pre-trained model***, functioning similarly to distillation, whereas MESA only enhances the model's generalization ability and ***doesn't use any pre-trained models.***

Just like in the early stages of studies on visual Transformer, currently vision Mamba research does not have a well-established and universally accepted training protocol. The conventional training setting for vision Transformer may not be optimal for vision Mamba and our Mamba-Like Linear Attention. Therefore, we additionally employ the overfitting prevention strategy MESA to alleviate the overfitting problem of MLLA model and fully demonstrate its potential. Our goal is to provide the community with more robust models.

We believe that excessive pursuit of strictly unchanged training setting could actually restrict the exploration of new architectures.

---

**3. Results without MESA.**

- To better address your request, we further provide the results without MESA.

- Actually, we already provided the result for ***MLLA-T trained without MESA*** in our previous response. Here, we offer a comparison with vision Mamba models based on this result. 

  |       Model       | #Params | FLOPs | Acc. |
  | :---------------: | :-----: | :---: | :--: |
  |       Vim-S       |   26M   | 5.1G  | 80.3 |
  |     VMamba-T      |   31M   | 4.9G  | 82.5 |
  |   LocalVMamba-T   |   26M   | 5.7G  | 82.7 |
  | MLLA-T (w/o MESA) |   25M   | 4.2G  | 83.3 |

  ***As you can see, without MESA, MLLA-T can also achieve comparable result and still significantly surpass various SOTA vision Mamba models.***

- We are currently in the process of training ***MLLA-S/B under the no MESA setting*** and will provide those results ***in a few days.***

- Furthermore, we will utilize the backbone trained without MESA to conduct downstream tasks.

- All these results will be included in the revised manuscript. We will provide more experimental results within our capabilities and try our best to benefit the community and follow-up works.

---

**3. Downstream tasks.**

- There seems to be a misunderstanding. In our previous response, we stated that Mask R-CNN 3x results for ***base level model, e.g. VMamba-B,*** is hardly reported by vision Mamba works. However, we did not claim that Mamba-based models have not been tested on 3x schedule. Indeed, we already provided comparison of ***tiny and small level models under 3x schedule in Table 9. of our paper.***
- The primary focus of our paper is on comparisons with vision Mamba models, rather than achieving SOTA results. Given that the works we compared with do not present 3x detection results for base level model, we also omit this experiment previously. Currently, we are working on this experiment to better address your request.
I appreciate the authors' response, which has addressed some of my concerns.  Although I understand that the purpose of this paper is to establish connections between various mechanisms in Mamba and Linear Attention, the main issue remains the unfair experimental implementation. The authors should provide results without MESA in Table 3 for tiny, small, and base models, as MESA can boost performance like token labeling, and this should be explicitly indicated. The current presentation may mislead follow-up works due to the unfair training conditions. The authors' claim that other models do not benefit from MESA lacks any supporting evidence.

Regarding downstream tasks, while the authors clarify that they did not use MESA, it’s important to note that the pre-trained models have incorporated MESA. This introduces an inherent unfairness in weight initialization, which may affect the fairness of the final performance comparisons.

The authors mention that other Mamba-based models have not been tested on 3x schedule-based detection tasks, which raises a question: are Mamba-based models unable to train successfully on 3x schedule-based detection tasks? If so, the potential applications of Mamba-based models would be significantly limited. Additionally, please note that the proposed method belongs to the vision transformer family, not Mamba-based models, and it is crucial to conduct related experiments as listed in many previous works. Reporting results with the 3x schedule is more meaningful, as this allows the models to converge to a more consistent outcome, providing a more reliable comparison.

The author should be aware that one of the most crucial aspects of a vision backbone paper is the experimental section, which must adhere to highly consistent training conditions with other competitors to ensure a fair comparison. Implicitly introducing additional tricks to improve performance is misleading and unfair to follow-up works. I must reiterate that the unfair experimental design remains a significant concern.