本文来源公众号“江大白”,仅用于学术分享,侵权删,干货满满。
原文链接:DefMamba新型视觉模型,多尺度Backbone与Mamba,创新多任务视觉算法!(附论文及源码)
导读
现有视觉 Mamba 多采用固定扫描顺序,限制了图像空间结构的建模。为此,作者提出 DefMamba,结合多尺度 Backbone 和可变形 Mamba 模块,动态调整扫描路径以聚焦关键区域。通过可变形扫描策略,DefMamba 有效提升了图像结构理解能力,并在多项视觉任务中取得领先性能。
论文链接:https://arxiv.org/pdf/2504.05794
代码链接:https://github.com/leiyeliu/DefMamba(即将开源)
近期,状态空间模型(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. 预备知识
3.2. 整体模型架构
与语言模型中使用的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)所示。
4. 实验
4.1. 图像分类
结果。表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. 目标检测
4.3. 语义分割
设置。为了评估DefMamba的语义分割性能,作者使用在ADE20K [45]上预训练的分类权重初始化的UperNet [37]训练yinqing-DefMamba_2504,迭代次数为160,000。作者采用了AdamW优化器 [26],学习率设置为 6x 10^-5。
作者的实验主要使用默认输入分辨率512x512进行。此外,作者引入了多尺度(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
THE END !
文章结束,感谢阅读。您的点赞,收藏,评论是我继续更新的动力。大家有推荐的公众号可以评论区留言,共同学习,一起进步。