Vision Mamba: 使用双向状态空间模型实现高效视觉表示学习.
0. TL;DR
本文提出了Vision Mamba（Vim），一种基于双向状态空间模型（SSM）的新型视觉骨干网络。Vim通过引入位置嵌入和双向SSM块，实现了对视觉数据的全局上下文建模，同时克服了传统SSM模型在视觉任务中的局限性。实验结果表明，Vim在ImageNet分类、COCO目标检测和ADE20K语义分割等任务上均优于现有的视觉Transformer模型（如DeiT），并且在处理高分辨率图像时具有显著的计算和内存效率优势。
1. 背景介绍
近年来，状态空间模型（SSMs）在序列建模领域逐渐崭露头角，成为Transformer架构的有力竞争者。SSMs的优势在于其线性时间推理、高度并行化的训练以及在长序列处理任务中的强大性能。Mamba作为一种基于SSM的模型，通过其选择性机制和硬件感知设计，实现了令人印象深刻的性能，成为挑战注意力机制的Transformer架构的有力候选者。
与此同时，视觉Transformer（ViTs）在视觉表示学习领域取得了巨大成功，特别是在大规模自监督预训练和下游任务的高性能方面。然而，Transformer的自注意力机制在处理长距离视觉依赖时，例如高分辨率图像，会带来速度和内存使用方面的挑战。
为了将Mamba的成功从语言建模领域转移到视觉领域，本文提出了Vim，一种基于纯SSM的通用视觉骨干网络。Vim通过引入双向SSM块和位置嵌入，实现了对视觉数据的全局上下文建模，同时保持了SSM的高效性。
2. Vim 模型
Vim模型结合了双向SSM，用于数据依赖的全局视觉上下文建模；并通过位置嵌入实现位置感知的视觉识别。
为了使用Mamba处理视觉任务，作者设计了如下数据处理流程：
- 使用16×16的核大小投影层将二维图像拆分成二维图像块；
- 将图像块展平为向量，进行线性投影，并添加位置编码；
- 引入类标记，使用Vim模型处理序列；
- 对输出类标记进行线性投影，得到分类结果。
原始的Mamba模块是为一维序列设计的，不适合需要空间感知理解的视觉任务。故Vim Block结合了双向序列建模以用于视觉任务。
- 输入序列首先通过归一化层；
- 将归一化的序列线性投影为$x,z$；
- $x$通过双向SSM块，第一个块以前向处理视觉序列，第二个块以后向处理视觉序列；
- 通过$z$对双向的输出序列进行门控；
- 相加并通过线性投影得到输出序列。
3. 实验分析
作者对Vim的双向设计进行了消融实验，结果表明，使用双向SSM，并且在 SSM 之前添加一个Conv1d的效果最好。
对于图像分类任务，Vim在ImageNet-1K数据集上的表现优于现有的卷积网络、Transformer网络和SSM网络。具体来说，Vim-Ti的Top-1准确率为76.1%，高于DeiT-Ti的72.2%。Vim-S的Top-1准确率为80.3%，高于DeiT-S的79.8%。此外，Vim在处理高分辨率图像时具有显著的速度和内存优势。
对于语义分割任务，Vim在ADE20K数据集上的表现也优于DeiT。具体来说，Vim-Ti的mIoU为41.0%，高于DeiT-Ti的39.2%。Vim-S的mIoU为44.9%，高于DeiT-S的44.0%。
对于目标检测和实例分割，Vim在COCO 2017数据集上的表现也优于DeiT。具体来说，Vim-Ti的目标检测AP为45.7%，高于DeiT-Ti的44.4%。实例分割的AP为39.2%，高于DeiT-Ti的38.1%。
