当前位置：首页 > article >正文

【论文阅读】：Mamba YOLO SSMs-Based YOLO For Object Detection

article 2024/11/15 18:35:08

摘要

Mamba架构已被证明可以有效地捕获长距离的地面依赖关系。

在深度学习技术的快速发展的推动下，YOLO系列为实时目标探测器设定了一个新的基准。在YOLO的基础上，不断探索再参数化、高效层聚合网络和无锚定技术的创新应用。为了进一步提高检测性能，引入了基于变压器的结构，显著扩展了模型的接受域，并取得了显著的性能提高。然而，因为自注意机制的二次复杂度增加了模型的计算负担。幸运的是，状态空间模型（SSM）作为一种创新技术的出现，有效地缓解了二次复杂性所引起的问题。基于这些进展，我们引入了一种基于SSM的Mamba-YOLO模型。

Mamba-YOLO不仅优化了SSM的基础，而且还特别适应于目标检测任务。考虑到SSM在序列建模中的潜在局限性，如接受域不足和图像局部性较弱，我们设计了LSBlock和RGBlock。这些模块能够更精确地捕获局部图像依赖性，并显著地增强了模型的鲁棒性。在公开的基准数据集COCO和VOC上的广泛实验结果表明，Mamba-YOLO在性能和竞争力方面都超过了现有的YOLO系列模型。

算法

前期研究和准备

传统的CNNs由于其局部感受野和层次结构设计，对图像中剧烈的尺度变化、复杂背景和多视角干扰的挑战具有一定的局限性，图像相关性较差。

因此，将Transformer引入了目标检测领域，如DETR系列，它依赖于强大的自注意全局建模能力来解决cnn的小接受域的问题。但近年来，更多的工作开始重新思考如何设计卷积神经网络来使模型更快，他们开始不满于Transformer结构的二次复杂度，开始使用混合结构来重构模型和降低复杂度，如MobileVit，EdgeVit，EfficientFormer。混合模型带来的问题就是性能的明显下降。因此，在性能和速度之间找到平衡一直是一个挑战。

YOLOv4中CSPNet被引入，YOLOv6中开始引入重新参数化，YOLOv7中使用ELAN重建模型，YOLOv8中采用解耦的头部和无锚点设计。YOLOv10 在其设计中集成了Transformer结构的元素，引入了部分自注意力( Partial Self-Attention，PSA )模块，旨在控制计算成本的同时增强模型的全局建模能力。 Gold Yolo引入了一种新的机制GD ( Gather-and-Distribute )，通过自注意力操作实现，解决了传统特征金字塔网络和Rep - PAN 的信息融合问题，并成功实现了SOTA。

Mamba 基于状态空间模型( SSMs )的方法，它具有对长距离依赖的强大建模能力和线性时间复杂度的优越性质。并且Mamba架构已经在图像分类中取得了成效。因此可以提出一个问题：能否将SSM结构引入目标检测领域，与当前的实时检测器相结合，利用SSM的优势，为YOLO系列带来新的性能提升?

介绍

该篇论文中，引入了ODSSBlock模块，将SSM结构应用于目标检测领域。与用于图像分类的VSSBlock不同，用于目标检测的图像具有更大的像素，并且由于SSM模型通常对文本序列进行建模，缺乏对图像的通道表达能力，我们提出了LSBlock对通道特征进行建模。得益于图像较大的像素和较多的通道维度，我们提出了RGBlock结构在SS2D输出后进一步解码，利用点乘的高维表达方式提高通道相关性。Mamba - YOLO是视觉识别和检测任务中的一个重要进展，旨在构建一种结合SSM和CNN优点的新型骨干网络。该架构将基于SSM的状态空间转换模型应用到YOLO的各层中，以有效捕获全局依赖关系，并利用局部卷积的强度来提高检测精度和模型对复杂场景的理解，同时保持实时性。这种混合架构有望突破现有视觉模型在处理大规模或高分辨率图像时的局限性，为下一代视觉基础模型提供强大而灵活的支持。我们在PASCAL VOC 和MS COCO 上进行了实验，结果表明Mamba - YOLO在一般的目标检测任务中非常有竞争力，在MSCOCO上取得了比基线YOLOv8高出8.1%的mAP。

算法

总体架构

Mamba YOLO架构概述如图2所示。我们的目标检测模型分为ODMamba主干部分和颈部部分。ODMamba由Simple Stem、减少采样Block组成。在颈部，我们遵循PAN - FPN的设计，使用ODSSBlock模块代替C2f ，以捕获更丰富的梯度信息流，其中Conv单独负责下采样。骨架首先通过Stem模块进行下采样，得到分辨率为H/4 W/4的2D特征图。随后，所有的模型都由ODSSBlock和一个用于进一步下采样的VisionClue Merge模块组成。

Mamba-YOLO目标检测模型由ODMamba骨干网络、PAFPN网络结构及Head检测头组成

ODMamba骨干网络：输入图片首先通过Simple Stem模块进行下采样，得到分辨率为H/4×W/4的2D特征图，随后经过串联的ODSSBlock和Vision Clue Merge模块，用于提取图像特征。

PANFPN网络结构：在Neck部分，采用了PAFPN的设计，并且用ODSSBlock代替原有结构的C2f模块，结构中还包含上采样的Upsample，单独负责下采样的Convolution（stride=2），该部分用于融合和细化不同尺度的特征，增强模型对于各种尺度物体的检测能力，尤其是在多尺度目标检测任务中非常有效。

Head检测头结构：检测头部结构分为3层，分别为80*80小尺度目标检测层、40*40中尺度目标检测层、20*20大尺度目标层。

ODSS Block

ODSSBlock是Mamba-YOLO的核心模块，该模块主要包含LSBlock、RGBlock、SS2D三个模块。在输入阶段经过一系列处理，可以使网络能够学习到更深入、更丰富的特征表示，同时通过批处理归一化保持训练推理过程的高效和稳定。ODSSBlock的批归一化、层归一化和残差连接设计允许模型在深层堆叠训练时有效流动。

LocalSpatial Block（LSBlock）

为了Mamba在解决处理复杂尺度变化的任务时，提取局部特征的挑战。提出LSBlock来增强对局部特征的捕捉。

对于给定的输入特征，它首先进行深度可分离卷积，在每个输入通道上单独操作。有效提取输入特征图的局部空间信息，同时减少计算成本和参数数量。然后进行Batch Normalization，提供一定程度的正则化效果的同时减少过拟合。

深度可分离卷积中通过1 × 1卷积混合通道信息，并通过激活函数更好地保留信息的分布，使模型能够学习到更复杂的特征表示，从而能够从输入的特征图中提取丰富的多尺度上下文信息。在LSBlock中，激活函数使用一个非线性的GeLU，在不改变空间维度的情况下，改变特征的通道数，从而增强特征表示。

最后，通过残差级联将原始输入与处理后的特征进行融合。使得模型能够理解和融合图像中不同维度的特征，从而增强对尺度变化的鲁棒性。

ResGated Block （RGBlock）

VMamba 架构中的MLP遵循Transformer设计，对输入序列进行非线性变换以增强模型的表达能力。最近的研究表明，Gated MLP在自然语言处理中表现出强大的性能，并且我们发现门控机制具有同样的视觉潜力。

在Xl - 2分支上使用深度可分离卷积作为位置编码模块，并通过残差级联的方式在训练过程中更有效地回流梯度，具有较低的计算成本，并通过保留和利用图像的空间结构信息显著地提高了性能。

RGBlock采用非线性GeLU作为激活函数，控制每一层信息的流动，然后通过元素乘法与Xl - 2的一个分支合并，再通过1x1卷积与全局特征进行细化，以融合通道信息，最后通过残差级联与原始输入Xl - 2和隐藏层中的特征相加。RGBlock可以捕获更多的全局特征，同时只带来计算成本的略微增加。

SS2D

State Space Two Dimensional（SS2D），是一种基于State Space Model的卷积操作。该操作通过对输入特征图进行扫描扩展，得到四个方向上的序列，然后将这些序列合并成全局特征表示。SS2D的主要作用是在不增加计算量的情况下，扩展模型的可学习范围，从而提高模型的性能。

当卷积神经网络( Convolutional Neural Networks，CNNs )和视觉变换( Vision Transformer，ViT )结构通常使用卷积进行下采样时，我们发现这种方法会干扰SS2D在不同信息流阶段的选择性操作。

为了解决这个问题，VMamba 对2D特征图进行分割，并使用1x1卷积进行降维。我们的研究结果表明，为状态空间模型( SSMs )保留更多的视觉线索有利于模型的训练。与传统的维度减半不同，我们通过：

去除范数来简化这个过程；
对维度图进行拆分；
将多余的特征图添加到通道维度上；
利用4倍压缩的1x1卷积进行下采样。与使用步长为2的3 × 3卷积不同，我们的方法保留了SS2D从上一层中选择的特征图。

实验

数据集：MS COCO和VOC数据集

COCO2017训练数据集上进行训练，并在COCO2017val数据集上进行验证。

实验结果表明，Mamba - YOLO相比现有的先进方法取得了显著的性能提升。

消融实验

消融实验中，每个提出的方法的有效性验证与VOC0712数据集，其中训练集包含约16551图像来自VOC2007和VOC2012训练集，和验证集由4952图像来自VOC2007测试集。采用从头训练策略进行训练，训练次数500轮。

Mamba YOLO的消融研究

这里对ODSS Block中的每个模块进行独立的检查，在没有Clue Merge的情况下，使用ViT的传统卷积方法进行下采样，以评估Vision Clue Merge对准确率的影响。

在VOC0712数据集上执行Mamba YOLO进行消融实验，测试模型为Mamba YOLO-T，结果如表2所示，线索合并为SSM保留了更多的视觉线索，也为ODSS Block是最优的论断提供了证据。

RGBlock结构的烧蚀研究

RGBlock通过逐像素获取全局依赖关系和全局特征来捕获逐像素的局部依赖关系。

RGBlock在多层感知机基础上的三种变体（图5展示）

1 )卷积MLP，在原MLP的基础上增加了DW - Conv；
2 )残差卷积MLP，即在原MLP的基础上加入残差级联的DW - Conv；
3 )门控MLP，一种在门控机制下设计的MLP变体。

表3展示了原始MLP、RG Block以及每个变体在VOC0712数据集中的性能，以验证我们对MLP的分析的有效性，测试模型为Mamba YOLO-T。

我们观察到单独引入卷积并没有带来性能的有效提升

在图5 中的Gated MLP中，其输出由两个元素相乘的线性投影组成，其中一个由残差连接的DWConv和门控激活函数组成，这实际上赋予了模型通过层次结构传播重要特征的能力，有效地提高了模型的准确性和鲁棒性。该实验表明，引入的卷积在处理复杂图像任务时性能的提升与门控聚合机制相当相关，前提是它们应用在残差连通性的背景下。

其他模型和实例分割的消融研究

采用MambaYOLO-T之上的v8-seg，并在COCOseg数据集上进行训练和测试，通过Bbox AP和Mask AP等指标评估模型性能。Mamba YOLO-T-seg在每个大小上都显著优于YOLOv5 和YOLOv8 的分割模型。RTMDet 基于一个包含深度卷积的基本构建块，在动态标签分配过程中引入软标签来计算匹配代价，在多个视觉任务中表现出优异的性能，并且Mamba YOLO-T-seg在Mask mAP上相比其Tiny仍然保持了2.3的优势。

总结

一种新的检测器模型Mamba-YOLO，将状态空间模型（SSM）引入到实时目标检测领域，并结合了当前最先进的YOLO系列算法的优点。具体来说，作者提出了ODSSBlock模块来应用SSM结构，同时设计了LSBlock和RGBlock两个新模块来增强模型的局部特征捕捉能力和通道表达能力。实验结果表明，Mamba-YOLO在PASCAL VOC和COCO数据集上表现出了比现有最先进算法更高的mAP值。此外，该论文还对该领域的其他相关研究进行了综述和比较，展示了Mamba-YOLO的独特性和优越性。未来，还可以运用到如自动驾驶、智能安防等领域中。