EfficientVMamba：軽量視覚MambaのためのAtrous選択的スキャン

EfficientVMamba: Atrous Selective Scan for Light Weight Visual Mamba

March 15, 2024

著者: Xiaohuan Pei, Tao Huang, Chang Xu

cs.AI

要旨

軽量モデル開発におけるこれまでの取り組みは主にCNNとTransformerベースの設計に焦点を当ててきたが、依然として課題が残っている。CNNは局所的特徴抽出に優れているが解像度を犠牲にし、Transformerはグローバルな範囲を提供するが計算量がO(N^2)と増大する。この精度と効率性の間のトレードオフは、依然として大きな障壁となっている。最近、Mambaなどの状態空間モデル（SSM）が、言語モデリングやコンピュータビジョンなどのさまざまなタスクで優れた性能と競争力を示し、グローバル情報抽出の時間計算量をO(N)に削減している。これに触発され、本研究では、軽量モデル設計における視覚状態空間モデルの可能性を探り、EfficientVMambaと呼ばれる新しい効率的なモデルバリアントを提案する。具体的には、EfficientVMambaは、効率的なスキップサンプリングによるatrousベースの選択的スキャンアプローチを統合し、グローバルおよび局所的な表現的特徴を活用するように設計されたビルディングブロックを構成する。さらに、SSMブロックと畳み込みの統合を調査し、追加の畳み込みブランチと組み合わせた効率的な視覚状態空間ブロックを導入し、モデルの性能をさらに向上させる。実験結果は、EfficientVMambaが計算複雑性を削減しながら、さまざまな視覚タスクで競争力のある結果を生み出すことを示している。例えば、1.3G FLOPsのEfficientVMamba-Sは、1.5G FLOPsのVim-TiをImageNetで5.6%の精度で大幅に改善する。コードはhttps://github.com/TerryPei/EfficientVMambaで公開されている。

English

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 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 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.3G FLOPs improves Vim-Ti with 1.5G FLOPs by a large margin of 5.6% accuracy on ImageNet. Code is available at: https://github.com/TerryPei/EfficientVMamba.