MambaVision: Un'architettura ibrida Mamba-Transformer per la visione artificiale

MambaVision: A Hybrid Mamba-Transformer Vision Backbone

July 10, 2024

Autori: Ali Hatamizadeh, Jan Kautz

cs.AI

Abstract

Proponiamo una nuova architettura ibrida Mamba-Transformer, denominata MambaVision, specificamente progettata per applicazioni visive. Il nostro contributo principale include la riprogettazione della formulazione Mamba per migliorarne la capacità di modellazione efficiente delle caratteristiche visive. Inoltre, conduciamo uno studio di ablazione completo sulla fattibilità di integrare Vision Transformers (ViT) con Mamba. I nostri risultati dimostrano che dotare l'architettura Mamba di diversi blocchi di self-attention negli strati finali migliora significativamente la capacità di modellazione per catturare dipendenze spaziali a lungo raggio. Sulla base delle nostre scoperte, introduciamo una famiglia di modelli MambaVision con un'architettura gerarchica per soddisfare vari criteri di progettazione. Per la classificazione di immagini sul dataset ImageNet-1K, le varianti del modello MambaVision raggiungono una nuova prestazione State-of-the-Art (SOTA) in termini di accuratezza Top-1 e throughput delle immagini. In task downstream come il rilevamento di oggetti, la segmentazione di istanze e la segmentazione semantica sui dataset MS COCO e ADE20K, MambaVision supera architetture di dimensioni comparabili e dimostra prestazioni più favorevoli. Codice: 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.