Vision Mamba: Efficiënt leren van visuele representaties met een bidirectioneel toestandsruimtemodel

Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

January 17, 2024

Auteurs: Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, Xinggang Wang

cs.AI

Samenvatting

Onlangs hebben state space models (SSMs) met efficiënte hardwarebewuste ontwerpen, zoals Mamba, groot potentieel getoond voor het modelleren van lange sequenties. Het bouwen van efficiënte en generieke visuele backbones uitsluitend gebaseerd op SSMs is een aantrekkelijke richting. Het representeren van visuele data is echter uitdagend voor SSMs vanwege de positiegevoeligheid van visuele data en de vereiste van globale context voor visueel begrip. In dit artikel tonen we aan dat de afhankelijkheid van visuele representatielearning van self-attention niet noodzakelijk is en stellen we een nieuwe generieke visuele backbone voor met bidirectionele Mamba-blokken (Vim), die de beeldsequenties markeert met position embeddings en de visuele representatie comprimeert met bidirectionele state space models. Bij ImageNet-classificatie, COCO-objectdetectie en ADE20k-semantische segmentatietaken behaalt Vim betere prestaties in vergelijking met gevestigde vision transformers zoals DeiT, terwijl het ook een aanzienlijk verbeterde rekenefficiëntie en geheugenefficiëntie demonstreert. Zo is Vim 2,8 keer sneller dan DeiT en bespaart het 86,8% GPU-geheugen bij batchinferentie om features te extraheren uit afbeeldingen met een resolutie van 1248×1248. De resultaten tonen aan dat Vim in staat is om de rekenefficiëntie- en geheugenbeperkingen te overwinnen bij het uitvoeren van Transformer-stijl begrip voor hoogresolutiebeelden en dat het groot potentieel heeft om de volgende generatie backbone te worden voor vision foundation models. Code is beschikbaar op https://github.com/hustvl/Vim.

English

Recently the state space models (SSMs) with efficient hardware-aware designs, i.e., Mamba, have shown great potential for long sequence modeling. Building efficient and generic vision backbones purely upon SSMs is an appealing direction. However, representing visual data is challenging for SSMs due to the position-sensitivity of visual data and the requirement of global context for visual understanding. In this paper, we show that the reliance of visual representation learning on self-attention is not necessary and propose a new generic vision backbone with bidirectional Mamba blocks (Vim), which marks the image sequences with position embeddings and compresses the visual representation with bidirectional state space models. On ImageNet classification, COCO object detection, and ADE20k semantic segmentation tasks, Vim achieves higher performance compared to well-established vision transformers like DeiT, while also demonstrating significantly improved computation & memory efficiency. For example, Vim is 2.8times faster than DeiT and saves 86.8% GPU memory when performing batch inference to extract features on images with a resolution of 1248times1248. The results demonstrate that Vim is capable of overcoming the computation & memory constraints on performing Transformer-style understanding for high-resolution images and it has great potential to become the next-generation backbone for vision foundation models. Code is available at https://github.com/hustvl/Vim.