MobileMamba: Lightweight Multi-Receptive Visual Mamba Network——轻量级多感受野视觉 Mamba 网络

这篇文章的核心内容是介绍了一种新型的轻量级视觉网络架构——MobileMamba,旨在平衡模型的效率和性能,特别是在移动设备等资源受限的环境中。文章详细阐述了MobileMamba的设计理念、架构细节、实验验证以及与其他现有模型的比较,以下是文章的主要研究内容总结:

研究背景与动机

  • 背景:随着移动设备的普及,对高效且准确的视觉处理需求增加。现有的轻量级模型主要分为基于CNN和基于Transformer的结构,但它们分别存在局部感受野和计算复杂度高的问题。

  • 动机:状态空间模型(SSMs)因其线性计算复杂度在视觉领域受到关注,但现有的轻量级Mamba基模型在吞吐量和性能上表现欠佳。因此,作者提出了MobileMamba框架,旨在通过设计优化提升轻量级模型的效率和性能。

MobileMamba框架设计

  • 三阶段网络结构:作者设计了一个三阶段网络,通过实验验证,与四阶段网络相比,三阶段网络在相同的吞吐量下能实现更高的准确率,并且推理速度更快。

  • 多感受野特征交互(MRFFI)模块:这是MobileMamba的核心模块,包含三个部分:

    • 长距离小波变换增强Mamba(WTE-Mamba):用于提取全局特征并增强高频边缘信息。

    • 高效多核深度可分离卷积(MK-DeConv):用于捕捉多尺度感受野。

    • 消除冗余恒等映射:减少高维空间中的特征冗余,降低计算复杂度。

  • 训练和测试策略:采用知识蒸馏和扩展训练周期来提升小模型的性能,并在测试阶段采用归一化层融合以提高推理速度。

实验验证

  • 图像分类:在ImageNet-1K数据集上,MobileMamba在不同模型规模下均展现出优越的性能,Top-1准确率最高达到83.6%,并且在GPU上比现有方法更快。

  • 目标检测和实例分割:在MS-COCO 2017数据集上,使用SSDLite、RetinaNet和Mask R-CNN进行实验,MobileMamba在mAP和吞吐量上均优于或接近现有方法。

  • 语义分割:在ADE20K数据集上,使用DeepLabv3、Semantic FPN和PSPNet进行实验,MobileMamba在mIoU上取得了优异的成绩,同时保持较低的FLOPs。

关键结论

  • MobileMamba通过其独特的三阶段网络结构和MRFFI模块,在保持低计算复杂度的同时,显著提高了模型的性能和推理速度。

  • 通过采用训练和测试策略,MobileMamba进一步提升了性能,尤其是在小模型上。

  • 在多个视觉任务上的广泛实验表明,MobileMamba在效率和性能之间取得了良好的平衡,优于现有的轻量级模型。

限制与未来工作

  • 尽管MobileMamba在GPU上表现出色,但在CPU和移动设备上的性能仍有提升空间。

  • 未来的工作将集中在进一步优化Mamba模型在各种设备上的推理能力,特别是提高其在边缘设备上的效率。

这篇文章的核心贡献在于提出了一种新的轻量级视觉网络架构,通过创新的设计和优化策略,在保持高效计算的同时,显著提升了模型的性能,为移动设备等资源受限环境中的视觉应用提供了新的解决方案。这里是自己的论文阅读记录,感兴趣的话可以参考一下,如果需要阅读原文的话可以看这里,如下所示:

官方项目地址在这里,如下所示:

摘要

以往轻量级模型的研究主要集中在基于卷积神经网络(CNN)和基于 Transformer 的设计上。CNN 由于其局部感受野,在捕捉长距离依赖关系方面存在困难,而 Transformer 尽管具有全局建模能力,但在高分辨率场景下受到二次计算复杂度的限制。最近,由于其线性计算复杂度,状态空间模型在视觉领域受到关注。尽管轻量级 Mamba 基模型的 FLOPs 很低,但现有的轻量级 Mamba 基模型在吞吐量方面表现欠佳。在本工作中,我们提出了 MobileMamba 框架,该框架在效率和性能之间取得了平衡。我们设计了一个三阶段网络,以显著提高推理速度。在细粒度层面,我们引入了多感受野特征交互(MRFFI)模块,包括长距离小波变换增强 Mamba(WTE-Mamba)、高效多核深度可分离卷积(MK-DeConv)和消除冗余恒等组件。该模块整合了多感受野信息,并增强了高频细节提取能力。此外,我们还采用了训练和测试策略,以进一步提高性能和效率。MobileMamba 在 Top1 上达到了 83.6% 的准确率,超越了现有的最先进方法,并且在 GPU 上比 LocalVim 快了最多 ×21↑。在高分辨率下游任务上的广泛实验表明,MobileMamba 超越了当前的高效模型,在速度和准确率之间实现了最佳平衡。

1. 引言

随着移动设备的普及,资源受限环境中对高效且准确的视觉处理的需求增加。轻量级模型显著降低了计算和存储成本,同时提高了推理速度。当前的轻量级模型主要分为基于 CNN 和基于 Transformer 的结构。基于 CNN 的 MobileNets [24, 26, 56] 使用深度可分离卷积来降低计算复杂度,为后续的基于 CNN 的工作 [3, 45, 61, 63, 79] 奠定了基础。然而,基于 CNN 的方法的主要缺点是其局部有效感受野(ERF),如图 1(i) 所示,它被限制在中心区域,缺乏长距离相关性。在具有高分辨率输入的下游任务(表 5)中,基于 CNN 的方法只能通过增加计算负载来实现性能提升。Vision Transformers(ViTs)展现出全局 ERF 和长距离建模能力,如图 1(ii) 所示。然而,它们的二次计算复杂度导致比 CNN 更高的开销。一些工作 [33, 34, 38, 46, 56, 66, 76] 通过降低分辨率或通道数来缓解这种复杂度,取得了显著成果。尽管如此,纯 ViT 仍然缺乏归纳偏差,促使研究人员开发混合 CNN-ViT 结构 [32, 74, 78],结合局部和全局 ERF 以提高性能,如图 1(iii) 所示。然而,基于 ViT 的方法仍然面临二次计算复杂度的问题,尤其是在具有高分辨率输入的下游任务(表 6)中。状态空间模型 [13–15, 58] 因其能够以线性计算复杂度捕捉长距离依赖关系而受到关注。研究人员已成功将这些模型应用于视觉领域 [39, 57, 83],取得了显著的效果和效率。最近的轻量级 Mamba 基模型 [29, 51] 引入了不同的高效扫描方法以降低复杂度。然而,这些工作仅报告了 FLOPs,这并不一定与快速推理速度相关。实验结果如图 2 所示,当前的 Mamba 基结构在推理速度和性能方面都存在问题。基于上述动机,我们提出了 MobileMamba,通过粗粒度、细粒度和训练/测试策略设计为高效的轻量级网络。首先,在第 3.1 节中,我们讨论了四阶段和三阶段网络之间的权衡,就准确率、速度和 FLOPs 而言。如图 3 所示,在相同的吞吐量下,三阶段网络实现了更高的准确率。同样,对于相同的性能,三阶段网络具有更高的吞吐量。因此,我们选择三阶段网络作为我们的粗粒度框架。在第 3.2 节的 MobileMamba 模块设计中,我们引入了高效的多感受野特征交互(MRFFI)模块。具体来说,输入特征沿通道维度被分为三个部分。第一部分使用长距离小波变换增强 Mamba(WTE-Mamba)模块来提取全局特征,同时增强边缘信息等细粒度细节的提取。第二部分采用多核深度可分离卷积(MK-DeConv)操作来捕捉多尺度感受野。最后部分使用消除冗余的恒等映射来减少高维空间中的通道冗余,降低计算复杂度并提高处理速度。通过 MRFFI 获得的特征整合了全局和多尺度局部感受野信息,增强了高频边缘细节的提取能力。最后,我们通过第 3.3 节中的两个训练阶段策略增强模型的学习能力,即知识蒸馏和扩展训练周期。此外,在测试阶段采用归一化层融合策略以提高模型的推理速度。在图 1(iv) 中,我们的方法利用全局 ERF,而多核局部卷积操作有助于提取邻近信息。图 1 底部的比较表明,MobileMamba†(采用训练策略)在 ImageNet1K [10] 上的 Top1 准确率分别为 76.9/78.9/80.7/82.2/83.3/83.6,适用于从 200M 到 4G FLOPs 的模型,超越了现有的 CNN、ViT 和 Mamba 基方法。与高效的 Mamba 基方法相比,如图 2 所示,MobileMamba 在 Top-1 上提高了 +0.7↑,同时比 LocalVim [30] 快了 ×21↑,并且在 Top-1 上提高了 +2.0↑,同时比 EfficientVMamba [51] 快了 ×3.3↑。这表明与现有的 Mamba 基轻量级模型设计相比具有显著优势。在下游任务上的广泛实验进一步验证了我们方法的有效性。在 Mask RCNN [22] 上,MobileMamba 将 mAPb 提高了 +1.3↑,mAPm 提高了 +1.0↑,吞吐量提高了 +56%↑,与 EMO [78] 相比。

在 RetinaNet [35] 上,它将 mAPb 提高了 +2.1↑,吞吐量提高了 ×4.3↑,与 EfficientVMamba [51] 相比。在 SSDLite [25] 上,它通过提高分辨率实现了 24.0/29.5 的 mAPs。在 DeepLabv3 [5]、Semantic FPN [31] 和 PSPNet [80] 上,它以较少的 FLOPs 实现了高达 37.4/42.5/36.9 的 mIoUs。与基于 CNN 的 MobileNetv2 [56] 和基于 ViT 的 MobileViTv2 [47] 相比,我们的方法在高分辨率 512x512 输入下游任务中分别实现了 +7.2↑ 和 +0.4↑ 的改进,而仅需 PSPNet [80] 的 8.5% 和 11.2% 的 FLOPs。总之,我们的贡献如下:

  • 我们提出了一个轻量级的三阶段 MobileMamba 框架,该框架在性能和效率之间取得了良好的平衡。MobileMamba 在分类任务以及三个高分辨率输入下游任务中的有效性和效率得到了验证。

  • 我们设计了一个高效的多感受野特征交互(MRFFI)模块,以增强具有更大 ERF 的多尺度感知能力,并提高细粒度高频边缘信息的提取能力。

  • MobileMamba 通过在不同 FLOPs 大小的模型中采用训练和测试策略,显著提高了性能和效率。

2. 相关工作

2.1 轻量级

视觉模型 研究最广泛的轻量级视觉网络可分为基于 CNN 和基于 Vision Transformer(ViT)的结构。基于 CNN 的 MobileNets [24, 26, 56] 从标准卷积过渡到深度可分离卷积,显著降低了计算复杂度。GhostNets [18, 40, 63] 用半数通道上的廉价运算替换了原始卷积。此外,许多基于 CNN 的工作 [3, 44, 61, 62, 79] 在移动设备上展现了卓越的性能和效率。这些方法的主要限制在于其局部感受野。相比之下,ViT 具备全局感受野,能够捕捉长距离依赖关系。然而,它们的二次计算复杂度导致比 CNN 更高的计算成本。因此,轻量级视觉 Transformer 被设计为保留全局感受野的同时减少计算开销。EfficientViT [38] 设计了一个三阶段网络,并提出了级联组注意力,显著提高了推理速度。SHViT [76] 引入了单头自注意力,仅选择部分通道使用 ViT,而通过恒等映射直接连接其余通道,极大地提高了操作效率。此外,许多混合方法 [33, 34, 46, 47, 50, 65, 66, 78] 也取得了杰出的成果。

2.2 状态空间模型

距离依赖关系方面的高效计算复杂度,SSMs [13, 15, 58] 受到了广泛关注。Mamba [14] 引入了 S6 模块,实现了一个简单结构,在长序列建模方面具有卓越的效率。由于这一优势,许多工作已将其应用于视觉任务 [19, 20, 39, 43, 54, 72, 83]。Vim [83] 提出了双向 Mamba 块,在高分辨率下展现了其速度和内存优势,优于 ViTs。VMamba [39] 引入了交叉扫描以增强建模能力。EfficientVMamba [51] 提出了高效扫描,通过跳过采样提高了扫描效率。LocalVim [29] 提出了局部窗口扫描以增强局部信息获取。尽管有各种设计,但没有任何轻量级 Mamba 基网络能够超越现有的 CNN 和 ViT 方法。本文探索轻量级 Mamba 基视觉网络,以实现更好的性能、更低的计算复杂度和更快的推理速度。

3. 方法论

3.1 MobileMamba 的粗粒度设计

在本节中,我们设计了高效的 MobileMamba 结构,包括一个三阶段网络,如图 3(B) 所示。大多数现有的网络 [3, 33, 74] 遵循图 3(A) 所示的四阶段框架。具体来说,在四阶段网络中,第一次下采样将输入图像 H×W×3 缩小到 H/4×W/4×C1,最终输出特征图是 H/32×W/32×C4。相比之下,三阶段网络在第一次下采样时将输入图像缩小到 H/16×W/16×C1,最终输出特征图是 H/64×W/64×C4。由于四阶段网络具有更大的特征图尺寸,因此需要更多的计算量,从而导致速度较慢。图 3 下方的表格比较了在类似吞吐量条件下,三阶段网络和各种四阶段网络在 ImageNet-1K [10] 数据集上的分类结果。在前两个实验中,四阶段网络的前两个阶段设计为纯 CNN 架构,以提高推理速度。第三个实验在整个网络的四个阶段中均采用 MobileMamba 块。结果表明,尽管四阶段网络在前两个阶段采用纯 CNN 结构时,推理速度和性能有所提高,但三阶段网络在 Top-1 和 Top-5 准确率方面均提高了 +0.4↑,并且推理速度更快。最终,我们选择三阶段网络结构以提高推理速度并改善分类结果。

3.2 MobileMamba 的细粒度设计

我们提出的高效的多感受野特征交互(MRFFI)模块放置在每个 MobileMamba 块的对称局部信息感知和前馈网络(FFN)之间。在 MRFFI 模块中,特征沿通道维度被分为三个部分。1)第一部分特征经过长距离小波变换增强 Mamba(WTE-Mamba),在进行全局建模的同时增强高频边缘信息的提取。2)第二部分通过多核深度可分离卷积(MK-DeConv)操作来增强不同感受野的感知能力。3)剩余特征经过恒等映射以减少高维空间中的特征冗余,降低计算复杂度,从而提高处理速度。长距离 WTE-Mamba。其目的是在全局建模的基础上增强提取细粒度信息的能力,例如高频边缘细节。此外,对 WT 特征图进行卷积操作相比于正常尺度具有更大的 ERF,并且计算复杂度更低。对于输入特

MobileMamba 设计了六种结构,如表 2 所示。在不同模型中,我们保持了相同的全局和局部通道比例。对于小模型,我们使用较小的输入分辨率以实现较低的计算复杂度和更快的运行时间。相反,对于大模型,我们使用较大的输入分辨率以获得更好的性能,如第 4.4 节所述。通过根据模型大小调整输入分辨率,我们在计算效率和性能之间取得了平衡。这种设计策略确保了 MobileMamba 能够有效地扩展以满足不同需求,同时保持一致的通道比例。

表 2. MobileMamba 模型变体的架构细节。

模型名称输入分辨率{C1, C2, C3}{D1, D2, D3}{ξ1, ξ2, ξ3}{µ1, µ2, µ3}
MobileMamba-T2192{144, 272, 368}{1, 2, 2}{0.8, 0.7, 0.6}{0.2, 0.2, 0.3}
MobileMamba-T4192{176, 368, 448}{1, 2, 2}{0.8, 0.7, 0.6}{0.2, 0.2, 0.3}
MobileMamba-S6224{192, 384, 448}{1, 2, 2}{0.8, 0.7, 0.6}{0.2, 0.2, 0.3}
MobileMamba-B1256{200, 376, 448}{2, 3, 2}{0.8, 0.7, 0.6}{0.2, 0.2, 0.3}
MobileMamba-B2384{200, 376, 448}{2, 3, 2}{0.8, 0.7, 0.6}{0.2, 0.2, 0.3}
MobileMamba-B4512{200, 376, 448}{2, 3, 2}{0.8, 0.7, 0.6}{0.2, 0.2, 0.3}

3.3 训练和测试策略

我们采用了两种训练策略,以在保持相同参数数量和计算复杂度的情况下,进一步提升小模型的性能和效率。此外,我们还使用了一种测试策略,以确保模型的有效性,同时提高推理速度。知识蒸馏。为了使轻量级学生模型 MobileMamba 从更强大的教师分类模型中学习,我们遵循了 DeiT [64] 中的软蒸馏设置。这涉及最小化教师模型和学生模型 softmax 输出之间的 Kullback-Leibler 散度。扩展训练周期。我们观察到,在常规的 300 个训练周期下,小模型 MobileMamba 的损失尚未完全收敛,Top-1 准确率也未达到其潜力。因此,为了提高轻量级模型的性能上限,我们将训练扩展到 1000 个周期。归一化层融合。卷积操作通常后面跟着批量归一化。在推理过程中,批量归一化可以与前面的卷积或线性层融合。重新计算新卷积层的权重和偏置,以确保其组合输出与原始层的输出一致。这种融合提高了计算效率,并通过减少层数加快了前向传播速度。

4. 实验

4.1 实现细节

我们在 ImageNet-1K [10] 数据集上进行图像分类。基线模型从头开始训练 300 个周期,分辨率为 224²。采用 AdamW [41] 优化器,其参数为 betas (0.9, 0.999),权重衰减为 5e-2,学习率为 1.5e-3,批量大小为 1024。我们使用余弦调度器 [42],带有 20 个热身周期,标签平滑 [60] 为 0.1,随机深度 [28] 和 RandAugment [9]。为了公平比较,我们遵循 [64] 中提出的数据增强技术,包括 Mixup [77]、随机擦除 [81] 和自动增强 [8]。对于增强模型†,我们训练了 1000 个周期,并遵循 DeiT [64] 中使用知识蒸馏的配方,以 TResNet-L [53] 作为教师模型。对于目标检测任务,我们在 MS-COCO 2017 [36] 数据集上使用轻量级 SSDLite [25] 和 RetinaNet [35] 进行验证。对于实例分割任务,我们在 COCO [36] 数据集上使用 Mask R-CNN [22] 进行实验。对于语义分割任务,我们在 ADE20K [82] 数据集上使用 DeepLabv3 [5]、PSPNet [80] 和 FPN [37] 进行实验。对于所有下游任务实验,我们使用标准的 MMDetection [4] 和 MMSegmentation [7] 库,并且仅将优化器替换为 AdamW [41],而没有调整其他参数。GPU 吞吐量是在单个 Nvidia L40S 上测量的,批量大小为 256。

4.2 MobileMamba

在 ImageNet-1K 分类上的表现 表 3 展示了 MobileMamba 在六个不同模型规模下与其他最先进方法在 ImageNet-1K 上的结果。不同的模型规模根据 FLOPs 进行分类。例如,与 MobileMamba-B1 模型相比,B2 和 B4 模型仅增加了输入分辨率,而没有增加网络深度或宽度。在前两个模型规模中,目前没有与 FLOPs 相当的 Mamba 基模型。MobileMamba-T2 在 Top-1 上比基于 Transformer 的 SHViTS1 [76] 高出 +0.8↑。MobileMamba-T4 在 Top-1 上比基于线性注意力的 VRWKV-T [11] 高出 +1↑,而其 FLOPs 仅为后者的 33%。对于 MobileMamba-S6 和 B1 模型,我们还观察到与其他 CNN、Transformer 和 Mamba 基模型相比有显著提升。MobileMamba-S6 在 Top-1 准确率上比 EfficientVMamba-T [51] 高出 1.5,同时减少了 18.5% 的 FLOPs。为了展示轻量级模型的扩展能力,我们保持 MobileMamba-B1 模型的网络架构不变,并将输入分辨率提高到 384² 和 512²,从而得到大约 2G 和 4G FLOPs 的模型。MobileMamba-B2 和 B4 模型在 FLOPs 较少的情况下实现了更高的分类结果。此外,使用训练策略†可以进一步提升模型性能。例如,应用训练策略的 MobileMamba-T2† 模型在 Top-1 上提高了 +3.3↑,在 Top-5 上提高了 +1.7↑。在所有模型规模中,训练策略始终能够显著提升性能。

表 3. 在 ImageNet-1K [10] 数据集上的分类性能。白色、灰色、黄色和蓝色背景分别表示基于 CNN、Transformer、Mamba/RWKV 和我们的 MobileMamba 的模型。这种显示方式在后续实验中继续使用。† 表示使用了训练策略。

模型名称FLOPs↓参数↓分辨率Top-1Top-5发表年份
EdgeNeXt-XXS [45]2601.322471.2-ECCVW’22
ShuffleNetV2×1.5 [44]3003.522472.6-ECCV’18
FasterNet-T0 [3]3403.922471.9-CVPR’23
MobileViTv3-0.5 [66]4811.425672.3-arXiv’2209
EfficientViT-M2 [38]2014.222470.8-CVPR’23
EMO-1M [78]2611.322471.5-ICCV’23
SHViT-S1 [76]2416.322472.8-CVPR’24
MobileMamba-T22558.819273.6--
MobileMamba-T2†2558.819276.9--
EdgeNeXt-XS [45]5402.322475.0-ECCVW’22
InceptionNeXt-A [75]5104.222475.3-CVPR’24
EfficientFormerV2-S0 [34]4003.522475.7-ICCV’23
EfficientViT-M4 [38]2998.822474.3-CVPR’23
EMO-2M [78]4392.322475.1-

4.3 MobileMamba

在下游任务上的表现 目标检测和实例分割。我们在 MS-COCO 2017 [36] 数据集上使用轻量级 SSDLite [25] 和 RetinaNet [35] 进行目标检测,以及使用 Mask R-CNN [22] 进行实例分割,以验证预训练的 MobileMamba 模型。对于 SSDLite(表 5),我们最初在 320² 分辨率下进行实验,然后将分辨率提高到 512²

,同时保持其他参数不变。MobileMamba-B1 在 320² 分辨率下比 EMO1M [78] 高出 +2↑。MobileMamba-B1 比 MViTv2-1.25 [49] 的 FLOPs 少 -0.3G↓,同时在 512² 分辨率下 mAP 提高了 +1.7↑。对于 RetinaNet(表 4),MobileMamba-B1 展示了比 EfficientVMamba-T [51] 高 ×4.3↑ 的 GPU 吞吐量,同时 mAPb 提高了 +2.1↑。与 EMO-5M [78] 相比,它在 GPU 吞吐量上提高了 +31%↑,mAPb 提高了 +0.7↑。对于 Mask R-CNN(表 4),MobileMamba-B1 比 EMO-5M [78] 的吞吐量提高了 +57%↑,同时 mAPb 和 mAPm 分别提高了 +1.3↑ 和 +1.0↑。与 SHViT-S4 [76] 相比,它在 mAPb 和 mAPm 上分别提高了 +1.6↑ 和 +1.5↑。

表 4. 使用 RetinaNet [35] 和 Mask RCNN [22] 在 MS-COCO 2017 [36] 数据集上的目标检测和实例分割结果。TP:单个 NVIDIA L40S 上的 GPU 吞吐量。

背骨网络mAPbmAPb@50mAPb@75mAPmmAPm@50mAPm@75TP
EfficientNet-B0 [61]31.951.034.529.447.931.271
ResNet-50 [21]38.058.641.434.455.136.741
FastViT-SA12 [65]38.960.542.235.957.638.136
EfficientViT-M4 [38]32.854.434.531.051.232.2121
PoolFormer-S12 [74]37.359.040.134.655.836.932
EfficientFormer-L1 [33]37.960.341.035.457.337.345
SHViT-S4 [76]39.061.241.935.957.937.9136
EMO-5M [78]39.361.742.436.458.438.797
MobileMamba-B140.661.843.837.458.939.9152

表 5. 使用 SSDLite [23] 在 MSCOCO 2017 [36] 数据集上进行目标检测的结果,分辨率为 320×320。

背骨网络分辨率FLOPs↓参数数量↓mAP
MobileNetv1 [27]3201.3G5.122.2
MobileNetv2 [55]3200.8G4.322.1
MobileNetv3 [23]3200.6G5.022.0
MobileViTv1-XXS [48]3200.9G1.719.9
MobileViTv2-0.5 [49]3200.9G2.021.2
EMO-1M [78]3200.6G2.322.0
MobileMamba-B13201.7G18.024.0
ResNet50 [21]5128.8G26.625.2
EdgeNeXt-S [45]5122.1G6.227.9
MobileViTv2-0.75 [49]5121.8G3.624.6
EMO-2M [78]5120.9G3.325.2
MobileViTv1-S [48]5123.4G5.727.7
MobileViTv2-1.25 [49]5124.7G8.227.8
EMO-5M [78]5121.8G6.027.9
MobileMamba-B15124.4G18.029.5

语义分割。我们在 ADE20K [82] 数据集上使用 DeepLabv3 [5]、Semantic FPN [31] 和 PSPNet [80] 对预训练的 MobileMamba 进行语义分割性能评估(表 6)。对于 DeepLabv3,MobileMamba-B4 比 EMO2M [78] 的 mIoU 高出 +1.3↑。对于 Semantic FPN,MobileMamba 展示了显著的优势。与 EMO-5M [78] 相比,它仅使用了 22% 的 FLOPs,同时 mIoU 提高了 +2.1↑。与 EdgeViT-XS [50] 相比,它在 mIoU 上提高了 +1.1↑,而 FLOPs 仅为后者的 20%。与具有类似结果的 PVTv2-B1 [69] 相比,我们的模型的 FLOPs 仅为后者的 16%。对于 PSPNet,MobileMamba-B4 比 EMO-2M [78] 的 mIoU 高出 +2.4↑。与 MobileViTv2-1.0 [49] 相比,它在 mIoU 上提高了 +0.4↑,而仅使用了 11% 的 FLOPs。

表 6. 在 ADE20K [82] 数据集上,使用 DeepLabv3 [5]、Semantic FPN [31] 和 PSPNet [80] 进行语义分割的结果,分辨率为 512×512。

背骨网络FLOPs↓参数数量↓mIoU
DeepLabv3 [5]---
MobileNetv2 [55]75.4G18.734.1
MobileNetv2-0.5 [49]26.1G6.331.9
MobileViTv3-0.5 [67]--33.5
EMO-1M [78]2.4G5.633.5
MobileViTv2-0.75 [49]40.0

4.4 额外的消融和解释性分析

逐步实验。图 5 展示了从基线 EfficientViT-M5 [38] 模型逐步推导出 MobileMamba 模型的过程,通过逐步实验。由于 FLOPs 并不总是完全反映模型的推理速度,因此我们包括了 GPU 吞吐量指标以展示模型的效率。在结构层面,级联组注意力逐渐被我们提出的 MRFFI 模块替换。这一过程整合了多尺度感受野,同时提高了模型的吞吐量。随后,应用细粒度设计以增强模型在 Mamba 的全局感受野和频率域细节方面的表示能力,从而提高模型性能。减少网络层数,同时增加维度和全局比例 ξ,可以在保持类似准确率的同时降低 FLOPs 并提高吞吐量。最后,通过减少 d 状态和扫描方向,同时增加 SSM 的扩展比率,模型性能得到进一步提升,吞吐量显著增加。最终,与基线相比,MobileMamba 在 Top-1 上提高了 +0.9↑,在 Top-5 上提高了 +0.6↑,同时吞吐量增加了 +729↑ 张图像每秒。效率比较。表 7 展示了与最先进方法在效率和有效性方面的比较。MobileMamba 在 GPU 吞吐量方面超越了所有方法。平均而言,三种不同大小的 MobileMamba 模型的 GPU 吞吐量比 EfficientVMamba [51] 快 ×3.5↑。然而,在 AMD EPYC 9K84 96-Core CPU 吞吐量和 iPhone15(ms)移动设备的延迟方面,MobileMamba 落后于基于 Transformer 的模型。这归因于 Mamba 方法在 CPU 上的当前工程实现仍然存在改进和优化的空间。尽管如此,与其他基于 Mamba 的方法相比,MobileMamba 在 CPU 上的延迟仅为 EfficientVMamba 的 15%-42%,同时在 Top-1 上平均提高了 +1.5↑。表 7. 与最先进方法在效率和有效性方面的比较。

模型名称FLOPs参数数量分辨率GPU 吞吐量CPU 吞吐量移动设备延迟Top-1 (%)
EfficientViT-M2 [38]2014.2224186932553.91.170.8
EMO-1M [78]2611.322483619110.95.171.5
MobileMamba-T22558.8192210718511.811.773.6
EfficientViT-M4 [38]2998.8224146122284.41.674.3
EMO-2M [78]4392.322463016715.07.675.1
MobileMamba-T441314.2192165718411.816.976.1
EfficientViT-M5 [38]52212.4224102711805.62.077.1
EfficientVMamba-T [51]8006.022462851470.0113.676.5
MobileMamba-S665215.0224110008012.519.778.0
EfficientVMamba-S [51]130011.022433277137.8287.978.7
EMO-6M [78]9616.122440384223.514.879.0
MobileMamba-B1108017.125669864920.347.079.9
EfficientViT-M5r512 [38]267012.451216943132.43.580.8
EfficientVMamba-B [51]400033.02246485198.2834.881.8
MobileMamba-B4431317.151218621284.2291.782.5

表 7. 与最先进方法在效率和有效性方面的比较。

在低分辨率下的小模型消融。为了提高小模型的性能并增加其吞吐量,我们研究了输入分辨率的影响。我们设置了三种输入分辨率:160²、192² 和 224²,并调整模型参数以确保每个分辨率的 FLOPs 大约为 250M。如表 8a 所示,尽管 FLOPs 相似,但较低的输入分辨率导致模型吞吐量更高,参数尺寸更大。综合考虑吞吐量、参数尺寸和性能,我们设计了输入分辨率为 192² 的小模型,实现了良好的平衡并取得了令人满意的结果。

在高分辨率下的大模型消融。我们探索了提高小模型扩展能力的方法。在表 8b 中,在标准分辨率 224² 下,增加模型的深度和宽度以实现 2G 和 4G FLOPs 并未显著提高性能,尽管计算负载增加。这是由于在当前三阶段框架中输入分辨率过低。因此,我们将输入分辨率提高到 384² 和 512²。在 FLOPs 相似且吞吐量略有下降的情况下,Top-1 分别提高了 +1.6↑ 和 +1.8↑。

训练策略的效果。表 8c 展示了使用训练和测试策略的逐步实验。应用知识蒸馏(KD)后,Top-1 和 Top-5 准确率在 ImageNet-1K 数据集上分别提高了 +2↑ 和 +0.7↑。将训练扩展到 1000 个周期进一步将这些指标分别提高了 +0.7↑ 和 +0.5↑。最终,具有 652M FLOPs 的模型在 Top1 上达到了 80.7,在 Top-5 上达到了 95.2,超过了未使用训练策略且具有 1080M FLOPs 的模型。此外,在测试阶段采用归一化融合可以进一步提高速度 ×1.2%↑。

Mamba 组件的消融。在 Mamba 模型的内部参数实验中,如表 8d 所示,S、R 和 D 分别代表扫描方向、扩展比率和 d 状态。减少 S 可以增加吞吐量,尽管性能略有下降。在相同的 S 数量下,使用 R=2 和 D=1 的组合比 R=1 和 D=16 的组合具有更高的吞吐量和更好的性能。因此,最终选择使用双向扫描,R=2 和 D=1。MK-DeConv 的影响。我们对高效 MK-DeConv 操作中的分组数 n 进行了实验(见表 8e)。对于 n = 1,所有通道都使用单个卷积模块,核大小为 3。对于 n = 3,通道被分为三组,分别使用 k = 3, 5, 7 的卷积,然后沿通道维度进行拼接。这些方法在参数、FLOPs 和吞吐量方面没有显著差异,结果相似。因此,我们为了简化采用 n = 1。然而,使用 k = 3 时,ERF 为 3。经过 WT 后,特征图大小减半,然后使用相同的 k = 3 进行卷积,随后进行 IWT 以恢复原始特征大小,有效地将感受野加倍至 6。通过结合单分支卷积和小波变换,这种方法实现了多核和多感受野特性。

小波变换组件的效果。小波变换生成一个低频和三个高频特征图。低频图保留了原始特征信息,而高频图捕获了边缘细节。经过小波变换后,减半的特征图进行卷积,然后进行逆小波变换,恢复原始大小,有效地将感受野加倍。尽管吞吐量可能会略有下降,但小波变换在扩大感受野和提取边缘信息方面的优势可以提高模型性能(见表 8f)。

5.结论

我们设计了 MobileMamba 框架,以平衡性能和效率,解决了现有 Mamba 基模型的局限性。提出的 MRFFI 模块增强了各种感受野的感知能力,同时保留了高频特征和推理效率。训练和测试策略进一步提高了性能和效率。在 ImageNet1K 数据集上的广泛实验验证了该方法的有效性、效率和在高分辨率下游任务中的可转移性。

限制和未来工作。尽管 Mamba 模型取得了进展,但它们在工程实现方面仍然存在不足,包括需要在 CPU 加速和边缘设备加速方面进行大量改进。在未来,我们将继续专注于提高 Mamba 模型在各种设备上的推理能力,特别关注效率。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

Together_CZ

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值