网站搜索

评估 Mamba 机制在视觉识别任务中的必要性-MambaOut


介绍

Transformer 是 BERT、GPT 系列和 ViT 等模型的支柱。然而,它的注意力机制具有二次复杂度,这使得它对于长序列具有挑战性。为了解决这个问题,人们开发了各种具有线性复杂度的令牌混合器。

最近,基于 RNN 的模型因其对长序列的高效训练和推理而受到关注,并显示出作为大型语言模型骨干的前景。

受到这些功能的启发,研究人员探索了在视觉识别任务中使用 Mamba,从而产生了 Vision Mamba、VMamba、LocalMamba 和 PlainMamba 等模型。尽管如此,实验表明,与最先进的卷积和基于注意力的模型相比,状态空间模型或基于 SSM 的视觉模型表现不佳。

最近的这篇论文并不关注设计新的视觉曼巴模型。相反,调查一个关键的研究问题:曼巴对于视觉识别任务来说是必要的吗?

什么是曼巴?

Mamba 是由卡内基梅隆大学和普林斯顿大学的研究人员开发的深度学习架构,旨在解决 Transformer 模型的局限性,特别是对于长序列。它使用结构化状态空间序列 (S4) 模型,结合连续时间、循环和卷积模型的优势,有效处理长依赖性和不规则采样数据。

最近,研究人员将 Mamba 应用于计算机视觉任务,类似于视觉变压器 (ViT) 的使用方式。 Vision Mamba (ViM) 通过利用双向状态空间模型 (SSM) 提高效率,解决传统 Transformer 的高计算需求,尤其是高分辨率图像。

曼巴架构

Mamba 通过引入独特的选择机制来增强 S4 模型,该选择机制根据输入调整参数,使其能够专注于序列中的相关信息。这种时变框架提高了计算效率。

Mamba 还采用硬件感知算法在 GPU 等现代硬件上进行高效计算,从而优化性能和内存使用。该架构将 SSM 设计与 MLP 模块集成在一起,使其适用于各种数据类型,包括语言、音频和基因组学。

曼巴变体

  • MambaByte:一种无标记语言模型,可处理原始字节序列,消除标记化及其相关偏差。
  • Mamba Mixture of Experts (MOE):将 Mixture of Experts 与 Mamba 集成,通过交替 Mamba 和 MOE 层来提高效率和可扩展性。
  • 曼巴愿景(ViM): ViM 采用 SSM 进行视觉数据处理,使用双向 Mamba 块进行视觉序列编码。这减少了计算需求,并提高了 ImageNet 分类、COCO 对象检测和 ADE20k 语义分割等任务的性能。
  • Jamba:Jamba 由 AI21 Labs 开发,是一种混合 Transformer 和 Mamba SSM 架构,拥有 520 亿个参数和 256k 个令牌的上下文窗口。

演示

在开始使用模型之前,我们将克隆存储库并安装一些必要的软件包,

!pip install timm==0.6.11
!git clone https://github.com/yuweihao/MambaOut.git
!pip install gradio

此外,我们添加了一个链接,可用于访问运行步骤并将使用 MambaOut 执行推理的笔记本。

cd /MambaOut

下面的单元格将帮助您运行 gradio Web 应用程序。

!python gradio_demo/app.py

类 RNN 模型和因果注意力

下图从记忆角度解释了因果注意力和类RNN模型的机制,其中_xi_表示第i步的输入标记。

(a) 因果注意力:将所有先前标记的键 (k) 和值 (v) 存储为内存。通过不断添加当前令牌的键和值来更新内存,使其无损。然而,随着序列的延长,将旧内存与当前标记集成的计算复杂性也会增加。因此,注意力对于短序列效果很好,但对于较长序列却很困难。

(b) 类 RNN 模型:将先前的标记压缩为固定大小的隐藏状态 (h),用作内存。这种固定大小意味着 RNN 内存本质上是有损的,无法与注意力模型的无损内存容量相匹配。然而,类似 RNN 的模型在处理长序列方面表现出色,因为无论序列长度如何,将旧内存与当前输入合并的复杂性保持不变。

由于其循环特性,Mamba 特别适合需要因果令牌混合的任务。具体来说,Mamba 擅长执行具有以下特征的任务:

1.任务涉及处理长序列。 2.任务需要因果令牌混合。

下一个问题是视觉识别任务是否有很长的序列?

对于 ImageNet 上的图像分类,典型的输入图像大小为 224x224,从而产生 196 个标记,块大小为 16x16。这个数字远小于长序列任务的阈值,因此 ImageNet 分类不被视为长序列任务。

在COCO上进行对象检测和实例分割,图像大小为800x1280,在ADE20K上进行语义分割(ADE20K是语义分割任务广泛使用的数据集,由150个语义类别组成。该数据集包括训练中的20,000张图像集和验证集中有 2,000 张图像),图像大小为 512x2048,token 数量约为4,000 个,补丁大小为 16x16。由于 4,000 个 token 超过了小序列的阈值并且接近基本阈值,因此 COCO 检测和 ADE20K 分割都被视为长序列任务。

MambaOut框架

MambaOut总体框架

图(a)表示MambaOut for Visual Recognition的整体框架: MambaOut 专为视觉识别而设计,遵循类似于 ResNet 的分层架构。它由四个阶段组成,每个阶段都有不同的通道维度,表示为 _Di_。这种层次结构允许模型在多个抽象级别处理视觉信息,从而增强其识别图像中复杂模式的能力。

(b) 门控 CNN 模块的架构: Gated CNN 块是 MambaOut 框架内的一个组件。它与 Mamba 块的不同之处在于它不包含状态空间模型 (SSM)。虽然这两个模块都使用具有门控机制的卷积神经网络 (CNN) 来调节信息流,但门控 CNN 模块中缺少 SSM,这意味着它不具有与 Mamba 模块相同的处理长序列和时间依赖性的能力,后者包含 SSM为了这些目的。

门控 CNN 和 Mamba 模块之间的主要区别在于状态空间模型 (SSM) 的存在。

在 MambaOut 中,使用 7x7 内核大小的深度卷积作为 Gated CNN 的令牌混合器,类似于 ConvNeXt。与 ResNet 类似,MambaOut 是使用 4 阶段框架构建的,通过在每个阶段堆叠门控 CNN 块来构建,如图所示。

在我们进一步讨论之前,我们先假设一下引入曼巴眼镜进行视觉识别的必要性。 假设1:在ImageNet上没有必要引入SSM进行图像分类,因为该任务不满足特征1或特征2。 假设2:仍然值得进一步探索SSM在视觉检测和分割方面的潜力,因为尽管不满足特征2,但这些任务符合特征1。

训练

ImageNet 上的图像分类

  • ImageNet用作图像分类的基准,有130万张训练图像和5万张验证图像。
  • 训练遵循 DeiT 方案,无需蒸馏,包括各种数据增强技术和正则化方法。
  • 使用 AdamW 优化器进行训练,学习率缩放规则为 lr=batchsize/1024 * 10^-3,导致学习率为 0.004,批量大小为 4096。
  • MambaOut 模型使用 PyTorch 和 timm 库实现,并在 TPU v3 上进行训练。

ImageNet 上的图像分类结果

  • MambaOut 模型不包含 SSM,在 ImageNet 上的所有模型大小上始终优于视觉 Mamba 模型。
  • 例如,MambaOut-Small 模型的 top-1 准确率达到 84.1%,比 LocalVMamba-S 好 0.4%,同时只需要 79% 的 MAC。
  • 这些结果支持假设 1,表明在 ImageNet 上引入 SSM 进行图像分类是不必要的。
  • Visual Mamba 模型目前明显落后于 ImageNet 上最先进的卷积和注意力模型。
  • 例如,CAFormer-M36 的准确度比所有同等尺寸的视觉 Mamba 模型高出 1% 以上。
  • 旨在挑战假设 1 的未来研究可能需要开发具有卷积和 SSM 令牌混合器的视觉 Mamba 模型,以在 ImageNet 上实现最先进的性能。

COCO 上的对象检测和实例分割

  • COCO 2017 用作对象检测和实例分割的基准。
  • MambaOut 用作 Mask R-CNN 中的骨干网,使用在 ImageNet 上预训练的权重进行初始化。
  • 训练遵循 12 个 epoch 的标准 1× 计划,训练图像调整大小,短边为 800 像素,长边不超过 1333 像素。
  • AdamW 优化器的学习率为 0.0001,总批量大小为 16。
  • 使用 PyTorch 和 mmdetection 库完成实施,并利用 FP16 精度来节省培训成本。
  • 实验在 4 个 NVIDIA 4090 GPU 上进行。

COCO 上的对象检测和实例分割结果

  • 虽然 MambaOut 在 COCO 上的对象检测和实例分割方面可以优于一些视觉 Mamba 模型,但它仍然落后于 VMamba 和 LocalVMamba 等最先进的视觉 Mamba。
  • 例如,作为 Mask R-CNN 骨干的 MambaOut-Tiny 落后 VMamba-T 1.4 APb 和 1.1 APm。
  • 这种性能差异凸显了将 Mamba 集成到长序列视觉任务中的好处,支持假设 2。
  • 然而,与 TransNeXt 等最先进的卷积注意力混合模型相比,视觉 Mamba 仍然显示出显着的性能差距。 Visual Mamba 需要通过在视觉检测任务中超越其他最先进的模型来证明其有效性。

ADE20K 上的语义分割

  • ADE20K用作语义分割任务的基准,包括150个语义类别,训练集中有20,000张图像,验证集中有2,000张图像。
  • Mamba 被用作 UperNet 的骨干网,并通过 ImageNet 预训练权重进行初始化。
  • 使用 AdamW 优化器进行训练,学习率为 0.0001,批量大小为 16,迭代次数为 160,000 次。
  • 使用 PyTorch 和 mmsegmentation 库完成实现,并在四个 NVIDIA 4090 GPU 上进行实验,利用 FP16 精度来提高训练速度。

ADE20K 上的语义分割结果

  • 与 COCO 上的对象检测类似,ADE20K 上语义分割的性能趋势表明,MambaOut 可以优于一些视觉 Mamba 模型,但无法与最先进的 Mamba 模型的结果相匹配。
  • 例如,LocalVMamba-T 在单尺度(SS)和多尺度(MS)评估中均超过 MambaOut-Tiny 0.5 mIoU,进一步从经验上支持假设 2。
  • 此外,与集成卷积和注意力机制的更先进的混合模型(例如 SG-Former 和 TransNeXt)相比,视觉 Mamba 模型继续表现出显着的性能缺陷。
  • Visual Mamba 需要通过在视觉分割任务中取得更强的性能来进一步展示其在长序列建模方面的优势。

结论

Mamba 机制最适合具有长序列和自回归特征的任务。 Mamba 显示了视觉检测和分割任务的潜力,这确实符合长序列特征。 MambaOut 模型超越了 ImageNet 上所有视觉 Mamba 模型,但仍落后于最先进的视觉 Mamba 模型。

然而,由于计算资源的限制,本文重点验证视觉任务的 Mamba 概念。未来的研究可以进一步探索 Mamba 和 RNN 概念,以及 RNN 和 Transformer 在大型语言模型 (LLM) 和大型多模态模型 (LMM) 中的集成,从而有可能在这些领域带来新的进步。

参考

  • 原创研究论文:- MambaOut:我们真的需要 Mamba 来实现愿景吗?
  • Vision Mamba:利用双向状态空间模型进行高效视觉表示学习
  • Mamba(深度学习架构)

相关文章: