Vision Mamba: Aprendizaje Eficiente de Representaciones Visuales con Modelo de Espacio de Estados Bidireccional

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

Resumen

Recientemente, los modelos de espacio de estados (SSMs) con diseños eficientes orientados al hardware, como Mamba, han demostrado un gran potencial para el modelado de secuencias largas. Construir backbones visuales eficientes y genéricos basados únicamente en SSMs es una dirección atractiva. Sin embargo, representar datos visuales es un desafío para los SSMs debido a la sensibilidad a la posición de los datos visuales y la necesidad de contexto global para la comprensión visual. En este artículo, mostramos que la dependencia del aprendizaje de representaciones visuales en la autoatención no es necesaria y proponemos un nuevo backbone visual genérico con bloques bidireccionales de Mamba (Vim), que marca las secuencias de imágenes con incrustaciones de posición y comprime la representación visual con modelos de espacio de estados bidireccionales. En tareas de clasificación en ImageNet, detección de objetos en COCO y segmentación semántica en ADE20k, Vim logra un rendimiento superior en comparación con transformadores visuales bien establecidos como DeiT, además de demostrar una eficiencia computacional y de memoria significativamente mejorada. Por ejemplo, Vim es 2.8 veces más rápido que DeiT y ahorra un 86.8% de memoria GPU al realizar inferencias por lotes para extraer características en imágenes con una resolución de 1248x1248. Los resultados demuestran que Vim es capaz de superar las limitaciones de computación y memoria al realizar una comprensión al estilo Transformer para imágenes de alta resolución y tiene un gran potencial para convertirse en el backbone de próxima generación para modelos fundamentales de visión. El código está disponible en 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.