Vision Mamba: Эффективное обучение визуальных представлений с использованием двунаправленной модели пространства состояний

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

January 17, 2024

Авторы: Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, Xinggang Wang

cs.AI

Аннотация

Недавно модели пространства состояний (SSM) с эффективными аппаратно-ориентированными решениями, такие как Mamba, продемонстрировали значительный потенциал для моделирования длинных последовательностей. Создание эффективных и универсальных базовых архитектур для задач компьютерного зрения, основанных исключительно на SSM, является перспективным направлением. Однако представление визуальных данных представляет сложность для SSM из-за чувствительности визуальных данных к положению и необходимости глобального контекста для понимания изображений. В данной работе мы показываем, что зависимость обучения визуальных представлений от механизма самовнимания не является обязательной, и предлагаем новую универсальную базовую архитектуру для задач зрения с двунаправленными блоками Mamba (Vim), которая маркирует последовательности изображений позиционными эмбеддингами и сжимает визуальное представление с помощью двунаправленных моделей пространства состояний. На задачах классификации ImageNet, детекции объектов COCO и семантической сегментации ADE20k Vim демонстрирует более высокую производительность по сравнению с хорошо зарекомендовавшими себя трансформерами для задач зрения, такими как DeiT, а также значительное улучшение вычислительной и памяти эффективности. Например, Vim работает в 2,8 раза быстрее, чем DeiT, и экономит 86,8% памяти GPU при пакетном выводе для извлечения признаков на изображениях с разрешением 1248×1248. Результаты показывают, что Vim способен преодолеть ограничения по вычислениям и памяти при выполнении трансформероподобного анализа высококачественных изображений и имеет большой потенциал стать базовой архитектурой следующего поколения для фундаментальных моделей компьютерного зрения. Код доступен по адресу 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.