Vision Mamba: Effizientes Lernen visueller Repräsentationen mit bidirektionalem State-Space-Modell

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

January 17, 2024

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

cs.AI

Zusammenfassung

Kürzlich haben State-Space-Modelle (SSMs) mit effizienten, hardwarebewussten Designs, wie z. B. Mamba, großes Potenzial für die Modellierung langer Sequenzen gezeigt. Die Entwicklung effizienter und generischer Vision-Backbones, die ausschließlich auf SSMs basieren, ist eine vielversprechende Richtung. Die Darstellung visueller Daten stellt jedoch eine Herausforderung für SSMs dar, da visuelle Daten positionssensitiv sind und ein globaler Kontext für das visuelle Verständnis erforderlich ist. In diesem Artikel zeigen wir, dass die Abhängigkeit des visuellen Repräsentationslernens von Self-Attention nicht notwendig ist, und schlagen einen neuen generischen Vision-Backbone mit bidirektionalen Mamba-Blöcken (Vim) vor, der die Bildsequenzen mit Positions-Einbettungen markiert und die visuelle Repräsentation mit bidirektionalen State-Space-Modellen komprimiert. Bei Aufgaben wie ImageNet-Klassifikation, COCO-Objekterkennung und ADE20k-Semantiksegmentierung erzielt Vim eine höhere Leistung im Vergleich zu etablierten Vision-Transformern wie DeiT und zeigt gleichzeitig eine deutlich verbesserte Rechen- und Speichereffizienz. Beispielsweise ist Vim 2,8-mal schneller als DeiT und spart 86,8 % GPU-Speicher, wenn Batch-Inferenz zur Merkmalsextraktion auf Bildern mit einer Auflösung von 1248×1248 durchgeführt wird. Die Ergebnisse zeigen, dass Vim in der Lage ist, die Rechen- und Speicherbeschränkungen bei der Durchführung von Transformer-artigem Verständnis für hochauflösende Bilder zu überwinden, und großes Potenzial hat, das nächste Backbone für Vision-Foundation-Modelle zu werden. Der Code ist verfügbar unter 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.