[논문 리뷰] LocalMamba: Visual State Space Model with Windowed Selective Scan

Tao Huang, Xiaohuan Pei외 4명|arXiv (Cornell University)|2024. 03. 14.

Data Visualization and Analytics인용 수 31

한 줄 요약

LocalMamba는 Vision Mamba에 윈도우드 로컬 스캐닝과 계층별 스캔 방향 탐색을 도입하여 ImageNet, COCO, ADE20K에서 우수한 성능과 효율적인 계산을 달성합니다.

ABSTRACT

Recent advancements in state space models, notably Mamba, have demonstrated significant progress in modeling long sequences for tasks like language understanding. Yet, their application in vision tasks has not markedly surpassed the performance of traditional Convolutional Neural Networks (CNNs) and Vision Transformers (ViTs). This paper posits that the key to enhancing Vision Mamba (ViM) lies in optimizing scan directions for sequence modeling. Traditional ViM approaches, which flatten spatial tokens, overlook the preservation of local 2D dependencies, thereby elongating the distance between adjacent tokens. We introduce a novel local scanning strategy that divides images into distinct windows, effectively capturing local dependencies while maintaining a global perspective. Additionally, acknowledging the varying preferences for scan patterns across different network layers, we propose a dynamic method to independently search for the optimal scan choices for each layer, substantially improving performance. Extensive experiments across both plain and hierarchical models underscore our approach's superiority in effectively capturing image representations. For example, our model significantly outperforms Vim-Ti by 3.1% on ImageNet with the same 1.5G FLOPs. Code is available at: https://github.com/hunto/LocalMamba.

연구 동기 및 목표

• 비전 상태 공간 모델에서 로컬 의존성을 보존하기 위해 로컬 2D 인식 스캐닝의 사용을 촉진한다.
• 이미지 내 로컬 구조를 더 잘 포착하면서도 글로벌 컨텍스트를 유지하기 위해 윈도우화된 로컬 스캔 메커니즘을 제안한다.
• 네트워크 깊이에 걸쳐 스캔 패턴을 최적화하기 위해 학습 가능한 계층별 스캔 방향 탐색을 도입한다.
• Plain 및 계층적 모델 변형(LocalVim 및 LocalVMamba)을 개발하여 확장성 및 효과를 검증한다.
• 분류, 탐지, 분할 작업에서 Vim, VMamba, CNNs, ViTs에 비해 개선을 보여준다.

제안 방법

• 입력 특징을 병렬 로컬 윈도우에서 처리하고 Spatial-Channel Attention (SCAttn)으로 집계하는 네 가지 분기 로컬 스캔 블록으로 LocalMamba를 도입한다.
• 여덟 개 후보 방향(수평, 수직, 2x2/7x7 로컬 윈도우 각각 표준 및 뒤집힌 형태)을 확장하고 미분 가능 탐색(DARTS 방식)으로 각 계층당 네 방향을 선택한다.
• 학습 중 여러 SSM을 계층당 결합하기 위해 연속 이완을 사용하고 추론 시 상위 4개 방향을 선택한다.
• Plain 변형인 LocalVim과 계층적 변형인 LocalVMamba 두 가지 아키텍처를 제공하고 Vim/VMamba 블록을 LocalMamba 블록으로 대체한다.
• ImageNet 분류, COCO 객체 탐지/분할 및 ADE20K 시맨틱 분할에서의 결과를 보고 이익을 나타낸다.

Figure 1 : Illustration of scan methods. (a) and (b): Previous methods Vim [ 60 ] and VMamba [ 32 ] traverse the entire row or column axis, resulting in significant distances for capturing dependencies between neighboring pixels within the same semantic region ( e.g. , the left eye in the image). (c

실험 결과

연구 질문

• RQ1 윈도우 기반의 로컬 스캐닝 전략이 글로벌 컨텍스트를 유지하면서 Vision Mamba 모델에서 로컬 2D 의존성을 보존하는 데 개선을 가져올 수 있는가?
• RQ2 계층별로 최적의 스캔 방향을 탐색하는 것이 고정되거나 단일 방향 스캔보다 측정 가능한 이점을 가져오는가?
• RQ3 Plain 및 계층적 LocalMamba 변형이 Vim, VMamba, CNNs, ViTs에 비해 분류, 탐지 및 분할 작업에서 어떻게 수행하는가?

주요 결과

Method	Image size	Params (M)	FLOPs (G)	Top-1 ACC (%)
LocalVim-T	224^2	8	1.5	76.2
LocalVim-S	224^2	28	4.8	81.2
VMamba-T	224^2	22	5.6	82.2
VMamba-S	224^2	44	11.2	83.5
LocalVMamba-T	224^2	26	5.7	82.7
LocalVMamba-S	224^2	50	11.4	83.7

• ImageNet-1K에서 LocalVim-T는 76.2% Top-1 정확도와 1.5G FLOPs로 DeiT-Ti(72.2%)를 능가한다.
• 계층적 LocalVMamba-T는 82.7% 정확도에 도달하여 Swin-T보다 1.4% 포인트 앞선다.
• LocalVim-S 및 LocalVMamba-S는 분류에서 Vim 및 VMamba 베이스라인보다 큰 이점을 보인다.
• COCO 객체 탐지에서 LocalVMamba-T는 46.7 APb와 42.2 APm로 Swin-T보다 여유 있게 앞선다.
• ADE20K 분할에서 LocalVim-S는 46.4 mIoU (SS)이고 LocalVMamba-S는 50.0 mIoU (SS) / 51.0 mIoU (MS) 다양한 설정에서.
• 절단 분석은 로컬 스캔이 Vim-T를 1.0% 개선하고 SCAttn이 약 0.6%의 이득을 부여한다는 것을 보여준다.

Figure 2 : By extending the original scan with our local scan mechanism, our method significantly improves the ImageNet accuracies of Vim [ 60 ] while keeping similar FLOPs.