人工智能
MambaOut: 我们真的需要 Mamba 来进行视觉任务吗?
在现代机器学习和人工智能框架中,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)样式的模型由于其能够并行训练和在长序列上高效地工作而受到关注。受 RNN 样式模型出色的表现的启发,研究人员正在尝试将 Mamba 模型家族引入视觉识别任务,因为 Mamba 模型的令牌混合器是具有循环神经网络精神的结构化状态空间模型。然而,实验结果表明,基于状态空间模型的框架在视觉任务上表现不佳,尤其是与注意力和最先进的卷积模型相比。
MambaOut 是一个尝试来调查 Mamba 模型家族的性质,并总结 Mamba 适合具有自回归或长序列特征的任务,因为状态空间模型具有内在的 RNN 机制。然而,大多数视觉任务不具有这些特征,基于一些实验,MambaOut 提出了两个假设。首先,状态空间模型不必要于图像分类,因为图像分类不符合自回归或长序列特征。第二,状态空间模型可能对实例分割和语义分割以及对象检测有潜力,因为它们表现出长序列特征,尽管它们不是自回归的。实验结果表明,Mamba 框架适合具有自回归或长序列特征的任务,并不必要于图像分类任务。至于 MambaOut 框架本身,它是一个基于门控卷积神经网络块的 Mamba 模型系列,移除了状态空间模型,实验结果表明 MambaOut 框架能够在图像分类任务上超越 Mamba 模型,但在检测和分割任务上无法复制最先进的 Mamba 模型的性能。
Mamba 适合哪些任务?
Mamba 框架的令牌混合器是一个具有四个输入依赖参数的选择性状态空间模型。框架的循环性质将 RNN 样式的状态空间模型与因果注意力区分开来。隐藏状态可以被视为一个固定大小的内存,存储历史信息。固定大小意味着内存是有损的,但它也确保了将内存与当前输入集成的计算复杂度保持不变。相反, 因果注意力层存储所有来自之前令牌的键和值,并通过添加当前令牌的键和值来扩展,并且这种内存是无损的,理论上讲。然而,内存大小会随着输入令牌的增加而增长,导致将内存与当前输入集成的复杂度增加。因果注意力和 RNN 样式模型之间的内存机制差异如以下图所示。

由于状态空间模型的内存是有损的,它不如因果注意力的无损内存。因此,Mamba 模型不能在短序列上表现出色,这是因果注意力机制容易处理的领域。然而,在涉及长序列的场景中,因果注意力方法由于二次复杂度而失败。在这种情况下,Mamba 框架展示了其在合并内存与当前输入方面的效率,并能够平滑地处理长序列,表明 Mamba 模型家族适合处理长序列。
值得注意的是,状态空间模型的循环性质使 Mamba 模型能够高效地处理长序列,但也引入了一个限制,因为它只能访问当前和之前的时间步的信息,这种令牌混合被称为因果模式,如以下图所示。由于其因果性质,这种方法适合 自回归生成任务。

完全可见模式适合理解任务,因为模型可以同时访问所有输入。另外,注意力默认处于完全可见模式,可以通过将因果掩码应用于注意力图来轻松转换为因果模式,而 RNN 样式模型由于其循环性质而固有地处于因果模式。总而言之,Mamba 框架适合具有长序列或需要因果令牌混合模式的任务。
视觉识别任务、因果令牌混合代码和非常长的序列
如前所述,完全可见的令牌混合模式允许无限制的混合范围,而因果模式则限制当前令牌只能访问之前的令牌的信息。另外,视觉识别被归类为理解任务,因为模型可以同时看到整个图像,这消除了对令牌混合的限制的需求,并且施加额外的限制可能会降低模型的性能。通常,完全可见模式适合理解任务,而因果模式更适合自回归任务。另外,这一说法得到了 BeRT 和 ViT 模型更常用于理解任务而不是 GPT 模型的事实的支持。
实验验证和结果
下一步是实验验证 MambaOut 框架的假设。如以下图所示,Mamba 块基于门控卷积神经网络块,Mamba 和门控 CNN 块的元架构可以被视为 MetaFormer 框架的令牌混合器和多层感知器的简化集成。

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

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

例如,MambaOut-Small 模型返回了超过 84% 的 top-1 准确率,高于其最接近的 Mamba 竞争对手 0.4%。这个结果强烈支持第一个假设,即引入状态空间模型对于图像分类任务是不必要的。
对象检测和实例分割任务
COCO 是对象检测和实例分割任务的基准。虽然 MambaOut 框架能够超越一些视觉 Mamba 模型的性能,但它仍然无法达到最先进的视觉 Mamba 模型(包括 LocalVMamba 和 VMamba)的性能。MambaOut 框架与最先进的视觉模型之间的性能差距强调了将 Mamba 模型家族集成到长序列视觉任务中的益处。然而,值得注意的是,状态先进的卷积-注意力-混合模型和视觉 Mamba 模型之间仍然存在显著的性能差距。

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












