作者：浙江师范大学
论文地址：https://arxiv.org/pdf/2406.05835
代码地址：https://github.com/HZAI-ZJNU/Mamba-YOLO

---

前言

计算机视觉的发展，从CNN开始火起来，到Transformer再将之推向高潮，最近火热的Mamba架构也开始突起爆发。如果将最近火热的Mamba模型与YOLO检测框架结合，能擦出什么火花呢？这篇文章将问你解读。

---

不同模型学习到的注意力图可视化

YOLOv5[46]、YOLOv6[27]、YOLOv7[28]、YOLOv8[42]和曼巴YOLO在主干中的随机初始权可视化结果输出的特征图。

图1:MSCOCO数据集上实时目标探测器的对比。基于ssm的目标检测方法在flops精度(左)方面具有显著优势。

一、摘要

在深度学习技术快速发展的推动下，YOLO系列为实时目标探测器树立了新的标杆。研究人员在YOLO的基础上不断探索重新参数化、高效层聚集网络和无锚技术的创新应用。为了进一步提高检测性能，引入了基于变压器的结构，大大扩展了模型的接受域，并取得了显着的性能提升。然而，这种改进是有代价的，因为自注意机制的二次复杂度增加了模型的计算负担。幸运的是，状态空间模型(SSM)作为一种创新技术的出现，有效地缓解了二次复杂度带来的问题。在此基础上，作者提出了一种新的基于SSM的目标检测模型Mamba-YOLO。Mamba-YOLO不仅优化了SSM基础，而且还专门适用于目标检测任务。考虑到SSM在序列建模中可能存在的局限性，如接收野不足和图像局域性弱，作者设计了LSBlock和RGBlock。这些模块能够更精确地捕获局部图像依赖关系，并显著增强模型的鲁棒性。在公开可用的基准数据集COCO和VOC上进行的大量实验结果表明，Mamba-YOLO在性能和竞争力方面都超过了现有的YOLO系列模型，展示了其巨大的潜力和竞争优势。PyTorch代码可在:https://github.com/HZAI-ZJNU/Mamba-YOLO获得

二、介绍

近年来，深度学习发展迅速，特别是在计算机视觉领域，一系列强大的结构已经取得了令人印象深刻的表现。从cnn[1, 2, 3, 4, 5]和Transformers[6, 7, 8, 9, 10]到Mamba架构[32,31,33,34]，各种结构的应用已经展示了它们在计算机视觉中的强大潜力。

在目标检测的下游任务中，主要使用CNN[15,16,17,20,21]和Transformer结构[11,13]。cnn及其一系列改进在确保准确性的同时提供了快速的执行速度。然而，由于图像相关性较差，研究人员将Transformers引入到目标检测领域，如DETR系列[11,12,14]，依靠强大的自关注全局建模能力来解决cnn接收域小的问题。幸运的是，随着硬件的进步，这种结构带来的内存计算的增加不会造成太大的问题。但是，近年来，更多的工作[5,59,60]开始重新思考如何设计cnn以使模型更快，更多的从业者开始对Transformer结构的二次复杂度感到不满，开始使用混合结构来重建模型并降低复杂性，如MobileVit [61]， EdgeVit [62]， effentformer[43]。然而，混合动力车型也带来了问题，性能的明显下降也是一个令人担忧的问题，因此在性能和速度之间寻找平衡一直是研究人员关注的问题。

YOLO系列一直是目标检测领域具有里程碑意义的实时探测器。从YOLOv4[25]开始，引入了CSPNet [63]， YOLOv6[27]开始纳入重新参数化，YOLOv7[28]使用ELAN重建模型，YOLOv8[42]采用了解耦的头部和无锚设计。最近提出的YOLOv10[48]将变压器结构的元素集成到其设计中，引入了部分自关注(PSA)模块，旨在增强模型的全局建模能力，同时控制计算成本。这证明了这个系列一直具有强大的生命力。此外，cnn的结构赋予了该模型强大的执行速度，许多从业者已经使用注意力机制来改进该模型，以达到他们所在领域期望的性能改进。

viti -YOLO[38]在YOLO中引入了MHSA-Darknet，以及增强的训练策略，如TTA和加权帧融合技术。然而，参数和FLOPs数量的增加并没有带来预期的性能提升，这表明Transformer在对象检测任务中的可扩展性存在局限性，特别是在YOLO中。YOLOS[37]采用了基于原有ViT架构的极简改造方案，将ViT中的CLS标记替换为DET标记，并在集成预测方法中采用二分类匹配损失。然而，它的性能令人失望，并且对预训练方案非常敏感，在不同的预训练策略下，YOLOS表现出很大的可变性。Gold-YOLO[29]提出了一种通过卷积基元和注意基元提取和融合特征信息来增强多尺度特征融合的方法。然而，这些方法在整合Transformer结构的同时，放弃了其强大的全局关注机制和长序列处理能力的核心优势，而寻求减少二次复杂度降低带来的计算量，这往往限制了模型的性能。

最近，基于状态空间模型(state space models, ssm)的方法，如Mamba[32]，由于其强大的远程依赖关系建模能力和优越的线性时间复杂度特性，为解决这些问题提供了新的思路。令人兴奋的是，研究人员已经成功地将曼巴结构引入视觉领域，并在图像分类方面取得了成功[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，其中LS Block有效地提取输入特征映射的局部空间信息，以补偿SSM的局部建模能力。通过对MLP层设计的重新思考，作者将门控聚合与有效卷积与残差连通性相结合，提出了RG块，有效地捕获了局部依赖关系，增强了模型的鲁棒性。
•设计了一组不同尺度的Mamba-YOLO (Tiny/Base/Large)模型来支持不同规模和尺度的任务部署，并在COCO和VOC两个数据集上进行了实验，如图1所示，结果表明作者的Mamba-YOLO与现有的最先进的方法相比，性能有了显著的提高。

二、相关工作

2.1 实时目标检测器

YOLOv1到YOLOv3[22,23,24]是yoloo系列模型的先行者，它们的性能提升都与骨干改进密切相关，使得暗网得到广泛应用。YOLOv4[25]引入了大量残差结构设计提出的CSPDarknet53骨干网，有效降低了计算冗余，实现了高性能的特征表达和高效的训练。YOLOv7[28]提出了E-ELAN结构，在不破坏原有模型的情况下增强模型能力。Yolov8[42]结合了前几代yolo8的特点，采用了具有更丰富梯度流的C2f结构，在兼顾精度的同时，轻量级、可适应不同场景。最近，Gold Yolo[29]提出了一种新的机制GD (gatherand - distribute)，通过自关注运算来解决传统特征金字塔网络[52]和Rep-PAN[27]的信息融合问题，并成功实现了SOTA。事实上，传统的cnn由于其局部接受域和分层结构设计，在面对图像尺度变化剧烈、背景复杂和多视点干扰等挑战时存在一定的局限性。

2.2 端到端的目标检测器

DETR[11]首次将Transformer引入对象检测，使用变压器编码器-解码器架构，绕过传统的手工制作组件，如锚点生成和非最大抑制，将检测视为直接的集成预测问题。Deformable DETR[12]引入了Deformable Attention，这是Transformer Attention的一种变体，用于对参考位置周围的稀疏关键点集进行采样，解决了DETR在处理高分辨率特征地图方面的局限性。DINO[13]集成了混合查询选择策略、可变形注意力和注入噪声的演示训练，并通过查询优化提高性能。RT-DETR[14]提出了一种混合编码器来解耦尺度内相互作用和跨尺度融合，以实现高效的多尺度特征处理。然而，deta在训练收敛性、计算成本和小目标检测方面存在挑战，而yolo系列在小模型领域仍然是精度和速度平衡的SOTA。

2.3 视觉状态空间模型

状态空间模型是近年来研究的热点。Mamba[32]在研究SSM[39,40,41]的基础上，显示了输入大小的线性复杂性，解决了Transformer在建模状态空间长序列上的计算效率问题。在广义视觉骨干领域，Vision Mamba[33]提出了一种基于SSM的纯视觉骨干模型，这是曼巴首次被引入视觉领域。vammba[31]引入了Cross-Scan模块，使模型能够对二维图像进行选择性扫描，增强了视觉处理能力，在图像分类任务上表现出优越性。LocalMamba[34]侧重于视觉空间模型的窗口扫描策略，优化视觉信息以捕获局部依赖关系，并引入动态扫描方法以搜索不同层的最优选择。MambaOut[49]探讨了Mamba架构在视觉任务中的必要性，它指出SSM对于图像分类任务不是必需的，但它对于遵循长序列特征的检测和分割任务的价值值得进一步探讨。在下游视觉任务中，曼巴也被广泛应用于医学图像分割[53,54,55]和遥感图像分割[56,57]的研究中。受VMamba[31]在视觉任务领域取得的显著成果的启发，本文首次提出了一种新的SSMs模型mamba YOLO，该模型旨在考虑全局感觉领域，同时展示其在目标检测任务中的潜力。

三 方法

3.1 基础知识Mamba

基于状态空间模型(SSM)的结构化状态空间序列模型S4[39]和Mamba[32]都源于一个连续系统，该系统通过隐式潜在中间状态h(t)∈rn将单变量序列x(t)∈R映射到输出序列y(t)。这种设计不仅连接了输入和输出之间的关系，而且还封装了时间动态。系统可以用数学方法定义为:
h ’ (t) =Ah(t) + Bx(t) (1)
y(t) =Ch(t)(2)
式(1)中，A∈R N×N表示状态转移矩阵，表示隐藏状态随时间的演变，B∈R N×1表示输入空间相对于隐藏状态的权值矩阵。其中，C∈R N×1是观测矩阵，它将隐藏的中间状态映射到输出。

有关Mamba讲解，可以参考其他资料。

3.2 总体框架

图2展示了Mamba YOLO体系结构的概述。作者的目标检测模型分为ODMamba脊柱和颈部部分。ODMamba由Simple Stem, Downsample Block组成。在颈部，遵循PAN-FPN的设计，使用ODSSBlock模块而不是C2f[42]来捕获更富梯度的信息流。主干首先通过Stem模块进行下采样，得到分辨率为h4, w4的二维特征图。

随后，所有模型都由ODSSBlock和VisionClue Merge模块组成，用于进一步下采样。在颈部部分， 采用了PAFPN的设计[42]，使用ODSSBlock代替C2f，其中Conv单独负责下采样。

简单干现代视觉变压器(ViTs)通常采用分割补丁作为其初始模块，将图像划分为不重叠的片段。这个分割过程是通过核大小为4、步长为4的卷积操作来实现的。然而，最近的研究，如来自EfficientFormerV2[43]的研究表明，这种方法可能会限制vit的优化能力，影响整体性能。为了在性能和效率之间取得平衡， 提出了一个流线型的阀杆层。 没有使用不重叠的补丁，而是使用两个卷积，步幅为2，内核大小为3。

虽然卷积神经网络(cnn)和视觉变压器(ViT)结构通常使用卷积进行下采样，但 发现这种方法会干扰SS2D在不同信息流阶段的选择性操作[31]。为了解决这个问题，vammba[31]分割了2D特征图，并使用1x1卷积降维。 研究结果表明，为状态空间模型(SSMs)保留更多的视觉线索有利于模型训练。与传统的尺寸减半相比， 通过以下方式简化了这个过程:1)去除规范;2)分割维度图;3)将多余的特征映射附加到通道维度;4)利用4倍压缩的点向卷积进行下采样。与使用步幅为2的3x3卷积不同，作者的方法保留了SS2D从上一层选择的特征图。

图三。SS2D操作说明。SS2D中的扫描扩展操作分为4个分支，分别在图像的4个不同方向路径上逐块扫描图像，得到4个序列。SS2D中的扫描合并操作将获得的序列作为S6块的输入，将不同方向的序列进行合并，从而将特征提取为全局特征。

3.3 ODSS块

如图4所示，ODSS Block是Mamba YOLO的核心模块，在输入阶段经过一系列处理，使网络能够学习到更深入、更丰富的特征表示，同时通过批处理归一化保持训练推理过程的高效和稳定。

扫描扩展、S6块和扫描合并是SS2D[31]算法的三个主要步骤，其主要流程如图3所示。扫描展开操作将输入图像展开为一系列子图像，每个子图像表示一个特定的方向，当从对角线视点观察时，扫描展开操作沿着四个对称方向进行，分别是自顶向下、自底向上、从左到右和从右到左。这种布局不仅全面覆盖了输入图像的所有区域，而且通过系统的方向变换，为后续的特征提取提供了丰富的多维信息库，提高了图像特征多维捕获的效率和全面性。然后对这些子图像进行S6块[32]操作的特征提取，最后通过扫描合并操作将这些子图像合并在一起，形成与输入图像大小相同的输出图像。

图4:(a) SS2D详细结构图(b) ODSSBlock架构图©残差门通块图(RGBlock)(d)局部空间块图(LSBlock)

3.3.1 LocalSpatial Block

曼巴架构已被证明在捕获远程地面依赖性方面是有效的。然而，在处理复杂尺度变化的任务时，局部特征的提取面临一定的挑战。在图4©中，本文提出了LocalSpatial Block来增强局部特征的捕获。具体来说，对于给定的输入特征F l−1∈R C×H×W，首先对其进行深度可分卷积，该卷积分别对每个输入通道进行操作，而不混合通道信息。有效提取输入特征图的局部空间信息，同时减少计算成本和参数数量，然后进行批处理归一化，在减少过拟合的同时提供一定程度的正则化效果;

四 实验

在本节中，作者对曼巴YOLO进行了对象检测任务和一些下游视觉任务的综合实验。作者使用MS COCO[36]数据集来验证所提出的Mamba YOLO比较与最先进的方法的优越性， 在COCO2017train数据集上进行训练，并在COCO2017val数据集上进行验证。在消融实验中，使用VOC0712[35]数据集验证了每种方法的有效性，其中训练集包含来自VOC2007和VOC2012训练集的约16,551幅图像，验证集由来自VOC2007测试集的4,952幅图像组成。作者提到的所有模型都是使用从头开始训练的策略进行训练的，总训练次数为500次。更多设置请参考附录。作者所有的模型都是在8个NVIDIA H800 gpu上训练的。

4.1 SOTA对比

图1和表1显示了MS-COCO2017val与其他最先进的目标探测器相比的结果，其中本文提出的方法在FLOPs和Params与精度之间具有最佳的综合权衡。具体来说，是曼巴YOLO-T与性能最佳的微型轻量级型号DAMO yoloo - t /YOLO MS-XS相比，AP显著增加3.4%/2.0%，与具有相同精度的基准YOLOv8-S相比，Params降低45.5%，FLOPs降低50%。Mamba YOLO-B与Params和FLOPs相似的Gold-YOLO-M相比，前者的AP增益比后者高4.5%。即使与具有相同精度的Gold-YOLO-M相比，Params也降低了47.2%，FLOPs降低了43.2%。在大型模型中，曼巴YOLO-L也实现了比每个先进的目标探测器更好或相似的性能。与表现最好的Gold-YOLO-L相比，Mamba YOLO-L的AP提高了0.3%，而Params降低了0.9%。上述对比结果表明，与现有的最先进的方法相比，作者提出的模型在不同尺度上对曼巴YOLO有显著的改进。

五 结论

在本文中，作者重新分析了CNN和Transformer架构在目标检测领域的优缺点，并指出了它们融合的局限性。在此基础上，提出了一种基于状态空间模型架构的检测器，并通过YOLO进行了扩展;作者重新分析了传统MLP的局限性，提出了RG块，设计了其门控机制和深度卷积残差连通性，使模型具有在分层结构中传播重要特征的能力。此外，为了解决Mamba体系结构在捕获本地依赖关系方面的局限性，LSBlock增强了捕获本地特征的能力，并将其与原始输入融合，增强了特征的表示，从而显著提高了模型的检测能力。作者的目标是建立一个新的YOLO基线，证明曼巴YOLO具有很强的竞争力。作者的工作是曼巴架构在实时目标检测任务中的首次探索，也希望为该领域的研究人员带来新的思路。