轻量化MobileMamba视觉模型来了

今天小编分享的科学经验：轻量化MobileMamba视觉模型来了，欢迎阅读。

浙大、腾讯优图、华中科技大学的团队，提出轻量化 MobileMamba！

既良好地平衡了效率与效果，推理速度远超现有基于 Mamba 的模型。

一直以来，轻量化模型研究的主阵地都在 CNN 和 Transformer 的设计。

但 CNN 的局部有效感受野在高分辨率输入时，难以获得长距离依赖；尽管 Transformer 有着全局建模能力，但是其平方级计算复杂度，限制了其在高分辨率下的轻量化应用。

最近的状态空间模型如 Mamba，因其线性计算复杂度和出色的效果被广泛用在视觉领網域。

然而，基于 Mamba 的轻量化模型虽然 FLOPs 低，但是实际的吞吐量极低。

△最近基于 CNN/Transformer/Mamba 方法的效果 vs. FLOPs 对比

团队首先在粗粒度上设计了三阶段网络显著提升推理速度。

随后在细粒度上提出了高效多感受野特征互動（MRFFI）模块包含长距离小波变换增强 Mamba （WTE-Mamba）、高效多核深度可分离卷积（MK-DeConv）和去冗余恒等映射三个部分——这有利于在长距离建模的特征上融合多尺度多感受野信息并加强高频细节特征提取。

最后，使用两个训练和一个推理策略，进一步提升模型的性能与效率。

大量实验验证，MobileMamba 在 ImageNet-1K 数据集上的 Top - 1 准确率最高可达 83.6，且速度是 LocalVim 的 21 倍、EfficientVMamba 的 3.3 倍。

同时，大量的下游任务实验也验证了该方法在高分辨率输入情况下，取得了效果与效率的最佳平衡。

现存缺陷：成本高、速度低

随着移动设备的普及，资源受限环境中对高效、快速且准确的视觉处理需求日益增长。

开发轻量化模型，有助于显著降低计算和存储成本，还能提升推理速度，从而拓展技术的应用范围。

现有被广泛研究的轻量化模型，主要被分为基于 CNN 和 Transformer 的结构。

基于 CNN 的 MobileNet，设计了深度可分离卷积大幅度减少了计算复杂度；GhostNet 提出将原本将原本全通道 1x1 卷积替换为半数通道进行廉价计算，另半数通道直接恒等映射。

这些方法给后续基于 CNN 的工作奠定了良好的基础。

但是基于 CNN 方法的主要缺陷在于其局部感受野，如图 ( i ) 所示，其 ERF 仅在中间区網域而缺少远距离的相关性。

并且在下游任务高分辨率输入下，基于 CNN 的方法仅能通过堆叠计算量来换取性能的少量提升。

如图 ( ii ) 所示，ViT 有着全局感受野和长距离建模能力。但由于其平方级别的计算复杂度，计算开销比 CNN 更大。

一些工作尝试从减少分辨率或者减少通道数上，来改减少所带来的计算复杂度的增长，也取得了出色的效果。

不过，基于纯 ViT 的结构缺少了归纳偏置，因此，越来越多的研究者将 CNN 与 Transformer 结合得到混合结构，获得更好的效果，并获得局部和全局的感受野（如图 ( iii ) ）。

不过，尤其在下游任务高分辨率输入下，基于 ViT 的方法仍然受到平方级别计算复杂度的问题。

提出 MobileMamba

最近，由于状态空间模型捕捉长距离依赖关系并且线性的计算复杂度表现出色，引起了广泛关注，大量研究者将其应用于视觉领網域，效果和效率都取得了出色的效果。

基于 Mamba 的轻量化模型 LocalMamba 提出了将影像划分为視窗并在視窗内局部扫描的方式减少计算复杂度，而 EfficientVMamba 设计了高效 2D 扫描方式来降低计算复杂度。

不过这两种模型都仅公布了 FLOPs，而FLOPs 低并不能代表推理速度快。

经实验发现（图 2），现有的基于 Mamba 结构的推理速度较慢并且效果较差。

MobileMamba 团队分别从粗粒度、细粒度和训练测试策略三个方面来设计高效轻量化网络。

首先，研究人员讨论了四阶段和三阶段在准确率、速度、FLOPs 上的权衡。

在同等吞吐量下，三阶段网络会取得更高的准确率；同样的相同效果下三阶段网络有着更高的吞吐量。

因此，团队选择三阶段网络作为 MobileMamba 的粗粒度设计框架。

在细粒度模块设计方面，研究人员提出了高效高效多感受野特征互動（MRFFI）模块。

具体来说，将输入特征根据通道维度划分三个部分。

第一部分将通过小波变换增强的 Mamba 模块提取全局特征的同时加强边缘细节等细粒度信息的提取能力。

第二部分通过高效多核深度可分离卷积操作获取多尺度感受野的感知能力。

然后部分通过去冗余恒等映射，减少高维空间下通道冗余的问题，并减少计算复杂度提高运算速度。

最终经过 MRFFI 得到的特征融合了全局和多尺度局部的多感受野信息，并且加强了边缘细节的高频信息提取能力。

最后，研究人员通过两个训练阶段策略知识蒸馏和延长训练轮数增强模型的学习能力，提升模型效果；以及一个归一化层融合的测试阶段策略提升模型的推理速度。

△MobileMamba 结构概述实验结果

实验表明，MobileMamba 有着全局感受野的同时，高效多核深度可分离卷积操作有助于提取相邻信息。

通过与 SoTA 方法的对比可知，MobileMamba 从 200M 到 4G FLOPs 的模型在使用训练策略后，在 ImageNet-1K 上的 Top-1，分别达到 76.9、78.9、80.7、82.2、83.3、83.6 效果，均超过现有基于 CNN、ViT 和 Mamba 的方法。

与同为 Mamba 的方法相比，MobileMamba 比 LocalVim 在 Top-1 上提升 0.7 ↑的同时，速度快 21 倍；比 EfficientVMamba 提升 2.0 ↑的同时速度快 3.3 ↑倍。

这均显著优于现有基于 Mamba 的轻量化模型设计。