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）を用いた新しい汎用視覚バックボーンを提案します。Vimは、画像シーケンスに位置埋め込みを施し、双方向状態空間モデルを用いて視覚表現を圧縮します。ImageNet分類、COCO物体検出、ADE20kセマンティックセグメンテーションのタスクにおいて、VimはDeiTのような確立された視覚トランスフォーマーと比較して高い性能を達成し、さらに計算効率とメモリ効率の大幅な改善を示しています。例えば、解像度1248×1248の画像に対してバッチ推論を行い特徴を抽出する場合、VimはDeiTよりも2.8倍高速で、GPUメモリを86.8%節約します。これらの結果は、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.