[논문 리뷰] Multi-Scale VMamba: Hierarchy in Hierarchy Visual State Space Model

Yuheng Shi, Minjing Dong외 1명|arXiv (Cornell University)|2024. 05. 23.

Image Retrieval and Classification Techniques인용 수 22

한 줄 요약

소개 Multi-Scale VMamba (MSVMamba), 다중 스케일 2D 스캐닝 (MS2D), 계층적 MS3 블록, 및 ConvFFN를 결합하여 매개변수가 제한된 상태에서 장거리 의존성 학습을 개선; ImageNet, COCO, ADE20K에서 SSM 기반 백본 중 최첨단 성능 달성.

ABSTRACT

Despite the significant achievements of Vision Transformers (ViTs) in various vision tasks, they are constrained by the quadratic complexity. Recently, State Space Models (SSMs) have garnered widespread attention due to their global receptive field and linear complexity with respect to the input length, demonstrating substantial potential across fields including natural language processing and computer vision. To improve the performance of SSMs in vision tasks, a multi-scan strategy is widely adopted, which leads to significant redundancy of SSMs. For a better trade-off between efficiency and performance, we analyze the underlying reasons behind the success of the multi-scan strategy, where long-range dependency plays an important role. Based on the analysis, we introduce Multi-Scale Vision Mamba (MSVMamba) to preserve the superiority of SSMs in vision tasks with limited parameters. It employs a multi-scale 2D scanning technique on both original and downsampled feature maps, which not only benefits long-range dependency learning but also reduces computational costs. Additionally, we integrate a Convolutional Feed-Forward Network (ConvFFN) to address the lack of channel mixing. Our experiments demonstrate that MSVMamba is highly competitive, with the MSVMamba-Tiny model achieving 82.8% top-1 accuracy on ImageNet, 46.9% box mAP, and 42.2% instance mAP with the Mask R-CNN framework, 1x training schedule on COCO, and 47.6% mIoU with single-scale testing on ADE20K.Code is available at \url{https://github.com/YuHengsss/MSVMamba}.

연구 동기 및 목표

• 매개변수로 제한된 비전 모델에서 SSM을 사용한 장거리 망각 문제를 해결한다.
• 중복 감소 및 미세한 정보 유지하기 위한 계층적 다중 스케일 스캐닝 전략 개발.
• SSM 기반 백본 내 채널 혼합 및 로컬 특성 추출 향상을 위해 ConvFFN 통합.

제안 방법

• VMamba의 SS2D를 다중 스케일 상태 공간(MS3) 블록으로 교체하여 MS2D 스캐닝과 ConvFFN 채널 믹서를 포함한다.
• 깊이wise 합성? 함께 설명된 MS2D를 strides 1 및 s로 사용해 다중 스케일 피처 맵 생성, 전체 해상도 맵과 다운샘플 맵을 S6 블록을 통해 처리하고 결과를 집계한다.
• MS2D 후 SE 블록을 Incorporate하고 ConvFFN (depthwise conv + two FC layers)으로 채널 간 정보 교환을 강화한다.
• Comparable FLOPs를 보장하기 위해 임베딩 차원 및 블록 수를 제어하여 LeViT-유사 예산과의 공정한 비교를 가능하게 한다.
• Nano, Micro, 및 Tiny 모델 변형을 제공하여 6.9M–33.0M 매개변수 및 0.9–4.6 GFLOPs로 확장 가능한 배치를 지원한다.

Figure 1 : Comparison on ImageNet.

실험 결과

연구 질문

• RQ1 다중 스케일 2D 스캐닝을 어떻게 설계하면 SSM 기반 비전 백본의 중복을 줄이고 장거리 의존성 학습을 개선할 수 있는가?
• RQ2 ConvFFN 및 SE 블록의 도입이 고정된 계산 예산 하에서 채널 간 정보 교환 및 전체 정확도에 미치는 영향은 무엇인가?
• RQ3 계층적 다중 스케일 VMamba 설계가 이미지넷, COCO, ADE20K에서 기존 VMamba 변형 및 다른 SOTA 백본보다 효율성을 유지하면서 더 나은 성능을 발휘할 수 있는가?

주요 결과

• MSVMamba-T는 ImageNet-1K에서 33M 매개변수 및 4.6 GFLOPs로 82.8%의 Top-1 정확도를 달성한다.
• MSVMamba-Nano, -Nano baseline은 대략 유사한 FLOPs에서 VMamba-Nano보다 최대 5.5 포인트 높은 Top-1 정확도를 달성한다.
• MSVMamba-T는 거의 같은 계산 비용에서 VMamba-T보다 0.6%의 Top-1 정확도를 상회한다.
• COCO 객체 검출에서 MSVMamba-T는 1x 일정에서 Swin-T보다 박스 AP +4.2, 마스크 AP +2.9를 상회한다.
• ADE20K 의미론적 분할에서 MSVMamba-T는 단일 스케일 테스트에서 47.6 mIoU(멀티스케일 테스트 시 48.5).
• MS2D와 SE 블록이 포함된 ConvFFN이 정확도 향상을 가져다주며(예: SE로 Top-1에 +0.5% 추가; Ablation에서 ConvFFN으로 +2.0%).

(a) Horizontal Scan