【深度学习】CVPR 2025|MambaOut:在视觉任务中,我们真的需要 Mamba 吗?
仅用于站内搜索,没有排版格式,具体信息请跳转上方微信公众号内链接
作者丨科技猛兽
编辑丨极市平台
导读
在ImageNet图像分类中有效地超过了视觉Mamba模型
1MambaOut:在视觉任务中,我们真的需要Mamba吗?(来自NUS,MetaFormer原作者)1MambaOut论文解读1.1在视觉任务中,我们真的需要Mamba吗?1.2本文有哪些新的发现?1.3概念讨论1.4视觉任务有长序列吗?1.5视觉任务需要Causal的tokenmixing模式吗?1.6关于Mamba在视觉的必要性的假设1.7GatedCNN和MambaOut1.8ImageNet实验结果1.9COCO目标检测和实例分割实验结果1.10ADE20K语义分割实验结果
在本文中,作者深入研究了Mamba的本质,并得出结论:Mamba非常适合具有长序列和自回归特征的任务。视觉任务中的图像分类任务与上面的特征都不匹配,因此作者假设Mamba在此任务中不是必需的;检测和分割任务也不是自回归的,但它们遵循长序列特征,因此作者认为探索Mamba在这些任务上的潜力仍然是值得的。
为了实证验证这一假设,作者通过在移除核心tokenmixerSSM的同时堆叠MambaBlock来构建一系列名为MambaOut的模型。实验结果有力地支持了我们的假设。具体来说,对于ImageNet图像分类,MambaOut超过了所有视觉Mamba模型,这表明该任务确实没有必要Mamba。对于检测和分割,MambaOut无法与最先进的视觉Mamba模型的性能相匹配,展示了Mamba对长序列视觉任务的潜力。
分析了SSM中类似RNN的机制,并在概念上得出结论,Mamba适用于具有长序列和自回归特征的任务。
检查了视觉任务的特征,认为SSM对于ImageNet图像分类任务是没必要的,因为此任务不满足自回归和长序列,但是探索SSM在检测分割任务的潜力仍然很有价值,因为这些任务尽管不是自回归的,但是符合长序列的特征。
开发了一系列名为MambaOut的模型,该模型基于GatedCNNBlock,但没有SSM。实验表明,MambaOut在ImageNet图像分类中有效地超过了视觉Mamba模型。
论文名称:MambaOut:DoWeReallyNeedMambaforVision?(Arxiv2024.05)
论文地址:
http ://arxiv.org/pdf/2405.07992
代码链接:
http ://github.com/yuweihao/MambaOut
近年来,Transformer[1]已经成为各种任务的主流骨干,支持BERT、GPT系列和ViT等突出模型。然而,Transformer的tokenmixerSelf-attention有关于序列长度的二次计算复杂度,对长序列任务提出了重大的挑战。为了解决这个问题,前人的工作引入了各种各样的具有线性复杂度的tokenmixer,比如动态卷积[2]、Linformer[3]、Longformer[4]、BigBird[5]和Performer[6]等。最近,出现了一种新的类似RNN的模型浪潮[7][8][9],因其可并行化训练以及对长序列执行有效推理的能力,引起了社区的极大兴趣。值得注意的是,RWKV[8]和Mamba[9]等模型被证明对大语言模型(LLM)的主干[10]有效。
受这些类似RNN的模型展示出的promising的能力的启发,各种研究工作试图将Mamba[9]引入到视觉工作中,例如VisionMamba[11],VMamba[12],LocalMamba[13]和PlainMamba[14]等工作。本着RNN的精神,Mamba的令牌混合器是结构化状态空间模型(StateSpaceModel,SSM)。然而,他们的实验表明,与最先进的基于卷积[15][16]和基于Attention[17][18]的模型相比,基于SSM的视觉模型的性能不足。这就产生了一个研究问题:在视觉任务中,我们真的是否需要Mamba?
在本文中,作者研究了Mamba的性质,并在概念上总结了Mamba非常适合具有两个关键特征的任务:长序列(long-sequence)和**自回归(autoregressive)**,原因是SSM固有的RNN机制[19][20]。但不是所有的视觉任务都具有这两个特征。例如,ImageNet图像分类两者都不符合,COCO目标检测和实例分割,ADE20K语义分割任务仅符合长序列这个特征。另一方面,自回归特性要求每个token仅从前一个token和当前token中聚合信息,也就是一种causalmode的tokenmixing策略。事实上,所有的视觉识别任务都属于是理解领域而非生成任务,意味着模型一次可以看到整个图像。因此,在视觉识别模型中对tokenmixing施加额外的causal约束可能会导致性能下降。虽然这个问题可以通过添加Bidirectional分支[21]来缓解,但这个问题不可避免地存在于每个分支中。
基于上述概念讨论,作者提出以下两个假设:
假设1:图像分类任务没必要使用SSM。因为这个任务既不属于长序列任务,也不属于自回归任务。
假设2:SSM可能对目标检测,实例分割和语义分割任务有益,因为它们尽管不是自回归的,但遵循长序列特征。
为了验证这两个假设,作者开发了一系列模型,称为MambaOut,通过堆叠GatedCNN[22]Block。GatedCNN和MambaBlock之间的主要区别在于SSM的存在,如图1所示。实验结果表明,更简单的MambaOut模型已经超过了视觉Mamba模型的性能,这反过来又验证了本文假设1。作者还通过实验证明,MambaOut在检测和分割任务中没有达到最先进的视觉Mamba模型的性能,表明SSM在这些任务上的潜力,并有效地验证了本文假设2。
在本节中,作者首先讨论Mamba模型适合的任务的特征。接下来,检查视觉识别任务是否符合这些特征。基于检查结果,作者提出了关于Mamba在视觉的必要性的假设。
Mamba模型的tokenmixer是selectiveSSM[9][20],其包含4个input-dependent的参数,并将它们通过下式变为:
然后,SSM的序列到序列变换可以表示为:
其中表示时间步长,表示输入,表示隐藏状态,表示输出,上式2的循环属性将SSM与CausalAttention区分开来。隐藏状态可以看作是存储所有历史信息的固定大小的memory。通过式2,更新此内存,同时其大小保持不变。固定大小意味着memory不可避免地是有损的,但它确保将memory与当前输入集成的计算复杂度保持不变。相反,CausalAttention将来自先前token的所有键和值存储为其memory,它通过为每个新输入添加当前token的键和值来扩展。这个memory在理论上是无损的。然而,随着输入更多的tokens,memory的大小也会增加,从而增加了将memory与当前输入集成的复杂性。图2进一步说明了类似RNN的模型和CausalAttention之间的记忆机制的差异。
SSM的memory在本质上是有损的,它在逻辑上落后于Attention的无损memory。因此,Mamba无法展示处理短序列方面的优势。而Attention在短序列处理上很容易处理好。然而,在涉及长序列的场景中,Attention会因其二次复杂度而发生变化。在这种情况下,Mamba可以突显其在将内存与当前输入合并时的效率,从而平滑地处理长序列。因此,Mamba特别适合处理长序列。
虽然SSM的循环性质(式2)允许Mamba有效地处理长序列,但它引入了一个显著的局限性:只能访问前一个时间步和当前时间步的信息。如图3所示,这种类型的tokenmixing称为CausalMode,可以表示为:
其中和分别表示第个token的输入和输出。由于其因果性质,这种模式非常适合自回归生成任务。
另一种模式称为完全可见模式,其中每个token都可以聚合来自所有先前和后续tokens的信息。这意味着每个token的输出取决于所有tokens的输入:
式中,代表总的token数。完全可见模式适用于理解任务,其中模型可以同时访问所有输入。
默认情况下,Attention处于完全可见的模式,但它可以通过将CausalMask应用于注意力图轻松转变为CausalMode。RNN-like模型由于其循环特性固有地以因果模式运行。由于这种固有特性,RNN-like模型无法转换为完全可见的模式。尽管RNN可以使用双向分支近似完全可见的模式,但每个分支仍然是CausalMode。因此,由于Mamba的循环特性的固有限制,导致它非常适合需要Causaltokenmixing的任务。
总之,Mamba非常适合显示以下特征的任务:
特征1:该任务涉及处理长序列。
特征2:该任务需要Causal的tokenmixing模式。
接下来,作者将讨论视觉识别任务是否表现出这两个特征。
由此,推导出二次项与线性项的比率为:
如果,则二次项的计算负载超过了线性项的计算量。这提供了一个简单的指标来确定任务是否涉及长序列。例如,在ViT-S中有384个channel,阈值,ViT-B中有768个channel,阈值。
对于ImageNet图像分类任务,典型的输入图像大小为,就会有个token,PatchSize为。显然,196远小于和,这表明ImageNet图像分类任务不属于长序列任务。
对于目标检测和实例分割任务,推理图像大小为,在ADE20K语义分割任务,推理图像大小为,给定PatchSize为,token的数量约为。由于以及,因此检测和ADE20K语义分割都可以被认为是长序列任务。
如上图3所示,完全可见模式进行tokenmixing时候不限制范围,而Causal模式限制仅仅访问之前token的信息。视觉识别被归类为理解任务,其中模型一次可以看到整个图像,消除了对tokenmixing的限制。对tokenmixing施加额外的约束可能会降低模型性能。如图3(b)所示,当Causal的限制应用于VisionTransformers(ViT)时,可以观察到性能的显着下降。一般来说,完全可见的模式适用于理解任务,而因果模式更适合自回归任务。这也是为什么BERT和ViT比GPT-1/2和ImageGPT更多用于理解任务。
基于之前的讨论,作者总结了假设,即为视觉识别任务引入Mamba的必要性:
假设1:没有必要在ImageNet上引入SSM进行图像分类,因为此任务不满足特征1或特征2。
假设2:仍然值得进一步探索SSM在检测和分割方面的潜力,因为这些任务与特征1一致,尽管不满足特征2。
式中,TokenMixer()代表tokenmixer,其在GatedCNN和Mamba中分别为:
ImageNet实验结果如下图7所示。本文不包含SSM的MambaOut模型始终优于包含所有模型大小的SSM的视觉Mamba模型。比如,MambaOut-Small模型实现了84.1%的top-1精度,比LocalVMamba-S高0.4%,同时只需要79%的MAC。这些结果有力地支持了本文假设1,即没有必要在ImageNet上引入SSM进行图像分类,这与Occam的razor的原理对齐。
此外,与最先进的卷积和注意力模型相比,视觉Mamba模型目前表现出显著的性能差距。例如,CAFormer-M36[18]采用了简单的可分离卷积和标准注意机制等传统tokenmixer,在准确率上优于所有大小相当的视觉Mamba模型。
作者使用MaskR-CNN[27]作为目标检测和实例分割框架。虽然MambaOut在COCO目标检测和实例分割任务中可以超越一些视觉Mamba模型,但它仍然落后于最先进的视觉Mamba,例如VMamba和LocalVMamba。这种性能差异强调了在长序列视觉任务中集成Mamba的好处,验证了本文假设2。然而,与最先进的卷积注意混合模型TransNeXt[28]相比,视觉Mamba仍然表现出显着的性能差距。视觉Mamba需要通过在视觉检测任务中优于其他最先进的模型来进一步验证其有效性。
作者使用UperNet[29]作为语义分割框架。ADE20K语义分割的性能趋势类似于COCO目标检测的结果。MambaOut可以胜过一些视觉Mamba模型,但不能匹配最先进的Mamba模型的结果。例如,LocalVMamba-T在单尺度(SS)和多尺度(MS)评估中都超过了MambaOut-Tiny0.5mIoU,进一步验证实了本文假设2。此外,与包含了卷积和注意力机制的更先进的混合模型(如SG-Former[30]和TransNeXt[28])相比,视觉Mamba表现出了显著的性能缺陷。因此,还应该通过在分割任务中提供更强的性能来进一步展示其长序列建模的优势。
参考
^AttentionisAllyouNeed
^LiteTransformerwithLong-ShortRangeAttention
^Linformer:Self-AttentionwithLinearComplexity
^Longformer:TheLong-DocumentTransformer
^BigBird:TransformersforLongerSequences
^MaskedLanguageModelingforProteinsviaLinearlyScalableLong-ContextTransformers
^TransformersareRNNs:FastAutoregressiveTransformerswithLinearAttention
^abRWKV:ReinventingRNNsfortheTransformerEra
^abcdMamba:Linear-TimeSequenceModelingwithSelectiveStateSpaces
^Jamba:Ahybridtransformer-mambalanguagemodel
^VisionMamba:EfficientVisualRepresentationLearningwithBidirectionalStateSpaceModel
^VMamba:VisualStateSpaceModel
^LocalMamba:VisualStateSpaceModelwithWindowedSelectiveScan
^Plainmamba:Improvingnon-hierarchicalmambainvisualrecognition
^AConvNetforthe2020s
^ScalingUpYourKernelsto31x31:RevisitingLargeKernelDesigninCNNs
^GoingdeeperwithImageTransformers
^abMetaFormerBaselinesforVision
^CombiningRecurrent,Convolutional,andContinuous-timeModelswithLinearState-SpaceLayers
^abEfficientlyModelingLongSequenceswithStructuredStateSpaces
^Bidirectionalrecurrentneuralnetworks
^LanguageModelingwithGatedConvolutionalNetworks
^LongShort-TermMemory
^LanguageModelingwithGatedConvolutionalNetworks
^MetaFormerIsActuallyWhatYouNeedforVision
^InceptionNeXt:WhenInceptionMeetsConvNeXt
^MaskR-CNN
^abTransNeXt:RobustFovealVisualPerceptionforVisionTransformers
^Unifiedperceptualparsingforsceneunderstanding
^SG-Former:Self-guidedTransformerwithEvolvingTokenReallocation