YOLO-Mamba

论文：《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的作用。