
极市导读
本文介绍一个种新型视觉基础模型DefMamba,通过多尺度 Backbone 结构和可变形 Mamba 模块,结合动态扫描策略(DS),显著提升了视觉任务的表现。该模型在图像分类、目标检测和语义分割等任务中均优于现有方法,展现了强大的特征提取和细节感知能力。 >>加入极市CV技术交流群,走在计算机视觉的最前沿

近期,状态空间模型(SsM),特别是Mamba,因其能有效平衡计算效率和性能,吸引了学者的广泛关注。
然而,大多数现有的视觉Mamba方法使用预定义的扫描顺序将图像展平为1D序列,导致模型在特征提取过程中对图像空间结构信息的利用能力减弱。为解决这一问题,作者提出了一种新型视觉基础模型——DefMamba。
该模型包含多尺度 Backbone 结构和可变形Mamba(DM)模块,能够动态调整扫描路径以优先处理重要信息,从而提升对相关输入特征的捕获和处理能力。
通过结合可变形扫描(DS)策略,该模型显著增强了学习图像结构的能力,并能检测物体细节的变化。大量实验表明,DefMamba在图像分类、目标检测、实例分割和语义分割等多种视觉任务中均取得了最先进的性能。
1. 引言
目前大多数视觉基础模型主要依赖卷积神经网络(CNNs)[25, 28, 32]和Transformer架构[9, 24, 33]。然而,CNNs受其滑动窗口结构的限制,这限制了感受野,并显著阻碍了输入数据全局信息的聚合。相比之下,Transformer由于注意力机制,在全局信息聚合方面表现出色,但其高计算复杂度在实现效率与性能之间的平衡上构成挑战。状态空间模型(SSMs)[12]为这一权衡提供了一个潜在的解决方案。SSMs通过隐藏状态矩阵聚合先前特征来更新当前特征,从而将计算复杂度降低为与序列长度呈线性关系。尽管SSMs以递归方式处理序列,但在简化后,SSMs可以并行计算序列。尽管具有这些优势,但由于状态矩阵更新过程中缺乏内容感知感知,SSMs难以捕获长程依赖。
近期,Mamba [11] 提出了一种改进的选择机制,旨在优化状态空间模型(SSMs)的训练过程。这种创新机制将内容感知引入特征提取流程,扩展了有效感受野,并在多种自然语言处理(NLP)任务中实现了显著的性能提升。因此,众多研究尝试将该方法扩展到更广泛的计算机视觉领域。这一过程中的主要挑战是如何在映射二维图像特征图到一维序列时,不丢失关键信息。大多数现有方法采用预定义的映射策略,如光栅扫描 [23, 47]、局部扫描 [19] 和连续扫描 [42]。然而,如图1所示,这些方法都依赖于固定的扫描路径。这导致在展平后相邻的token不再相邻。因此,它们忽略了图像的固有空间结构,导致结构信息的丢失。为解决这一问题,QuadMamba [39] 根据图像不同区域包含的信息量确定扫描窗口大小。然而,每个窗口内的扫描顺序是固定的,这导致上述问题未能得到完整解决。GrootV [38] 基于输入特征自适应构建树状拓扑,并从该拓扑中提取特征。但它在构建拓扑时仅使用相邻特征,并在块上均匀分配注意力。上述方法要么基于固定的扫描顺序,导致结构信息丢失;要么平等对待感知区域内的信息,导致对物体细节变化的敏感性不足。

为解决这一问题,作者提出了一种名为DefMamba的新颖框架,该框架受可变形机制[6, 36]的启发。然而,直观地将可变形机制应用于状态空间模型(SSMs)仍会导致结构信息丢失并增加计算复杂度。因此,作者提出了一种可变形状态空间模型和可变形扫描策略(DS),该策略基于重要信息优先考虑可变形 Token ,并将参考点滑动至重要区域。这种方法使SSMs能够更有效地捕获和处理与输入相关的相关特征。具体而言,作者将参考点从固定位置移动到可调节位置,该位置提供了更多有用信息,从而有助于感知目标细节的变化。另一方面,作者还通过偏移向量动态调整扫描顺序,以获得结构感知序列。
通过这种方式,yinqing-DefMamba_2504自适应地感知目标细节的变化,以找到最合适的特征点,并根据输入图像特征确定与目标结构一致的优化扫描顺序。
作者进行了广泛的实验,以验证DefMamba在多个视觉基准上的有效性,包括在ImageNet [8]上的图像分类、在COCO [21]上的目标检测和实例分割,以及在ADE20K [45]上的语义分割。
这些结果表明,yinqing-DefMamba_2504在所有基准上都优于现有的基于SSM的方法,并且与基于CNN和transformer的方法保持竞争力。
2. 相关工作
2.1. Mamba用于视觉应用
大量研究已成功将Mamba [11]应用于视觉任务 ,并取得初步成果。ViM [47]引入双向扫描方法将2D图像转换为1D序列,随后将这些序列输入SSM进行全局上下文建模,标志着Mamba首次应用于视觉任务。VMamba [23]采用四向扫描算法将2D图像转换为1D序列。PlainMamba [42]将扫描方法从栅格改为连续,保留了图像的空间依赖性。
MSVMamba [31]基于四次扫描对序列进行下采样,以减少计算冗余并缓解信息丢失问题。GrootV [38]在四向平面图上构建最小生成树,利用相邻特征之间的差异,根据不同输入动态调整扫描顺序。QuadMamba [39]在扫描过程中根据图像的信息内容自适应调整窗口粒度,以更好地聚合局部信息。虽然GrootV和QuadMamba可以根据输入数据调整扫描方法,但GrootV在生成最小生成树时仅考虑相邻元素之间的关系,忽略了全局信息。另一方面,QuadMamba仍然依赖预定义的扫描方法,并未实现真正的动态扫描。相比之下,作者的DefMamba引入了内容感知可变形扫描策略,允许网络动态学习扫描顺序和参考点位置。
2.2. 可变形卷积神经网络和注意力机制
可变形卷积[6, 20, 41, 48]采用一种能够适应输入特征图几何变化的卷积核,从而克服了传统卷积在处理复杂目标时表现不佳的局限性。近期,可变形机制已被扩展到视觉Transformer[9],以增强其捕捉局部特征和适应几何变化的能力。DPT[5]提出了一种自适应块嵌入方法,该方法在保留块语义信息的同时动态调整块的位置和大小。PSViT[44]在ViT[9]之前引入了一个渐进式采样模块,该模块迭代地识别当前图像中最合适的可变形点位置。
DAT[36]首次将可变形机制与Transformer中的自注意力机制相结合,将可变形注意力集成到视觉主干网络中。该方法学习一组与全局关键点对应的特征,并适应空间变化。先前的方法探索了将可变形机制有效融入Transformer架构的各种方式。随着Mamba[11]的近期流行,作者尝试将可变形机制引入Mamba。然而,直接应用这些机制导致了结构信息丢失和需要额外模块的问题。在此背景下,作者设计的DS策略通过有效优先处理可变形 Token 并将参考点引导至关键区域而脱颖而出。
3. 方法
在本节中,作者首先总结了第3.1节中的状态空间模型(SSMs)。然后在第3.2节中,作者描述了所提出网络的整体结构。第3.3节介绍了一种可变形状态空间模型(DSSM)。最后,第3.4节展示了跨多个尺度的模型配置设计。
3.1. 预备知识
SSMs,包括S4[12]和Mamba[11]等显著实现,是结构化序列架构,结合了循环神经网络 (RNNs)和卷积神经网络(CNNs)的元素,能够实现关于序列长度的线性或近线性扩展。这些源自连续系统的模型,通过隐藏状态 ,定义了输入 到输出 的一维函数到函数映射。其中 表示时间。更正式地,SSMs的特征由[11]中给出的连续时间常微分方程(ODE)呈现:
其中 是当前隐藏状态。 是更新后的隐藏状态。 是当前输入。 是输出。 是SSM的演化矩阵, 和 分别是输入和输出投影矩阵。
为了使状态空间模型(SSMs)在深度学习的序列建模任务中得以应用,必须对其进行离散化处理,将SSM从连续时间函数到函数的映射转换为离散时间的序列到序列映射。S4[12]和 Mamba[11]是连续系统的离散化示例,它们引入了时间尺度参数 ,将连续参数 转换为其离散对应物 。这种离散化通常通过零阶保持(ZOH)[11]方法实现:

S4[12]和 Mamba[11]都采用了与方程2中概述的类似的离散化过程,但Mamba通过S6选择性扫描机制对参数 和 进行了条件化,该条件化基于输入 。其中, 表示批大小, 表示序列长度, 表示特征维度。
3.2. 整体模型架构
DefMamba采用了一种与许多CNN[25,28]和Transformer[24]相似的通用多尺度 Backbone 结构。如图2(a)所示,图像 首先通过一个patch嵌入层被分割成块,该层生成一个具有空间维度 和通道维度 的2D特征图。随后,使用多个网络阶段来创建维度为 的层次化表示。每个阶段由一个Deformable Mamba(DM)块堆栈和一个下采样层(第四阶段除外)组成。最后,特征通过平均池化后送入分类头以获得预测结果。特别地,作者遵循[38]并在patch嵌入层和下采样层中使用了重叠形式。具体结构细节请参考附录。

与语言模型中使用的Mamba结构[11]不同,DM模块遵循Transformer模块的流行结构[9, 24],该结构包含两个Layer Norm(LN)层、一个FFN、一个DSSM(将在下一小节介绍)以及残差连接,如图2(b)所示。
3.3. 可变形状态空间模型
可变形状态空间模型的总体架构如图3(a)所示。受[23, 31, 38]的启发,作者在视觉Mamba模块[47]中用深度卷积替换了原有的1D卷积,并引入了一个包含可变形扫描和可变形状态空间模型(DSSM)的可变形分支。作者保留了标准的前向和后向分支以确保模型稳定收敛,因为与之前的扫描方法相比,yinqing-DefMamba_2504引入了更多的空间 Token 跳跃,这可能会使模型训练和学习变得复杂。后续实验验证了这一点(表4)。


可变形扫描。考虑到多个可变形点之间的相互干扰问题,作者将可变形点的偏移量限制在特定范围内。这确保了可变形点与参考点之间的关系保持不变,使作者能够利用变形前的关系来近似变形后的关系。此外,考虑到计算的简便性,作者采用并行方法同时调整参考点和扫描顺序,从而降低计算负担。具体结构如图3(b)所示。
给定输入特征 ,其中 表示通道维度, 表示空间分辨率。作者首先使用一个子网络生成偏移量 ,该子网络利用 输出参考点和参考 Token 的偏移量 。
作者最初实现的子网络如图3(c)所示。输入特征首先通过一个 的深度卷积进行处理,以捕获局部特征。随后,采用GELU,层归一化(LN)和一个 卷积来推导偏移值,这些偏移值总共有三个维度。前两个维度表示参考点在二维空间中的偏移,而第三个维度表示参考 Token索引在整个 Patch 中的偏移。
在作者的实验中,作者观察到 Token 的偏移需要全局感知 Patch 内的特征,而这一点无法仅通过卷积实现。基于这一点,并考虑到先前方法的研究发现[29,31]中Mamba在通道维度存在冗余,作者在深度卷积层之后引入了通道注意力(CA)机制[18]。该机制减少了通道冗余,并促进了全局上下文信息的整合。值得注意的是,根据DAT[36]的配置,作者在四个阶段中设置 为 , ,并在 卷积中省略了偏差。
为稳定训练过程,作者采用 函数来减轻 中极端值的影响,其中 。随后,作者将 沿通道维度分为两部分,一部分包含 2 个通道,另一部分包含 1 个通道,以获得点偏移 和 Token 索引偏移 。如前所述,作者需要进一步约束 ,以稳定训练并简化结构。作者将 的水平维度和垂直维度分别除以 和 ,从而将偏移限制在单个 Token 的范围内。详细过程如下:

然后,作者将点偏移量 ,Token 索引偏移量 和输入特征 分别发送到点偏移分支和索引分支,以获得最终输出。
点偏移。为了获得对物体变化更敏感的特征表示,作者动态地调整网络的参考点,使其变形为包含更多相关信息的点,这些点基于输入进行确定。首先,作者生成了参考点 中的值对应于从 到 的二维坐标点。为了简化网络计算,作者将 从其原始范围归一化到 ,其中 表示左上角的点,[1,1]表示右下角的点。然后,作者将参考点 和偏移量相加,得到变形点 。由于偏移量 包含小数部分,不能直接使用。因此,作者使用双线性插值从输入 中提取与偏移点 对应空间位置的特征。
当yinqing-DefMamba_2504执行点偏移时,特征会在空间中移动,这可能导致初始阶段添加的位置编码失效,从而降低模型性能。为了解决这个问题,作者基于 Swin Transformer 中的相对位置编码设计了偏移偏差。具体来说,给定大小为 的特征图,点的相对坐标位移在两个维度上分别位于 和 的范围内。因此,作者设置了一个可学习的相对偏移偏差矩阵 。然而,考虑到这样一个矩阵会导致参数数量显著增加,作者对这个矩阵进行了下采样操作,得到 。同时,作者将点位移除以 2 以适应这一变化。然后,作者使用点位移在 上进行插值,计算相应的补偿。最后,将这个补偿加到插值特征上。具体过程如下:

其中 表示使用双线性插值函数提取位置 在 索引偏移上的特征。作者修改了扫描顺序,通过改变参考 Token 索引和可变形 Token 索引,使模型能够有效感知输入目标的结构。作者最初生成了参考 Token 索引 ,其中 。 中的值表示当前 Patch 内的 Token 位置,范围从 0 到 。为简化网络计算,作者将 从原始范围归一化到[ -1 , 1]。这使作者能够计算可变形 Token 索引 。
由于推导出的 包含小数成分,它们不能直接使用。因此,对于 ,作者应用了一种排序算法,根据其值的幅度来确定偏移后的索引。最后,作者根据这些索引将偏移特征 转换为 1 D序列,从而获得内容自适应的图像特征序列。需要注意的是,排序算法会截断梯度,导致网络无法训练。为了解决这个问题,作者对最终图像序列的梯度在维度上进行平均,并将它们复制到 中,以近似扫描顺序偏移的梯度。
4. 实验
4.1. 图像分类
设置。图像分类实验使用ImageNet-1K [8]数据集进行,该数据集包含超过128万张训练图像和5万张验证图像,涵盖1000个类别。作者的训练设置严格遵循先前实践的方法 [23, 31, 38],并结合了多种数据增强技术,如随机裁剪、随机水平翻转、标签平滑正则化、mixup、autoaugment和随机擦除。模型使用AdamW [26]优化器进行训练,采用余弦退火学习率调度器,包括20个周期的预热阶段。总批处理大小设置为 1024 ,模型在 A 800 GPU上进行训练。优化器参数配置为beta设置为 ,动量为 0.9 ,初始学习率为 ,权重衰减为 0.05 ,并采用指数移动平均(EMA)。
结果。表1展示了作者提出的DefMamba模型(T、S、B)与各种当前最优(SOTA)方法之间的比较。具体而言,DefMamba-T实现了78.6%的Top-1准确率,分别比基于CNN的RegNetY800M[28]和基于transformer的DeiT-Ti[33]高2.3%和6.4%。此外,DefMamba-T优于最近引入的SSMs模型,在参数和计算复杂度方面,分别比ViM-T[47]、LocalViM-T[19]和MSVMamba-N[31]高2.5%、2.4%和1.3%。此外,它在实现0.7%性能提升的同时,将计算负担降低了60%,优于PlainMamba-L1[42]。DefMamba-S实现了83.5%的Top-1准确率,超过GrootV-T[38]和EfficientVMamba-B[27]。此外,DefMamba-B达到了84.2%的准确率,比VMamba-S[23]高0.6%,展示了yinqing-DefMamba_2504的有效性。

4.2. 目标检测
设置。作者使用Mask R-CNN框架[16]在MSCOCO 2017数据集[21]上评估了DefMamba,用于目标检测和实例分割任务。遵循先前工作[24,31,38],作者使用在ImageNet-1K上预训练的 Backbone 网络进行初始化。作者采用了标准的训练策略,包括 ( 12 个epoch)的MultiScale (MS)训练,以确保公平的比较。
结果。如表2所示,yinqing-DefMamba_2504在大多数评估指标上优于现有方法。具体而言,在 调度下,DefMamba-S在框mAP( )上达到 47.5 ,在 Mask mAP( )上达到 42.8。yinqing-DefMamba_2504超越了ResNet-50[15],Swin-T[24]和ConvNeXt-T[25]。同时,yinqing-DefMamba_2504与LocalVMamba-T[19]和QuadMamba-S[39]相比,将 提升了 0.6 个百分点。此外,yinqing-DefMamba_2504在性能上与之前的SOTA方法 VMamba-T[23]相当,同时将计算负载降低了 。

4.3. 语义分割
设置。为了评估DefMamba的语义分割性能,作者使用在ADE2OK[45]上预训练的分类权重初始化的UperNet[37]i川练yinqing-DefMamba_2504,迭代次数为160,000。作者采用了AdamW优化器[26],学习率设置为 。
作者的实验主要使用默认输入分辨率 进行。此外,作者引入了多尺度(MS)测试来评估性能变化。
结果。DefMamba-S模型在语义分割方面相较于多种SOTA方法表现出色,如表3所示。DefMamba-S在单尺度评估中达到48.8的mIOU,在多尺度评估中达到49.6。这超越了ResNet-50 [15]、Swin-T [24]和ConvNeXt-T [25]的性能。此外,DefMamba-S的性能超过了近期SSM方法,包括GrootV-T [38]、QuadMamba-S [39]和MSVMamba-T [31]。yinqing-DefMamba_2504在单尺度mIoU指标上分别提升了0.3分、1.6分和1.2分。

4.4. 消融研究
DSSM结构的影响。为证明所提出的可变形分支的有效性,作者在表4中针对不同的分支设置进行了一系列实验。如表4所示,FB-BB指的是用于特征提取的前向和后向分支,CB代表连续扫描[42]分支,LB是局部扫描[19]分支,DB表示作者在图2(a)中提出的可变形分支。通过比较表4中的结果(1)和(5),作者观察到,所提出的可变形分支在合理的计算预算内显著提高了ImageNet数据集的准确率1.7%,从而证明了所提出可变形分支的有效性。此外,与表4中的(3)、(4)和(5)相比,yinqing-DefMamba_2504仅增加了0.1G的计算成本,同时显著提高了ImageNet数据集的准确率1.4%。这进一步证明,与其它固定扫描方法相比,作者的可变形扫描方法更能够捕捉物体的结构信息并提升模型性能。从表4中的(1)和(2)中,作者注意到,仅加入所提出的可变形分支会导致性能下降,这是由于空间token跳跃的增加。为实现更稳定的训练和更高的模型性能,作者在所提出的模型中遵循了之前的范式,将FB-BB和DB结合起来。
可变形扫描组件的影响。为了全面评估作者提出的可变形扫描方法中每个组件的贡献,作者在表5中进行了消融研究。如表5所示,DP、DT、OB和CA分别代表图2(b)中可变形扫描的组件。DP(可变形点)涉及生成可变形点的过程,包括初始化参考点、偏移网络和双线性插值。DT(可变形 Token )指的是根据预测的偏移动态改变 Token 顺序的过程,包括初始化 Token 索引和可变形 Token 索引,以及偏移网络。OB表示生成偏移偏差的操作。CA表示偏移网络中的通道注意力操作。通过比较表5中的结果(1)、(2)和(3),作者观察到添加DP操作或DT操作均可将 Baseline (1)的性能提升0.2%-0.4%,且计算成本仅略微增加(0.1G),这证明了DP和DT的有效性。此外,当结合DP和DT操作(4)时,与 Baseline 相比,性能进一步提升了1%,如表5中的(1)和(4)所示。这些结果有力地证明,无论单独添加DT或DP,还是结合两者,模型都能更好地学习图像结构,促进对物体细节变化的感知,从而提升性能。此外,作者通过比较表5中的(4)和(5),以及(5)和(6)分别运行消融实验,以验证OB和CA的有效性。这些实验验证了yinqing-DefMamba_2504的有效性。

可视化结果。为了更好地展示作者可变形扫描策略的优势,作者在图4中展示了不同方法在各个位置的图像激活图,并用红色和橙色点在图像中清晰 Token 。具体来说,作者使用[23]中yinqing-DefMamba_2504可视化了第二阶段最终层的激活图。如图4最后一行笔的激活图所示,yinqing-DefMamba_2504在处理包含多个重叠物体的复杂场景时,表现出更强的聚焦于物体结构和形状信息的能力。这种能力使得识别和分割更加精确,进一步突显了yinqing-DefMamba_2504在捕捉关键细节方面的有效性。

作者还可视化了可变形点和可变形 Token 索引,以直观展示yinqing-DefMamba_2504的表现,如图5所示。在图5(a)的红色框中,作者可以观察到一些位于物体外的焦点点向物体移动,这些移动使网络能够关注更多物体信息。与图5(b)相比,yinqing-DefMamba_2504(c)调整了扫描顺序以强调重要 Token 。例如,如图5第一行所示,对应蛇头的 Token 从栅格扫描中的中间位置移动到yinqing-DefMamba_2504中的第一个位置。这种偏移有助于网络学习相关特征。

5. 局限性
尽管yinqing-DefMamba_2504取得了显著成果,但仍存在一定局限性。当图像包含不完整的物体结构或按照一定规则排列的多个物体时,可变形扫描策略可能效果不佳。如图6所示,当图像仅显示棒球的一部分时,可变形机制无法捕捉完整的结构信息,导致偏移量过小并收敛于预定义的扫描方法。同时,当多个物体按照一定规则排列时,相邻 Token 之间的信息变化极小,模型因此处于惰性学习状态。

6. 结论
在本工作中,作者致力于解决现有基于Mamba的方法所面临的挑战,这些方法依赖于固定的扫描技术来提取特征。这些方法未能充分利用图像中固有的空间结构信息。为了克服这一局限性,作者提出了一种基于深度扫描(DS)的新型基础模型,命名为DefMamba。该创新模型旨在增强网络学习和表示复杂图像结构的能力,以及检测物体细节中的微妙变化。
深度扫描主要包含两个关键操作:焦点点的移动和扫描顺序的调整。第一个操作有效地将参考点重新定位到重要的感兴趣区域,从而增强模型对物体细节变化的敏感性。第二个操作修改扫描顺序,以创建一个与输入特征所基于的底层物体结构更匹配的结构感知序列。在ImageNet、COCO和ADE20K等基准数据集上的大量实验评估有力地表明,yinqing-DefMamba_2504优于现有的扫描自注意力模型(SSMs),并且在与其他卷积神经网络(CNNs)和基于Transformer的方法相比时仍具有竞争力。
参考
[1]. DefMamba: Deformable Visual State Space Model
(文:极市干货)