EfficientVMamba: Atrous Selective Scan for Light Weight Visual Mamba

2024年03月15日
先前轻量级模型的开发主要集中在CNN和Transformer设计上,但仍面临着持续的挑战。CNN擅长局部特征提取,但会降低分辨率,而Transformer在提供全局信息的同时也会增加计算需求,达到$\mathcal{O}(N^2)$。这种精度和效率之间的权衡仍然是一个重要的障碍。最近,状态空间模型(SSMs)如Mamba在各种任务中表现出色,例如语言建模和计算机视觉,同时将全局信息提取的时间复杂度降低到$\mathcal{O}(N)$。受此启发,本文提出探索视觉状态空间模型在轻量级模型设计中的潜力,并引入一种新的高效模型变体,称为EfficientVMamba。具体而言,我们的EfficientVMamba集成了一种基于孔卡扫描的选择性扫描方法,通过高效的跳跃采样构成了旨在利用全局和局部表示特征的构建块。此外,我们研究了SSM块和卷积之间的集成,并引入了一个高效的视觉状态空间块,结合一个额外的卷积分支,进一步提高了模型性能。实验结果表明,EfficientVMamba在缩小计算复杂度的同时,在各种视觉任务中产生了竞争力的结果。例如,我们的EfficientVMamba-S使用1.3G FLOPs,在ImageNet上的Vim-Ti使用1.5G FLOPs,精度提高了5.6%。代码可在以下网址找到:\url{https://github.com/TerryPei/EfficientVMamba}。
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}.
许愿