MambaVision: ハイブリッドMamba-Transformerビジョンバックボーン

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アーキテクチャの最終層にいくつかのセルフアテンションブロックを追加することで、長距離空間依存性を捉えるモデリング能力が大幅に向上することが示された。これらの知見に基づき、様々な設計基準を満たす階層型アーキテクチャを持つMambaVisionモデルファミリーを導入した。ImageNet-1Kデータセットにおける画像分類では、MambaVisionモデルのバリエーションがTop-1精度と画像スループットの両面で新たなState-of-the-Art (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.