轻量级视觉骨干网络 MobileMamba: Lightweight Multi-Receptive Visual Mamba Network

快速链接

原文 : https://arxiv.org/pdf/2411.15941
代码 : https://github.com/lewandofskee/MobileMamba

解决问题：视觉模型在移动设备端性能和效果的平衡性

因为移动设备的算力受限，移动设备上常使用轻量级网络。在图像领域，常用的轻量级网络主要分基于CNN的网络，以及基于vision-transformer的网络，但都有其局限性。基于CNN的网络，速度快，效果好，但无法捕捉长距的关系。基于vision-transformer的网络可捕捉长距离间关系，但其计算复杂度与输入图尺寸的平方成正比。state-space模型可以用线性复杂度的计算量捕获长距关系。（state-space和Mamba模型应存在某些关系）。已有的基于Mamba的网络flops低，但推理速度慢，效果差。本文提出的MobileMamba可以平衡效果和性能, imagenet-1K Top1 83.6% ，超最优。GPU上速度比localVim快21倍。

解决方法：改进网络结构+训练和测试策略

本文通过改进Mamba网络结构，使用一些训练和测试策略，达到对Mamba模型的轻量化，并且保证其优异的效果。

网络结构改进

粗颗粒上：本文对比了常用的4阶段网络与提出的3阶段网络在输出分辨率上的差异。用imagenet-1k的图像分类实验结果说明，提出的3阶段网络在推理速度，准确性上都更优。感觉存粹是一种基于实验经验判断的结构改进。

细颗粒度上：本文提出了MRFFI模块（Multi-Receptive Field Feature Interaction， 多感受野特征交互）。输入特征在channel维度分配到以下三部分。
（1）long-range wavelet trasform-enhanced mamba（长程基于小波转换的mamba），用于提取全局特征，增强细颗粒度细节提取。
（2）multi-kernel depthwise convolution（多核深度可分离卷积），用于捕获多尺度的感受野。
（3）identity mapping（恒等映射，就是输出和输入相同），用于减少高分辨率空间通道冗余，降低计算复杂度，提高处理速度。

训练和测试策略

训练方面，使用了2个技巧。
（1）知识蒸馏：最小化教师模型和学生模型的 Softmax 输出之间概率分布差异。
（2）延长训练epochs: 训练的epoch数量从传统的300延长到1000，让此模型收敛。
测试方面，将批归一化层融合到前面卷积或线性层的权重、偏置中，并且保证融合前后计算结果一致。这样，通过减少计算层数，提高计算速度。

实验支撑：图像分类、分割，目标检测等

图像分类结果对比

目标检测和实例分割结果对比

语义分割结果对比