[论文解读] EfficientVMamba: Atrous Selective Scan for Light Weight Visual Mamba

Xiaohuan Pei, Tao Huang+1 more|arXiv (Cornell University)|Mar 15, 2024

Advanced Vision and Imaging被引用 12

一句话总结

EfficientVMamba 引入基于膨胀卷积的选择性扫描策略（ES2D）和双路径 EVSS 块，将全局状态空间建模与局部卷积融合，在保持竞争力准确度的同时实现更低 FLOPs 的视觉任务。

ABSTRACT

Prior efforts in light-weight model development mainly centered on CNN and Transformer-based designs yet faced persistent challenges. CNNs adept at local feature extraction compromise resolution while Transformers offer global reach but escalate computational demands $\mathcal{O}(N^2)$. This ongoing trade-off between accuracy and efficiency remains a significant hurdle. Recently, state space models (SSMs), such as Mamba, have shown outstanding performance and competitiveness in various tasks such as language modeling and computer vision, while reducing the time complexity of global information extraction to $\mathcal{O}(N)$. Inspired by this, this work proposes to explore the potential of visual state space models in light-weight model design and introduce a novel efficient model variant dubbed EfficientVMamba. Concretely, our EfficientVMamba integrates a atrous-based selective scan approach by efficient skip sampling, constituting building blocks designed to harness both global and local representational features. Additionally, we investigate the integration between SSM blocks and convolutions, and introduce an efficient visual state space block combined with an additional convolution branch, which further elevate the model performance. Experimental results show that, EfficientVMamba scales down the computational complexity while yields competitive results across a variety of vision tasks. For example, our EfficientVMamba-S with $1.3$G FLOPs improves Vim-Ti with $1.5$G FLOPs by a large margin of $5.6\%$ accuracy on ImageNet. Code is available at: \url{https://github.com/TerryPei/EfficientVMamba}.

研究动机与目标

• 激励在不高计算成本的前提下保持全局上下文的轻量级视觉模型。
• 提出 ES2D，以在降低扫描复杂度的同时保持全局感受野。
• 引入通过 SE 融合将全局状态空间表示与局部卷积相融合的 EVSS 块。
• 探索倒置残差插入以优化各阶段的块放置。
• 展示在图像分类、目标检测和语义分割等任务中的有效性。

提出的方法

• 引入基于膨胀的选择性扫描（ES2D），通过跳跃采样将扫描的令牌数量从 N 降低到 N/p^2。
• 开发 Efficient Visual State Space（EVSS）块，将基于 ES2D 的全局特征与 3x3 卷积分支以及 SE 重校准融合。
• 通过在 SE 之后的逐元素求和融合全局与局部路径，得到 X^{l+1} = SE(ES2D(X^l)) + SE(Conv(X^l))。
• 采用倒置插入：在早期阶段放置 EVSS 块以获取全局表示，在更深的阶段放置 InRes 块以获取局部特征。
• 提供三种模型变体（EfficientVMamba-T、-S、-B），FLOPs 和参数逐步提高。

实验结果

研究问题

• RQ1 ES2D 能否在保持全局上下文的同时降低全局扫描的计算成本？
• RQ2 将全局 ES2D 路径与局部卷积分支耦合是否在资源受限的情况下提升准确性？
• RQ3 在轻量级模型中，倒置残差插入是否有利于将基于 SSM 的块与 CNN 块结合？
• RQ4 与现有轻量级主干网相比，EfficientVMamba 变体在 ImageNet 分类、COCO 目标检测和 ADE20K 语义分割上的表现如何？

主要发现

• EfficientVMamba-T/S/B 在低 FLOPs（分别为 0.8/1.3/4.0 GFLOPs）下实现具有竞争力的 ImageNet 准确率。
• EfficientVMamba-S 在 ImageNet 上以 1.3 GFLOPs 达到 78.7% 的 top-1，优于若干更大尺寸的骨干网。
• EfficientVMamba-B 在 ImageNet 上以 4.0 GFLOPs 和 3300万参数达到 81.8% 的 top-1。
• 在 COCO RetinaNet 实验中，EfficientVMamba-T 达到 37.5 AP，EfficientVMamba-B 达到 42.8 AP，参数数量比某些基线更小。
• 在 ADE20K 语义分割中，EfficientVMamba 变体超越了更重的模型，并具备具有竞争力的 mIoU 分数（例如，跨变体的 SS 测试为 46.5%–46.5%+）。
• 消融实验表明 ES2D 在降低 FLOPs 的同时保持准确性，SE 融合提升了性能；在早期阶段的倒置插入更好地利用全局特征。