Mamba YOLO: SSMs-Based YOLO For Object Detection
- 总结
- 前言
- 感受野
- 为什么Transformer 的结构被引入,显著扩展了模型的感受野?
- 状态空间模型SSM
- 介绍
- 相关工作
- 实时目标检测
- 端到端目标检测器
- 视觉状态空间模型
- 方法
- 预处理
- 整体架构
- ODSS Block
- LocalSpatial Block
- ResGated Block
- 实验
- 与最先进技术的比较
- 消融研究
- Mamba YOLO 的消融研究
- RG Block结构的消融研究
- 其他模型和实例分割的消融研究
- 结论
总结
Mamba-YOLO的Mamba主要体现在SS2D这个模块,并且这个SS2D模块使用的是VMamba中提出来的SS2D模块。
由于Mamba以前主要用于处理文本等一维数据,为了将其应用于图像、视频等领域,从而创建了这个SS2D模块。
Mamba是一种选择性的状态空间模型/S6模型,用于处理一维数据。在SS2D中的体现就是S6模块。
SS2D通过在不同方向上采样,得到一维数据,然后输入到S6模块处理,最后进行合并。
关键就是S6 Block,S6 Block 是一种基于状态空间模型的架构,它能够捕捉序列中长距离依赖关系。
S6 Block 可以使用并行化技术对序列数据进行高效计算。
由于采用了状态空间表示和选择性扫描(Selective Scan)等优化技术,S6 Block 处理序列数据时在内存和计算资源上都相对节省。
(1)状态空间
状态空间包含完全描述一个系统所需的最少变量数。它是一种通过定义系统可能状态来数学表示问题的方法。
让我们简化一下。想象我们正在穿越迷宫。“状态空间”就是所有可能位置(状态)的地图。每个点代表迷宫中的一个独特位置,并带有特定的细节,比如你距离出口有多远。
“状态空间表示”是对这张地图的简化描述。它显示了你当前的位置(当前状态)、你可以前往的下一个位置(可能的未来状态)、以及你如何到达下一个位置(向右或向左移动)。
在我们的例子中,X 和 Y 坐标,以及到出口的距离,可以表示为“状态向量”。
(2)SSM(状态空间模型)
SSM(状态空间模型)是一类用于描述这些状态表示并预测其下一个状态的模型,预测基于某些输入。
SSM接受连续序列作为输入并预测输出序列。
SSM 假设动态系统(如在 3D 空间中移动的物体)可以通过时间 t 的状态和两个方程来预测。
通过求解这些方程,我们假设可以根据观测到的数据(输入序列和先前状态)揭示预测系统状态的统计原则。
其目标是找到这个状态表示 h(t),使我们能够从输入到输出序列。
由于我们通常是离散输入,所以需要对模型离散化。
可以使用零阶保持技术:首先,每当我们接收到离散信号时,我们保持其值,直到我们接收到新的离散信号为止。这个过程会生成一个连续信号,供 SSM 使用:
我们保持数值的时间由一个可学习的新参数 步长 ∆ 表示。
现在我们为输入生成了一个连续信号,接下来我们可以生成连续输出,并根据输入的时间步长对这些值进行采样。
这些采样的值就是我们离散化的输出!
(3)S6 Block注意力机制体现在哪里
S6 Block 的注意力机制并非像传统 Transformer 那样通过自注意力来实现,而是通过状态空间模型的设计,间接实现了对序列数据中重要特征的选择性关注
- 状态空间建模:S6 Block使用状态空间模型来代替直接的注意力计算。状态空间模型通过构建一个状态向量来表示序列的历史信息,能在序列的长依赖关系上保持信息
- 选择性扫描机制:选择性的对序列中不同部分进行扫描,而不是逐点处理整个序列。这样可以跳过不重要特征,仅在序列中信息密集的部分投入计算资源
- 通过状态向量关注长距离信息:S6 Block通过维护一个隐含的状态向量来表示序列中长程依赖关系。每个状态向量都可以视为一种聚合特征,通过持续的更新和叠加,能够关注到不同时间步的关联信息
前言
在深度学习技术快速发展的推动下,YOLO 系列在实时目标检测器领域设立了新的标杆。研究人员在 YOLO 的基础上不断探索创新的重参数化、有效的层聚合网络和无锚技术的应用。
为了进一步提升检测性能,基于 Transformer 的结构被引入,显著扩展了模型的感受野,并取得了显著的性能提升。然而,这些改进也带来了代价,因为自注意力机制的二次复杂性增加了模型的计算负担。幸运的是,状态空间模型 (SSM) 的出现有效地缓解了由二次复杂性引发的问题。
基于这些进展,我们提出了 Mamba-YOLO,这是一种基于 SSM 的新型目标检测模型。Mamba-YOLO 不仅优化了 SSM 基础,还专门针对目标检测任务进行了适配。考虑到 SSM 在序列建模中的潜在局限性,如感受野不足和图像局部性较弱,我们设计了 LSBlock 和 RGBlock。这些模块能够更精确地捕捉图像的局部依赖关系,并显著增强了模型的鲁棒性。
在 COCO 和 VOC 这些公开的基准数据集上的大量实验结果表明,Mamba-YOLO 在性能和竞争力方面均优于现有的 YOLO 系列模型,展现出其强大的潜力和竞争优势。
感受野
模型的感受野(Receptive Field)是指模型在处理输入数据时,能够感知到的输入区域的大小。对于图像处理任务来说,感受野代表了在输入图像中,一个神经元或特征所能“看到”或“感知”的像素范围。更大的感受野可以帮助模型捕捉全局信息,比如目标的大致位置或整个形状,而较小的感受野则更专注于局部细节。
为什么Transformer 的结构被引入,显著扩展了模型的感受野?
状态空间模型SSM
Mamba是一种状态空间模型(SSM)架构,它改进了S4架构。它有时也被称为S6或者selective SSM
介绍
近年来,深度学习迅速发展,尤其是在计算机视觉领域,一系列强大的结构取得了令人瞩目的性能表现。从CNN[1, 2, 3, 4, 5]和Transformer[6, 7, 8, 9, 10]到Mamba架构[32, 31, 33, 34],各种结构的应用展现了其在计算机视觉中的强大潜力。
在目标检测这一下游任务中,CNN[15, 16, 17, 20, 21]和Transformer结构[11, 13]被广泛使用。CNN及其一系列改进提供了快速的执行速度,同时确保了准确性。然而,由于CNN的图像关联性较弱,研究人员引入了Transformer到目标检测领域,例如DETR系列[11, 12, 14],通过自注意力的强大全局建模能力来解决CNN感受野较小的问题。随着硬件的进步,这种结构带来的内存计算增加不再是大问题。然而,近年来,越来越多的工作[5, 59, 60]开始重新思考如何设计CNN以提高模型的速度,同时越来越多的从业者对Transformer结构的二次复杂度感到不满,开始使用混合结构来重构模型并降低复杂度,例如MobileVit[61]、EdgeVit[62]、EfficientFormer[43]。然而,混合模型也带来了问题,性能的明显下降也是一个令人担忧的问题,因此如何在性能和速度之间找到平衡一直是研究人员关注的重点。
YOLO系列在目标检测领域一直是实时检测器的里程碑。自YOLOv4[25]起,CSPNet[63]被引入,YOLOv6[27]开始结合重参数化技术,YOLOv7[28]使用ELAN重构模型,而YOLOv8[42]采用了解耦头和无锚设计。最近提出的YOLOv10[48]将Transformer结构的元素集成到设计中,引入了部分自注意力(PSA)模块,旨在提升模型的全局建模能力的同时控制计算成本。这表明该系列始终充满活力。此外,CNN的结构为模型提供了强大的执行速度,许多从业者使用注意力机制来改进此模型,以在其领域实现预期的性能提升。
ViT-YOLO[38]引入了MHSA-Darknet到YOLO中,并结合了TTA和加权帧融合等增强训练策略。然而,参数数量和FLOPs的增加并未带来预期的性能提升,显示出Transformer在YOLO等目标检测任务中扩展性的局限性。 YOLOS[37]基于原始ViT架构采用极简改造方案,将ViT中的CLS标记替换为DET标记,并在集成预测方法中使用二分匹配损失。然而,其性能表现令人失望,对预训练方案非常敏感,YOLOS在不同预训练策略下表现出极大的波动性。**Gold-YOLO[29]通过卷积和注意力原语提取和融合特征信息,提出了一种增强多尺度特征融合的方法。**然而,这些方法虽然集成了Transformer结构,却放弃了其核心优势,即强大的全局注意力机制和长序列处理能力,同时试图通过减少二次复杂度带来的计算激增来提升性能,这往往限制了模型的表现。
最近,基于状态空间模型(SSM)的方法,如Mamba[32],由于其强大的远距离依赖建模能力和线性时间复杂度的优越特性,为解决这些问题提供了新的思路。令人兴奋的是,研究人员已成功地将Mamba架构引入视觉领域,并在图像分类中取得成功[31, 33]。受此启发,我们提出一个问题:是否可以将SSM结构引入目标检测领域,结合当前的实时检测器,以利用SSM的优势并为YOLO系列带来新的性能提升?
本文提出了一种名为Mamba-YOLO的检测器模型。我们引入了ODSSBlock模块(如图4所示),将SSM结构应用于目标检测领域。与用于图像分类的VSSBlock不同,目标检测的输入图像具有更大的像素尺寸,而由于SSM模型通常用于建模文本序列,缺乏图像的通道表达能力,我们提出了LSBlock来对通道特征进行建模。基于图像中更大的像素和更多的通道维度,我们提出了RGBlock结构,在SS2D输出后进一步解码,使用高维点乘表达来增强通道相关性。Mamba-YOLO是视觉识别和检测任务中的一项重要进展,旨在构建一种结合SSM和CNN优势的新型主干网络。该架构将基于SSM的状态空间转换模型应用于YOLO的各层,有效捕捉全局依赖性,并利用局部卷积的优势提升检测精度和模型对复杂场景的理解能力,同时保持实时性能。这样的混合架构有望突破现有视觉模型在处理大规模或高分辨率图像方面的限制,为下一代视觉基础模型提供强大且灵活的支持。我们在PASCAL VOC[35]和COCO[36]上进行了详尽的实验,结果显示,Mamba-YOLO在通用目标检测任务中非常具有竞争力,在MSCOCO上获得了比基线模型YOLOv8高8.1%的mAP。
本文的主要贡献如下:
- 我们提出了基于SSM的Mamba-YOLO,为YOLO系列在目标检测领域设立了新的基线,并为基于SSM的更高效和更有效的检测器的未来发展奠定了坚实基础。
- 我们提出了ODSSBlock模块,其中LSBlock有效提取了输入特征图的局部空间信息,以弥补SSM的局部建模能力不足。通过重新设计MLP层,我们结合门控聚合和有效卷积及残差连接的思想,提出了RGBlock,有效捕捉局部依赖性并增强模型的鲁棒性。
我们设计了一组具有不同规模的Mamba-YOLO模型(Tiny/Base/Large),以支持不同大小和规模的任务部署,并在COCO和VOC两个数据集上进行了实验(如图1所示),结果表明我们的Mamba-YOLO相比现有的最先进方法取得了显著的性能提升。
相关工作
实时目标检测
从 YOLOv1 到 YOLOv3 [22, 23, 24] 是 YOLO 系列模型的开创者,它们的性能提升都与主干网络的改进密切相关,使 DarkNet 得到了广泛应用。YOLOv4 [25] 引入了大量残差结构设计,并提出了 CSPDarknet53 主干网络,有效地减少了计算冗余,实现了高性能的特征表达和高效训练。YOLOv7 [28] 提出了 E-ELAN 结构,在不破坏原有结构的基础上增强了模型能力。YOLOv8 [42] 结合了之前 YOLO 版本的特点,采用了具有更丰富梯度流的 C2f 结构,在保证准确性的同时具备轻量化特性,适用于不同场景。最近,Gold YOLO [29] 引入了一种名为 GD(Gather-and-Distribute)的新机制,通过自注意力操作实现,用于解决传统特征金字塔网络 [52] 和 Rep-PAN [27] 的信息融合问题,并成功实现了 SOTA(当前最优)。实际上,传统的 CNNs 因为它们的局部感受野和层次结构设计,在应对图像中剧烈尺度变化、复杂背景和多视角干扰方面存在局限性。
端到端目标检测器
DETR [11] 首次将 Transformer 引入目标检测,使用 Transformer 编码-解码架构,绕过了传统的手工设计组件,如锚点生成和非极大值抑制,将检测视为一种直接的集合预测问题。Deformable DETR [12] 引入了 Deformable Attention,一种 Transformer 注意力的变体,通过在参考位置周围采样稀疏的关键点集,解决了 DETR 在处理高分辨率特征图时的局限性。DINO [13] 集成了混合查询选择策略、可变形注意力,并通过注入噪声的训练和查询优化来提升性能。RT-DETR [14] 提出了混合编码器,以解耦尺度内的交互和跨尺度融合,实现高效的多尺度特征处理。然而,DETRs 在训练收敛、计算成本和小目标检测方面仍存在挑战,而 YOLO 系列在小模型领域仍然是平衡准确性和速度的 SOTA(当前最优)。
视觉状态空间模型
最近,状态空间模型(State Space Models, SSM)成为研究热点。基于 SSM 的研究 [39, 40, 41],Mamba [32] 展示了输入规模上的线性复杂度,解决了 Transformer 在长序列建模中的计算效率问题。在通用视觉主干领域,Vision Mamba [33] 提出了基于 SSM 的纯视觉主干模型,首次将 Mamba 引入视觉领域。VMamba [31] 引入了 Cross-Scan 模块,使模型能够进行二维图像选择性扫描,增强了视觉处理能力,并在图像分类任务中展示了其优越性。LocalMamba [34] 专注于视空间模型的窗口扫描策略,优化视觉信息以捕获局部依赖关系,并引入动态扫描方法为不同层次寻找最佳选择。MambaOut [49] 探讨了 Mamba 架构在视觉任务中的必要性,指出 SSM 对图像分类任务非必需,但其在具有长序列特征的检测和分割任务中的价值值得进一步探索。在下游视觉任务中,Mamba 也广泛应用于医学图像分割 [53, 54, 55] 和遥感图像分割 [56, 57] 的研究中。受 VMamba [31] 在视觉任务领域取得的显著成果启发,本文首次提出 Mamba YOLO,这是一种新的 SSM 模型,旨在兼顾全局感知场,并展示其在目标检测任务中的潜力。
方法
预处理
整个序列变换也可以用卷积形式表示,其定义如下:
整体架构
【解释】:
这里的 PAN-FPN 指的是一种特定的网络结构设计,即 Path Aggregation Network - Feature Pyramid Network。这种结构在目标检测任务中广泛应用,尤其在 YOLO 等网络中,用于增强多尺度特征的融合和信息传递。
PAN-FPN设计的核心思想:
- FPN(Feature Pyramid Network):主要用于捕获不同尺度的特征。通过多层次的特征提取,FPN可以更好地识别不同大小的物体。传统FPN通过自底向上的特征金字塔(如从高分辨率到低分辨率)以及自顶向下的路径来融合特征,确保模型能够获取到丰富的语义信息。
- PAN(Path Aggregation Network):在FPN的基础上,PAN添加了一个自底向上的路径,这样不仅从上层向下融合语义特征,还能够在低层次特征中加入更多的语义信息。简单来说,PAN通过自底向上的路径进一步增强了低层特征的表达能力,有助于提高模型的检测效果,特别是在小目标检测方面效果显著。
图3:SS2D操作示意图。SS2D中的扫描扩展操作分为四个分支,在图像的四个不同方向路径上逐块扫描图像并获得四个序列。SS2D中的扫描合并操作将获得的序列 作为 S6 块的输入,并合并来自不同方向的序列,以便将特征提取为全局特征。
SS2D的结构见下图4
ODSS Block
如图 4 所示,ODSS Block 是 Mamba YOLO 的核心模块,它在输入阶段经历一系列处理,以使网络能够学习更深层次和更丰富的特征表示,同时通过批量归一化保持训练推理过程的高效性和稳定性。
LocalSpatial Block
ResGated Block
实验
在本节中,我们对 Mamba YOLO 进行了全面的实验,测试其在目标检测任务及一些下游视觉任务中的表现。我们使用 MS COCO【36】数据集来验证所提出的 Mamba YOLO 的优越性。与当前先进方法的比较中,我们在 COCO2017train 数据集上进行训练,并在 COCO2017val 数据集上进行验证。在消融实验中,我们采用 VOC0712【35】数据集来验证每个所提方法的有效性,其中训练集包含约 16,551 张来自 VOC2007 和 VOC2012 训练集的图像,验证集包含 4,952 张来自 VOC2007 测试集的图像。所有提到的模型均采用从头训练策略,总训练轮数为 500。更多设置详见附录。我们所有的模型均在 8 块 NVIDIA H800 GPU 上进行训练。
与最先进技术的比较
图 1 和表 1 展示了 MS-COCO2017val 数据集上与其他最先进目标检测器的比较结果,本文提出的方法在 FLOPs、参数量和准确性之间实现了最佳综合权衡。具体而言,Mamba YOLO-T 相比于性能最佳的轻量级模型 DAMO YOLO-T/YOLO MS-XS,在 AP 上分别有 3.4% 和 2.0% 的显著提升,且相比于具有相似准确率的基线 YOLOv8-S,参数量减少了 45.5%,FLOPs 减少了 50%。将 Mamba YOLO-B 与具有相似参数和 FLOPs 的 Gold-YOLO-M 相比,前者的 AP 增益比后者高 4.5%。即使与具有相同准确率的 Gold-YOLO-M 比较,Mamba YOLO-B 的参数量减少了 47.2%,FLOPs 减少了 43.2%。在大型模型中,Mamba YOLO-L 与各先进目标检测器相比也取得了更好或相似的性能。与表现最佳的 Gold-YOLO-L 相比,Mamba YOLO-L 的 AP 提高了 0.3%,参数量减少了 0.9%。上述比较结果表明,所提出的模型在不同规模的 Mamba YOLO 上相比现有最先进方法具有显著改进。
消融研究
Mamba YOLO 的消融研究
在本节中,我们对 ODSS Block 中的每个模块进行了独立测试,在不使用 Clue Merge 的情况下,采用 Vision Transformer 的传统卷积方法进行下采样,以评估 Vision Clue Merge 对准确性的影响。我们在 VOC0712 数据集上对 Mamba YOLO-T 进行了消融实验。表 2 的结果显示,线索融合为状态空间模型(SSM)保留了更多的视觉线索,并证明了 ODSS Block 结构的确是最优的。
RG Block结构的消融研究
RGBlock 通过逐像素捕捉局部依赖关系,将全局依赖关系和全局特征整合到逐像素的处理过程中。关于 RG Block 的设计细节,我们在多层感知器(MLP)基础上还考虑了三种变体:1) 卷积 MLP,在原始 MLP 中添加 DW-Conv;2) 残差卷积 MLP,以残差连接方式添加 DW-Conv 到原始 MLP;3) 门控 MLP,这是一种基于门控机制设计的 MLP 变体。图 5 展示了这些变体,表 3 则展示了原始 MLP、RG Block 及每个变体在 VOC0712 数据集上的表现,以验证我们对 MLP 分析的有效性,测试模型为 Mamba YOLO-T。我们观察到,单独引入卷积并未显著提升性能,而在变体图 5(d) 的门控 MLP 中,其输出由两个线性投影的元素乘法构成,其中一个包含残差连接的 DWConv 和门控激活函数,这实际上赋予了模型通过分层结构传播重要特征的能力,有效提高了模型的准确性和鲁棒性。此实验表明,在处理复杂图像任务时,引入卷积对性能的提升与门控聚合机制密切相关,前提是在残差连接的上下文中应用它们。
其他模型和实例分割的消融研究
为了评估我们提出的基于 SSM 的 Mamba YOLO 架构的优越性和良好的可扩展性,我们将其应用于实例分割任务,除了目标检测领域之外。在 Mamba YOLO-T 的基础上添加 v8 分割头【42】,并在 COCOseg 数据集上进行训练和测试,通过 Bbox AP 和 Mask AP 等指标评估模型性能。Mamba YOLO-T-seg 在每种尺寸上显著优于 YOLOv5【46】和 YOLOv8【42】的分割模型。基于包含大核深度卷积的基础构建块的 RTMDet【58】在动态标签分配过程中引入软标签以计算匹配成本,在多个视觉任务中表现出色,而 Mamba YOLO-T-seg 仍比其 Tiny 型在 Mask mAP 上保持 2.3 的优势。结果如表 4 和图 8 所示。
结论
在本文中,我们重新分析了 CNN 和 Transformer 架构在目标检测领域的优缺点,并指出它们融合的局限性。基于此,我们提出了一种基于状态空间模型架构并扩展自 YOLO 的检测器,重新分析了传统 MLP 的局限性,并提出了 RG Block,其门控机制和深度卷积残差连接旨在赋予模型在分层结构中传播重要特征的能力。此外,为了解决 Mamba 架构在捕捉局部依赖关系方面的局限性,LSBlock 增强了捕捉局部特征的能力,并将其与原始输入融合,从而增强特征的表示能力,显著提高了模型的检测能力。我们的目标是建立 YOLO 的新基准,证明 Mamba YOLO 具有很强的竞争力。我们的工作是对 Mamba 架构在实时目标检测任务中的首次探索,同时也希望为该领域的研究者带来新的思路。