Connect with us

人工智能

MambaOut: 我们真的需要 Mamba 来进行视觉任务吗?

mm

在现代机器学习和人工智能框架中,Transformer 是各个领域中最广泛使用的组件,包括 GPT 系列、BERT 在自然语言处理中以及 Vision Transformer 在计算机视觉任务中。虽然在模型架构中包含 Transformer 可以显著提高模型性能,但 Transformer 中的注意力模块的计算复杂度随着序列长度的增加而呈二次方增长,从而导致高计算挑战。多年来,不同的模型探索了不同的策略来解决计算挑战,包括核化、历史内存压缩、令牌混合范围限制和低秩方法。最近,循环神经网络(RNN)如 Mamba 和 RWKV 等方法因其在大型语言模型中的出色表现而受到关注。

Mamba,一种具有循环神经网络类似令牌混合器的状态空间模型家族,最近被提出以解决注意力机制的二次复杂性,并随后被应用于视觉任务。研究人员已经探索了将 Mamba 和状态空间模型(SSM)或状态空间模型集成到视觉识别任务中的方法,Vision Mamba 是一个将 Mamba 集成到视觉模型中的例子,类似于 Vision Transformer。另一方面,LocalMamba 将局部感知偏差集成到视觉 Mamba 模型中,以增强其性能,而 VMamba 框架使用基本的 Mamba 模型来构建类似 ResNet 和 AlexNet 的层次模型。然而,Mamba 框架是否真的必要于视觉识别任务?这个问题是因为 Mamba 模型家族在视觉任务中的性能一直不如传统的注意力和卷积模型。

MambaOut 尝试回答 Mamba 是否适合具有自回归和长序列特征的任务。MambaOut 框架假设 Mamba 不适合视觉任务,因为图像分类不符合长序列或自回归特征。虽然分割和检测任务也不是自回归的,但它们表现出长序列特征,导致 MambaOut 框架假设 Mamba 对这些任务可能有潜力。MambaOut 框架通过堆叠 Mamba 块并移除状态空间模型(其核心令牌混合器)构建。实验结果支持 MambaOut 框架的假设,因为它能够在 ImageNet 图像分类框架上超越所有视觉 Mamba 模型,表明 Mamba 不必要于视觉任务。另一方面,对于检测和分割任务,MambaOut 框架无法复制最先进的 Mamba 模型的性能,表明 Mamba 模型家族对于长序列视觉任务可能有潜力。

本文旨在深入探讨 MambaOut 框架,我们将探讨其机制、方法、架构以及与最先进框架的比较。让我们开始吧。

MambaOut: 我们真的需要 Mamba 来进行视觉任务吗?

随着机器学习应用和能力的进步,Transformer 已经成为各类任务的主流骨架,包括 Vision Transformer、GPT 系列、BERT 等。然而,Transformer 的令牌混合器会随着序列长度的增加而呈二次方增长,从而导致计算挑战。为了解决这个问题,已经提出了许多令牌混合器,其计算复杂度与令牌长度呈线性关系,例如 Linformer、Longformer、Performer、Dynamic Convolution 和 Big Bird。然而,近年来,循环神经网络类似模型由于其能够并行训练和在长序列上表现出色的能力而受到关注。受 RNN 类似模型出色的表现的启发,研究人员正在尝试将 Mamba 模型家族应用于视觉识别任务,因为 Mamba 模型的令牌混合器是具有状态空间模型的循环神经网络。然而,实验结果表明,状态空间模型基于框架在视觉任务中的表现不如传统注意力和卷积模型。

MambaOut 是一个尝试来研究 Mamba 模型家族的性质,总结出 Mamba 适合具有自回归或长序列特征的任务,因为状态空间模型具有内在的 RNN 机制。然而,大多数视觉任务不具有这些特征,基于一些实验,MambaOut 提出了两个假设。首先,状态空间模型对于图像分类任务不是必要的,因为图像分类不符合长序列或自回归特征。第二,状态空间模型可能对于实例分割、语义分割和目标检测任务有潜力,因为这些任务表现出长序列特征,尽管它们不是自回归的。实验结果表明,Mamba 框架适合具有自回归或长序列特征的任务,对于图像分类任务不是必要的。至于 MambaOut 框架本身,它是一个基于门控卷积神经网络块的 Mamba 模型,实验结果表明 MambaOut 框架能够在图像分类任务中超越 Mamba 模型,但在检测和分割任务中无法复制最先进的 Mamba 模型的性能。

Mamba 适合哪些任务?

Mamba 框架的令牌混合器是一个具有四个输入依赖参数的选择性状态空间模型。框架的递归特性使其与因果注意力区别开来。隐藏状态可以看作是一个固定大小的内存,存储历史信息。固定大小意味着内存是有损的,但也确保了将内存与当前输入集成的计算复杂度保持不变。相反,因果注意力层存储所有来自之前令牌的键和值,并通过添加当前令牌的键和值来扩展,并且这种内存是无损的,理论上讲。然而,内存大小会随着输入令牌的增加而增长,从而增加了将内存与当前输入集成的复杂度。因果注意力和 RNN 类似模型之间的内存机制差异如以下图所示。

由于状态空间模型的内存是有损的,因此它不如因果注意力的无损内存。因此,Mamba 模型无法在短序列上表现出其优势,而因果注意力机制可以轻松地处理短序列。然而,在涉及长序列的场景中,因果注意力方法会由于二次复杂度而失败。在这种情况下,Mamba 框架展示了其在合并内存与当前输入方面的效率,并能够平滑地处理长序列,表明 Mamba 模型家族适合处理长序列。

值得注意的是,状态空间模型的递归性质使 Mamba 模型能够高效地处理长序列,但也引入了一个限制,因为它只能访问当前和之前的时间步的信息,这种令牌混合方式被称为因果模式,如以下图所示。由于其因果性质,这种方法适合自回归生成任务。

完全可见模式适合理解任务,其中模型可以同时访问所有输入。另外,注意力默认处于完全可见模式,可以通过将因果掩码应用于注意力图来轻松转换为因果模式,而 RNN 类似模型由于其递归特性,天然地处于因果模式。总之,Mamba 框架适合具有长序列或因果令牌混合模式的任务。

视觉识别任务、因果令牌混合代码和非常长的序列

如前所述,完全可见的令牌混合模式允许无限制的混合范围,而因果模式限制了当前令牌只能访问之前的令牌的信息。另外,视觉识别被归类为理解任务,其中模型可以同时看到整个图像,这消除了对令牌混合的限制的需求,并且施加额外的约束可能会降低模型的性能。通常,完全可见模式适合理解任务,而因果模式更适合自回归任务。此外,这一说法得到了 BeRT 和 ViT 模型更常用于理解任务而不是 GPT 模型的事实的支持。

实验验证和结果

下一步是实验验证 MambaOut 框架提出的假设。如以下图所示,Mamba 块基于门控卷积神经网络块,Mamba 和门控 CNN 块的元架构可以看作是 MetaFormer 框架的令牌混合器和多层感知器的简化集成。

Mamba 块通过添加一个状态空间模型来扩展门控卷积神经网络块,状态空间模型的存在是区分门控 CNN 和 Mamba 块的关键。另外,为了提高实际速度,MambaOut 框架仅在部分通道上进行深度卷积,如以下算法所示,门控 CNN 块的实现简单而有效。

图像分类任务

ImageNet 是图像分类任务的基准,包含超过 1000 个常见类别、130 万张训练图像和 5 万张验证图像。实验中使用的数据增强包括随机裁剪、Mixup、颜色抖动、随机擦除、CutMix 和 Rand Augment。以下表格总结了 Mamba 模型家族、MambaOut 模型和其他注意力和卷积模型在 ImageNet 数据集上的性能。如图所示,MambaOut 框架在没有状态空间模型的情况下,能够在所有模型大小上一致地超越视觉 Mamba 模型。

例如,MambaOut-Small 模型返回的 top-1 准确率超过 84%,比其最接近的 Mamba 竞争者高出 0.4%。这个结果强烈支持第一个假设,即引入状态空间模型对于图像分类任务不是必要的。

目标检测和实例分割任务

COCO 是目标检测和实例分割任务的基准。虽然 MambaOut 框架能够超越一些视觉 Mamba 模型的性能,但它仍然不如最先进的视觉 Mamba 模型,包括 LocalVMamba 和 VMamba。MambaOut 框架与最先进视觉模型之间的性能差距强调了将 Mamba 模型家族集成到长序列视觉任务中的潜在益处。然而,值得注意的是,仍然存在显著的性能差距 zwischen 最先进的卷积-注意力-混合模型和视觉 Mamba 模型。

最后的思考

Mamba 模型家族似乎适合具有自回归和长序列特征的任务。MambaOut 框架假设 Mamba 不适合视觉任务,因为图像分类不符合长序列或自回归特征。虽然分割和检测任务也不是自回归的,但它们表现出长序列特征,导致 MambaOut 框架假设 Mamba 对这些任务可能有潜力。MambaOut 框架通过堆叠 Mamba 块并移除状态空间模型构建,其核心令牌混合器。实验结果支持 MambaOut 框架的假设,因为它能够在 ImageNet 图像分类框架上超越所有视觉 Mamba 模型,表明 Mamba 不必要于视觉任务。另一方面,对于检测和分割任务,MambaOut 框架无法复制最先进的 Mamba 模型的性能,表明 Mamba 模型家族对于长序列视觉任务可能有潜力。

专业为工程师,心为作家。 Kunal是一名技术作家,对AI和ML有着深厚的热爱和理解,致力于通过其引人入胜和信息丰富的文档来简化这些领域中的复杂概念。