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.8$\times$ faster than DeiT and saves 86.8% GPU memory when performing batch inference to extract features on images with a resolution of 1248$\times$1248. 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.
翻译:近年来,具有高效硬件感知设计的状态空间模型(SSMs),即Mamba,在长序列建模中展现出巨大潜力。基于SSMs构建高效且通用的视觉基础模型是一个极具吸引力的方向。然而,由于视觉数据的位置敏感性以及视觉理解需要全局上下文,SSMs在表征视觉数据时面临挑战。本文表明,视觉表征学习对自注意力的依赖并非必要,并提出了一种基于双向Mamba模块(Vim)的新型通用视觉基础模型。该模型通过位置嵌入标注图像序列,并利用双向状态空间模型压缩视觉表征。在ImageNet分类、COCO目标检测和ADE20k语义分割任务中,Vim相比DeiT等成熟的视觉Transformer取得了更高性能,同时在计算与内存效率上也有显著提升。例如,在1248×1248分辨率图像上进行批推理提取特征时,Vim比DeiT快2.8倍,并节省了86.8%的GPU内存。结果表明,Vim能够克服对高分辨率图像进行Transformer式理解时的计算与内存限制,并具有成为下一代视觉基础模型骨干的巨大潜力。代码已在https://github.com/hustvl/Vim开源。