MobileMamba: Lightweight Multi-Receptive Visual Mamba Network——论文笔记

论文链接：https://arxiv.org/pdf/2411.15941

项目代码：https://github.com/lewandofskee/MobileMamba

一、主要内容

        本文提出了MobileMamba框架，它平衡了效率和性能。设计了一个三级网络，显著提高了推理速度。在细粒度层面，本文引入了多接受场特征交互(MRFFI)模块，包括远程小波变换增强曼巴(WTE-Mamba)、高效多核深度卷积(MK-DeConv)和消除冗余同一性组件。该模块集成了多感受野信息，增强了高频细节提取。

二、提出背景

        目前的轻量化模型主要分为基于cnn的结构和基于transformer的结构。基于cnn的MobileNets使用深度可分离卷积来降低计算复杂度，为后续基于cnn的工作奠定基础。然而，基于cnn的方法的主要缺点是它们的局部有效感受野(ERF)，如图1(i)所示，它局限于中心区域，缺乏远程相关性。在具有高分辨率输入的下游任务中，基于cnn的方法只能通过增加计算负荷来实现性能提升。

        Vision Transformers（ViTs）具有全局的ERF和 图(ii)中的远程建模能力。然而，与cnn相比它们的二次计算复杂度较高 。 这促使研究人员开发混合CNN-ViT结构，结合局部和全局ERF以提高图(iii)中的性能。然而，基于vit的方法仍然面临二次计算复杂度的问题，特别是下游任务的高分辨率输入。

        文中2.2节部分表明，尽管有各种各样的设计，基于mamba的轻量级网络都无法超越现有的CNN和ViT方法。图 2 的实验结果表明，目前基于mamba的结构存在推理速度慢、性能差的问题。
        因此，本文提出了MobileMamba，它通过粗粒度、细粒度和训练/测试策略被设计成一个高效的轻量级网络。如图1(iv)所示，MobileMamba有着全局感受野的同时，高效多核深度可分离卷积操作有助于提取相邻信息。

图 1 不同架构的有效接受域(ERF)的可视化

图 2 基于mambaba方法的准确性与速度

三、贡献

        （1）作者提出了一个轻量级的三阶段MobileMamba框架，该框架在性能和效率之间实现了良好的平衡。

        （2）作者设计了一个高效的多感受野特征交互（MRFFI）模块，以通过更大的有效感受野增强多尺度感知能力，并改进细粒度高频边缘信息的提取。

        （3）MobileMamba通过在不同FLOPs大小的模型上采用训练和测试策略，显著提升了性能和效率。

四、模型架构介绍

图 4 MobileMamba概述。(a) MobileMamba的架构。(b) 16 ×16 DownSample patchchembed。(c) MobileMamba块结构。(d)细粒度设计。提出了高效的多感受野特征交互(MRFFI)模块。

1、MobileMamba的粗粒度设计

 

图 3 粗粒度设计

         根据图3可得，在前两个实验中，四阶段网络的前两个阶段采用纯CNN架构设计，提高了推理速度。第三个实验在网络的所有四个阶段使用MobileMamba块。但实验结果表明，在相同吞吐量下，三级网络的准确率更高。同样，对于相同的性能，三级网络具有更高的吞吐量。因此，我们选择一个三阶段网络作为我们的粗粒度框架。

2、MobileMamba的细粒度设计

         在细粒度模块设计方面在3.2节，作者提出了高效多感受野特征交互 (MRFFI)模块。具体来说，输入特征沿着通道维度分为三个部分。

        第一部分使用远程小波变换增强曼巴(WTE-Mamba)模块提取全局特征，同时增强边缘信息等细粒度细节的提取。

        第二部分通过高效多核深度可分离卷积(Multi-Kernel depth Convolution, MK-DeConv)运算来增强不同感受野的感知能力。

        最后部分通过去冗余恒等映射（Identity mapping），减少高维空间下通道冗余的问题，降低计算复杂度，提高处理速度。

        通过MRFFI获得的特征融合了全局和多尺度局部感受野信息，增强了高频边缘细节的提取。最后，我们通过两种训练阶段策略(知识蒸馏和扩展训练时期)来增强模型的学习能力。此外，测试阶段的归一化层融合策略提高了模型的推理速度。

3、训练和测试策略

        本文采用两种训练策略来进一步提高小模型的性能和效率，同时保持相同的参数数量和计算复杂度。此外，本文使用测试策略来确保模型的有效性，同时提高推理速度。

        知识蒸馏。为了使轻量级学生模型MobileMamba能够从更鲁棒的的教师分类模型中学习，本文遵循了DeiT的软蒸馏设置。这涉及到最小化教师模型和学生模型的softmax输出之间的Kullback-Leibler分歧。

        { 蒸馏，特别是知识蒸馏，是机器学习中的一种技术，通常用于将大型模型（教师模型）的知识转移到小型模型（学生模型）中。这样做的目的是让学生在保持较小规模的同时，尽可能接近甚至超越教师的性能。常见的蒸馏方法包括硬蒸馏和软蒸馏。硬蒸馏可能是指学生直接模仿教师的输出标签，而软蒸馏可能涉及到概率分布层面的知识转移。

        软蒸馏（Soft Distillation）​ 是一种知识蒸馏（Knowledge Distillation）​技术，旨在将复杂模型（教师模型）的知识迁移到简单模型（学生模型）中，通过让学生模型学习教师模型的“软目标”（概率分布）而非仅依赖真实标签（硬目标），从而提升学生模型的泛化能力和性能。）

        （通过Kullback-Leibler（KL）散度​（一种衡量两个概率分布差异的指标），迫使学生模型的输出概率分布（经过Softmax）尽可能接近教师模型的输出概率分布。}

        扩展训练时期。本文观察到 常规300次，损耗小模MobileMamba还没有完全融合，Top-1的准确率还没有达到它的潜力。因此，为了提高轻量级模型的性能上限，本文将训练扩展到1000次。

        归一化层融合。卷积操作之后通常是批处理规范化。在推理过程中，批归一化可以与之前的卷积层或线性层融合。重新计算新的卷积层的权重和偏差确保其组合输出与原始层的输出匹配。这种融合通过减少层数提高了计算效率，加快了向前通过的速度。