在深度学习和计算机视觉领域,YoloV9以其卓越的性能和高效的检测能力赢得了广泛认可。为了进一步提升YoloV9的性能,我们创新性地引入了GroupMambaLayer作为其RepNCSP模块的核心改进。这一策略不仅显著增强了模型的性能,还优化了参数效率和计算资源的使用,为YoloV9的应用带来了全新的活力。
一、GroupMambaLayer的核心优势
GroupMambaLayer,作为参数高效且准确的组视觉状态空间模型的关键组件,其核心优势在于其独特的调制组曼巴设计。该设计通过将输入通道分为四组,并独立地在每个方向上进行视觉单选择扫描(VSSS),有效提升了空间依赖性的建模能力。同时,通道亲和力调制(CAM)算子的引入,进一步增强了跨通道的信息交换,解决了分组操作中可能存在的交互限制问题。
如下图:
图中这部分,就是GroupMambaLayer的操作,将输入分成四个相等的份,然后,每份都经过VSSS Block,这个四个VSSS Block代表不同的方向,分别是:
二、YoloV9 RepNCSP模块的革新
在YoloV9的架构中,RepNCSP模块扮演着至关重要的角色,其性能直接影响到整个模型的检测精度和效率。传统的RepNCSP模块使用RepNBottleneck结构,虽然在一定程度上保证了性能,但在面对复杂场景和更高精度的需求时,仍显不足。
我们创新性地将GroupMambaLayer引入RepNCSP模块,替换原有的RepNBottleneck结构。这一改进不仅保留了原有模块的优势,还通过GroupMambaLayer的高效参数利用和强大建模能力,实现了性能的大幅提升。具体来说,GroupMambaLayer的分组扫描和跨通道调制机制,使得RepNCSP模块在处理图像特征时,能够更加精准地捕捉空间信息和跨通道信息,从而提高了检测的准确性和鲁棒性。
编译环境
需要使用Ubuntu,我用的系统是Ubuntu 22.04,也可以使用WSL,使用方式:
https://blog.csdn.net/hhhhhhhhhhwwwwwwwwww/article/details/136147987
安装g++和gcc
sudo apt update
sudo apt install build-essential
下载官方链接,链接:
https://github.com/Amshaker/GroupMamba
解压后,进入kernels/selective_scan/文件夹下面,然后安装,执行命令:
python setup.py install
编译时间比较长,等待一段时间就行了。
论文翻译:《GroupMamba:参数高效且准确的组视觉状态空间模型》
https://arxiv.org/html/2407.13772v1
状态空间模型(SSM)的最新进展展示了在具有次二次复杂性的长距离依赖建模中的有效性能。然而,纯基于SSM的模型在稳定性和在计算机视觉任务中实现最佳性能方面仍面临挑战。本文解决了将基于SSM的模型扩展到计算机视觉领域的挑战,特别是大型模型尺寸的不稳定性和低效性。为此,我们引入了一个调制组曼巴层(Modulated Group Mamba layer),该层将输入通道分为四组,并独立地对每组应用我们提出的基于SSM的高效视觉单选择扫描(VSSS)块,每个VSSS块在四个空间方向之一中进行扫描。调制组曼巴层还将四个VSSS块封装成一个通道调制算子,以改善跨通道通信。此外,我们引入了一个基于蒸馏的训练目标,以稳定大型模型的训练,从而获得一致的性能提升。我们的综合实验证明了所提出贡献的优点,在ImageNet-1K的图像分类、MS-COCO的目标检测和实例分割以及ADE2OK的语义分割方面,相比现有方法取得了更优的性能。我们包含2300万个参数的微型变体在ImageNet-1K上实现了83.3%的分类Top-1准确率,达到了最先进的性能,同时与相同模型尺寸的最佳现有曼巴设计相比,在参数方面提高了26%的效率。我们的代码和模型可在以下网址获取:https://github.com/Amshaker/GroupMamba。
1、引言
在语言和视觉理解领域,出现了各种上下文建模方法。这些方法包括卷积[22, 66]、注意力[59.],以及最近出现的状态空间模型[17, 16]。具有多头自注意力机制的Transformer[59.]在语言模型(如GPT-3[2])和视觉模型(如Vision Transformers[10,36])中都处于核心地位。然而,由于注意力机制的计算复杂度为二次方,特别是对于较长的序列,这导致了一些挑战,进而促使了状态空间模型(如 S 4 \mathrm{S}_{4} S4[17.])的出现。
虽然 S 4 \mathrm{S}_{4} S4[17.]在处理较长的输入序列时由于其线性复杂度而有效,但在信息密集数据的全局上下文处理方面遇到了限制,特别是在计算机视觉等领域,因为该模型具有数据无关的特性。作为替代方案,提出了诸如基于全局卷积的状态空间模型[14]和Liquid S4[20]等方法来缓解上述限制。最近的Mamba[16]引入了S6架构,旨在提高状态空间模型处理长距离依赖的能力。Mamba引入的选择性扫描算法使用了依赖于输入的状态空间参数,这使得它在上下文学习方面表现更好,同时与自注意力相比,计算效率仍然较高。
然而,Mamba,特别是其S6算法,在例如图像分类等任务中,特别是在扩展到较大规模时,存在不稳定的问题[45]。此外,在图像分类中使用的Mamba模型变体,通常称为VSS(Visual State Space)块,在参数和计算需求方面可以根据通道数量进行更高效的设计。VSS块包括广泛的输入和输出投影以及深度卷积,其参数和计算复杂度与输入中的通道数直接成正比。为了解决这个问题,我们提出了一种调制组曼巴层(Modulated Group Mamba layer),它以计算和参数高效的方式缓解了上述问题。本文的主要贡献如下:
-
我们受组卷积的启发,引入了调制组曼巴层,该层通过使用多方向扫描方法来实现全面的空间覆盖和局部及全局信息的有效建模,从而提高了状态空间模型的计算效率和交互性。
-
我们引入了通道亲和力调制(Channel Affinity Modulation,CAM)算子,该算子通过增强跨通道通信来改善特征聚合,解决了分组操作中固有的有限交互问题。
-
我们采用基于蒸馏的训练目标来进一步稳定具有大量参数的模型的训练。
-
我们基于提出的调制组曼巴层构建了一系列参数高效的通用分类模型,称为“GroupMamba”。我们的微型变体在ImageNet-1k[9]上实现了83.3%的Top-1准确率,仅包含2300万个参数和4.6G浮点运算次数(FLOPs)。此外,我们的基础变体在包含5700万个参数和14G FLOPs的情况下,实现了84.5%的Top-1准确率,优于所有最近的状态空间模型方法(见图1)。
2、相关工作
自AlexNet[29]问世以来,卷积神经网络(ConvNets)一直是计算机视觉任务中的热门选择。该领域迅速发展,涌现出了多种具有里程碑意义的ConvNet架构[51, 55, 22, 25, 56]。随着这些架构的进步,人们还投入了大量精力来优化单个卷积层,包括深度卷积[65]、组卷积[7]和可变形卷积[8]。最近,ConvNeXt变体[37, 63]通过结合宏观设计与高级设置和训练策略,朝着现代化传统2D ConvNets迈出了实质性步伐,实现了与最先进模型相当的性能。
近年来,开创性的视觉转换器(Vision Transformer,ViT)[10]对计算机视觉领域产生了重大影响,包括图像分类[57, 36, 35, 12]、目标检测[3, 71, 43, 68]和分割[5, 50, 27]等任务。ViT[10]引入了一种整体设计,将图像视为一系列扁平化的二维块(patch),而不带有特定于图像的归纳偏置。ViT在计算机视觉任务中的卓越表现及其可扩展性,激发了后续众多旨在设计更好架构的努力。早期的基于ViT的模型通常需要大规模数据集(如JFT-300M[54])进行预训练。随后,DeiT[57]除了将蒸馏标记(distillation token)集成到架构中外,还提出了先进的训练技术,使得在较小数据集(如ImageNet-1K[9])上的有效训练成为可能。自此以后,后续研究通过结合卷积神经网络(CNN)和ViT模块,设计了分层和混合架构,以提高不同视觉任务的性能[53, 40, 11, 49, 12]。另一项工作是缓解自注意力固有的二次复杂度,这是ViT的主要瓶颈。这一努力带来了显著的改进,并产生了更高效和近似的变体[61, 49, 44, 42, 28, 6, 58],在降低复杂度的同时保持了有效性。
最近,状态空间模型(State Space Models,SSMs)已成为ViT(Vision Transformer)的替代方案[59],用于捕捉语言序列中复杂的动态和相互依赖性[17]。该领域的一个显著方法是结构化状态空间序列模型(S4)[17],它旨在解决长距离依赖问题,同时保持线性复杂度。沿着这一方向,已经提出了几种模型,包括S5[52]、
H
3
\mathrm{H}_{3}
H3[13]和GSS[41]。最近,Mamba[16]引入了一种依赖于输入的SSM层,并利用并行选择性扫描机制(S6)。
在视觉领域,各种工作已将SSM应用于不同的任务。特别是在图像分类方面,VMamba[34]在层次化Swin-Transformer[36]风格的设计中,使用Mamba在两个空间维度上进行双向扫描,以有效地构建全局感受野。同时,Vision Mamba(Vim)[70]则提出了一种整体设计,对整个图像进行单次双向扫描,性能优于传统的视觉转换器,如DeiT。LocalVMamba[26]通过在不同窗口内引入扫描方法(受SwinTransformer[36]启发),并结合网络层之间的动态扫描方向,解决了捕捉详细局部信息的挑战。EfficientVMamba[46]集成了基于空洞的选择性扫描和双路径模块,以实现高效的全局和局部特征提取,在减少计算复杂度的同时取得了具有竞争力的结果。这些模型已被应用于图像分类以及图像分割[33,39,48,15]、视频理解[67,30,4]和其他各种任务[19,23,60,18,31]。它们的广泛应用表明了SSM(状态空间模型)[17, 52, 13, 41],特别是Mamba[16]在视觉领域的有效性。在本文中,我们提出了一种调制组Mamba层,该层缓解了默认视觉Mamba块的缺点,如稳定性不足[45]以及相对于通道数量的参数增加。
3、方法
动机:我们的方法基于对现有视觉状态空间模型局限性的观察。
-
大型模型稳定性不足。我们从[45]中观察到,当扩展到大量参数时,基于Mamba[16]和MLP通道混合器的图像分类模型会变得不稳定。这种不稳定性可以在SiMBA-L(MLP)[45]中看到,其分类结果次优,准确率为49%。我们通过引入调制组Mamba设计以及蒸馏目标(如第3.4节所述)来解决这个问题,从而在不修改通道混合器的情况下稳定Mamba SSM训练。
-
高效/改进的交互:考虑到基于Mamba的设计对通道数量的计算影响,所提出的调制组Mamba层在计算上更为经济,且参数效率高于默认Mamba,并能够通过多方向扫描对输入标记的局部和全局信息进行建模。此外,本文还提出了一种额外的通道亲和力调制算子,以补偿分组操作导致的有限通道交互。
3.1、预备知识
状态空间模型:状态空间模型(SSMs),如S4 [17] 和Mamba [16],是受循环神经网络(RNNs)和卷积神经网络(CNNs)组合启发的结构化序列架构,其序列长度具有线性或接近线性的缩放特性。SSMs来源于连续系统,通过隐藏状态 h ( t ) ∈ R N h(t) \in \mathbb{R}^{N} h(t)∈RN 定义了一个输入 x ( t ) ∈ R L → y ( t ) ∈ R L x(t) \in \mathbb{R}^{L} \rightarrow y(t) \in \mathbb{R}^{L} x(t)∈RL→y(t)∈RL 的一维函数到函数的映射。更正式地说,SSMs由方程1中的连续时间常微分方程(ODE)描述。
h ′ ( t ) = A h ( t ) + B x ( t ) , y ( t ) = C h ( t ) , \begin{array}{l} h^{\prime}(t)=\mathbf{A} h(t)+\mathbf{B} x(t), \\ y(t)=\mathbf{C} h(t), \end{array} h′(t)=Ah(t)+Bx(t),y(t)=Ch(t),
其中
h
(
t
)
h(t)
h(t) 是当前隐藏状态,
h
′
(
t
)
h^{\prime}(t)
h′(t) 是更新后的隐藏状态,
x
(
t
)
x(t)
x(t) 是当前输入,
y
(
t
)
y(t)
y(t) 是输出,
A
∈
R
N
×
N
\mathbf{A} \in \mathbb{R}^{N \times N}
A∈RN×N 是SSM的演化矩阵,而
B
∈
R
N
×
1
\mathbf{B} \in \mathbb{R}^{N \times 1}
B∈RN×1 和
C
∈
R
N
×
1
\mathbf{C} \in \mathbb{R}^{N \times 1}
C∈RN×1 分别是输入和输出的投影矩阵。
离散状态空间模型:为了使这些模型能够在深度学习的序列建模任务中使用,需要将它们离散化,即将SSM从连续时间函数到函数的映射转换为离散时间序列到序列的映射。S4 [17] 和 Mamba [16] 是连续系统的离散化适应版本之一,它们引入了一个时间尺度参数
Δ
\boldsymbol{\Delta}
Δ 来将连续参数
A
,
B
\mathbf{A}, \mathbf{B}
A,B 转换为它们的离散等价物
A
,
B
\mathbf{A}, \mathbf{B}
A,B。这种离散化通常通过方程2中给出的零阶保持(ZeroOrder Hold,ZOH)方法来完成。
A = exp ( Δ A ) , B = ( Δ A ) − 1 ( exp ( Δ A ) − I ) ⋅ Δ B h t = A h t − 1 + B x t , y t = C h t . \begin{array}{l} \mathbf{A}=\exp (\Delta \mathbf{A}), \\ \mathbf{B}=(\Delta \mathbf{A})^{-1}(\exp (\Delta \mathbf{A})-\mathbf{I}) \cdot \Delta \mathbf{B} \\ h_{t}=\mathbf{A} h_{t-1}+\mathbf{B} x_{t}, \\ y_{t}=\mathbf{C} h_{t} . \end{array} A=exp(ΔA),B=(ΔA)−1(exp(ΔA)−I)⋅ΔBht=Aht−1+Bxt,yt=Cht.
虽然S4 [17] 和 Mamba [16] 都使用了上述方程2中所述的类似离散化步骤,但Mamba通过S6选择性扫描机制将参数 Δ ∈ R B × L × D \Delta \in \mathbb{R}^{B \times L \times D} Δ∈RB×L×D、 B ∈ R B × L × N \mathbf{B} \in \mathbb{R}^{B \times L \times N} B∈RB×L×N 和 C ∈ R B × L × N \mathbf{C} \in \mathbb{R}^{B \times L \times N} C∈RB×L×N 根据输入 x ∈ R B × L × D x \in \mathbb{R}^{B \times L \times D} x∈RB×L×D 进行条件化,从而与S4区分开来。其中, B B B 是批量大小, L L L 是序列长度, D D D 是特征维度。
3.2、总体架构
如图2(a)所示,我们的模型采用了一种分层架构,类似于Swin Transformer [36],包含四个阶段,以在不同分辨率下高效地处理图像。假设输入图像为
I
∈
R
H
×
W
×
3
\mathbf{I} \in \mathbb{R}^{H \times W \times 3}
I∈RH×W×3,我们首先应用一个Patch Embedding层将图像划分为大小为
4
×
4
4 \times 4
4×4 的不重叠块,并将每个块嵌入到一个
C
1
C_{1}
C1 维的特征向量中。Patch Embedding层通过两个步长为2的
3
×
3
3 \times 3
3×3 卷积实现。这在第一阶段产生了大小为
H
4
×
W
4
×
C
1
\frac{H}{4} \times \frac{W}{4} \times C_{1}
4H×4W×C1 的特征图。这些特征图被传递给一组我们的Modulated Grouped Mamba块(详见第3.3节)。在每个后续阶段,一个下采样层将
2
×
2
2 \times 2
2×2 区域内的块合并,随后是另一组Modulated Grouped Mamba块。因此,第二、三、四阶段的特征大小分别为
H
8
×
W
8
×
C
2
\frac{H}{8} \times \frac{W}{8} \times C_{2}
8H×8W×C2、
H
16
×
W
16
×
C
3
\frac{H}{16} \times \frac{W}{16} \times C_{3}
16H×16W×C3 和
H
32
×
W
32
×
C
4
\frac{H}{32} \times \frac{W}{32} \times C_{4}
32H×32W×C4。
3.3、调制组Mamba层
我们提出了针对输入序列
X
in
\mathbf{X}_{\text {in }}
Xin 的调制组Mamba层的总体操作(如图2(b)所示),其维度为
(
B
,
H
,
W
,
C
)
(B, H, W, C)
(B,H,W,C),其中
B
B
B 是批量大小,
C
C
C 是输入通道数,
H
/
W
H/W
H/W 是特征图的宽度和高度,具体见方程3。
x
G
M
=
GroupedMamba
(
x
i
n
,
Θ
)
x
CAM
=
CAM
(
x
G
M
,
Affinity
(
x
i
n
)
)
x
out
=
x
in
+
FFN
(
LN
(
x
CAM
)
)
\begin{aligned} \mathbf{x}_{\mathrm{GM}} & =\operatorname{GroupedMamba}\left(\mathbf{x}_{\mathrm{in}}, \Theta\right) \\ \mathbf{x}_{\text {CAM }} & =\operatorname{CAM}\left(\mathbf{x}_{\mathrm{GM}}, \operatorname{Affinity}\left(\mathbf{x}_{\mathrm{in}}\right)\right) \\ \mathbf{x}_{\text {out }} & =\mathbf{x}_{\text {in }}+\operatorname{FFN}\left(\operatorname{LN}\left(\mathbf{x}_{\text {CAM }}\right)\right) \end{aligned}
xGMxCAM xout =GroupedMamba(xin,Θ)=CAM(xGM,Affinity(xin))=xin +FFN(LN(xCAM ))
其中,
X
G
M
\mathbf{X}_{\mathrm{GM}}
XGM 是方程6的输出,
x
C
A
M
\mathbf{x}_{\mathrm{CAM}}
xCAM 是方程9的输出,
L
N
\mathrm{LN}
LN 是层归一化[1]操作,FFN是方程5中描述的前馈网络,而
X
out
\mathbf{X}_{\text {out }}
Xout 是调制组Mamba块的最终输出。各个操作,即GroupedMamba操作符、GroupedMamba操作符内部使用的VSSS块以及CAM操作符,分别在3.3.1节、3.3.2节和3.3.3节中介绍。
3.3.1、视觉单选择扫描(VSSS)块
VSSS块(图2(c))是一个基于Mamba操作符的标记和通道混合器。在数学上,对于输入标记序列
Z
in
\mathbf{Z}_{\text {in }}
Zin ,VSSS块执行如方程4所述的操作。
z
out
′
=
z
in
+
Mamba
(
LN
(
z
in
)
)
z
out
=
z
out
′
+
FFN
(
LN
(
z
out
′
)
)
\begin{array}{l} \mathbf{z}_{\text {out }}^{\prime}=\mathbf{z}_{\text {in }}+\operatorname{Mamba}\left(\operatorname{LN}\left(\mathbf{z}_{\text {in }}\right)\right) \\ \mathbf{z}_{\text {out }}=\mathbf{z}_{\text {out }}^{\prime}+\operatorname{FFN}\left(\operatorname{LN}\left(\mathbf{z}_{\text {out }}^{\prime}\right)\right) \end{array}
zout ′=zin +Mamba(LN(zin ))zout =zout ′+FFN(LN(zout ′))
其中,
Z
out
\mathbf{Z}_{\text {out }}
Zout 是输出序列,Mamba是方程2中描述的Mamba SSM操作符的离散化版本。
FFN
(
L
N
(
z
out
′
)
)
=
GELU
(
L
N
(
Z
out
′
)
W
1
+
b
1
)
W
2
+
b
2
\operatorname{FFN}\left(L N\left(\mathbf{z}_{\text {out }}^{\prime}\right)\right)=\operatorname{GELU}\left(L N\left(\mathbf{Z}_{\text {out }}^{\prime}\right) \mathbf{W}_{1}+\mathbf{b}_{1}\right) \mathbf{W}_{2}+\mathbf{b}_{2}
FFN(LN(zout ′))=GELU(LN(Zout ′)W1+b1)W2+b2
其中,GELU [24] 是激活函数,而
W
1
,
W
2
,
b
1
\mathbf{W}_{1}, \mathbf{W}_{2}, \mathbf{b}_{1}
W1,W2,b1 和
b
2
\mathbf{b}_{2}
b2 是线性投影的权重和偏置。
3.3.2、分组Mamba操作符
考虑到第3节中之前提出的动机,我们旨在设计一种Mamba[16]的变体,该变体既计算高效又能有效建模输入序列的空间依赖性。鉴于在输入序列的大量通道 C C C上,Mamba的计算效率不高,我们受分组卷积的启发,提出了该操作符的分组变体。分组Mamba操作是第3.3.1节中介绍的VSSS块的一种变体,其中输入通道被分为若干组,并且VSSS操作符被分别应用于每组。具体来说,我们将输入通道分为四组,每组大小为 C 4 \frac{C}{4} 4C,并对每组应用一个独立的VSSS块。为了更好地建模输入中的空间依赖性,这四个组中的每一个都沿着标记序列的四个方向之一进行扫描:从左到右、从右到左、从上到下和从下到上,如图2(d)所示。
设 G = 4 G=4 G=4为表示四个扫描方向(从左到右、从右到左、从上到下和从下到上)的组数。我们从输入序列 X in \mathbf{X}_{\text {in }} Xin 中形成了四个序列,即 X L R \mathbf{X}_{\mathrm{LR}} XLR、 X R L \mathbf{X}_{\mathrm{RL}} XRL、 x T B \mathbf{x}_{\mathrm{TB}} xTB和 x B T \mathbf{x}_{\mathrm{BT}} xBT,每个序列的形状为( B , H , W , C 4 B, H, W, \frac{C}{4} B,H,W,4C),分别代表前面指定的四个方向之一。然后,我们将它们展平以形成一个形状为( B , N , C 4 B, N, \frac{C}{4} B,N,4C)的单一标记序列,其中 N = W × H N=W \times H N=W×H是序列中的标记数。对于这四个组中的每一个,可以通过 θ L R \theta_{\mathrm{LR}} θLR、 θ R L \theta_{\mathrm{RL}} θRL、 θ T B \theta_{\mathrm{TB}} θTB和 θ B T \theta_{\mathrm{BT}} θBT分别指定参数,这些参数代表VSSS块的参数。
基于上述定义,分组Mamba操作符的总体关系可以写为方程6所示。
x
G
M
=
GroupedMamba
(
x
in
,
Θ
)
=
Concat
(
VSSS
(
x
L
R
,
Θ
L
R
)
,
VSSS
(
x
R
L
,
Θ
R
L
)
,
VSSS
(
x
T
B
,
Θ
T
B
)
,
VSSSS
(
x
B
T
,
Θ
B
T
)
)
\begin{aligned} \mathbf{x}_{G M}=\text { GroupedMamba }\left(\mathbf{x}_{\text {in }}, \Theta\right)=\operatorname{Concat}\left(\operatorname{VSSS}\left(\mathbf{x}_{\mathrm{LR}}, \Theta_{\mathrm{LR}}\right),\right. \\ \operatorname{VSSS}\left(\mathbf{x}_{\mathrm{RL}}, \Theta_{\mathrm{RL}}\right), \\ \operatorname{VSSS}\left(\mathbf{x}_{\mathrm{TB}}, \Theta_{\mathrm{TB}}\right), \\ \left.\operatorname{VSSSS}\left(\mathbf{x}_{\mathrm{BT}}, \Theta_{\mathrm{BT}}\right)\right) \end{aligned}
xGM= GroupedMamba (xin ,Θ)=Concat(VSSS(xLR,ΘLR),VSSS(xRL,ΘRL),VSSS(xTB,ΘTB),VSSSS(xBT,ΘBT))
其中:
- X L R , X R L , X T B \mathbf{X}_{\mathrm{LR}}, \mathbf{X}_{\mathrm{RL}}, \mathbf{X}_{\mathrm{TB}} XLR,XRL,XTB 和 X B T \mathbf{X}_{\mathrm{BT}} XBT 分别表示在各自方向上扫描的输入张量。
- Θ L R , Θ R L , Θ T B \Theta_{\mathrm{LR}}, \Theta_{\mathrm{RL}}, \Theta_{\mathrm{TB}} ΘLR,ΘRL,ΘTB 和 Θ B T \Theta_{\mathrm{BT}} ΘBT 分别表示每个方向上VSSS块的参数。
- 每个Mamba操作符的输出被重新整形为 ( B , H , W , C 4 ) \left(B, H, W, \frac{C}{4}\right) (B,H,W,4C),并再次连接回来以形成标记序列 X G M \mathbf{X}_{\mathrm{GM}} XGM,其大小再次为 ( B , H , W , C ) (B, H, W, C) (B,H,W,C)。
3.3.3、通道亲和力调制(CAM)
由于分组Mamba操作符中每个操作符仅操作 C 4 \frac{C}{4} 4C个通道,因此它可能在通道间的信息交换方面存在局限性。为了鼓励跨通道的信息交换,我们提出了通道亲和力调制(Channel Affinity Modulation, CAM)操作符,该操作符通过重新校准通道特征响应来增强网络的表示能力。在此模块中,我们首先对输入进行平均池化以计算通道统计信息,如方程7所示。
ChannelStat ( x in ) = AvgPool ( x in ) \text {ChannelStat }\left(\mathbf{x}_{\text {in }}\right)=\operatorname{AvgPool}\left(\mathbf{x}_{\text {in }}\right) ChannelStat (xin )=AvgPool(xin )
其中, X in \mathbf{X}_{\text {in }} Xin 是输入张量,AvgPool表示全局平均池化操作。接下来是亲和力计算操作,如方程8所示。
Affinity ( x in ) = σ ( W 2 δ ( W 1 ChannelStat ( x in ) ) ) \operatorname{Affinity}\left(\mathbf{x}_{\text{in}}\right)=\sigma\left(W_{2} \delta\left(W_{1} \text {ChannelStat }\left(\mathbf{x}_{\text {in }}\right)\right)\right) Affinity(xin)=σ(W2δ(W1ChannelStat (xin )))
其中, δ \delta δ和 σ \sigma σ表示非线性函数, W 1 W_{1} W1和 W 2 W_{2} W2是可学习的权重。 σ \sigma σ的作用是为每个通道分配一个重要性权重以计算亲和力。亲和力计算的结果用于重新校准分组Mamba操作符的输出,如方程9所示。
x CAM = CAM ( x GM , Affinity ( x in ) ) = x GM ⋅ Affinity ( x in ) \mathbf{x}_{\text{CAM}}=\operatorname{CAM}\left(\mathbf{x}_{\text{GM}}, \operatorname{Affinity}\left(\mathbf{x}_{\text{in}}\right)\right)=\mathbf{x}_{\text{GM}} \cdot \operatorname{Affinity}\left(\mathbf{x}_{\text{in}}\right) xCAM=CAM(xGM,Affinity(xin))=xGM⋅Affinity(xin)
其中, X CAM \mathbf{X}_{\text{CAM}} XCAM是重新校准后的输出, X GM \mathbf{X}_{\text{GM}} XGM是方程6中四个VSSS组拼接后的输出, X in \mathbf{X}_{\text{in}} Xin是输入张量, Affinity ( x in ) \operatorname{Affinity}\left(\mathbf{x}_{\text {in }}\right) Affinity(xin )是通过通道亲和力计算操作获得的通道注意力分数。
其中, X C A M \mathbf{X}_{CAM} XCAM是重新校准后的输出, X G M \mathbf{X}_{GM} XGM是方程6中四个VSSS组拼接后的输出, X i n \mathbf{X}_{in} Xin是输入张量,而Affinity ( x i n ) \left(\mathbf{x}_{in}\right) (xin)则是从方程8的通道亲和力计算操作中获得的通道注意力分数。
3.4、蒸馏损失函数
如第3节中动机部分所述,当扩展到大型模型时,Mamba训练会变得不稳定[45]。为了缓解这个问题,我们提出在标准交叉熵目标旁边使用蒸馏目标。知识蒸馏涉及训练一个学生模型,通过最小化分类损失和蒸馏损失的组合来学习教师模型的行为。蒸馏损失是通过计算教师和学生模型logits之间的交叉熵目标来获得的。给定来自学生模型的logits ( Z s ) \left(Z_{s}\right) (Zs),来自教师模型(在我们的情况下是RegNetY-16G[47])的logits ( Z t ) \left(Z_{t}\right) (Zt),真实标签 y y y,以及教师模型的硬决策 y t = argmax c Z t ( c ) y_{t}=\operatorname{argmax}_{c} Z_{t}(c) yt=argmaxcZt(c),联合损失函数如方程10所示。
L total = α L C E ( Z s , y ) + ( 1 − α ) L C E ( Z s , y t ) \mathcal{L}_{\text {total }}=\alpha \mathcal{L}_{\mathrm{CE}}\left(Z_{s}, y\right)+(1-\alpha) \mathcal{L}_{\mathrm{CE}}\left(Z_{s}, y_{t}\right) Ltotal =αLCE(Zs,y)+(1−α)LCE(Zs,yt)
其中, L C E \mathcal{L}_{\mathrm{CE}} LCE是交叉熵目标, α \alpha α是加权参数。我们在第4节中通过实验表明,使用这种蒸馏目标进行训练可以稳定训练过程,从而在更大的模型变体上带来一致的性能提升。
4、实验
4.1、图像分类
设置:图像分类实验基于ImageNet-1K数据集[9],该数据集包含超过128万张训练图像和5万张验证图像,涵盖1000个类别。我们遵循[35]的设置,使用AdamW[38]优化器和余弦衰减学习率调度器训练我们的模型,共训练300个周期,包括20个周期的预热。总批量大小设置为1024,模型在8个A100 GPU上进行训练,每个GPU具有80GB的CUDA内存。优化器的beta值设置为(0.9, 0.999);动量设置为0.9;初始学习率设置为
1
×
1
0
−
3
1 \times 10^{-3}
1×10−3,权重衰减为0.05。在与蒸馏目标(见第3-4节)结合使用时,还使用了0.1的标签平滑。
结果:表1展示了我们所提出的GroupMamba模型(T、S、B)与各种最先进方法的比较。GroupMamba模型在准确性和计算效率之间表现出了显著的平衡。GroupMamba-T以2300万个参数和4.5 GFLOPs的计算量实现了83.3%的top-1准确率,分别比ConvNeXt-T[37]和Swin-T[36]高出1.2%和2.0%,同时参数更少。此外,GroupMamba-T还超越了最近引入的SSM模型,分别比VMamba-T[34]和LocalVMamba-T[26]高出0.8%和0.6%,同时参数比VMamba-T少26%。GroupMamba-S以3400万个参数和7.0 GFLOPs的计算量实现了83.9%的准确率,超越了VMamba-S[34]、Swin-S[36]和EfficientVMamba-B[46]。与LocalVMamba-S[26]相比,性能高出0.2%,同时参数减少了32%。此外,GroupMamba-B以仅5700万个参数和14 GFLOPs的计算量实现了84.5%的准确率,比VMamba B[34]高出0.6%,同时参数减少了36%。
4.2、目标检测和实例分割
设置:我们在MS-COCO 2017数据集[32]上评估了GroupMamba-T在目标检测方面的性能。我们的方法基于Mask-RCNN[21]检测器,其超参数设置与Swin[36]相同。我们使用AdamW[38]优化器,并训练带有GroupMamba-T主干的Mask-RCNN模型,共训练12个周期。主干网络使用ImageNet-1K[9]进行初始化和微调。我们使用的初始学习率为
1
×
1
0
−
4
1 \times 10^{-4}
1×10−4,并在第9和第11个周期时将学习率衰减10倍。
结果:表2展示了GroupMamba-T在MS-COCO数据集上使用Mask R-CNN框架进行目标检测和实例分割的结果,并将其与各种最先进的模型进行了比较。我们的模型在边界框AP(AP)上达到了47.6,在掩码AP( A P m AP^{\mathrm{m}} APm)上达到了42.9。它超越了ResNet-50[22]、Swin-T[35]和ConvNeXt-T[37]。此外,与VMamba-T[34]和LocalVMamba-T[26]相比,GroupMamba-T具有竞争力的性能,并且与VMamba-T相比,参数减少了约20%。在图3中,我们展示了GroupMamba-T在COCO验证集[32]样本上的定性结果,展示了其在实例分割和目标检测方面的性能。我们的模型能够准确地定位不同场景和尺度下的对象,并正确地进行分割。
4.3、语义分割
设置:我们还在ADE20K[69]数据集上评估了GroupMamba-T在语义分割方面的性能。该框架基于UperNet[64]架构,我们遵循与Swin[36]主干相同的超参数设置。更具体地说,我们使用AdamW[38]优化器进行了总共160k次迭代,初始学习率设置为 6 × 1 0 − 5 6 \times 10^{-5} 6×10−5。我们实验中使用的默认输入分辨率为 512 × 512 512 \times 512 512×512。
结果:如表3所示,与各种最先进的方法相比,GroupMamba-T模型在语义分割方面表现出优异的性能。GroupMamba-T在单尺度评估下达到了48.6的mIoU,在多尺度评估下达到了49.2的mIoU,同时拥有49M参数和955G FLOPs。这优于ResNet-50[22]、Swin-T[36]和ConvNeXt-T[37]。此外,GroupMamba-T还超越了最新的SSM方法,包括ViM-S[70]、VMamba-T[34]和LocalVMamba[26],同时参数数量更少。在图4中,我们展示了GroupMamba-T在ADE20K[69]验证集上用于语义分割的样本的定性结果。第一行显示了真实标签掩码,第二行显示了预测的掩码。值得注意的是,我们的模型能够准确地描绘出掩码,突出了语义分割的有效性。GroupMamba的定量和定性结果证明了我们的GroupMamba主干在包括语义分割、目标检测和实例分割在内的多种下游任务中的强大泛化能力。
4.4、消融研究
表4展示了与VMamba-T相比,我们每项贡献的效果。在第二行中,没有CAM模块的GroupMamba-T实现了
82.2
%
82.2\%
82.2%的Top-1准确率。如第三行所示,集成了我们提出的CAM模块后,性能提高了
0.3
%
0.3\%
0.3%,同时参数和FLOPs相当。在第四行中,我们增加了通道数以匹配VMamba-T的参数数量(未使用蒸馏)。我们的GroupMamba-T*在参数数量相同的情况下,比VMamba-T高出
0.6
%
0.6\%
0.6%。在最后一行中,我们在不增加通道数的情况下引入了提出的蒸馏损失,与VMamba-T相比,性能提升了
0.8
%
0.8\%
0.8%,同时使用的参数减少了
26
%
26\%
26%。
5、结论
在本文中,我们通过引入一种名为调制组曼巴(Modulated Group Mamba)的新型层,解决了计算机视觉任务中视觉SSM的计算效率低下和稳定性挑战。我们还提出了一种多方向扫描方法,该方法通过在四个空间方向上扫描并利用通道亲和力调制(CAM)算子来增强跨通道的特征聚合,从而提高参数效率。为了稳定训练,特别是对于较大的模型,我们采用了基于蒸馏的训练目标。我们的实验结果表明,所提出的GroupMamba模型在需要更少参数的情况下,性能优于最近的SSM。
代码
由于这次的代码比较多,所以我只贴出GroupMambaLayer的代码,我对GroupMambaLayer做了一些修改,否则会出现nan的问题。代码如下: