Vision Mamba: Apprendimento Efficiente delle Rappresentazioni Visive con Modelli Spaziali di Stato Bidirezionali

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

January 17, 2024

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

cs.AI

Abstract

Recentemente, i modelli a spazio di stato (SSM) con design efficienti orientati all'hardware, come Mamba, hanno dimostrato un grande potenziale per la modellazione di sequenze lunghe. Costruire backbone visivi efficienti e generici basati esclusivamente su SSM rappresenta una direzione promettente. Tuttavia, rappresentare dati visivi è una sfida per gli SSM a causa della sensibilità alla posizione dei dati visivi e della necessità di un contesto globale per la comprensione visiva. In questo articolo, dimostriamo che la dipendenza dell'apprendimento della rappresentazione visiva dall'auto-attenzione non è necessaria e proponiamo un nuovo backbone visivo generico con blocchi Mamba bidirezionali (Vim), che contrassegna le sequenze di immagini con incorporamenti di posizione e comprime la rappresentazione visiva con modelli a spazio di stato bidirezionali. Nelle attività di classificazione su ImageNet, rilevamento di oggetti su COCO e segmentazione semantica su ADE20k, Vim raggiunge prestazioni superiori rispetto a transformer visivi consolidati come DeiT, dimostrando anche un significativo miglioramento dell'efficienza computazionale e di memoria. Ad esempio, Vim è 2,8 volte più veloce di DeiT e risparmia l'86,8% della memoria GPU durante l'inferenza batch per estrarre caratteristiche da immagini con una risoluzione di 1248x1248. I risultati dimostrano che Vim è in grado di superare i vincoli computazionali e di memoria nell'eseguire una comprensione in stile Transformer per immagini ad alta risoluzione e ha un grande potenziale per diventare il backbone di prossima generazione per i modelli di fondazione visivi. Il codice è disponibile all'indirizzo 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.