人工智能

MambaOut：我们真的需要 Mamba 来实现愿景吗？

发布时间

3星期前

2024 年 5 月 24 日

在现代机器学习和人工智能框架中，变压器是各个领域中使用最广泛的组件之一，包括自然语言处理中的 GPT 系列和 BERT，以及计算机视觉任务中的视觉变压器。尽管在模型架构中包含 Transformer 可以显着提高模型性能，但 Transformers 中的注意力模块会随序列长度呈二次方扩展，从而导致计算挑战。多年来，不同的模型探索了不同的策略来应对计算挑战，包括内核化、历史内存压缩、令牌混合范围限制和低秩方法等方法。最近，Mamba 和 RWKV 等递归神经网络方法因其在大型语言模型中的良好结果而引起了广泛关注。

Mamba 模型系列具有递归神经网络的架构，例如状态空间模型的令牌混合器，最近被引入以解决注意力机制的二次复杂性，并随后应用于视觉任务。研究人员已经探索了将 Mamba 和 SSM 或状态空间模型结合到视觉识别任务中的方法，而 Vision Mamba 结合了 Mamba 来开发类似于 Vision Transformer 的各向同性视觉模型，就是一个很好的例子。另一方面，LocalMamba 结合了局部归纳偏差来增强视觉 Mamba 模型，VMamba 框架采用基础 Mamba 模型来构建类似于 ResNet 和 AlexNet 的分层模型。然而，Mamba 框架对于视觉识别上下文任务真的至关重要吗？之所以出现这个问题，是因为与传统的基于注意力的模型和卷积模型相比，Mamba 系列视觉任务模型的性能迄今为止一直表现平平。

MambaOut 是一部试图深入探讨 Mamba 框架本质，并回答 Mamba 是否非常适合具有自回归和长序列特征的任务的作品。 MambaOut 框架假设 Mamba 对于视觉任务来说不是必需的，因为图像分类不符合长序列或自回归特征。尽管分割和检测任务也不是自回归的，但它们显示出长序列特征，导致 MambaOut 框架假设 Mamba 在这些任务中的潜力。 MambaOut 框架是通过将 Mamba 块堆叠在一起构建的，同时删除了状态空间模型（其核心令牌混合器）。实验结果支持了 MambaOut 框架提出的假设，因为它能够超越 ImageNet 图像分类框架上的所有视觉 Mamba 模型，表明 Mamba 对于视觉任务来说并不是必需的。另一方面，对于检测和分割任务，MambaOut 框架无法复制最先进的 Mamba 模型提供的性能，这展示了 Mamba 系列模型在长序列视觉任务中的潜力。

本文旨在深入介绍 MambaOut 框架，我们探讨了该框架的机制、方法、架构，以及它与最先进框架的比较。那么让我们开始吧。

MambaOut：愿景真的需要曼巴吗？

随着机器学习应用和能力的进步，Transformers 已成为一系列任务的主流骨干，为包括视觉变形金刚、GPT系列模型、BERT等等。然而，变压器的令牌混合器会导致序列长度的二次复杂度，并且对较长的序列提出了重大挑战。为了解决这个问题，引入了许多具有线性复杂度的令牌长度的令牌混合器，如 Linformer、Longformer、Performer、Dynamic Convolution 和 Big Bird。然而，近年来，类似循环神经网络的模型由于其并行训练的能力以及在较长序列上提供高效性能的能力而受到重视。在类 RNN 模型所提供的卓越性能的指导下，研究人员正在尝试将 Mamba 系列模型引入并利用到视觉识别任务中，因为 Mamba 模型的令牌混合器是遵循循环神经网络精神的结构化状态空间模型。然而，实验结果表明，与基于注意力的最先进的卷积模型相比，基于状态空间模型的视觉框架在现实世界的视觉任务中表现不佳。

MambaOut 是一次尝试调查 MambaOut 的本质曼巴模型族，并总结 Mamba 适合自回归或长序列的任务，因为状态空间模型具有固有的 RNN 机制。然而，大多数视觉任务并不具备这两个特征，在一些实验的基础上，MambaOut 提出了以下两个假设。首先，状态空间模型对于图像分类来说不是必需的，因为图像分类任务既不符合自回归也不符合长序列特征。其次，状态空间模型可能对实例分割和语义分割以及对象检测有利，因为它们遵循长序列特征，尽管它们不是自回归的。分析循环神经网络类状态空间模型机制的实验结果表明，Mamba 框架适合具有自回归或长序列特征的任务，而对于图像分类任务则不需要。来到 MambaOut 框架本身，它是一系列基于门控卷积神经网络块、没有状态空间模型的 Mamba 模型，实验结果表明 MambaOut 框架能够在图像分类任务中优于 Mamba 模型，但无法复制图像检测和分割任务的性能。

曼巴适合做什么任务？

Mamba 框架的令牌混合器是一个选择性状态空间模型，定义了四个与输入相关的参数。该框架的循环特性将类似 RNN 的状态空间模型与因果注意力模型区分开来。隐藏状态可以看作是存储历史信息的固定大小的存储器。固定大小意味着内存是有损的，但它也确保了将内存与当前输入集成的计算复杂度保持不变。相反，因果注意力层存储先前标记中的所有键和值，并通过将当前标记的键和值与每个新输入相加来进行扩展，并且从理论上讲，这种记忆是无损的。然而，存储器大小随着输入更多标记而增加，增加了将存储器与当前输入集成的复杂性。因果注意力模型和类 RNN 模型的记忆机制之间的差异如下图所示。

由于状态空间模型的记忆本质上是有损的，它达不到因果注意力的无损记忆，因此，曼巴模型无法展示其在处理短序列方面的优势，而因果注意机制在这个领域可以轻松表现良好。然而，在涉及长序列的场景中，因果注意力方法由于二次复杂度而变得不稳定。在这种情况下，Mamba 框架展示了其将内存与当前输入合并的效率，并且能够顺利处理长序列，这表明 Mamba 系列模型非常适合处理长序列。

还值得注意的是，一方面，状态空间模型的循环性质允许 Mamba 模型有效地处理长序列，但它引入了一定的限制，因为它只能从当前和之前的时间步访问信息，而这种类型的令牌混合称为因果模式，如下图所示。由于其因果性质，该方法适用于自回归生成任务.

完全可见模式适合理解模型可以一次性访问所有输入的任务。此外，默认情况下，注意力处于完全可见模式，并且可以通过将因果掩模应用于注意力图来轻松地将其转变为因果模式，并且类似 RNN 的模型由于其循环特性，本质上以因果模式运行。总而言之，Mamba 框架适合涉及处理长序列的任务，或需要因果令牌混合模式的任务。

视觉识别任务、因果标记混合代码和非常大的序列

如前所述，完全可见的令牌混合模式允许不受限制的混合范围，而因果模式则限制当前令牌只能访问先前令牌的信息。此外，视觉识别被归类为一种理解任务，其中模型可以立即看到整个图像，这消除了对令牌混合的限制，并且对令牌混合施加额外的约束可能会降低模型的性能。一般来说，完全可见模式适合理解任务，而休闲模式更适合自回归任务。此外，BeRT 和 ViT 模型比 GPT 模型更能用于理解任务，这一事实进一步支持了这一说法。

实验验证及结果

下一步是通过实验验证 MambaOut 框架提出的假设。如下图所示，Mamba 模块基于门控卷积神经网络模块，Mamba 和门控 CNN 模块的元架构可以视为 MetaFormer 框架的令牌混合器和 MLP 的简化集成。

Mamba 模块通过额外的状态空间模型扩展了门控卷积神经网络，SSm 的存在是门控 CNN 和 Mamba 模块的区别。此外，为了提高实际速度，MambaOut 框架仅在部分通道上进行深度卷积，正如下面的算法所示，门控 CNN 块的实现简单、有效且优雅。

图像分类任务

ImageNet 是图像分类任务的基准，因为它由一千多个公共类、超过 1.3 万张训练图像和超过 50,000 张验证图像组成。用于实验的数据增强包括随机调整大小的裁剪、混合、颜色抖动、随机擦除、CutMix 和 Rand 增强。下表总结了 Mamba 系列模型、MambaOut 模型以及其他基于注意力和卷积模型在 ImageNet 数据集上的性能。可以看出，不带状态空间模型的 MambaOut 框架在所有模型大小上均优于具有 SSM 的视觉 Mamba 模型。

例如，MambaOut-Small 模型返回超过 1% 的 top-84 准确度分数，比最接近的 Mamba 竞争对手高出 0.4%。这一结果有力地支持了第一个假设，即不需要为图像分类任务引入状态空间模型。

对象检测和实例分割任务

COCO 充当对象检测和实例分割任务的基准。尽管 MambaOut 框架能够超越一些视觉 Mamba 模型的性能，但它仍然低于最先进的视觉 Mamba 模型，包括 LocalVMamba 和 VMamba。 MambaOut 与最先进的视觉模型的性能差异强调了将 Mamba 系列模型集成到长序列视觉任务中的好处。然而，值得注意的是，最先进的卷积注意力混合模型和视觉 Mamba 模型之间仍然存在显着的性能差距。

总结

在本文中，我们讨论了 Mamba 系列模型的概念，并得出结论：它适合涉及自回归和长序列特征的任务。 MambaOut 是一部试图深入探讨 Mamba 框架本质，并回答 Mamba 是否非常适合具有自回归和长序列特征的任务的作品。 MambaOut 框架假设 Mamba 对于视觉任务来说不是必需的，因为图像分类不符合长序列或自回归特征。尽管分割和检测任务也不是自回归的，但它们显示出长序列特征，导致 MambaOut 框架假设 Mamba 在这些任务中的潜力。 MambaOut 框架是通过将 Mamba 块堆叠在一起构建的，同时删除了状态空间模型（其核心令牌混合器）。实验结果支持了 MambaOut 框架提出的假设，因为它能够超越 ImageNet 图像分类框架上的所有视觉 Mamba 模型，表明 Mamba 对于视觉任务来说并不是必需的。另一方面，对于检测和分割任务，MambaOut 框架无法复制最先进的 Mamba 模型提供的性能，这展示了 Mamba 系列模型在长序列视觉任务中的潜力。