极市导读
本文提出VCMamba,先用CNN卷积提局部细节,再换四向Mamba扫全局,全程线性复杂度;ImageNet上82.6% Top-1、ADE20K 47.1 mIoU,参数比同级模型少37%仍领先,为高精度视觉任务提供轻量级新骨干。>>加入极市CV技术交流群,走在计算机视觉的最前沿
精简阅读版本
本文主要解决了什么问题
-
视觉Transformer (ViTs) 和状态空间模型 (SSMs) 在捕捉细粒度局部特征方面不如卷积神经网络 (CNNs) 有效,而CNNs缺乏ViTs和Mamba等模型的全局推理能力。 -
现有视觉SSM模型主要依赖patch embeddings,无法充分利用CNN在视觉处理早期阶段擅长捕捉的丰富、空间密集的局部特征。 -
标准非重叠patch embeddings可能导致信息损失,特别是对密集预测任务至关重要的细粒度细节。
本文的核心创新是什么
-
提出了VCMamba,一种新颖的分层视觉架构,有效结合了多阶段卷积特征提取和多方向Mamba SSMs,实现高效且强大的全局上下文建模。 -
采用卷积茎和分层结构,在早期和中间阶段使用卷积块提取丰富的多尺度局部特征,后期阶段使用多向Mamba块建模长距离依赖关系和全局上下文。 -
设计了多向Mamba块,利用四向扫描机制和方向感知更新,有效捕获2D空间关系,同时保持相对于图像分辨率的线性复杂度。
结果相较于以前的方法有哪些提升
-
在ImageNet-1K分类任务上,VCMamba-B实现了82.6%的top-1准确率,以37%较少参数超过PlainMamba-L3的0.3%,以64%较少参数优于Vision GNN-B的0.3%。 -
在ADE20K语义分割任务上,VCMamba-B获得了47.1 mIoU,比EfficientFormer-L7高出2.0 mIoU,同时使用的参数减少了62%。 -
VCMamba-S在ImageNet-1K上达到78.7%的top-1准确率,超过了ViM-Ti(76.1%)和Pyramid ViG-Ti(78.2%),展示了在轻量级模型上的优势。
局限性总结
-
论文未明确讨论VCMamba在更高分辨率图像上的计算效率和性能表现,尽管提到Mamba具有线性复杂度优势。 -
未提供VCMamba在其他视觉任务(如目标检测、实例分割)上的实验结果,限制了对其通用性的全面评估。 -
未详细分析多向Mamba块的计算开销及其对整体推理速度的影响,这在实际应用中可能是一个重要考量因素。
导读
Vision Transformers(ViTs)和 State Space Models(SSMs)的最新进展已经挑战了 Convolutional Neural Networks(CNNs)在计算机视觉领域的主导地位。ViTs 擅长捕捉全局上下文,而像 Mamba 这样的 SSMs 为长序列提供了线性复杂度,但它们在捕捉细粒度局部特征方面不如 CNNs 有效。相反,CNNs 对局部特征具有强大的归纳偏置,但缺乏 transformers 和 Mamba 的全局推理能力。为了弥合这一差距,作者引入了 VCMamba,这是一种新颗的视觉 Backbone 网络,整合了 CNNs 和多向 Mamba SSMs 的优势。 VCMamba 采用卷积茎和分层结构,在其早期阶段使用卷积块来提取丰富的局部特征。这些卷积块随后由包含多向 Mamba 块的后期阶段处理,这些块旨在有效地建模长距离依赖关系和全局上下文。这种混合设计允许优越的特征表示,同时保持相对于图像分辨率的线性复杂度。作者通过在 ImageNet-1K 分类和 ADE20K 语义分割上的广泛实验证明了 VCMamba 的有效性。作者的 VCMambaB 在 Image Net-1K 上实现了 的 top-1 准确率,以 的较少参数超过了 PlainMamba-L3 的 ,并以 的较少参数优于 Vision GNN-B 的 。此外,VCMamba-B 在 ADE20K 上获得了 47.1 mloU ,比 EfficientFormer-L7 高出 2.0 mloU ,同时使用的参数减少了 。
代码 https://github.com/Wertyuui345/VCMamba
01 引言
深度学习视觉架构的发展主要由卷积神经网络(CNNs)所塑造。其固有的归纳偏置,如局部性和权重共享,使得分层特征的高效学习成为可能。然而,卷积的固定感受野可能会限制其有效捕获长距离依赖关系的能力。视觉Transformer(ViTs)作为一种强大的替代方案出现,它将图像视为 Patch 序列,并利用自注意力机制来建模全局关系,通常能够实现卓越的性能,尽管其复杂度相对于 Patch 数量呈二次方增长。
最近,状态空间模型(SSMs)作为一种有前景的架构引起了关注。Mamba 是一个显著的SSM变体,引入了选择机制和硬件感知算法,实现了线性时间序列建模和在语言任务上的强大性能。这激发了将SSMs适配于视觉领域的兴趣,催生了如VMamba、Vision Mamba (Vim) 和PlainMamba等模型。VMamba使用2D选择性扫描模块将一维SSMs适配于2D视觉数据。Vim利用双向Mamba块进行视觉表示,而PlainMamba提出了一种简单的、非层次化的SSM,并具有特定的2D扫描适配用于视觉识别。这些模型展示了SSMs在视觉领域的潜力,在性能和效率之间提供了有吸引力的平衡,特别是对于高分辨率输入。
尽管取得了这些进展,许多当代视觉SSM(类似于ViTs)主要依赖于patch embeddings,这可能无法充分利用CNN在视觉处理早期阶段擅长捕捉的丰富、空间密集的局部特征。虽然PlainMamba专注于非分层SSM结构,而其他视觉SSM探索不同的扫描策略或双向机制,但仍有机会创建一种视觉 Backbone 网络,明确融合分层卷积阶段的强大局部特征提取与后期High-Level多向Mamba模块的高效全局建模。标准的非重叠patch embeddings可能导致信息损失,特别是对于密集预测任务至关重要的细粒度细节,而这一挑战可以通过为更高分辨率阶段设计的分层卷积结构来缓解。
为了弥合这一差距,作者引入了VCMamba,一种整合了CNNs和多向Mamba SSMs优势的新型视觉 Backbone 网络。VCMamba采用了一个卷积茎和分层结构,在其早期和中间阶段使用卷积块(实现为带有 和 卷积的 FFN ),使能够提取丰富的多尺度局部特征。这些特征随后由包含多向Mamba块的后期阶段处理,特别利用了中详述的四向扫描机制,以有效地建模长距离依赖和全局上下文。这种混合的分层设计允许卓越的特征表示,同时在其Mamba阶段保持相对于图像分辨率的线性复杂度,这是高分辨率图像处理的一个关键优势。
作者通过在ImageNet-1K分类和ADE20K语义分割上的大量实验证明了VCMamba的有效性。VCMambaB在ImageNet-1K上实现了 的top-1准确率,以少 的参数量超越了 PlainMamba-L3( 准确率) ,并在使用少 参数量的情况下优于Vision GNN-B( 准确率) 。此外,VCMamba-B在ADE20K语义分割上获得了 47.1 的平均 ,在使用少 参数量的情况下超过了EfficientFormer-L7( 45.1 mloU ) 2.0 mloU 。这些结果展示了VCMamba的强劲性能和高效率,特别是在受益于其混合特征提取能力的任务中。
作者的贡献是:
-
作者提出了VCMamba,一种新颖的分层视觉架构,它有效地结合了多阶段卷积特征提取和多方向Mamba SSMs,以实现高效且强大的全局上下文建模。
-
作者提出了几种VCMamba变体,并证明了它们相对于领先的CNN、ViT、视觉GNN和视觉SSM架构的优越性能。
-
作者表明,作者的混合和分层方法在图像分类和语义分割任务中都取得了优异的性能,同时保留了Mamba在高分辨率输入方面的效率优势。
本文的组织结构如下。第2节涵盖相关工作。第3节解释了state space models的初步信息。第4节描述了作者的hierarchical feature extraction、multi-directional mamba blocks以及VCMamba架构。第5节描述了作者在ImageNet-1K图像分类和ADE20K语义分割方面的实验设置和结果。最后,第6节总结了作者的主要贡献。
02 相关工作
2.1 视觉架构
卷积神经网络(CNNs)长期以来一直是计算机视觉领域的主导架构。其成功源于固有的归纳偏置,如局部性和权重共享,这些特性使得层次化视觉特征的高效学习成为可能。诸如ResNet 、EfficientNet 和ConvNeXt 等架构不断推动性能边界。对于移动应用,轻量级CNN如MobileNet 引入了深度可分离卷积等高效操作。近期研究进一步探索了通过使用高效扩展感受野的技术来增强用于移动视觉的CNN。尽管CNN在局部特征提取方面表现出色,但与基于注意力的模型相比,其固定的感受野在捕获全局、长距离依赖关系方面可能存在局限性。
Vision Transformers (ViTs)通过将最初为自然语言处理设计的Transformer架构应用于图像数据,标志着一种范式转变。ViTs 将图像视为 Patch 序列,并使用自注意力机制对它们之间的全局关系进行建模,通常能获得最先进的结果。然而,自注意力相对于 Patch 数量(图像分辨率)的二次复杂度为高分辨率输入和密集预测任务带来了重大的计算挑战。为缓解这一问题所做的努力包括MobileViT 和MobileViTv2 ,它们旨在通过融入卷积原理为移动设备创建更高效的ViT变体。
视觉图神经网络(ViGs)通过将图像建模为相互连接的 Patch 或节点的图提供了另一种视角。ViG 是早期提倡使用GNN作为通用视觉 Backbone 网络的工作,它利用K-最近邻算法连接图中的相似节点。对于移动应用,MobileViG 引入了一种基于静态图的连接机制。进一步的优化如GreedyViG 、WiGNet 和ClusterViG 专注于引入新的动态和高效的图构建算法。虽然ViGs提供了一种灵活的建模关系的方式,但图构建和传播步骤可能会引入计算开销。
2.2 视觉中的状态空间模型(SSMs)
状态空间模型(SSMs)最近作为一种引人注目的序列建模替代方案而出现,其提供了与序列长度成线性关系的复杂度。Mamba 作为一个突出的SSM,引入了一种选择性扫描机制,该机制实现了高效的、依赖于输入的处理,并在语言建模中表现出强大的性能。这种在语言建模方面的成功激发了将SSMs,特别是Mamba,适应于计算机视觉任务的浓厚兴趣。将Mamba的能力转换到视觉领域的初步努力包括:
-
VMamba 提出了2D选择性扫描模块(SS2D)来为SSMs提供2D空间感知能力,将图像转换为通过四向扫描处理的有序块序列。 -
Vim 引入了一个采用双向Mamba块和位置嵌入的视觉 Backbone 网络。 -
PlainMamba 专注于用于视觉的非分层Mamba架构,引入连续2D扫描以更好地将Mamba的选择性扫描适应到2D图像。
其他最近的研究方向包括领域特定适应如VideoMamba,专门任务如使用QMamba进行质量评估,以及效率提升如PTQ4VM 中的训练后量化。同时,分析性工作如MambaOut 和批判性地检查了Mamba组件在视觉领域的作用和功效。这些模型强调了SSMs为各种视觉任务提供高效强大 Backbone 的潜力,尽管活跃的研究仍在继续优化它们在2D数据上的应用。
2.3 混合架构
不同架构的互补优势导致了各种混合模型的出现。像CoAtNet 和MobileFormer 这样的CNN-ViT混合模型旨在结合CNN的局部特征提取与ViT的全局上下文建模能力。EfficientFormer 进一步优化了这种融合,以提高移动设备上的速度和效率。MambaVision 代表了一种最近的Mamba-Transformer混合模型,它将Transformer块集成到其后期阶段,以改善对长距离空间依赖性的捕获。
作者提出的VCMamba通过创建一个独特的CNN-SSM混合架构为这一研究方向做出了贡献。与将Transformer块与Mamba集成的MambaVision不同,VCMamba采用了一种分层结构,在其早期和中间阶段使用卷积块进行强大的多尺度特征提取,随后在后期阶段过渡到多向Mamba块以进行高效的全局建模。VCMamba旨在利用CNNs的局部特征丰富性和Mamba的序列建模能力,而不依赖Transformer块。
03 预备知识:状态空间模型 (Mamba)
状态空间模型(SSMs)通过状态变量描述系统。一个连续线性SSM定义为:
其中 是输入, 是潜在状态, 是输出,( )是系统矩阵。对于深度学习应用,这些通过使用时间尺度参数 进行离散化,将 A 和 B 转换为它们的离散对应物 和 。 Mamba[9]通过使参数 和 依赖于输入(选择性),显著增强了传统的 SSM。这使得模型能够基于当前 token 动态调整其行为,沿着序列选择性地传播或遗忘信息。离散化递归为:
其中下标 表示Mamba[9]中参数的输入依赖性。Mamba采用硬件高效的并行扫描算法进行训练和推理,实现了线性复杂度。
将Mamba适配到2D视觉数据,如PlainMamba[47]中所示,涉及将图像块展平为一维序列,然后应用选择性扫描。为了捕获2D空间上下文,PlainMamba采用了连续2D扫描等技术,该技术以多种(例如四种)预定义顺序处理视觉token,确保空间邻接性,以及方向感知更新,其中表示扫描方向的可学习参数 被整合到SSM的更新规则中。例如,对于第 k 个扫描方向和第 i 个 token的隐藏状态 的更新可以增强为:
这些多向扫描的输出然后通常会被聚合。这种多向方法对VCMamba的后续阶段至关重要。
04 VCMamba 架构
作者将VCMamba设计为一个分层视觉 Backbone 网络,在其早期阶段集成了CNN的强大局部特征提取能力,在后期阶段结合了多方向Mamba SSMs的高效全局上下文建模能力。这种混合方法使VCMamba能够在保持计算效率的同时,有效地处理多尺度的视觉信息。如图1所示的整体架构,包括一个卷积茎和随后的四个阶段的特征提取块,阶段之间有下采样层以创建特征金字塔。
4.1 层次化特征提取
VCMamba架构以卷积主干(图1(b))开始处理,这是现代视觉主干网络中用于初始特征提取和空间下采样的常见策略。该主干包含两个连续的 卷积层,每个卷积层使用步长为 2。每个卷积后都跟随批归一化和ReLU激活函数。总体而言,该主干将输入图像下采样4倍,在降低的空间分辨率上高效生成Low-Level特征表示,为后续阶段做准备。
在主干网络之后,VCMamba采用四阶段分层结构(图1(a))以逐步细化特征并构建多尺度表示。每个阶段在不同的空间分辨率上运行。采用下采样层将空间尺寸减半并扩展通道容量。这些下采样层通过一个带步长的 卷积后接批量归一化来实现(图1(c))。这种金字塔架构能够在不同尺度上学习特征,这对于在多样化视觉任务上实现稳健性能至关重要。该架构具有基于模型大小的不同模型变体(VCMamba-S、VCMamba-M和VCMamba-B),如第4.3节所讨论。
在初始三个阶段以及第四阶段的早期部分,VCMamba主要利用卷积 FFN(FFN)块进行特征转换和细化。这些FFN块在架构上类似于移动CNN中突出的高效倒置残差块 。在结构上,每个FFN块(图1(d))包含一个核心MLP模块,该模块首先使用 卷积扩展通道维度,然后使用 深度卷积进行空间混合,随后使用另一个 卷积将特征投影回来。在这些MLP模块中应用了批归一化和GeLU激活函数,以确保训练稳定并引入非线性。这种FFN设计优先考虑在网络高分辨率的早期阶段进行高效的局部特征提取和表示学习。
4.2 用于全局上下文的多向Mamba块
在其最后阶段,VCMamba从卷积FFN块转换以整合多向Mamba块(图1(e))。这种转变使模型能够利用前面的卷积阶段提取的丰富的多尺度局部特征,并有效地建模特征图上的长距离依赖和全局上下文。
VCMamba中的每个基于Mamba的块都是一个复合结构。它接收来自前一层的2D特征图,应用批归一化(Batch Normalization),然后通过一个核心的2D适应Mamba模块进行处理。该Mamba模块的输出通过残差连接与输入相结合,接着是另一个批归一化(Batch Normalization)和一个卷积MLP,其结构与4.1节中描述的FFN块中使用的结构相同。
基于Mamba的块通过利用Plain Mamba中建立的扫描原理,将Mamba的选择性扫描机制适配于2D视觉数据。作者的Mamba块还利用位置嵌入来帮助理解空间关系。正如作者的消融研究中所详述的,该块的设计通过用 Shortcut 替代乘法分支、交错Mamba和FFN以及层归一化进行了优化,以最大化性能。对于特征图 的处理过程如下:
-
位置和局部上下文编码:输入特征图 首先通过带有位置嵌入的 卷积,将其投影到内部维度 并进行归一化。然后将得到的特征展平成一个序列。为了在主要的SSM操作之前用局部空间上下文丰富这些token,应用了一个 深度卷积,然后接一个SiLU激活函数。这一步骤对于准备视觉token以供SSM进行有效的顺序处理至关重要。
-
多向选择性扫描:为了全面捕捉2D空间关系,该模块采用多向扫描策略。与单一的单向扫描不同,视觉 Token 沿着四个不同的、空间连续的路径进行处理(例如,行和列的"蛇形"模式),如图 2中的概念性说明。这种方法类似于连续2D扫描方法,确保当 Token 被SSM顺序处理时,一维序列中的相邻性对应于原始2D特征图中的空间相邻性,从而保留了关键的语义和空间连续性。沿着这 个扫描方向中的每一个,Mamba选择性扫描机制利用其输入相关参数( ),根据公式(2)中的递归关系更新每个 Token 的隐藏状态 。
-
方向感知更新:为了在每次1D扫描期间明确告知模型2D空间遍历的性质,引入了一种方向感知更新机制,如中所述。这涉及一组可学习参数 ,每个参数对应一个扫描方向(以及一个初始"开始"方向)。这些方向参数被整合到SSM的状态更新方程中,通常通过增强输入相关矩阵 来实现。因此,对于第 个扫描方向和第 个 token 的修改后状态更新可以表示为:
其中 表示针对当前token和扫描路径的离散化方向参数。
-
聚合与输出投影:由四个方向扫描中的每一个产生的特征序列通过求和进行聚合。这种结合了多方向上下文的组合表示,使用LayerNorm进行归一化。然后接着是一个 卷积和批归一化。
这个多向Mamba模块使VCMamba能够在其深层特征提取阶段高效地建模全局交互和长距离依赖。
4.3 VCMamba网络架构
如图1所示,整体VCMamba架构集成了卷积茎、带有卷积FFN块的分层阶段、阶段间下采样,以及最后阶段交替使用基于多方向Mamba的块与卷积FFN块。
作者通过调整深度(每个阶段的块数)和宽度(通道维度)来定义几种VCMamba变体,从而在性能和计算成本之间进行权衡。VCMamba-S、VCMamba-M和VCMamba-B的配置在表1中详细说明,展示了从最小模型尺寸(VCMamba-S)到最大模型尺寸(VCMamba-B)的宽度和深度调整。这些配置指定了卷积FFN块的数量和多向
四个阶段中各自的基于Mamba的模块(MDM Blocks)及其相应的通道维度 。
05 实验
在本节中,作者详细介绍了实验设置,并对作者提出的VCMamba架构进行了全面的性能评估。作者将VCMamba与现有的用于图像分类和语义分割任务的突出架构进行了基准测试。作者的结果表明,VCMamba实现了出色的准确性和计算效率,优于几种最先进的CNN、ViT、ViG和其他基于Mamba的视觉模型。
5.1 在ImageNet-1K上的图像分类
作者在广泛使用的ImageNet-1K数据集上评估VCMamba,该数据集包含约 130 万张训练图像和 5 万张验证图像,涵盖 1000 个物体类别。所有VCMamba模型都从零开始训练 300 个 epoch,使用标准输入分辨率 。作者的实现利用了PyTorch和Timm库。遵循训练现代视觉 Backbone 网络的常见做法,作者的训练方案包括AdamW优化器 ,学习率为 和余弦退火调度,以及数据增强技术,如RandAugment、Mixup 、CutMix和随机擦除。
作者的中等大小VCMamba-M模型,拥有 21.0 M 参数,达到了 的top-1准确率,优于 DeiT-Small(拥有 22 M 参数,准确率为 )和PVT-Small(拥有 24.5 M 参数,准确率为 )。轻量级VCMamba-S模型仅用 10.5 M 参数就达到了 的top-1准确率,超过了ViM-Ti(拥有 7 M 参数,准确率为 )和Pyramid ViG-Ti(拥有 10.7 M 参数,准确率为 )。这些结果突显了VCMamba的分层混合设计在各种模型规模上实现强大分类性能的有效性。
这些不同规模的结果突显了VCMamba的混合CNN-Mamba设计的有效性,该设计在早期阶段利用卷积的优势,在后期阶段利用多向Mamba能力,实现了良好的精度-参数权衡。
5.2. ADE20K上的语义分割
为评估VCMamba在密集预测任务上的能力,作者在ADE20K数据集上进行了语义分割实验。ADE20K包含 20 K张训练图像和 2 K张验证图像,涵盖 150 个语义类别。作者遵循既定方法,构建了以Semantic FPN作为分割解码器的VCMamba。
VCMamba Backbone 网络使用其在ImageNet-1K上预训练的权重进行初始化。随后,模型进行了40K次迭代的微调。作者使用AdamW优化器,初始学习率为 ,并使用幂为 0.9的多项式调度进行衰减。训练输入分辨率为 。
如表3所示,VCMamba在语义分割任务中表现出强大的性能。作者的VCMamba-S(1050万 Backbone 网络参数)达到了 42.0 mIoU ,超越了其他轻量级模型,如EfficientFormerL1 ( )。作者更大的模型VCMamba-B(3150万 Backbone 网络参数)达到了令人印象深刻的 47.1 mloU 。这超越了FastViT-SA36 4.2 mIoU 以及参数量大得多的 EfficientFormer-L7(拥有 82.1 M Backbone 网络参数, 45.1 mloU ) 2.0 mloU ,尽管VCMamba-B的 Backbone 网络参数减少了约 。这些结果突显了VCMamba架构在为密集预测任务学习强大表示方面的有效性,有效利用了其混合卷积和多向Mamba设计。
5.3. 消融研究
为了验证作者的架构设计选择,作者在ImageNet-1K上进行了一系列消融研究,从使用PlainMamba 层作为作者最终阶段的 Baseline 开始,并逐步集成作者提出的VCMamba的关键组件。该 Baseline 模型在拥有参数的情况下达到了的top-1准确率。逐步改进的细节在表4中详细说明。
首先,作者通过用简单的 Shortcut 替代乘法分支来优化Mamba块的内部结构,这使准确率提高到 ,同时略微减少了参数量至 31.0 M 。接下来,为了更好地融合特征表示,作者交错排列Mamba和FFN层,并在其后应用批量归一化,进一步将性能提升至 。然后,作者观察到在Mamba层内添加一个额外的LayerNorm带来了显著提升,准确率达到 。接下来,作者用卷积替换线性层头,使准确率达到 。最后,为确保模型层次结构中稳定的特征分布,作者用批量归一化层包豪每个阶段。这个最终模型是作者的VCMamba-B,它的top-1准确率达到 ,比作者的CNN-Mamba Baseline 模型总共提高了 。这一系列消融实验验证了最终VCMamba架构的有效性。
06 结论
在本文中,作者介绍了VCMamba,一种新颗的分层视觉主干网络,旨在协同地桥接卷积神经网络(CNNs)的强大局部特征提取能力与状态空间模型(SSMs)的高效全局上下文建模。VCMamba采用多阶段架构,在其早期和中间阶段利用卷积块来构建丰富的多尺度特征层次结构。在其更深、更低分辨率的阶段,它转向多向Mamba块以有效捕获长距离依赖关系。这种混合设计使 VCMamba在其Mamba阶段保持线性复杂度,为高分辨率视觉理解提供了可扩展的解决方案。
作者的广泛评估表明VCMamba在ImageNet-1K分类和ADE20K语义分割任务上具有强大的性能和效率。值得注意的是,VCMamba-B在ImageNet-1K上达到了 的top-1准确率,以减少 参数量的情况下,超过了PlainMamba-L3模型 ,并在ADE20K上获得了 47.1 mloU ,在使用减少 参数量的同时,超过了EfficientFormer-L7 2.0 mIoU 。这些结果验证了VCMamba作为一个引人注目且高效的主干网络,适用于多种计算机视觉任务。
参考
[1]. VCMamba: Bridging Convolutions with Multi-Directional Mamba for Efficient Visual Representation

公众号后台回复“数据集”获取100+深度学习各方向资源整理
极市干货

点击阅读原文进入CV社区
收获更多技术干货

