[轻量化]MobileMamba: Lightweight Multi-Receptive Visual Mamba Network, CVPR 2025

MobileMamba: Lightweight Multi-Receptive Visual Mamba Network, CVPR 2025

MobileMamba：轻量级多感受野 Visual Mamba 网络
2024 年 11 月 24 日
paper&code

提出了 MobileMamba 框架，它平衡了效率和性能。我们设计了一个三阶段网络，以显著提高推理速度。在细粒度层面上，引入了多感受野特征交互 （MRFFI） 模块，包括长距离小波变换增强 Mamba （WTE-Mamba）、高效多核深度卷积 （MK-DeConv） 和消除冗余身份组件。该模块集成了多感受野信息，并增强了高频细节提取。

粗粒度下的pipeline

图中(d)MRFFI模块三个分支的实现（由下而上）
1）第一部分经历了远程小波变换增强的Mamba（WTE-MAMBA），在执行全局建模时，它可以提取高频边缘信息。
2）第二部分是通过多内核深度卷积（MK-DECONV）操作来处理的，以增强不同感受野的感知能力。
3）其余的效果受到identity映射的影响，减少高维空间中的特征延长，并降低组合复杂性，从而提高处理速度。

具体来说

1. WTE-Mamba分支
   增强根据全局建模提取细粒度信息（例如高频边缘细节）的能力。WT特征图上的卷积操作具有更大的ERF，和较低的计算复杂性。对于输入特征$x^I\in {\mathbb{R}}^{h\times w\times c}$，特征$x_G^I\in {\mathbb{R}}^{h\times w\times \xi c}$通过双向扫描Mamba模块处理以学习全局信息，并且$0\le \xi \le 1$。
   
   同时，相同的特征映射经历了HAAR小波转换以获得特征表示$x_w^I\in {\mathbb{R}}^{\frac{h}{2}\times \frac{w}{2}\times 4\xi c}$，然后执行本地卷积信息提取和逆小波转换（IWT）以恢复原始特征映射大小
   
   将全局信息提取的Mamba模块的输出功能映射添加到从小波转换和卷积的本地信息的输出特征图中
   
2. Efficient MK-DeConv
   通过不同的ERF提取局部信息，从而实现了多感受野。对于其余特征，选择$x_L^I\in {\mathbb{R}}^{h\times w\times \mu c}$，其中局部通道比例表示为$\mu \le 1-\xi$。然后将这些通道分为n∈N部分。每个部分$n\in \mathbb{N}$。每个部分$x_{Lj}^I\in {\mathbb{R}}^{h\times w\times \frac{\mu c}{n}}$进行不同核大小的局部卷积操作。最后，将不同的卷积操作的结果组合，以形成输出特征$x_L^O\in {\mathbb{R}}^{h\times w\times \mu c}$。
   
3. 消除冗余Identity
   最后，为了降低高维空间中特征冗余，我们将identity mapping应用于其余$(1-\xi -\mu )c$的通道。

最终的输出

Training and Testing Strategies

知识蒸馏为了使轻量化的学生模型MobileMamba从更强大的教师分类模型中学习，我们遵循DEIT [64]的软蒸馏设置。这涉及最大程度地减少教师模型和学生模型的软性输出之间的kullback-leibler差异。
扩展训练epoch我们观察到，在常规的300个时期下，小型模型的损失尚未完全融合，并且Top-1的准确性尚未达到其潜力。因此，为了提高轻量级模型的功能上限，我们将训练扩展到1000个时代。
归一化层融合卷积操作通常紧随其后进行批准。在推断过程中，可以将批准化与先前的汇合或线性层融合。重新计算新的卷积层的重量和偏见可确保其组合输出与原始层的输出相匹配。这种融合可以通过减少层数来提高计算效率，并加快前向过程的速度。