Vision Mamba : Apprentissage efficace des représentations visuelles avec un modèle d'espace d'état bidirectionnel

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

Résumé

Récemment, les modèles d'espace d'état (SSMs) avec des conceptions efficaces adaptées au matériel, tels que Mamba, ont montré un grand potentiel pour la modélisation de longues séquences. Construire des architectures visuelles efficaces et génériques basées uniquement sur les SSMs est une direction prometteuse. Cependant, la représentation des données visuelles est un défi pour les SSMs en raison de la sensibilité à la position des données visuelles et de la nécessité d'un contexte global pour la compréhension visuelle. Dans cet article, nous montrons que la dépendance de l'apprentissage de représentation visuelle sur l'auto-attention n'est pas nécessaire et proposons une nouvelle architecture visuelle générique avec des blocs Mamba bidirectionnels (Vim), qui marque les séquences d'images avec des embeddings de position et compresse la représentation visuelle avec des modèles d'espace d'état bidirectionnels. Sur les tâches de classification ImageNet, de détection d'objets COCO et de segmentation sémantique ADE20k, Vim atteint des performances supérieures par rapport à des vision transformers bien établis comme DeiT, tout en démontrant une efficacité de calcul et de mémoire significativement améliorée. Par exemple, Vim est 2,8 fois plus rapide que DeiT et économise 86,8 % de mémoire GPU lors de l'inférence par lots pour extraire des caractéristiques sur des images avec une résolution de 1248x1248. Les résultats démontrent que Vim est capable de surmonter les contraintes de calcul et de mémoire pour effectuer une compréhension de style Transformer sur des images haute résolution et qu'il a un grand potentiel pour devenir l'architecture de base de la prochaine génération pour les modèles de fondation visuels. Le code est disponible à l'adresse 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.