点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
状态空间模型（SSM）的最新进展已显示出在以次平方复杂度建模长距离依赖关系方面的有效性能。然而，基于纯SSM的模型仍然面临稳定性问题，并且在计算机视觉任务上实现最优性能的挑战。作者的论文针对基于SSM的模型在计算机视觉中的扩展挑战，特别是大型模型在稳定性和效率方面的问题。
为了解决这一点，作者引入了一种调制组Mamba层，该层将输入通道分为四组，并分别对每组应用作者提出的基于SSM的高效视觉单选扫描（VSSS）块，每个VSSS块在一个四空间方向上进行扫描。调制组Mamba层还将四个VSSS块包装到一个通道调制算子中，以改进跨通道通信。
此外，作者引入了一种基于蒸馏的训练目标来稳定大型模型的训练，从而带来一致的性能提升。作者全面的实验验证了所提出贡献的优势，在ImageNet-1K图像分类、目标检测、MS-COCO上的实例分割以及ADE20K上的语义分割等任务上均优于现有方法。
作者具有2300万参数的小型变体在ImageNet-1K上以83.3%的分类top-1准确率达到了最先进性能，同时与同尺寸的最佳现有Mamba设计相比，参数效率提高了26%。
作者的代码和模型可在以下链接找到：https://github.com/Amshaker/GroupMamba。
1 Introduction
在语言和视觉理解的领域中，涌现了多种上下文建模方法。这些方法包括卷积、注意力[59]以及最近的状态空间模型。以多头自注意力机制[59]为核心的 Transformer 模型，在语言模型（如GPT-3[2]）和视觉模型中都占据了中心地位。然而，由于注意力机制的二次计算复杂度，特别是在处理较长序列时，带来了挑战，这也促成了如S4[17]等状态空间模型的近期出现。
尽管S4[17]在处理扩展输入序列方面表现出色，因其序列长度具有线性复杂度，但在处理信息密集型数据中的全局上下文时遇到限制，尤其是在计算机视觉领域，这归因于模型的与数据无关的性质。作为替代方案，基于全局卷积的状态空间模型[14]和液态S4[20]等方法被提出以减轻上述限制。最近的Mamba[16]引入了S6架构，旨在提高状态空间模型有效处理长距离依赖的能力。Mamba引入的选择性扫描算法使用了输入相关的状态空间参数，这允许了更好的上下文学习，同时与自注意力相比，在计算上仍然高效。
然而，Mamba特别是S6算法，在例如图像分类的任务中，当扩展到较大规模时[45]，被认为是不稳定的。此外，用于图像分类的Mamba模型变体，通常称为VSS（视觉状态空间）块，在通道数量方面可以更高效地减少参数和计算需求。VSS块包括广泛的输入和输出投影以及深度卷积，其参数和计算复杂度与输入通道的数量成正比。为了解决这个问题，作者提出了一个调制组Mamba层，它以计算和参数高效的方式减轻了上述问题。
本文的主要贡献包括：
- 作者引入了一种受群卷积启发的 Modulated Group Mamba，通过使用多方向扫描方法，增强了状态空间模型中的计算效率和交互作用，实现了全面的空间覆盖以及对局部和全局信息的有效建模。
- 作者提出了一种通道亲和力调制（_Channel Affinity Modulation (CAM)_）算子，通过增强通道间的通信来改善特征聚合，解决了分组操作中交互作用有限的固有局限。
- 作者采用基于知识蒸馏的训练目标，进一步稳定具有大量参数的模型的训练。
- 作者构建了一系列参数高效的通用分类模型系列“GroupMamba”，基于所提出的调制群玛巴层。作者的“微型”变体在ImageNet-1k [9]上以个参数和 FLOPs达到了83.3%的top-1准确率。
- 此外，作者的“基础”变体以个参数和 FLOPs达到了84.5%的top-1准确率，超过了所有最近的SSM方法（见图1）。
2 相关工作
自AlexNet [29]问世以来，卷积神经网络（ConvNets）一直是计算机视觉任务中的首选。该领域通过几个具有里程碑意义的ConvNet架构快速发展。在这些架构进步的同时，也付出了大量努力来完善单个卷积层，包括深度可分离卷积[65]、群卷积[7]和可变形卷积[8]。最近，ConvNeXt变体[37; 63]通过将宏观设计与先进的设置和训练配方相结合，使传统2D ConvNets现代化，以达到与最先进模型相媲美的性能。
近年来，开创性的视觉 Transformer [10]对计算机视觉领域产生了重大影响，包括图像分类、目标检测和分割等任务。ViT [10]引入了一种单片式设计，将图像视为一系列扁平化的2D图像块，没有图像特定的归纳偏置。ViT在计算机视觉任务上的出色性能及其可扩展性启发后续研究设计了更好的架构。早期的基于ViT的模型通常需要大规模数据集（例如JFT-300M [54]）进行预训练。后来，DeiT [57]在引入知识蒸馏 Token 的同时，提出了先进的训练技术，使得在小数据集（例如ImageNet-1K [9]）上也能有效训练。从那时起，后续研究通过结合CNN和ViT模块设计分层和混合架构，以提高不同视觉任务的性能。另一项工作是减轻自注意力中固有的二次复杂度，这是ViT的主要瓶颈。这一努力导致了显著的改进和更高效、近似的变体，在保持有效性的同时降低了复杂度。
近期，状态空间模型（SSMs）已成为ViTs [59]的一种替代方案，它们捕捉语言序列中的复杂动态和相互依赖关系 [17]。这一领域中的一个显著方法是结构化状态空间序列模型（S4）[17]，旨在解决长距离依赖问题，同时保持线性复杂性。沿着这个方向，已经提出了几种模型，包括S5 [52]、H3 [13]和GSS [41]。最近，Mamba [16] 引入了一个依赖于输入的SSM层，并利用了一种并行选择扫描机制。
在视觉领域，各种研究将SSMs应用于不同的任务。特别是对于图像分类，VMamba [34] 使用了在分层Swin-Transformer [36] 风格设计中沿两个空间维度进行双向扫描的Mamba，高效地构建了一个全局感受野。
与此同时，Vision Mamba (Vim) [70] 提出了一种单片式设计，仅对整个图像进行单向双向扫描，其性能超越了传统的视觉 Transformer ，如DeiT。LocalVMamba [26] 通过引入在特定窗口内的扫描方法（灵感来源于Swin-Transformer [36]）以及跨网络层的动态扫描方向，解决了捕捉详细局部信息的挑战。EfficientVMamba [46] 集成了基于扩张的选择性扫描和双路径模块，用于高效地提取全局和局部特征，在降低计算复杂性的同时取得了具有竞争力的结果。这些模型已经应用于图像分类，以及图像分割 ，视频理解 ，和各种其他任务 。它们广泛的应用证明了SSMs，尤其是Mamba [16]，在视觉领域的有效性。在本文中，作者提出了一个“调制组Mamba”层，它减轻了默认视觉Mamba块的缺点，例如稳定性不足 [45] 以及随着通道数量增加的参数数量问题。
3 Method
动机：作者的方法是基于对现有视觉状态空间模型局限性的观察。
- 大型模型的稳定性不足：作者从文献[45]中观察到，基于Mamba[16]的图像分类模型，在参数数量扩展到很大时，如果使用MLP通道混合器，模型会变得不稳定。这种不稳定性可以在SiMBA-L (MLP)[45]中看到，导致分类结果的准确率仅为，并不理想。作者通过引入一种名为_Modulated Group Mamba_的设计，并配合蒸馏目标（第3.4节中介绍）来稳定Mamba SSM的训练，而无需修改通道混合器。
- 高效的改进交互：鉴于基于Mamba的设计在通道数量上的计算影响，作者提出的_Modulated Group Mamba_层在计算成本和参数效率上都比默认的Mamba低，并且能够通过多方向扫描，从输入标记中建模局部和全局信息。本文还提出了一种新增的_Channel Affinity Modulation_算子，以补偿由于分组操作导致的通道交互限制。
Preliminaries
状态空间模型： 状态空间模型（SSM）如S4 [17]和Mamba [16]，是受循环神经网络（RNNs）和卷积神经网络（CNNs）启发的一维序列结构，其序列长度具有线性或近线性扩展。
状态空间模型： 状态空间模型（SSM），如S4 [17]和Mamba [16]，这些一维序列结构受到了循环神经网络（RNNs）和卷积神经网络（CNNs）的影响，其序列长度可以线性或近线性地扩展。源自于连续系统，SSM通过隐藏状态定义了一个从输入到输出的函数到函数的映射。更正式地说，SSM可以通过方程1中的常微分方程（ODE）来描述。
其中是当前的隐藏状态，是更新后的隐藏状态，是当前输入，是输出，是SSM的演化矩阵，和分别是输入和输出投影矩阵。
离散状态空间模型： 为了让这些模型能够用于深度学习中的序列建模任务，需要对它们进行离散化处理，将SSM从连续时间的函数到函数映射转换成离散时间的序列到序列映射。S4 [17]和Mamba [16]是连续系统离散化改编的例子，它们引入了时间尺度参数，以将连续参数转换为它们的离散等价形式。这种离散化通常通过方程2中的零阶保持（ZOH）方法完成。
尽管S4 [17]和Mamba [16]都使用了如方程2所述的相似离散化步骤，但Mamba通过其S6选择性扫描机制，在输入上调节参数，和，从而与S4区分开来。其中是批大小，是序列长度，是特征维度。
Overall Architecture
如图2（a）所示，作者的模型采用类似于Swin Transformer [36]的分层架构，共四个阶段以有效地处理不同分辨率的图像。假设输入图像为 ，作者首先应用一个Patch Embedding层将图像划分为大小为的不重叠 Patch ，并将每个 Patch 嵌入到一个维度为的特征向量中。Patch Embedding层通过使用两个步长为2的卷积来实现。在第一阶段，这会产生大小为的特征图。这些特征图被传递到一组作者设计的调制分组Mamba块中（细节见第3.3节）。在随后的每个阶段，一个下采样层将区域内的 Patch 合并，之后是另一组作者的调制分组Mamba块。因此，第二、三和四个阶段特征的大小分别为，和。
Modulated Group Mamba Layer
作者提出了所提出的调制群魔斑层（Modulated Group Mamba layer，图2（b））对输入序列 的整体操作，其维度为 ，其中 是批处理大小， 是输入通道数， 分别是特征图的宽度和高度，如方程3所示。
在这里， 是方程6的输出， 是方程9的输出，LN 是层归一化[1]操作，FFN 是如方程5所描述的前馈网络， 是调制群魔 Patch 的最终输出。各个操作，即群魔斑算子（GroupedMamba operator）、在群魔斑算子内部使用的VSSS块，以及CAM算子，分别在3.3.1节、3.3.2节和3.3.3节中介绍。
3.3.1 Visual Single Selective Scan (VSSS) Block
VSSS块（图2（c））是基于Mamba操作符的 Token 和通道混合器。在数学上，对于输入 Token 序列，VSSS块执行如方程4所描述的操作。
其中是输出序列，Mamba是如方程2所描述的Mamba SSM操作符的离散版本。
其中GELU[24]是激活函数，、、和是线性投影的权重和偏置。
VSSS块（图2（c）部分）是一种基于Mamba操作符的 Token 和通道混合模块。在数学表达上，对于给定的输入 Token 序列，VSSS块执行如下方程4所示的运算：
在此，表示输出序列，而Mamba是如方程2所定义的Mamba SSM操作符的离散化形式。
其中，GELU[24]是指激活函数，、、和分别代表线性投影的权重和偏置项。
3.3.2 Grouped Mamba Operator
鉴于第3节中提出的动机，作者旨在设计一种计算效率高且能有效模拟输入序列空间依赖性的Mamba变体[16]。考虑到Mamba在大数量通道的输入序列上计算效率低下，作者提出了一种受分组卷积启发的分组运算符变体。分组Mamba操作是第3.3.1节中提出的VSSS块的变体，其中输入通道被分成若干组，VSSS算子分别应用于每组。具体来说，作者将输入通道分为四组，每组大小为，一个独立的VSSS块应用于每组。为了更好地模拟输入中的空间依赖性，四个组分别沿如图2（d）所示的四条方向扫描标记序列：从左到右、从右到左、从上到下以及从下到上。
设表示代表四种扫描方向的组数：从左到右、从右到左、从上到下以及从下到上。作者从输入序列形成四个序列，分别为、、和，每个序列的形状为，代表前面指定的四种方向之一。这些序列随后被展平，形成一个形状为的单个标记序列，其中是序列中的标记数量。四个组的参数分别由、、和指定，分别代表5SSS块的参数。
根据以上定义，分组Mamba算子的整体关系可以如方程6所示编写。
其中：
- 、、和表示沿各自方向扫描的输入张量。
- 、、和表示每个方向的VSSS块的参数。
- 每个Mamba算子的输出被重新调整为的形状，并再次连接形成大小为的标记序列。
3.3.3 Channel Affinity Modulation (CAM)
独立来看，分组Mamba算子可能在通道间信息交换方面存在局限，因为组内的每个算子仅操作个通道。为了促进通道间的信息交换，作者提出了一种通道亲和性调制算子，它重新校准通道的特征响应，以增强网络的表示能力。在这一模块中，作者首先对输入进行平均池化，以计算通道统计量，如方程7所示。
其中，是输入张量，代表全局平均池化操作。接下来是亲和性计算操作，如方程8所示。
其中，和代表非线性函数，和是可学习的权重。的作用是为每个通道分配重要性权重以计算亲和性。亲和性计算的结果用于重新校准分组Mamba算子的输出，如方程9所示。
其中，是重新校准后的输出，是来自方程6的四个VSSS组的拼接输出，是输入张量，而是从方程8中的通道亲和性计算操作获得的通道注意力得分。**
Distilled Loss Function
正如在第3节的动机部分提到的，当Mamba训练扩展到大型模型时会出现不稳定的情况[45]。为了缓解这个问题，作者 Proposal 在标准的交叉熵目标基础上，利用知识蒸馏目标。知识蒸馏包括训练一个学生模型，通过最小化分类损失和蒸馏损失的组合来学习教师模型的行为。蒸馏损失是通过计算教师和学生模型的逻辑值之间的交叉熵目标来计算的。给定来自学生模型的逻辑值（），来自教师模型（在作者的情况下是RegNetY-16G[47]）的逻辑值（），真实标签，以及教师的硬决策，联合损失函数定义如方程10所示。
其中是交叉熵目标，是权重参数。作者在第4节实验性地展示了，使用这种蒸馏目标进行训练可以稳定训练过程，并在更大模型变体上带来一致的性能提升。
4 Experiments
在本节中，作者将展示一系列实验，以验证作者提出的方法的有效性。所有实验均在配置为Intel Xeon CPU E5-2690和NVIDIA Tesla K40m GPU的服务器上进行。作者使用了由三个数据集组成的综合基准测试，涵盖了多种类型的视觉任务。以下是作者实验的详细设置和结果。
Image Classification
设置： 图像分类实验基于ImageNet-1K [9]，包含超过128万张训练图像和5万张验证图像，涵盖1000个类别。遵循[35]，作者使用AdamW [38]优化器和余弦退火学习率调度器训练模型，共300个周期，包括20个周期的 Warm up 。总批处理大小设置为1024，模型在8x A100 GPU上训练，每个GPU具有80GB的CUDA内存。优化器的betas设置为(0.9, 0.999)；动量设置为0.9，初始学习率为，权重衰减为0.05。标签平滑设置为0.1，与蒸馏目标（见第3.4节）一起使用。
结果： 表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%。具有3400万参数和7.0 GFLOPs的GroupMamba-S达到了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%。
Object Detection and Instance Segmentation
设置： 作者在MS-COCO 2017数据集[32]上评估了GroupMamba-T在目标检测任务中的性能。作者的方法基于Mask-RCNN[21]检测器，并使用了与Swin[36]相同的超参数。作者采用AdamW[38]优化器，用GroupMamba-T Backbone 网络训练Mask-RCNN共12个周期。该 Backbone 网络从ImageNet-1K[9]初始化并微调。作者使用的初始学习率为，在第9和第11个周期时学习率衰减10倍。
结果： 表2展示了GroupMamba-T的结果，与MS-COCO数据集上使用Mask R-CNN框架的各种最先进的目标检测和实例分割模型进行了比较。作者的模型在边界框AP（AP）上达到了47.6，在 Mask AP（AP）上达到了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]样本上的定性结果，证明了其在实例分割和目标检测方面的性能。作者的模型能够准确地对物体进行定位，并在不同的场景和尺度下正确地进行分割。
Semantic Segmentation
设置：作者还评估了GroupMamba-T在ADE20K [69] 数据集上进行语义分割的性能。该框架基于UperNet [64] 架构，并且作者使用了与Swin [36] 主干网络相同的超参数。更具体地说，作者使用AdamW [38] 优化器进行总共次迭代，初始学习率为。作者实验中使用的默认输入分辨率为。
结果：如表3所示，与各种最先进的方法相比，GroupMamba-T模型在语义分割方面表现出优越的性能。GroupMamba-T在单尺度评估中达到mIoU为，在多尺度评估中达到，参数量为，FLOPs为。这超越了ResNet-50 [22]，Swin-T [36] 和ConvNeXt-T [37]。此外，GroupMamba-T在参数更少的情况下，超过了包括ViM-S [70]，VMamba-T [34] 和LocalVMamba [26]在内的近期SSM方法的性能。在图4中，作者展示了GroupMamba-T在ADE20K [69] 验证集样本上进行语义分割的定性结果。第一行展示了 GT 值 Mask ，而第二行显示了预测的 Mask 。值得注意的是，作者的模型准确地勾画出了 Mask ，突显了在语义分割方面的有效性。GroupMamba的定量和定性结果展示了作者的GroupMamba主干网络在包括语义分割、目标检测和实例分割在内的各种下游任务中的强大泛化能力。### 消融研究
表4说明了与VMamba-T相比，作者各项贡献的影响。去除了CAM模块的提出的GroupMamba-T在第二行中达到82.2%的top-1准确率。如第三行所示，融合了所提出的CAM模块，性能提升了0.3%，参数量和FLOPs可比较。在第四行，作者将通道数量扩展到与VMamba-T（不使用蒸馏）的参数数量相匹配。作者的GroupMamba-T*在参数数量相等的情况下，比VMamba-T高出0.6%。在最后一行，作者没有扩展通道，而是融入了所提出的蒸馏损失，与VMamba-T相比，性能提升了0.8%，同时使用了26%的较少参数。
5 结论
在本文中，作者针对计算机视觉任务中视觉SSM的计算效率低下和稳定性挑战问题，引入了一种名为 Modulated Group Mamba 的新层。同时，作者还提出了一种多方向扫描方法，通过在四个空间方向上进行扫描并利用通道亲和力调制（CAM）算子来增强跨通道的特征聚合，从而提高了参数效率。
为了稳定训练，特别是对于大型模型，作者采用了一种基于蒸馏的训练目标。实验结果表明，所提出的组曼巴模型在性能上超过了最近的SSM，同时所需的参数更少。
参考
[1].GroupMamba: Parameter-Efficient and Accurate Group Visual State Space Model.
点击上方卡片，关注 「AI视界引擎」 公众号