在过去十年中,计算机视觉领域一直由两大架构巨头主导:** 卷积神经网络 (CNN)**,以及近来兴起的 视觉 Transformer (ViT)。CNN 以其高效性和对局部模式的强大归纳偏置而著称,而 ViT 则凭借自注意力机制,在捕捉图像全局关系方面表现出色。
然而,这种强大的能力是有代价的——自注意力机制的**复杂度与图像块数量呈二次关系 **(\(O(N^2)\)) ,这使得 ViT 的计算成本高昂且运行缓慢,尤其是在目标检测和分割等任务中常见的高分辨率图像场景下。
如果存在第三种途径呢?一种既具备 Transformer 的全局上下文建模能力,又拥有 CNN 的线性时间高效性的架构?
自然语言处理 (NLP) 领域的研究人员一直在探索一种新的模型类别——状态空间模型 (SSM)。其中,近期最具潜力的进展之一是 Mamba,它在长序列语言任务中表现卓越,同时保持了线性复杂度。
受 Mamba 突破性成果的启发,论文《VMamba: Visual State Space Model》将这种强大的 NLP 架构迁移到计算机视觉领域。研究人员提出了 VMamba,一种旨在融合两者优势的新型视觉骨干网络:
- 像 Transformer 一样具备全局感受野和动态、内容感知的权重。
- 像 CNN 一样具有高效率和良好可扩展性。
本文将深入解析 VMamba 的工作原理,剖析其核心理念、创新架构以及令人印象深刻的基准测试成果。
背景: 从卷积到状态空间
要理解 VMamba 的创新意义,首先需要回顾当前视觉模型的格局,并简介状态空间的概念。
CNN 与 Vision Transformer
**CNN **(如 ResNet、ConvNeXt) 通过在整幅图像上应用可学习滤波器 (卷积核) 构建分层特征表示。它们能高效捕捉边缘、纹理等局部特征,但需要依赖多层堆叠来实现长程依赖建模。
ViT 将图像视作一系列图像块,并采用自注意力机制,使所有图像块从第一层开始即可直接交互。这带来了全局感受野,但其二次方的计算复杂度对大尺寸图像而言十分昂贵。
图 1: 自注意力机制 (a) 在所有图像块之间建立全连接,实现了全局上下文但计算成本高昂。SS2D (b) 则通过结构化扫描路径连接图像块,更高效地实现全局覆盖。
状态空间模型 (SSM) 简介
SSM 源于控制理论,用于建模随时间动态演化的系统。连续时间的 SSM 表示为:
\[ \mathbf{h}'(t) = \mathbf{A}\mathbf{h}(t) + \mathbf{B}u(t) \]\[ y(t) = \mathbf{C}\mathbf{h}(t) + D u(t) \]其中:
- \(u(t)\): 输入信号
- \(\mathbf{h}(t)\): 总结过去输入的隐藏状态
- \(y(t)\): 输出信号
- \(\mathbf{A}, \mathbf{B}, \mathbf{C}, \mathbf{D}\): 控制系统演化的可学习参数
在图像或文本处理任务中,这类系统会被离散化为递归形式:
\[ \mathbf{h}_k = \overline{\mathbf{A}}\,\mathbf{h}_{k-1} + \overline{\mathbf{B}}\,u_k \]\[ y_k = \overline{\mathbf{C}}\,\mathbf{h}_k + \overline{\mathbf{D}}\,u_k \]这种形式在推理阶段极其高效——类似 RNN——并且在训练阶段也可并行化。
然而,传统 SSM 使用固定矩阵 (线性时不变系统) ,无法灵活适应不同上下文。Mamba 通过引入选择性扫描机制 (S6) 解决了这一限制,其参数可根据输入内容动态生成,从而让模型智能地决定保留或丢弃哪些信息,实现更丰富的表征能力。
核心创新: 构建 VMamba
将 Mamba 从处理一维序列数据 (文本) 扩展到处理二维空间数据 (图像) 的关键挑战在于,图像没有天然的序列顺序。VMamba 通过引入二维选择性扫描 (SS2D) 模块来解决这一问题。
VMamba 架构
VMamba 采用了类似 Swin Transformer 和 ConvNeXt 的分层架构:
- Stem 模块: 将输入图像划分为小块。
- 网络阶段: 每个阶段包含多个视觉状态空间 (VSS) 模块。
- 阶段间通过下采样层降低空间分辨率并提升通道深度。
图 3: (a) VMamba 的分层架构。(b-d) 从 Mamba 模块到基础 VSS 模块,再到最终优化版 VSS 模块的演进。
最终的 VSS 模块结构类似 Transformer 模块:
- 归一化 → SS2D 模块 → 残差连接 → 归一化 → MLP。
SS2D 将自注意力机制替换为针对二维数据优化的状态空间机制。
二维选择性扫描 (SS2D)
SS2D 流程 (图 2) 会将二维特征图展开为多个一维序列,使用 Mamba 风格的 S6 模块进行处理,再合并回二维布局:
图 2: 步骤 1 — 交叉扫描: 从二维特征图提取四个方向的序列。步骤 2 — 每个序列通过 S6 模块处理。步骤 3 — 交叉合并为二维特征图。
步骤说明:
交叉扫描 (Cross-Scan):
- 正向: 左上 → 右下
- 反向: 右下 → 左上
- 翻转正向: 右上 → 左下
- 翻转反向: 左下 → 右上
选择性扫描 (S6): 每条路径独立处理,在遍历过程中选择性地累积上下文信息。
交叉合并 (Cross-Merge): 将四个输出融合,并还原成原始空间布局。
这种结构化扫描方法以线性复杂度实现了全局上下文获取。
加速优化
初版 VMamba 虽然准确率高,但速度不足以部署。作者进行了多项优化:
- 使用 Triton 重写交叉扫描/合并,提高 GPU 效率。
- 用优化型线性层替代
einsum
操作。 - 调整超参数 (
d_state
、ssm-ratio
) 以提升吞吐量。 - 移除乘法分支,简化计算流程。
实验结果
图像分类 — ImageNet-1K
VMamba 与主流 CNN 和 Transformer 进行了对比测试。
表 1: 在相似模型规模下,VMamba 的性能优于 Swin 和 ConvNeXt 的同类模型,且吞吐量更高。
例如: VMamba-T 的 Top-1 准确率达 82.6%,比 Swin-T 高 1.3%,运行速度快 35%。
下游任务
VMamba 的优势在密集预测任务中同样明显:
表 2: 在 COCO 目标检测和 ADE20K 语义分割任务中,VMamba 对比 Swin 和 ConvNeXt 有显著提升。
在 COCO 目标检测中:
- VMamba-T 比 Swin-T 高 +4.6 mAP。
在 ADE20K 语义分割中:
- VMamba-T 比 Swin-T 高 +3.4 mIoU。
性能曲线在图 4(a) 中显示了其稳定领先优势:
图 4: (a) 在相同的 ImageNet 准确率下,VMamba 在密集任务上表现更优。(b) 随输入分辨率提高,准确率下降更缓和。
线性扩展优势
视觉 Transformer 的二次方扩展性限制了其在高分辨率下的性能,而 VMamba 凭借 SS2D,可以像 CNN 一样实现线性扩展:
图 5: FLOPs/内存随分辨率线性增长,即使输入尺寸很大也能保持高吞吐量。
深度解析: 可解释性
有效感受野 (ERF)
大多数模型在训练前的有效感受野较为局部,训练后 VMamba 则形成了类似 ViT 的全局有效感受野,且二维覆盖更均匀:
图 7: 训练后,VMamba 具备全局有效感受野,更好地聚合远距离特征。
激活图
SS2D 可类比为注意力图。
例如: 在一张有绵羊的图像中,对某个查询块,模型会高亮其扫描路径上的相关区域,即便距离较远的部分也能被记住。
图 6: VMamba 能有效关注物体所有相关区域,利用扫描路径积累上下文信息。
结论与展望
VMamba 通过将状态空间建模引入二维空间数据,实现了视觉骨干网络的范式转型:
- 线性扩展: 高效处理高分辨率图像,性能可与 CNN 相媲美。
- 全局上下文: 从初始层即可获得,类似 Transformer。
- 顶级表现: 在分类、检测、分割等任务中持续领先。
未来方向包括:
- 对基于 SSM 的架构进行无监督预训练。
- 探索适用于多模态的通用扫描模式。
- 将 VMamba 扩展至更大规模的架构。
随着对于高效、强大视觉模型的需求不断上升,VMamba 等架构或将重塑行业格局——为 CNN 与 Transformer 之外的第三大巨头腾出空间。