Mamba YOLO World

论文地址：https://arxiv.org/pdf/2409.08513v1

代码地址： GitHub - Xuan-World/Mamba-YOLO-World: Mamba-YOLO-World: Marrying YOLO-World with Mamba for Open-Vocabulary Detection

开集检测（OVD）旨在检测预定义类别之外的物体。作为将YOLO系列集成到OVD的SOTA模型，YOLO-World非常适合注重速度和效率的场景。然而，其性能受到其Neck特征融合机制的限制，导致二次复杂度和受限制的引导感受野。为了应对这些限制，作者提出了Mamba-YOLO-World，这是一种基于YOLO的OVD模型，采用MambaFusion Path Aggregation Network（MambaFusion-PAN）作为其Neck架构。具体而言，作者引入了一种创新的状态空间模型特征融合机制，包括并行引导选择扫描(Parallel-Guided Selective Scan)算法和串行引导选择扫描(Serial-Guided Selective Scan)算法，具有线性复杂度和全局引导感受野。它利用多模态输入序列和mamba隐藏状态来指导选择扫描过程。实验表明，Mamba-YOLO-World在COCO和LVIS基准测试中均优于原始YOLO-World，同时保持可比的参数和FLOPs。此外，它还超越了现有最先进的OVD方法，在更少的参数和FLOPs下。

INTRODUCTION

目标检测，作为计算机视觉的一个基本任务，在自动驾驶车辆、个人电子设备、医疗健康和安全等领域中起着至关重要的作用。传统方法在目标检测方面取得了巨大的进步。然而，这些模型都是在闭合的数据集上进行训练的，限制了它们在预定义类别（例如，COCO数据集中的80个类别）之外进行检测的能力。为了克服这些限制，开集检测（OVD）作为一种新任务应运而生，要求模型在预定义类别之外检测物体。

一些早期的OVD工作试图利用预训练的视觉语言模型（VLMs）的固有图像-文本对齐能力。然而，这些VLMs主要在图像-文本 Level 进行训练，因此在区域-文本 Level 的对齐能力不足。最近的工作，如MDETR，GLIP，DetClip，Grounding DINO，mm-Grounding-DINO 和YOLO-World重新定义OVD为视觉语言预训练任务，使用传统的目标检测器直接在大规模数据集上学习区域-文本 Level 的开放式词汇对齐能力。

根据上述相关工作，将传统目标检测器转变为 OVD 模型的关键在于实现一个适应模型现有 Neck 结构的视觉-语言特征融合机制，例如 YOLO-World 中的 VL-PAN 和 Grounding-DINO 中的 Feature-Enhancer。作为将 YOLO 系列集成到 OVD 的先驱模型，YOLO-World 非常适合在注重速度和效率的场景中部署。尽管如此，其性能受到其 VL-PAN 特征融合机制的制约。

具体来说，VL-PAN在文本到图像特征融合流中使用了最大Sigmoid视觉通道注意力机制，在图像到文本融合流中使用了多头交叉注意力机制，这导致了几个局限性。

首先，由于交叉模态注意力机制，两种融合流的复杂性会随图像大小和文本长度的乘积呈平方增长。

其次，VL-PAN缺乏全局引导的感受野。一方面，文本到图像融合流仅生成一个视觉通道权重向量，在像素级缺乏空间指导。另一方面，图像到文本融合流仅允许图像信息逐个指导每个单词，无法利用文本描述中的上下文信息。

为了克服上述局限性，作者提出了一个基于YOLO的OVD模型Mamba-YOLO-World，其Neck结构采用了提出的MambaFusion Path Aggregation Network（MambaFusion-PAN）。最近，Mamba作为一种新兴的态空间模型（SSM），已经展示了避免二次复杂度和捕捉全局感受野的能力。然而，在Mamba中，将多个模态的特征简单拼接在一起会导致复杂度为 $O(N+M)$ ，与拼接序列的长度成正比。这对OVD中的大型词汇特别有影响。因此，提出了一种基于MambaFusion-PAN的MambaFusion特征融合机制。作者使用Mamba隐藏状态作为不同模态之间特征融合的中介，这需要 $O(N+1)$ 的复杂度，并提供全局引导的感受野。图1中的可视化结果表明，作者的Mamba-YOLO-World在所有尺寸变体方面的准确性和泛化性都显著优于YOLO-World。

贡献可以概括为以下几点：

提出了Mamba-YOLO-World，这是一个基于YOLO的OVD模型，其Neck架构采用了MambaFusion-PAN。
提出了一种基于状态空间模型的特征融合机制，包括并行引导选择扫描（Parallel-Guided Selective Scan）算法和串行引导选择扫描（Serial-Guided Selective Scan）算法，具有 $O(N+1)$ 的复杂度和全局引导感受野。
实验证明，Mamba-YOLO-World在保持可比参数和FLOPs的同时，优于原始YOLO-World。此外，它在更少的参数和FLOPs下超越了现有的最先进OVD方法。

METHOD

Mamba-YOLO-World 主要是基于 YOLOv8 开发的，包括一个 Darknet Backbone 和一个 CLIP 文本编码器作为模型的 Backbone ，MambaFusion-PAN 作为模型的Neck，以及一个文本对比分类头和边界框回归头作为模型的 Head ，如图2 所示。

Mamba Preliminaries

对于一个连续输入信号 $u(t) \in \mathbb{R}$ ，通过隐藏状态 $h(t) \in \mathbb{R}^E$ ，SSM将它映射到连续输出信号 $y(t) \in \mathbb{R}$ 。

其中E是SSM的状态扩展因子， $A \in \mathbb{R}^{E\times E}$ 是状态转移矩阵， $B \in \mathbb{R}^{E\times 1}$ 和 $C \in \mathbb{R}^{1\times E}$ 对应输入和输出的映射矩阵。基于SSM，Mamba引入了选择扫描算法，使得A、B和C成为输入序列的函数。

MambaFusion-PAN

MambaFusion-PAN是作者提出的特征融合网络，用于替代YOLO中的Path Aggregation Feature Pyramid Network。如图2(a)所示，MambaFusion-PAN利用提出的基于SSM的并行和串行特征融合机制，通过视觉和语言分支之间的三个阶段特征融合流程，同时聚合多尺度图像特征和增强文本特征：文本到图像，图像到文本，最后文本到图像。具体组件如下部分的详细说明。

Mamba Hidden State

目前，基于Transformer和Mamba的VLMs仅将多模态特征ConCat ，随着文本序列长度和图像分辨率的增长，不可避免地增加了复杂性。尽管YOLO-World中的VL-PAN采用单向融合且无需特征 ConCat ，但它仍然导致 $O(N^{2})$ 的复杂性。这是因为文本到图像融合流中的视觉通道注意力机制和图像到文本融合流中的多头交叉注意力机制。

为了应对这些问题，作者提出通过mamba隐状态 $h(t) \in \mathbb{R}^{D\times E}$ 来提取压缩序列信息，作为不同模态之间特征融合的中间媒介，其中D是输入序列的维度，E是SSM状态扩展因子。由于D和E都是常数，不受序列长度的影响，因此作者的特征融合机制复杂度为 $O(N+1)$ ，其中N来自一个模态的输入序列，1来自另一个模态的mamba隐状态。

TextMambaBlock

TextMambaBlock由堆叠的 Mamba 层组成。给定由 CLIP 文本编码器生成的文本嵌入 $w_{0} \in \mathbb{R}^{L_{t} \times D_{t}}$ ，使用图 2(b)所示的 TextMambaBlock 不仅提取输出文本特征 $w_{1} \in \mathbb{R}^{L_{t} \times D_{t}}$ ，还提取文本隐藏状态 $THS \in \mathbb{R}^{D_{t} \times E_{t}}$ ，这将用于后续的文本到图像特征融合。

MF-CSPLayer

如图2(c)所示，作者通过MambaFusion CSPLayer（MF-CSPLayer）将THS与多尺度图像特征相结合。MF-CSPLayer将提出的并行引导选择扫描算法集成到YOLO CSPLayer风格的网络中。经过MF-CSPLayer处理后，作者可以不仅得到输出图像特征，还可以得到图像隐藏状态 $IHS \in \mathbb{R}^{D_{t} \times E_{t}}$ ，该状态将用于后续图像到文本特征融合。

Parallel-Guided Selective Scan

SSM选择扫描算法根据输入序列动态调整内部参数。受此启发，作者创新性地提出了并行引导选择扫描（PGSS）算法，该算法在扫描过程中，根据输入图像序列和THS的值动态调整Mamba内部参数（A，B和C）的值，如图2（d）和算法1所示。

因此，压缩文本信息与整个视觉选择扫描过程并行注入到Mamba中，使多尺度图像特征在像素 Level 而不是通道 Level 得到引导。其生成的输出传递给MF-CSPLayer的后续层。以下作者称这部分为文本到图像特征融合流程。

Serial-Guided Selective Scan

Mamba Selective Scan算法根据输入序列持续压缩信息到 $h(t)$ 。受此启发，作者提出了串行引导选择扫描（SGSS）算法，并将其与TextMambaBlock结合，如图2（e）所示。SGSS旨在将前序序列中的先验知识压缩到 $h(t)$ ，并将其作为后序序列的指导。具体来说，SGSS-TextMambaBlock在提取文本特征之前，通过扫描压缩的视觉信息IHS调整Mamba内部参数（A，B和C）的值serial。接下来，作者将这部分称为图像到文本特征融合流程。

EXPERIMENT

实施细节

Mamba-YOLO-World是基于MMYOLO工具箱和MMDetection工具箱开发的。作者提供了三种尺寸变体，即小型（S）、中型（M）和大型（L）。实验包括预训练阶段和微调阶段。在预训练阶段，作者采用包括Objects365（V1），GQA和Flickr30k的检测和定位数据集。与其他OVD方法一样，在排除COCO中的图像后，将GQA和Flickr30k数据集合并为GoldG数据集。在微调阶段，作者使用预训练的Mamba-YOLO-World并在下游任务数据集上进行微调。除非另有说明，作者将实验按照YOLO-World的设置进行。

零样本结果

在预训练之后，作者在零样本的方式下，直接在LVIS和COCO基准上评估所提出的Mamba-YOLO-World，并与YOLO-World和其他现有的最先进方法进行全面比较。

在 LVIS 上的零样本评估

LVIS数据集包含了1203个长尾物体类别。遵循之前的工作，作者使用固定AP指标，并在LVIS最小验证集上报告每个图像的1000个预测，进行公平比较。

根据表1，Mamba-YOLO-World在小变体上实现了+1.5%的AP提高，在中变体上实现了+1.8%的AP提高，同时保持可比参数和FLOPs。此外，它在与YOLO-World相比，通过+0.4%~+3.2%的APr和+1.4%~+3.4%的APc在所有尺寸变体上都优于YOLO-World。此外，Mamba-YOLO-World-L在较少的参数和FLOPs下，与先前的最先进方法相比，获得了更好的结果。

在 COCO 上零样本评估

COCO 数据集包含 80 个类别，是最常用的目标检测数据集。如表2所示，Mamba-YOLO-World 显示出整体优势，在所有尺寸变体中，其平均精度（AP）比 YOLO-World 高 +0.4%～+1%。

微调结果

在表3中，作者对COCO基准上的微调结果进行了进一步评估。在COCO train2017上进行微调后，Mamba-YOLO-World在所有大小变体上的AP准确性均高于微调后的YOLO-World，优势在+0.2%至+0.8%之间。

消融研究

在表4中，作者对基于Mamba-YOLO-World-S的MambaFusion-PAN文本到图像和图像到文本特征融合流程进行了消融实验，以分析其对性能的影响。在COCO基准测试上的零样本评估结果表明，作者的并行（文本到图像）和串行（图像到文本）特征融合方法在提高性能的同时，并未增加参数或FLOPs。

此外，作者分析了随着输入图像分辨率的增加，计算成本的变化。如图3所示，MambaFusion-PAN（Mamba-YOLO-World的Neck）在所有尺寸变体中，与VL-PAN（YOLO-World的Neck）相比，消耗了高达15%的FLOPs，表明MambaFusion-PAN的模型复杂度较低。