论文:《Mamba YOLO: SSMs-Based YOLO For Object Detection》
地址:https://arxiv.longhoe.net/abs/2405.01828
一、论文概述
基于CNN和Transformer目前发展的局限性,CNN的局部特征提取局限性以及Transformer的二次计算复杂度。然而,目前很多学者对于两者的结合工作,例如:MobileViT、EfficientFormer以及EdgeViT等等模型,在性能与速度上的平衡也是一个很大的问题。
当下的问题就是利用SSM的优点,既要做到保留全局信息捕捉的能力以弥补CNN的劣势,还要做到具备快速的推理速度以满足检测目标的实时性。本文开发了一个基于SSM结构应用于YOLO框架(具体是YOLOv8)的目标检测模型。
二、模型结构
以YOLOv8为框架,在Backbone中使用新的下采样模块和特征提取模块替代旧模型;在Neck结构中保持使用Conv作为下采样模块同时使用新的特征提取模块。最后,保留原始检测头结构。整体结构如下所示:
2.1 Simple Stem(下采样模块)
[在ViT中提出通过卷积操作将原始图像划分为一个个独立的patch,然而EfficientFormerV2表明,这种方式限制了ViT的优化能力,影响了模型整体性能。]因此,本文提出一种简化的下采样层—Stem层,通过使用两次卷积操作而不使用单独的补丁,对输入特征图进行下采样。具体结构
2.2 ODSSBlock(特征提取模块)
2.2.1 LocalSpatial Block(局部空间特征提取模块)
Mamba架构已被证明在捕获远程地面依赖方面是有效的。然而,在处理复杂尺度变化的任务时,局部特征的提取面临一定的挑战。本文提出了LocalSpatial Block来增强局部特征的捕获。LSBlock能够有效地捕获和表示输入特征图的局部空间信息(通过DWconv和MLP结合),并将其与原始输入融合(残差连接),增强特征的表示,为Mamba YOLO提供了强大的处理视觉数据中尺度变化和上下文信息的能力。
2.2.2 ResGated Block(残差结构的门控MLP模块)
通过残差拼接的方式对梯度进行更有效的回流,计算成本更低,并且通过保留和利用图像的空间结构信息,显著提高了性能。与传统的MLP相比,RG Block将全局依赖关系和全局特征转移到每个像素上(元素乘法),捕捉相邻特征的依赖关系,使上下文信息更加丰富,进一步增强了模型的表达能力。
2.2.3 最终结构
其中Zl−3和Zl分别表示输入和输出特征,LS和RG分别表示LocalSpatial Block和ResGated Block, Zl−1表示2D-Selective-Scan后的中间状态[31]。扫描扩展、S6块和扫描合并是SS2D[31]算法的三个主要步骤。
(扫描展开操作将输入图像展开为一系列子图像,每个子图像表示一个特定的方向,当从对角线视点观察时,扫描展开操作沿着四个对称方向进行,分别是自顶向下、自底向上、从左到右和从右到左。这种布局不仅全面覆盖了输入图像的所有区域,而且通过系统的方向变换,为后续的特征提取提供了丰富的多维信息库,提高了图像特征多维捕获的效率和全面性。然后对这些子图像进行S6块[32]操作的特征提取,最后通过扫描合并操作将这些子图像合并在一起,形成与输入图像大小相同的输出图像)
三、实验部分
3.1 对比实验
对比实验选择的模型为YOLOv5、YOLOv6、YOLOv7、YOLOv8、DAMO-YOLO、GLOD-YOLO、YOLO MS模型。
在小模型的检测精度上,Mamba-YOLO达到了SOTA水平。但是在大参数模型下,性能不如YOLOv7与YOLOv8。
3.2 消融实验
作者做了两个消融实验,一个是从整体上验证所提出的ODSSBlock模块的作用,另一个是验证所提出的门控MLPRGBlock的作用。