MambaYOLO:用状态空间模型做目标检测,简单且高效
1. 思路与背景
YOLO 系列长期主导实时目标检测,Transformer 架构的引入进一步扩大了感受野,但自注意力的二次复杂度让计算开销偏高。状态空间模型(SSM)最近在视觉骨干上展现了线性复杂度的优势,却很少被直接用在检测任务里。MambaYOLO 试着把这两条线搭起来:用 Mamba 风格的 SSM 搭建检测器,同时不依赖大规模预训练,结构保持简洁。
相关工作脉络大致是这样:
- YOLO 一路迭代:从 DarkNet 到 CSPDarknet53 + C2f,再到 Gold-YOLO 用 Gather-and-Distribute 机制做信息融合,靠骨干网络和特征金字塔的改进不断推高精度。
- DETR 系列:把 Transformer 搬进检测,靠编码器-解码器绕过锚点和 NMS,但训练收敛慢,小目标吃力,通常需要大量预训练。
- 视觉 SSM:Vision Mamba、VMamba 等提出 Cross-Scan、窗口扫描等技巧,在图像分类上效果不错,可一直没人拿它做实时检测的基线。
MambaYOLO 就从这个缺口入手,希望拿出一个训练简单、推理快、精度能打的 SSM-YOLO 组合。
2. 怎么做的
2.1 状态空间模型回顾
SSM 通过隐状态 (h(t)) 把输入 (x(t)) 映射到输出 (y(t)),连续形式为:
$$h'(t) = \mathbf{A}h(t) + \mathbf{B}x(t) \tag{1}$$ $$y(t) = \mathbf{C}h(t) \tag{2}$$
Mamba 用零阶保持(ZOH)离散化后得到:
$$\overline{\mathbf{A}}=\operatorname{exp}(\mathbf{\Delta A}) \tag{3}$$ $$\overline{\mathbf{B}}=(\mathbf{\Delta A})^{-1}(\operatorname{exp}(\mathbf{\Delta A})-\mathbf{I})\mathbf{\Delta B} \tag{4}$$
于是可以写成线性递归:
$$h'(t)=\overline{\mathbf{A}}h_{t-1}+\overline{\mathbf{B}}x_{t} \tag{5}$$ $$y_{t}=\mathbf{C}h_{t} \tag{6}$$
或者变成更利于训练的卷积形式:
$$\overline{\mathbf{K}}=(\mathbf{C}\overline{\mathbf{B}},\mathbf{C}\overline{\mathbf{A}}\overline{\mathbf{B}},...,\mathbf{C}\overline{\mathbf{A}}^{L-1}\overline{\mathbf{B}}) \tag{7}$$ $$y=x*\overline{\mathbf{K}} \tag{8}$$
训练时用卷积并行计算,推理时切回递归模式自回归生成,这是 SSM 高效的关键。
2.2 整体结构
MambaYOLO 架构分成 ODMamba 主干和 PAFPN 脖子。跟多数检测器一样,先由 Stem 下采样,再接若干 ODSSBlock 和 VisionClueMerge 模块逐级压缩特征图。脖子部分把 C2f 换成 ODSSBlock,保留多尺度特征融合。
2.3 Simple Stem
主流 ViTs 喜欢用 4×4 不重叠分块,但这可能影响优化。这里直接用两个步长为 2 的 3×3 卷积,简单省事,性能和效率平衡得更好。
2.4 VisionClueMerge
传统下采样(比如步长为 2 的 3×3 卷积)容易打乱 SS2D 按方向选择特征的过程。VMamba 的做法是拿 1×1 卷积分割降维,再拼回去。MambaYOLO 则更进一步:去掉归一化,把维度映射拆开,往通道里塞一些冗余特征图,最后用 4 倍压缩的逐点卷积做下采样。这样 SS2D 在前一层选出的重要线索不会被提前冲淡。
2.5 ODSSBlock
这是核心模块,一个 ODSSBlock 分两部分:SS2D 负责全局空间关系,RGBlock 负责通道融合。输入先过一层 ConvModule + BN + SiLU:
$$Z^{l-2}=\sigma\left(BatchNorm\left(ConvModule(Z^{l-3})\right)\right) \tag{9}$$
然后像 Transformer 那样用残差连接和层归一化:
$$Z^{l-1}=SS2D\left(LayerNorm(Z^{l-2})\right)+Z^{l-2} \tag{10}$$ $$Z^{l}=RGBlock\left(LayerNorm(Z^{l-1})\right)+Z^{l-1} \tag{11}$$

