Demystify Mamba in Vision: A Linear Attention Perspective

Dongchen Han; Ziyi Wang; Zhuofan Xia; Yizeng Han; Yifan Pu; Chunjiang Ge; Jun Song; Shiji Song; Bo Zheng; Gao Huang

Source-linked AI summary

Demystify Mamba in Vision: A Linear Attention Perspective

Dongchen Han, Ziyi Wang, Zhuofan Xia, Yizeng Han, Yifan Pu, Chunjiang Ge, Jun Song, Shiji Song, Bo Zheng, Gao Huang

arXiv:2405.16605v2May 26, 2024cs.CV

TL;DR

The paper addresses why Mamba succeeds despite the historically weak practical performance of linear attention. It unifies their formulations and analyzes six design differences, finding that forget gates and block design are central; MILA incorporates these merits and outperforms various vision Mamba models across image classification and high-resolution dense prediction tasks.

Problem
Linear attention offers linear complexity but has been reported to lack expressive power and deliver unsatisfactory practical results, whereas Mamba performs effectively with linear complexity.
Method
The paper unifies selective state space models and linear attention, analyzes six Mamba-specific distinctions, and incorporates the key findings into MILA.
Results
MILA outperforms various vision Mamba models across image classification and high-resolution dense prediction tasks.
Takeaways & Limitations
Forget gates and block design largely enhance performance, while the other distinctions contribute marginally or impair performance.
Takeaways & Limitations
The forget gate requires recurrent computation, reducing throughput and limiting its suitability for non-causal vision models.

Abstract

from arXiv · show

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.

1 Introduction

Mamba offers linear-complexity sequence modeling, but linear attention has historically struggled with expressive power and practical performance. This paper unifies their formulations, identifies six differences, and finds forget gates and block design central to Mamba’s effectiveness.

Motivation: Mamba provides effective sequence modeling with linear complexity, enabling manageable computation for extremely long sequences and high-resolution vision inputs.This contrasts with the quadratic complexity of prevailing Transformer models.
Motivation: Linear attention reduces attention complexity from O(N^2) to O(N) by changing computation order, but previous work found insufficient expressive power.It replaces Softmax with linear normalization and computes Q(K^⊤V) instead of (QK^⊤)V.
Approach: The paper reformulates selective state space models and linear attention in a unified framework, portraying Mamba as a variant with six distinct designs.The distinctions are input gate, forget gate, shortcut, no attention normalization, single-head operation, and modified block design.
Findings: Empirical analyses identify the forget gate and block design as the two core contributors to Mamba’s superiority.The study evaluates the effects of each special design on vision tasks.
Outcome: The authors introduce MILA by incorporating these core contributors or their alternatives into linear attention Transformer models for vision.The contribution list characterizes MILA as more suitable for vision tasks than the original Mamba model.

2 Related Works

Vision Transformers offer strong recognition performance, but Softmax attention becomes costly for high-resolution images. Linear attention provides linear complexity yet has generally failed to deliver satisfactory results, motivating this paper’s comparison with Mamba.

Vision Transformer and attention: Softmax attention’s quadratic complexity makes it challenging to apply Transformers to high-resolution images.Prior approaches reduce cost through local attention windows or sparsity.
Vision Transformer and attention: Linear attention inherently provides O(N) complexity and can model long sequences, but previous studies report unsatisfactory results that limit its applicability.Its efficiency does not consistently translate into effective practical performance.
Mamba: Prior vision adaptations of Mamba include cross-scan, local-window, and atrous-based selective-scan designs for processing images efficiently.These examples include VMamba, LocalMamba, and EfficientVMamba.
Mamba: This paper connects inferior linear attention Transformers with powerful Mamba models to investigate the factors underlying Mamba’s success.The comparison is framed as a way to demystify Mamba’s key advantages.

3 Preliminaries

This section introduces Softmax and linear attention, classical and discrete state space models, and Mamba’s selective state space formulation. It then presents Mamba’s input-dependent recurrence and the equivalent forms used for comparison.

Attention Mechanism: Softmax attention computes query-key similarities with O(N^2) complexity, creating high costs for long sequences.Linear attention addresses this computational cost by replacing Softmax normalization with linear normalization.
Attention Mechanism: Linear attention rearranges computation from (QK⊤)V to Q(K⊤V), reducing complexity to O(N) while retaining a global receptive field.In causal settings, the i-th token instead aggregates only from preceding tokens.
State Space Model: A state space model maps an input through a hidden state to an output, and its continuous dynamics are h′(t) = Ah(t) + Bx(t).For neural networks, the continuous model is discretized using a timescale parameter ∆, yielding hi = Ahi−1 + Bxi.
Selective State Space Model: Mamba makes the state space parameters input-dependent, producing the selective recurrence hi = Aihi−1 + Bixi.The parameters B, C, and ∆ depend on xi, and the discretized parameters Ai and Bi therefore vary with the input.
Selective State Space Model: Mamba uses diagonal state-transition matrices so Aihi−1 becomes an element-wise product e Ai ⊙ hi−1.For sequence inputs, Mamba applies the selective formulation independently across channels, with xi and ∆i in R1×C.
Unified Formulation: The unified formulation interprets ∆i as an input gate and e Ai as a forget gate controlling current-input admission and previous-state attenuation.Mamba derives these quantities from the input, and its channel-wise recurrence operates over sequence features.

4 Connecting Mamba and Linear Attention Transformer

The paper unifies selective state space models and linear attention, showing Mamba as a linear-attention variant with six differences spanning core operations and macro design. It analyzes the roles and trade-offs of these designs, especially gating, normalization, recurrence, and block architecture.

Unified formulation: Selective state space models and linear attention can be expressed in a unified formulation, with corresponding state, key, value, and query terms.The paper maps h_i to S_i, B_i to K_i^T, x_i to V_i, and C_i to Q_i.
Core-operation differences: Mamba’s selective state space operation adds an input gate, forget gate, and shortcut while omitting attention normalization and multi-head design.The forget gate attenuates the previous hidden state; the shortcut directly connects input to output.
Input gate: The input gate modestly improves accuracy by assigning larger values to informative foreground tokens and smaller values to less useful background tokens.Its values are predicted from the current token through a Softplus-based transformation, without considering overall semantic context.
Trade-offs: The forget gate requires recurrent computation because its values depend on inputs, reducing throughput relative to parallel linear attention and limiting image-token receptive fields to preceding sequence positions.This causal formulation may suit autoregressive data better than non-causal images; positional encodings are suggested as a possible substitute.
Normalization: Removing attention normalization can let longer tokens dominate feature maps, potentially causing training instability and reducing expressiveness.Normalization keeps attention weights summing to 1 and mitigates scale-dependent attention scores.
Macro block design: Mamba replaces the conventional linear-attention-plus-MLP block with a modified block combining selective SSM, convolution, mappings, activation, and gating operations.The paper evaluates replacing either the entire Transformer block or only its attention sub-block with Mamba’s block design while using linear attention.

5 Empirical Study

The empirical study isolates six design differences between Mamba and linear attention, finding that the forget gate and block design are the main contributors to performance. MILA incorporates these strengths while replacing recurrent forget-gate behavior with positional encodings, achieving strong accuracy, dense-prediction performance, and faster inference.

Study design: The study replaces Swin’s Softmax attention with linear attention, then adds Mamba’s six distinctions separately before evaluating MILA.The six distinctions are input gate, forget gate, shortcut, no attention normalization, single-head design, and block design.
Ablation results: Input gating improves accuracy by 0.2 but decreases throughput by 7%.Its gate values depend only on the current token, limiting image-level semantic adaptation.
Ablation results: Forget gating raises accuracy from 77.6 to 78.4 while reducing throughput from 1152 to 743 because it requires recurrent computation.The study therefore investigates parallelizable substitutes for its local-bias and positional-information properties.
Ablation results: Removing attention normalization drops accuracy from 77.6 to 72.4, while removing multi-head design significantly reduces performance despite lower computational cost.Without normalization, token-length variation increases and some long tokens dominate the model.
Comparison with vision Mamba: MILA outperforms vision Mamba models across ImageNet classification, COCO detection, and ADE-20K segmentation, while delivering 4.5x faster inference than Mamba2D.Compared with VMamba, it provides a 1.5x speedup and a 0.5 accuracy gain; MILA-B reaches 85.3 accuracy on ImageNet.

6 Conclusion

The paper reframes Mamba as a linear-attention variant and identifies forget gating and block design as its key advantages. MILA incorporates these designs and surpasses various vision Mamba models while retaining parallel computation and high inference speed.

Conclusion: Mamba and linear attention Transformer fit within a unified formulation with six distinguishing designs.These are input gate, forget gate, shortcut, no attention normalization, single-head processing, and modified block design.
Conclusion: Forget gate and block design largely enhance performance, whereas the other distinctions contribute marginally or impair performance.
Conclusion: MILA incorporates the merits of the forget gate and block design into linear attention.
Conclusion: MILA surpasses various vision Mamba models across multiple tasks while maintaining parallel computation and high inference speed.

A Datasets and Experiment Details

The experiments cover ImageNet-1K classification, COCO object detection and instance segmentation, and ADE20K semantic segmentation. Training follows established Swin Transformer settings, with dataset-specific frameworks and evaluation protocols.

ImageNet classification: ImageNet-1K contains 1.28 million training images, 50,000 validation images, and 1,000 classes.Models are trained from scratch for 300 epochs using AdamW, cosine decay, 20-epoch warm-up, weight decay 0.05, and total batch size 4096.
COCO object detection: COCO provides 118K training and 5K validation images for object detection and instance segmentation.Experiments follow the standard 1x and 3x Mask R-CNN settings, using pretrained MILA models as backbones.
ADE20K semantic segmentation: ADE20K contains 25K images across 150 semantic categories, with 20K training, 2K validation, and 3K testing images.UPerNet is used with Swin Transformer training settings, and both single-scale and multi-scale testing results are reported.

B Additional Experimental Results

Additional experiments compare MILA with advanced linear-attention designs and examine the effect of the MESA training strategy. These results extend the evaluation beyond the primary comparisons while noting training-protocol considerations.

Additional comparison with advanced linear attention designs: MILA outperforms various advanced linear attention designs without additional bells and whistles.The comparison results are reported in Table 6.
Additional comparison with advanced linear attention designs: Table 6 presents the comparison with advanced linear attention designs.
Ablation on the impact of MESA: Because vision Mamba lacks a universally accepted training protocol, MILA training additionally uses MESA to demonstrate its potential.The paper notes that conventional vision Transformer settings may not be optimal for vision Mamba and MILA.

C Model Architectures

MILA uses a common four-stage architecture built by stacking MILA blocks at each stage. The architecture is illustrated in Figure 7 and detailed in Table 8.

Model architecture: MILA adopts a common 4-stage framework.
Model architecture: MILA blocks are stacked at each stage to construct the model.
Model architecture: Figure 7 illustrates the MILA architecture, while Table 8 summarizes its detailed structure.

D Limitations

The paper acknowledges that its comparison between Mamba and linear attention Transformer is not exhaustive. Other implementation-level differences may remain beyond the six major distinctions analyzed.

Scope: The analysis is not exhaustive and may omit smaller implementation differences between Mamba and linear attention Transformer.The authors explicitly note that additional implementation-level differences may exist.
Scope: The six identified design distinctions should therefore be understood as major differences rather than a complete implementation correspondence.The paper frames its comparison around six major distinctions while acknowledging other possible differences.
Scope: The limitation concerns the completeness of the Mamba–linear-attention comparison, not the paper’s reported proposal of MILA.The stated boundary is limited to potentially omitted implementation details.