视觉Mamba来了：速度提升2.8倍，内存能省87%

原文来源：机器之心

图片来源：由无界 AI生成

Vision Mamba 不是个普通模型。

号称「全面包围 Transformer」的 Mamba，推出不到两个月就有了高性能的视觉版。

本周四，来自华中科技大学、地平线、智源人工智能研究院等机构的研究者提出了 Vision Mamba（Vim）。

论文地址：https://arxiv.org/pdf/2401.09417.pdf

项目地址：https://github.com/hustvl/Vim

论文标题：Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model

效果如何呢？在 ImageNet 分类任务、COCO 对象检测任务和 ADE20k 语义分割任务上，与  DeiT 等成熟的视觉 Transformers 相比，Vim 实现了更高的性能，同时还显著提高了计算和内存效率。例如，在对分辨率为 1248×1248 的图像进行批量推理提取特征时，Vim 比 DeiT 快 2.8 倍，并节省 86.8% 的 GPU 内存。结果表明，Vim 能够克服对高分辨率图像执行 Transformer 式理解时的计算和内存限制，并且具有成为视觉基础模型的下一代骨干的巨大潜力。

接下来我们看看论文内容。

Mamba 的提出带动了研究者对状态空间模型（state space model，SSM）兴趣的增加，不同于 Transformer 中自注意力机制的计算量会随着上下文长度的增加呈平方级增长，由于 SSM 擅长捕捉远程依赖关系，因而开始受到大家追捧。

在此期间，一些基于 SSM 的方法如线性状态空间层（LSSL）、结构化状态空间序列模型（S4）、对角状态空间（DSS）和 S4D 都被研究者提出来，用于处理各种序列数据，特别是在建模远程依赖关系方面。

Mamba 将时变参数纳入 SSM 中，并提出了一种硬件感知算法来实现高效的训练和推理。Mamba 卓越的扩展性能表明它在语言建模方面是 Transformer 有前途的替代品。

然而，到目前为止，研究者还尚未在视觉任务中探索出通用的基于纯 SSM 的骨干网络。

受 Mamba 在语言建模方面成功的激励，研究者开始设想能否将这种成功从语言转移到视觉，即用先进的 SSM 方法设计通用且高效的视觉主干。然而，由于 Mamba 特有的架构，需要解决两个挑战，即单向建模和缺乏位置感知。

为了应对这些问题，研究者提出了 Vision Mamba (Vim) 块，它结合了用于数据依赖的全局视觉上下文建模的双向 SSM 和用于位置感知视觉识别的位置嵌入。 

与其他基于 SSM 的视觉任务模型相比，Vim 是一种基于纯 SSM 的方法，并以序列方式对图像进行建模。与基于 Transformer 的 DeiT 相比，Vim 在 ImageNet 分类上取得了优越的性能。此外，Vim 在 GPU 内存和高分辨率图像的推理时间方面更加高效。

方法介绍

Vision Mamba (Vim) 的目标是将先进的状态空间模型 (SSM)，即 Mamba 引入到计算机视觉。 

Vim 的概述如图 2 所示，标准的 Mamba 是为 1-D 序列设计的。为了处理视觉任务，首先需要将二维图像

转换成展开的 2-D patch 

。式中 (H, W) 为输入图像的大小，C 为通道数，P 为图像 patch 的大小。接下来，需要将 x_p 线性投影到大小为 D 的向量上，并添加位置嵌入

得到如下公式：

Vim 块

原始的 Mamba 块是为一维序列设计的，不适合需要空间感知理解的视觉任务。Vim 块集成了用于视觉任务的双向序列建模，Vim 块如上图 2 所示。

Vim 块的操作算法如下所示。

架构细节

架构的超参数如下：

L：块数

D：隐藏状态维度

E：扩展状态维度

N：SSM 维度