Vision Mamba: Aprendizado Eficiente de Representação Visual com Modelo de Espaço de Estados Bidirecional

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

January 17, 2024

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

cs.AI

Resumo

Recentemente, os modelos de espaço de estados (SSMs) com designs eficientes e conscientes de hardware, como o Mamba, têm demonstrado grande potencial para a modelagem de sequências longas. Construir backbones visuais eficientes e genéricos baseados exclusivamente em SSMs é uma direção atraente. No entanto, representar dados visuais é desafiador para os SSMs devido à sensibilidade à posição dos dados visuais e à necessidade de contexto global para a compreensão visual. Neste artigo, mostramos que a dependência do aprendizado de representação visual na auto-atenção não é necessária e propomos um novo backbone visual genérico com blocos bidirecionais de Mamba (Vim), que marca as sequências de imagens com embeddings de posição e comprime a representação visual com modelos de espaço de estados bidirecionais. Em tarefas de classificação no ImageNet, detecção de objetos no COCO e segmentação semântica no ADE20k, o Vim alcança desempenho superior comparado a transformers visuais bem estabelecidos como o DeiT, além de demonstrar uma eficiência significativamente melhorada em termos de computação e memória. Por exemplo, o Vim é 2,8 vezes mais rápido que o DeiT e economiza 86,8% da memória da GPU ao realizar inferência em lote para extrair características de imagens com resolução de 1248x1248. Os resultados demonstram que o Vim é capaz de superar as restrições de computação e memória ao realizar compreensão no estilo Transformer para imagens de alta resolução e tem grande potencial para se tornar o backbone de próxima geração para modelos de fundação visual. O código está disponível em 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.