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 и пропускной способности изображения. В задачах наследования, таких как обнаружение объектов, сегментация экземпляров и семантическая сегментация на наборах данных 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.