Metaphor

视觉Mamba模型

10分钟阅读

引言

视觉Transformer（ViT）自2017年提出以来，已成为计算机视觉领域的主导架构。然而，其自注意力机制的计算复杂度随图像尺寸呈二次方增长，这给高分辨率图像处理带来了显著的计算和内存压力。

Mamba是一种具有选择性机制和硬件感知架构的状态空间模型（SSM），在长序列建模方面展现出巨大潜力。2024年初，研究者开始探索将Mamba应用于计算机视觉任务，提出了Vision Mamba（Vim）和VMamba等视觉骨干网络。12

本文综述视觉Mamba的核心架构、关键技术及其在各类视觉任务中的应用。

核心架构

Vision Mamba（Vim）

Vim是首个将双向Mamba块应用于视觉表征学习的通用骨干网络。其核心创新在于：

双向状态空间模型

Vim的核心是双向Mamba块（Bidirectional Mamba Block）：

┌─────────────────────────────────────────────────────────┐
│                    Vim 整体架构                           │
├─────────────────────────────────────────────────────────┤
│  输入图像 (H×W×3)                                        │
│       ↓                                                  │
│  Patch Embedding: 划分 16×16 patches                     │
│       ↓                                                  │
│  序列长度 N = (H/16) × (W/16)                             │
│       ↓                                                  │
│  ┌─────────────────────────────────────────────────┐    │
│  │           VSS Block × N                          │    │
│  │  ┌─────────────┐  ┌─────────────┐               │    │
│  │  │ 前向 SSM    │  │ 反向 SSM    │ → Concat      │    │
│  │  │ Forward SSM │  │ Backward SSM│               │    │
│  │  └─────────────┘  └─────────────┘               │    │
│  │       ↓                ↓                        │    │
│  │  LayerNorm          LayerNorm                    │    │
│  │       ↓                ↓                        │    │
│  │  ┌─────────────────────────────────────────┐    │    │
│  │  │         门控线性单元 (GLU)                │    │    │
│  │  └─────────────────────────────────────────┘    │    │
│  └─────────────────────────────────────────────────┘    │
│       ↓                                                  │
│  分类头 / 检测头 / 分割头                                  │
└─────────────────────────────────────────────────────────┘

关键设计：

  • 将图像划分为序列后，同时进行前向和反向两个方向的SSM扫描
  • 拼接双向输出以获得完整的全局上下文信息
  • 无需依赖自注意力机制即可捕获长距离依赖1

位置编码

视觉数据具有位置敏感性，这与文本序列不同。Vim采用以下策略：

  1. 绝对位置编码（Absolute Position Embedding）：在输入序列前添加可学习的位置嵌入
  2. 残差连接：保持空间位置信息不丢失
  3. 位置感知的选择性扫描：使模型能够区分不同位置的视觉特征

VMamba 与 2D 扫描策略

VMamba的核心贡献是提出**2D选择性扫描（2D Selective Scan, SS2D）**模块，解决了一维序列扫描难以捕获二维空间信息的问题。2

十字扫描模式（Cross-Scan Pattern）

SS2D通过四条扫描路径将2D图像展平为1D序列：

┌────────────────────────────────────────────────────────────┐
│                  2D选择性扫描 (SS2D)                        │
├────────────────────────────────────────────────────────────┤
│                                                            │
│   原始特征图 (H×W)                                          │
│   ┌───┬───┬───┬───┐                                       │
│   │ 0 │ 1 │ 2 │ 3 │  ← 行主序扫描                          │
│   ├───┼───┼───┼───┤                                       │
│   │ 4 │ 5 │ 6 │ 7 │                                       │
│   ├───┼───┼───┼───┤                                       │
│   │ 8 │ 9 │10 │11 │                                       │
│   └───┴───┴───┴───┘                                       │
│                                                            │
│   四条扫描路径:                                            │
│   ┌────────────────────────────────────────────────────┐  │
│   │  →  右向扫描 (→)                                    │  │
│   │  ←  左向扫描 (←)                                    │  │
│   │  ↓  下向扫描 (↓)                                    │  │
│   │  ↑  上向扫描 (↑)                                    │  │
│   └────────────────────────────────────────────────────┘  │
│                                                            │
│   扫描后 → 独立SSM处理 → 特征融合                          │
│                                                            │
└────────────────────────────────────────────────────────────┘

工作原理：

  1. 将每个像素作为起点，沿四个方向进行选择性扫描
  2. 每个方向的扫描结果由独立的SSM模块处理
  3. 四路输出通过求和或拼接融合
  4. 恢复2D特征图结构

这种设计使VMamba能够从不同方向和视角收集上下文信息，弥补了1D扫描与2D视觉数据之间的结构差异。2

VSS Block 结构

┌─────────────────────────────────────────┐
│              VSS Block                   │
├─────────────────────────────────────────┤
│  输入 X                                 │
│     ↓                                   │
│  ┌───────────────┐                      │
│  │  2D → 1D      │  (SS2D扫描)          │
│  │  展开         │                      │
│  └───────────────┘                      │
│     ↓                                   │
│  ┌───────────────┐                      │
│  │  选择性SSM    │  (状态更新)          │
│  │  (×4 方向)    │                      │
│  └───────────────┘                      │
│     ↓                                   │
│  ┌───────────────┐                      │
│  │  1D → 2D      │  (SS2D合并)          │
│  │  恢复         │                      │
│  └───────────────┘                      │
│     ↓                                   │
│  ┌───────────────┐                      │
│  │  门控MLP      │  (特征变换)          │
│  └───────────────┘                      │
│     ↓                                   │
│  输出 + 残差                            │
└─────────────────────────────────────────┘

复杂度与效率对比

线性复杂度 vs 二次复杂度

视觉Mamba相较于Vision Transformer的核心优势在于计算复杂度：

架构自注意力复杂度SSM复杂度序列长度依赖
ViT/DeiT-二次方增长
Vim/VMamba-线性增长

其中 为序列长度（与图像尺寸相关）。

具体效率提升（Vim vs DeiT）：

  • 推理速度：2.8倍提升
  • GPU内存：节省86.8%（1248×1248分辨率下）
  • 高分辨率图像：Vim可有效处理，而DeiT面临内存瓶颈1

全局感受野

虽然SSM本身是线性复杂度的，但它通过状态传递能够有效捕获全局感受野：

传统CNN:        局部感受野逐步扩展
      [□] → [□□] → [□□□□] → [□□□□□□□] → ...
      
ViT:           全局注意力（两两交互）
      [□] ←→ [□] ←→ [□] ←→ ...
      
Mamba/Vim:     选择性状态传递
      [□] → ■ → [□] → ■ → [□] → ■ → [□]
                 ↓
            隐藏状态携带
            全局信息

状态空间模型通过隐藏状态向量在序列中传递信息，每个时间步的输出都隐含地编码了之前所有输入的信息。

视觉任务基准性能

ImageNet 分类

模型Top-1 Acc参数量GFLOPs
DeiT-S79.8%22M4.5
Vim-S80.6%26M4.9
Vim-M81.9%51M9.8
VMamba-S80.2%26M4.8
VMamba-B82.1%88M15.6

Vim和VMamba在ImageNet上达到了与DeiT相当甚至更好的性能，同时保持更低的内存消耗。

语义分割（ADE20K）

模型mIoU参数量
DeiT-S (UperNet)45.8%39M
Vim-S (UperNet)48.1%43M
VMamba-S (UperNet)48.6%44M

目标检测与实例分割（COCO）

模型box mAPmask mAP
DeiT-S (RetinaNet)43.1%-
Vim-S (RetinaNet)46.2%-
VMamba-S (Mask-RCNN)48.5%43.8%

实验表明，Vim和VMamba在下游检测和分割任务中展现出竞争力，尤其在高分辨率输入场景下优势更为明显。12

其他视觉SSM变体

视觉Mamba领域发展迅速，已涌现出多种变体和扩展：3

按扫描方式分类

类型代表方法特点
十字扫描VMamba四方向独立扫描
双向扫描Vim前向+后向双通道
螺旋扫描LocalMamba局部窗口内螺旋遍历
分支扫描SwinMamba结合窗口机制

按任务分类

  • 图像分类：Vim、VMamba、PlainMamba
  • 语义分割：VMamba-Seg、SegMamba
  • 目标检测：Vim-Det、MambaDet
  • 视频理解：VideoMamba、ActionMamba
  • 点云处理：PointMamba、Point-Bert V2
  • 多模态：MambaLM、MambaCOD

多尺度视觉Mamba

为适应密集预测任务，研究者提出了多尺度架构：

┌─────────────────────────────────────┐
│        多尺度VMamba架构              │
├─────────────────────────────────────┤
│  Stage 1: 高分辨率, 窄通道           │
│     ↓ (下采样)                      │
│  Stage 2: 中分辨率, 中等通道         │
│     ↓ (下采样)                      │
│  Stage 3: 低分辨率, 宽通道           │
│     ↓ (上采样+融合)                 │
│  多尺度特征金字塔                    │
└─────────────────────────────────────┘

未来方向

当前挑战

  1. 扫描策略设计：如何为不同任务设计最优的扫描路径
  2. 状态容量：隐藏状态维度的选择对模型性能的影响
  3. 训练稳定性：SSM在大规模视觉预训练中的收敛问题

研究展望

方向潜在突破
视觉基础模型构建类似语言模型规模的视觉骨干
高分辨率处理医学影像、遥感图像的超大尺寸输入
多模态融合与CLIP式对齐的视觉-语言模型
高效微调适配器式方法在Vim上的应用
动态架构根据输入内容自适应调整扫描路径

视觉Mamba正逐步成为下一代视觉基础模型的有力候选者，其线性复杂度和全局建模能力为处理高分辨率视觉数据提供了新的可能性。123

参考文献

Footnotes

  1. Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model. arXiv:2401.09417 (ICML 2024) ↩ ↩2 ↩3 ↩4 ↩5

  2. VMamba: Visual State Space Model. arXiv:2401.10166 (2024) ↩ ↩2 ↩3 ↩4 ↩5

  3. Visual Mamba: A Survey and New Outlooks. arXiv:2404.18861 (2024) ↩ ↩2

关系图谱

反向链接