MambaVision: 하이브리드 맘바-트랜스포머 비전 백본

MambaVision: A Hybrid Mamba-Transformer Vision Backbone

July 10, 2024

저자: Ali Hatamizadeh, Jan Kautz

cs.AI

초록

우리는 비전 애플리케이션에 특화된 새로운 하이브리드 Mamba-Transformer 백본을 제안하며, 이를 MambaVision이라고 명명합니다. 우리의 핵심 기여는 시각적 특징을 효율적으로 모델링하기 위해 Mamba 공식을 재설계한 것입니다. 또한, Vision Transformers(ViT)와 Mamba를 통합하는 것의 타당성에 대한 포괄적인 절제 연구를 수행했습니다. 우리의 결과는 Mamba 아키텍처의 최종 레이어에 여러 개의 self-attention 블록을 추가함으로써 장거리 공간 의존성을 포착하는 모델링 능력이 크게 향상됨을 보여줍니다. 이러한 발견을 바탕으로, 다양한 설계 기준을 충족하기 위해 계층적 아키텍처를 가진 MambaVision 모델군을 소개합니다. ImageNet-1K 데이터셋에서의 이미지 분류 작업에서 MambaVision 모델 변종들은 Top-1 정확도와 이미지 처리량 측면에서 새로운 최첨단(SOTA) 성능을 달성했습니다. MS COCO 및 ADE20K 데이터셋에서의 객체 탐지, 인스턴스 분할, 의미론적 분할과 같은 다운스트림 작업에서 MambaVision은 비슷한 크기의 백본들을 능가하며 더 우수한 성능을 보여줍니다. 코드: https://github.com/NVlabs/MambaVision.

English

We propose a novel hybrid Mamba-Transformer backbone, denoted as MambaVision, which is specifically tailored for vision applications. Our core contribution includes redesigning the Mamba formulation to enhance its capability for efficient modeling of visual features. In addition, we conduct a comprehensive ablation study on the feasibility of integrating Vision Transformers (ViT) with Mamba. Our results demonstrate that equipping the Mamba architecture with several self-attention blocks at the final layers greatly improves the modeling capacity to capture long-range spatial dependencies. Based on our findings, we introduce a family of MambaVision models with a hierarchical architecture to meet various design criteria. For Image classification on ImageNet-1K dataset, MambaVision model variants achieve a new State-of-the-Art (SOTA) performance in terms of Top-1 accuracy and image throughput. In downstream tasks such as object detection, instance segmentation and semantic segmentation on MS COCO and ADE20K datasets, MambaVision outperforms comparably-sized backbones and demonstrates more favorable performance. Code: https://github.com/NVlabs/MambaVision.