VMamba:状态空间模型如何超越Transformer,重塑计算机视觉格局

GitHub April 2026
⭐ 3133
来源:GitHub归档:April 2026
当Transformer在视觉领域的统治地位看似不可撼动时,VMamba携状态空间模型(SSM)强势登场。其核心创新2D选择性扫描(SS2D)模块,以线性复杂度实现全局感受野,在ImageNet分类、目标检测与分割任务上不仅超越Swin Transformer,推理速度更是快出一倍。

Transformer在计算机视觉领域的主导地位正面临一位真正的挑战者。VMamba,一个基于Mamba状态空间模型(SSM)架构的全新视觉骨干网络,证明了线性复杂度的序列模型能够与ViT、Swin Transformer等模型所依赖的二次复杂度注意力机制相抗衡——甚至在某些指标上实现超越。其核心创新在于2D选择性扫描(SS2D)模块,该模块沿四个方向轴(左上到右下、右下到左上、右上到左下、左下到右上)遍历图像块,将空间上下文压缩进一个随图像尺寸线性更新的隐藏状态中。这一设计使得VMamba能够在无需自注意力O(n²)内存与计算开销的前提下,获得全局感受野。在ImageNet-1K分类任务上,VMamba-Tiny以22M参数和4.3G FLOPs达到82.5%的Top-1准确率,超越Swin-Tiny的81.3%(28M参数,4.5G FLOPs)。在COCO目标检测和ADE20K语义分割上,VMamba同样展现出更强的性能与更快的推理速度。这一突破意味着,高分辨率图像处理、实时视频理解以及医疗影像分析等场景,将迎来效率与精度并重的新范式。

技术深度解析

VMamba的架构是对Mamba状态空间模型(Gu & Dao, 2024)在二维图像领域的一次精心适配。核心挑战在于:Mamba原本处理一维序列,通过一个压缩上下文的隐藏状态逐个处理token;而图像本质上是具有复杂空间依赖关系的二维网格。VMamba团队通过2D选择性扫描(SS2D)模块解决了这一问题。

SS2D机制: 该模块接收一个二维特征图,并沿四条不同的遍历路径将其展平:
- 路径1:行主序从左到右,从上到下
- 路径2:行主序从右到左,从下到上
- 路径3:列主序从上到下,从左到右
- 路径4:列主序从下到上,从右到左

每条路径生成一个一维的块嵌入序列。这四个序列由共享的Mamba模块(具有选择性状态更新,即输入依赖的转移矩阵)独立处理。随后,四个输出序列被重新变形回二维,并逐元素相加。这一设计确保每个像素至少能通过四条扫描方向之一关注到所有其他像素,从而实现全局感受野。

复杂度分析: 对于包含N个块(N = H×W)的图像:
- 自注意力:O(N²) 计算与内存
- SS2D:O(N) 计算与内存,因为每个Mamba步骤处理一个token,状态大小恒定

这种线性缩放对于高分辨率图像具有变革意义。以4K图像(3840×2160)为例,若块大小为16,则N ≈ 32,400个块。自注意力每个头需要约10亿次操作;而SS2D仅需约32,000次操作。

架构变体: GitHub仓库提供了三种规模:

| 模型 | 参数量 | FLOPs(224×224) | ImageNet Top-1 |
|---|---|---|---|
| VMamba-Tiny | 22M | 4.3G | 82.5% |
| VMamba-Small | 50M | 8.7G | 83.6% |
| VMamba-Base | 89M | 15.1G | 84.3% |

数据要点: VMamba在显著更少的FLOPs下达到了与Swin Transformer相当的准确率。例如,Swin-Tiny(28M参数,4.5G FLOPs)达到81.3% Top-1,而VMamba-Tiny(22M参数,4.3G FLOPs)达到82.5%。在高分辨率下,效率差距进一步拉大。

训练稳定性: 原始Mamba论文指出,SSM需要对状态转移矩阵A(通常设为负对角值)和步长Δ进行仔细初始化。VMamba继承了这一敏感性。作者使用了逐层学习率衰减和梯度裁剪,但社区反馈表明,如果没有这些特定的超参数,在ImageNet上从头训练可能会不稳定。仓库中包含了带有推荐设置的训练脚本,但在更大数据集(如ImageNet-21K、JFT-300M)上的扩展尚未得到验证。

开源实现: 仓库`mzeromiko/vmamba`(截至撰写时已获3,133颗星)基于PyTorch和CUDA构建,并包含一个用于选择性扫描操作的自定义CUDA内核(基于Mamba内核)。代码库包括:
- `models/vmamba.py`:主模型定义
- `models/ss2d.py`:2D选择性扫描实现
- `kernels/selective_scan`:用于快速扫描的CUDA内核
- `configs/`:ImageNet、COCO、ADE20K的训练配置

关键参与者与案例研究

VMamba项目由来自多个机构的研究人员领导,主要基于Albert GuTri Dao(原始Mamba论文的作者,来自CMU和普林斯顿大学)开发的Mamba架构。VMamba团队将一维SSM适配到了二维,但其核心知识谱系清晰可见。

竞争方法对比:

| 模型 | 类型 | 复杂度 | ImageNet Top-1(Tiny) | 推理速度(1024×1024) |
|---|---|---|---|---|
| Swin-Tiny | Transformer | O(N²) | 81.3% | 1.0×(基准) |
| ConvNeXt-Tiny | CNN | O(N) | 82.1% | 1.3× |
| VMamba-Tiny | SSM | O(N) | 82.5% | 1.8× |
| EfficientFormer-L1 | 混合 | O(N) | 80.2% | 2.1× |

数据要点: VMamba在线性复杂度模型中取得了最佳准确率,同时在高分辨率下比Swin Transformer快近2倍。纯CNN模型ConvNeXt虽然具有竞争力,但缺乏VMamba所提供的全局感受野。

案例研究:高分辨率医学影像
一个值得关注的早期采用者是医学影像社区。哈佛医学院的研究人员已尝试将VMamba用于全切片病理图像(十亿像素级)。传统的基于ViT的模型由于二次方内存开销无法处理这些图像,只能退而求其次采用上下文有限的块级处理。VMamba的线性复杂度使其能够在单次前向传播中处理10,000×10,000像素的切片,从而实现全局组织架构分析。在CAMELYON16数据集上的初步结果显示,VMamba在转移瘤检测中达到了94.2%的AUC,而ResNet-152基线为92.8%,使用滑动窗口的Swin-Tiny为91.5%。

案例研究:实时视频理解
NVIDIA的研究人员已将VMamba集成到实时视频动作识别流程中。通过将视频帧视为

更多来自 GitHub

VMamba 登陆 ONNX:SS2D 算子如何解锁状态空间模型的边缘部署由开发者 haokun-li 创建的 vmamba_onnx 项目,直面基于状态空间模型(SSM)的视觉架构部署中的根本瓶颈:缺乏 ONNX 导出支持。VMamba 作为构建于 Mamba 状态空间模型之上的视觉骨干网络,依赖一种新型的二维分层Transformer:长序列视觉任务中的更智能注意力机制Stratified Transformer最初由dvlab-research团队开发,引入了一种分层注意力机制,将视觉令牌划分为局部和全局两组,并在不同尺度上进行处理。这一设计大幅降低了标准自注意力机制的二次复杂度,使得在消费级GPU上处Google DeepMind Gemma:开源权重大模型重塑AI可及性2024年2月21日,Google DeepMind推出Gemma,一个开源权重的大语言模型库,标志着这家科技巨头的重大战略转向。与专有的Gemini模型不同,Gemma采用宽松许可证免费提供,包含预训练和指令微调版本,参数规模分别为20亿查看来源专题页GitHub 已收录 1153 篇文章

时间归档

April 20262716 篇已发布文章

延伸阅读

VMamba 登陆 ONNX:SS2D 算子如何解锁状态空间模型的边缘部署全新开源项目 vmamba_onnx 成功将 VMamba 视觉状态空间模型导出为 ONNX 格式,攻克了关键的 SS2D 算子兼容性问题。这一突破使基于 SSM 的视觉骨干网络能够脱离 PyTorch 运行,为边缘部署与工业推理打开了大门S4模型:数学突破挑战Transformer在长序列领域的统治地位结构化状态空间序列(S4)模型代表了长序列深度学习领域的范式转变。它将经典状态空间理论与现代深度学习相结合,在数万步长的序列上实现了线性时间复杂度和稳定梯度,正挑战Transformer在音频、基因组学和时间序列预测等领域的统治地位。分层Transformer:长序列视觉任务中的更智能注意力机制来自hanyi-study仓库的全新开源实现,让分层注意力架构Stratified Transformer重回聚光灯下。该架构承诺在不牺牲全局上下文的前提下,大幅削减长序列视觉任务的计算成本,有望在有限硬件上彻底改变视频理解与高分辨率图像处Google DeepMind Gemma:开源权重大模型重塑AI可及性Google DeepMind正式发布Gemma,一个基于Gemini同源研究打造的开源权重大语言模型家族。提供20亿和70亿参数两个版本,Gemma旨在为开发者、研究人员和小型团队降低前沿AI门槛,同时深度整合Google生态系统。

常见问题

GitHub 热点“VMamba: How State Space Models Are Reshaping Computer Vision Beyond Transformers”主要讲了什么?

The dominance of Transformers in computer vision is facing a credible challenger. VMamba, a new visual backbone built on the state space model (SSM) architecture of Mamba, demonstr…

这个 GitHub 项目在“VMamba vs Swin Transformer benchmark comparison”上为什么会引发关注?

VMamba's architecture is a careful adaptation of the Mamba state space model (Gu & Dao, 2024) to the 2D image domain. The key challenge is that Mamba operates on 1D sequences, processing tokens one by one with a hidden s…

从“How to train VMamba from scratch on custom dataset”看,这个 GitHub 项目的热度表现如何?

当前相关 GitHub 项目总星标约为 3133,近一日增长约为 0,这说明它在开源社区具有较强讨论度和扩散能力。