AI 模型与平台
MambaOut:我们真的需要Mamba进行视觉任务吗?
在现代机器学习和人工智能框架中,Transformer是各个领域中最广泛使用的组件,包括自然语言处理中的GPT系列和BERT,以及计算机视觉任务中的Vision Transformers。虽然将Transformer包含在模型架构中可以显著提高模型性能,但Transformer中的注意力模块的计算复杂度随着序列长度的增加而呈二次增长,从而导致高计算挑战。多年来,各种模型都在探索不同的策略来解决计算挑战,包括核化、历史内存压缩、令牌混合范围限制和低秩方法。最近,像Mamba和RWKV这样的循环神经网络方法因其在大型语言模型中的出色表现而受到关注。
Mamba是一种具有循环神经网络结构的模型家族,其令牌混合器是状态空间模型。Mamba最近被引入视觉任务中,研究人员已经探索了将Mamba和状态空间模型(SSM)集成到视觉识别任务中的方法。例如,Vision Mamba将Mamba集成到视觉Transformer中,而LocalMamba则将局部感知偏差引入视觉Mamba模型中。然而,Mamba框架是否真正必要 для视觉识别任务仍然是一个问题,因为Mamba家族模型在视觉任务中的性能尚未令人满意。
MambaOut是一个尝试回答Mamba是否适合视觉任务的框架。MambaOut假设Mamba不适合视觉任务,因为图像分类不具有长序列或自回归特性。然而,分割和检测任务虽然不具有自回归特性,但具有长序列特性,因此MambaOut框架假设Mamba可能适合这些任务。实验结果支持MambaOut框架的假设,因为它能够在图像分类任务中超越所有视觉Mamba模型,表明Mamba不必要 для视觉任务。然而,对于检测和分割任务,MambaOut框架无法复制状态-of-the-art Mamba模型的性能,表明Mamba家族模型可能适合长序列视觉任务。
本文旨在深入介绍MambaOut框架,探讨其机制、方法、架构以及与状态-of-the-art 框架的比较。让我们开始吧。
MambaOut:我们真的需要Mamba进行视觉任务吗?
随着机器学习应用和能力的进步,Transformer已成为主流的骨架,支持包括Vision Transformers、GPT系列、BERT在内的众多模型。然而,Transformer中的令牌混合器会随着序列长度的增加而呈二次增长,从而导致计算挑战。为了解决这个问题,已经引入了多种令牌混合器,包括Linformer、Longformer、Performer、Dynamic Convolution和Big Bird。然而,最近,像Mamba和RWKV这样的循环神经网络方法因其并行训练能力和在长序列上的高效性能而受到关注。受这些方法的出色表现的启发,研究人员正在尝试将Mamba家族模型引入视觉识别任务中,因为Mamba模型的令牌混合器是状态空间模型。
MambaOut是一个尝试探索Mamba家族模型性质的框架,总结Mamba适合具有自回归和长序列特性的任务。然而,大多数视觉任务不具有这些特性,因此MambaOut提出两个假设。首先,状态空间模型不必要 для图像分类,因为图像分类不具有长序列或自回归特性。其次,状态空间模型可能适合实例分割、语义分割和物体检测,因为这些任务具有长序列特性。实验结果支持MambaOut框架的假设,因为它能够在图像分类任务中超越所有视觉Mamba模型,表明Mamba不必要 для视觉任务。然而,对于检测和分割任务,MambaOut框架无法复制状态-of-the-art Mamba模型的性能,表明Mamba家族模型可能适合长序列视觉任务。
什么任务适合Mamba?
Mamba框架的令牌混合器是具有四个输入依赖参数的选择性状态空间模型。框架的循环性质使其与因果注意力区别开来。隐藏状态可以被视为一个固定大小的内存,存储历史信息。固定大小意味着内存是有损的,但也确保了将内存与当前输入集成的计算复杂度保持不变。相反,因果注意力层存储所有来自前一个令牌的键和值,并通过添加当前令牌的键和值来扩展,每次输入都会增加复杂度。这种内存机制的差异在下图中进行了说明。

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

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

Mamba块通过添加状态空间模型来扩展门控卷积神经网络,而状态空间模型的存在是Mamba块和门控CNN块之间的区别。为了提高实际速度,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模型,但它仍然无法达到状态-of-the-art视觉Mamba模型(包括LocalVMamba和VMamba)的性能。MambaOut框架与状态-of-the-art视觉模型之间的性能差异凸显了将Mamba家族模型集成到长序列视觉任务中的益处。然而,值得注意的是,状态-of-the-art卷积注意力混合模型和视觉Mamba模型之间仍然存在显著的性能差距。

最终思考
Mamba家族模型似乎适合具有自回归和长序列特性的任务。MambaOut框架假设Mamba不适合视觉任务,因为图像分类不具有长序列或自回归特性。然而,分割和检测任务虽然不具有自回归特性,但具有长序列特性,因此MambaOut框架假设Mamba可能适合这些任务。MambaOut框架通过堆叠Mamba块并移除状态空间模型来构建,其实验结果支持MambaOut框架的假设,因为它能够在ImageNet图像分类框架上超越所有视觉Mamba模型,表明Mamba不必要 для视觉任务。然而,对于检测和分割任务,MambaOut框架无法复制状态-of-the-art Mamba模型的性能,表明Mamba家族模型可能适合长序列视觉任务。












