Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model 论文总结

最新推荐文章于 2024-10-14 20:01:31 发布

向岸看

最新推荐文章于 2024-10-14 20:01:31 发布

阅读量962

点赞数 20

文章标签： Vision Mamba 计算机视觉

本文链接：https://blog.csdn.net/qq_45981086/article/details/142133577

版权

题目：Vision Mamba（视觉Mamba）: Efficient Visual Representation Learning （有效的视觉表示学习） with Bidirectional State Space Model （双向状态空间模型）

论文（ICML 2024）：[2401.09417] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (arxiv.org)

源码：hustvl/Vim: [ICML 2024] Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (github.com)

一、摘要

二、引言

三、方法

3.1. Vision Mamba（整体架构）

一、摘要

研究背景：近年来，具有高效硬件感知设计的状态空间模型 (ssm)，即Mamba，在长序列建模中显示出巨大的潜力。同时，纯粹基于ssm构建高效和通用的视觉主干是一个吸引人的方向。

研究问题：然而，由于视觉数据的位置敏感性和视觉理解对全局上下文的要求，对于ssm来说，表示视觉数据是一个挑战。

主要工作：本文提出了一种新的基于双向Mamba模块 (bidirectional Mamba blocks, Vim) 的通用视觉主干，该主干利用 位置嵌入 对图像序列进行 标记（序列化），并利用双向状态空间模型对视觉表征进行压缩。

研究成果：在ImageNet分类、COCO对象检测和ADE20k语义分割任务上，Vim 比 DeiT 等成熟的视觉转换器实现了更高的性能，同时也显示出显著提高的计算和内存效率。Vim 比 DeiT 快2.8倍，在对分辨率为 1248×1248 的图像进行批量推理提取特征时，比 DeiT 节省了 86.8% 的 GPU 内存。

结论：结果表明，Vim能够克服对高分辨率图像进行 Transformer 式理解时的计算和内存限制，具有成为下一代视觉基础模型的潜力。

二、引言

状态空间模型 (SSM) 研究现状概括：

1. 早期方法：线性状态空间层（LSSL）、结构化状态空间序列模型（S4）、对角状态空间（DSS）和 S4D。
2. 联合 SSM 与 CNN 或 Transformer 架构：2-D SSM、SGConvNeXt 和 ConvSSM。
3. Mamba（最新）：将时变参数纳入SSM，并提出了一种硬件感知算法。

Vision Transformers概述：

与卷积神经网络相比，ViT 的核心优势在于它能够通过自注意为每个图像 patch 提供依赖于数据/patch 的全局上下文。
另一个优点是通过将图像视为没有 2D 归纳偏差的片序列来进行模态不可知建模，这使其成为多模态应用的优选架构。

研究问题：

Mamba 面临着两个挑战，单向建模和缺乏位置感知。

解决方法（主要工作）：

为了解决这些问题，提出了视觉 Mamba（Vision Mamba，Vim）模型，该模型结合了双向SSM和位置嵌入技术，前者用于数据相关的全局视觉上下文建模，后者用于位置感知的视觉识别。在当前阶段，我们使用 ImageNet 数据集在监督图像分类任务上训练 Vim 模型，然后使用预训练的 Vim 作为主干来执行下游密集预测任务的顺序视觉表示学习，即，语义分段、对象检测和实例分段。

主要贡献：

我们提出Vision mamba (Vim)，它结合了双向SSM，用于数据依赖的全局视觉上下文建模和位置嵌入，用于位置感知视觉理解。
无需注意，所提出的 Vim 具有与 ViT 相同的建模能力，而它只有次二次时间计算和线性内存复杂度。具体来说，Vim 比 DeiT 快2.8倍，在执行批量推理以提取分辨率为 1248×1248 的图像上的特征时节省了86.8%的GPU内存。
我们在 ImageNet 分类和密集预测下游任务上进行了大量的实验。结果表明，与成熟且高度优化的 plain vision Transformer (即 DeiT )相比，Vim具有优越的性能。

Q：单向建模是什么？有什么坏处？

A：单向建模（Unidirectional Modeling）是指在处理序列数据时，模型只考虑数据序列中的一个方向，即从开始到结束。例如，对于一个文本序列，模型从句子的第一个词开始，逐步处理到最后一个词。在每个时间步，模型的输出仅依赖于当前时间步及之前的时间步的信息，而不依赖于未来时间步的信息。（SSM建模是这样的，依赖于前向时间步长，不相邻步长之间的信息没有交互）

Q：为什么会缺乏位置感知？怎么避免？

A：缺乏位置感知主要是因为单向建模仅在一个方向上处理序列数据，而不能同时利用序列中的前后信息。例如，在文本生成任务中，单向RNN只能基于当前时间步及之前时间步的词汇生成新的词汇，而不能考虑当前时间步之后的词汇。一般的处理方法，使用双向建模或引入位置编码。

三、方法

3.1. Vision Mamba（整体架构）

问题：标准 Mamba 是为 1-D 序列设计的。

解决方法（序列化）：为了处理视觉任务，首先将二维图像 $\mathrm{t \in R^{H\times W \times C}}$ 变换为被展平的二维 patch $\mathrm{x_p \in R^{J\times (P^2 \cdot C)}}$ ，其中 (H, W) 为输入图像的大小，C 为通道数，P 为图像 patch 的大小。接下来，将 $x_p$ 线性投影到大小为 D 的向量上，并添加位置嵌入 $\mathrm{E_{pos} \in R^{(J+1) \times D}}$ ，如下所示：

其中 $\mathrm{t^j_p}$ 是 t 的第 j 个补丁， $\mathrm{W \in R^{(P^2 \cdot C) \times D}}$ 是可学习的投影矩阵。受 ViT 和 BERT 的启发，使用类标记来表示整个补丁序列，表示为 $\mathrm{t_{cls}}$ 。然后，将令牌序列 $\mathrm{T_{l - 1}}$ 发送到 Vim 编码器的第 $\mathrm{l}$ 层，并获得输出 $\mathrm{T_1}$ 。最后，将输出类令牌 $\mathrm{T^0_L}$ 归一化，并将其馈送到多层感知器（MLP），以获得最终预测 $\mathrm{\hat{p}}$ ，如下所示：

其中，Vim是所提出的 Vision Mamba 模块，L 是层数，Norm 是归一化层。（仿照 ViT 架构）

3.2. Vim Block

Vim block 中的操作，如算法1所示：

过程：

首先，输入 token 序列 $\mathrm{T_{l - 1}}$ 由归一化层进行归一化。
接下来，将归一化序列 线性投影 到维度大小为 $\mathrm{E}$ 的 $\mathrm{x}$ 和 $\mathrm{z}$ 上。
然后，从 forward 和 backward 方向处理 $\mathrm{x}$ 。
------------------------------- SSM 内部 --------------------------------
对于每个方向，首先，将 1-D 卷积应用于 $\mathrm{x}$ 并得到 $\mathrm{x'_o}$ 。然后，分别将 $\mathrm{x'_o}$ 线性投影到 $B_o, C_o, \Delta_o$ 。（时变操作）
然后，分别用 $\Delta_o$ 生成 $\bar{A}_o, \bar{B}_o$ 。（离散化）
------------------------------- SSM 内部 --------------------------------
最后，通过 SSM 计算 $y_{forward}$ 和 $y_{backward}$ 。然后， $y_{forward}$ 和 $y_{backward}$ 由 $\mathrm{z}$ 进行门控，并将它们加在一起以得到输出 token序列 $\mathrm{T_l}$ 。

3.3. 架构细节

L：模块数，D：隐藏状态维度，E：扩展状态维度，N：SSM维度。

在 ViT 和 DeiT 之后，首先使用 16×16 核大小投影层得到 1-D 序列的非重叠 patch 嵌入。随后，直接堆叠 L 个 Vim 模块 。默认情况下，将 模块数 L 设置为 24，SSM 维度 N 设置为16。为了与 DeiT系列的模型尺寸保持一致，将 隐藏状态维度 D 设置为192，并将小尺寸变体的 状态维度 E 设置为384 。对于小尺寸的变体，将 D 设置为384，将 E 设置为768。

3.4. 效率分析

1) I/O 效率

Mamba时间效率分析：高带宽内存 (HBM) 和 SRAM 是 gpu 的两个重要组成部分。其中 SRAM 具有更大的带宽，HBM 具有更大的内存大小。使用 HBM 的 Vim 的 SSM 操作的标准实现要求内存 IO 的数量达到 O(BMEN) 的数量级。

内核融合操作改进：受 Mamba 的启发，Vim 首先将 O(BME + EN) 字节的内存 $\mathrm{(\Delta _o,A_o,B_o,C_o)}$ 从慢速 HBM 读取到快速SRAM。然后，Vim 在 SRAM 中得到大小为 (B, M, E, N) 的离散的 $\mathrm{\bar{A}_o, \bar{B}_o}$ 。最后，Vim 在 SRAM 中执行 SSM 操作，并将大小为 (B, M, E) 的输出写回 HBM。这种方法可以帮助将 IOs 从 O(BMEN) 减少到 O(BME + EN) 。

Q：这样做，GPU的快速SRAM内存够吗？

2) Memory 效率（沿用Mamba的重计算策略，减小内存占用）

为了在处理长序列时避免内存不足问题并实现更低的内存使用，Vim选择了与Mamba相同的重新计算方法。对于计算梯度大小为(B, M, E, N) 的中间状态，Vim在网络反向传递时重新计算它们。对于中间激活，如激活函数和卷积的输出，Vim还重新计算它们以优化GPU内存需求，因为激活值占用大量内存，但重新计算速度很快。

3) Computation 效率

与自注意力相比：

其中自注意力与序列长度 M 成二次关系，SSM与序列长度 M 成线性关系（N是固定参数，默认设置为16）。

四、实验

4.1 图像分类

数据集：在 ImageNet-1K 数据集上对 Vim 进行基准测试，该数据集包含来自1000个类别的1.28M训练图像和50K验证图像。

公平性声明：为了公平比较，训练集主要遵循DeiT。

数据增强：应用随机裁剪、随机水平翻转、标签平滑正则化、混合和随机擦除作为数据增强。

实验细节：在 224×224 的输入图像上训练时，采用AdamW，动量为0.9，总批量大小为1024，权重衰减为 0.05 来优化模型。初始学习率为 $1\times 10^{-3}$ 训练了300 个 epoch 的 Vim 模型。在测试过程中，在验证集上应用中心裁剪。实验在 8块 A800 gpu 上进行。

Results：如上表所示，在参数大致相似的情况下，VimSmall 的 top-1 准确率达到 80.5，比 ResNet50 提高了 4.3 个点。与传统的基于自注意力的 ViT 相比，Vim 在参数数量和分类精度方面都有相当大的优势。与高度优化的 vit 变体(即DeiT)相比，Vim在不同规模上都超过了它，在参数数量相当的情况下：VimTiny 比 DeiT- tiny 高3.9个点，Vim- small 比 DeiT- small 高 0.7 个点。

图 (b) 和 (c) 比较了 VimTiny 和 DeiT 的 FPS 和 GPU 内存。具体来说，当图像大小为512×512时，Vim实现了与DeiT相似的FPS和内存。当图像大小增加到1248×1248时，Vim比DeiT快2.8倍，节省了86.8%的显存。Vim在序列长度上线性缩放的明显优势，使其可以用于高分辨率的下游视觉应用和长序列多模态应用。

4.2.语义分割

数据集：在 ADE20K 上进行语义分割实验。

Results：如上表所示。Vim在不同尺度上始终优于DeiT：Vim-Ti 比 DeiT-Ti 高 1.8 mIoU，Vim-S 比 DeiT-S 高 0.9 mIoU。与ResNet-101主干网络相比，Vim-S 以几乎少 2 倍的参数实现了相同的分割性能。

4.3. 消融研究

在 ADE20K 上使用 ImageNet-1K 分类和 Segmenter 语义分割框架，消除了 Vim 的关键双向设计。研究这些双向策略：

没有。直接采用曼巴块来处理只包含前向方向的视觉序列。
双向序列。在训练过程中，随机翻转视觉序列。这就像数据增强一样。
双向区块。将堆叠的块配对。每对的第一个块在正向处理视觉序列，每对的第二块在反向处理视觉序列。
双向SSM。为每个块添加额外的SSM，以反向处理视觉序列。
双向SSM + Conv1d。在backward SSM的基础上，在backward SSM之前进一步添加一个backward Conv1d。

Results：如上表所示，在语义分割方面，双向Mamba模块比普通的单向Mamba模块性能提高1.3 mIoU。通过添加额外的后向SSM和Conv1d，实现了卓越的分类精度(73.9 top-1 acc vs. 73.2 top-1 acc)和异常的分割优势(35.9 mIoU vs. 32.3 mIoU)。我们使用双向SSM + Conv1d策略作为我们的Vim块中的默认设置。