Demystify Mamba in Vision: A Linear Attention Perspective (2405.16605v2)

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.

• The paper reveals that reinterpreting Mamba as a linear attention Transformer enables linear computational complexity for high-resolution vision tasks.
• It analyzes key architectural differences—such as input and forget gates and a modified block design—that drive improved performance in vision models.
• The study introduces the MILA model, demonstrating that streamlined linear attention frameworks can match or exceed complex architectures in tasks like image classification and segmentation.

Demystifying Mamba in Vision: Linear Attention Perspective

The paper presents a comprehensive paper on the Mamba model, a state space model that achieves efficient linear complexity for sequence modeling, specifically targeting high-resolution vision tasks. It highlights the surprising equivalencies between the Mamba model and linear attention Transformers, exploring key distinctions that account for Mamba’s superior practical performance.

Summary

Mamba’s state space model offers linear computational complexity, a crucial advantage over traditional Transformer architectures, whose quadratic complexity limits their efficiency on large-scale inputs. Interestingly, despite its simpler computation, Mamba performs robustly across vision tasks, previously considered challenging for linear complexity models. This paper reevaluates Mamba by positioning it as a linear attention Transformer variant, elucidating six core differences: input gate, forget gate, shortcut, absence of attention normalization, single-head attention, and a modified block design.

Core Findings:

1. Unified Formulation:
   • By unifying selective state space and linear attention models’ formulations, the paper identifies key conceptual overlaps and modifications.
2. Analysis of Distinct Elements:
   • Input Gate: Functions as a moderating factor, selectively emphasizing relevant inputs based on their information content.
   • Forget Gate: Crucial for maintaining local bias and encoding positional information, though it demands sequential computation, reducing operational efficiency.
   • Shortcut and Absence of Normalization: While offering structural advantages, the lack of normalization suppresses potential performance benefits.
   • Single-head and Modified Block Design: The absence of multi-head design is supplemented by a unique block architecture, enhancing data processing effectiveness.
3. Empirical Evaluation:
   • Extensive experiments reveal the forget gate and macro block design as primary contributors to the model’s success. Moreover, the ability to replace the forget gate with positional encodings in non-autoregressive tasks like vision is underscored, providing similar advantages without compromising performance.
4. Mamba-Inspired Linear Attention (MILA):
   • Building on these insights, the MILA model integrates these core principles into a linear attention framework, demonstrating advancements over existing Mamba variants in tasks such as image classification, object detection, and semantic segmentation.

Implications and Future Outlook

The results achieved by MILA challenge the prevailing assumption that linear attention inherently lacks the sophistication required for vision-based tasks. By reforming aspects of the attention mechanism and macro architecture, it's evident that linear attention can match or outperform more complex counterparts like Mamba. This finding could inspire further research into refining linear attention models across different data domains, possibly leading to more computationally efficient architectures for broader applications.

Future work could optimize these elements further, seeking to balance the trade-offs between model complexity, operational throughput, and accuracy, especially in the context of training non-causal models. Additionally, the potential incorporation of more advanced positional encoding schemes may offer new avenues for research, enhancing model flexibility and scalability.

In summary, the paper provides significant insights into Mamba’s functionality, opening new perspectives on how leveraging linear attention's foundational simplicity, with strategic enhancements, can lead to high-performance, efficient vision models. This nuanced understanding invites more exploratory research into linear-complexity models, potentially redefining best practices for machine learning in high-dimensional data spaces.