超越Transformer的长序列建模:全面解析Mamba架构

前言与个人理解

Mamba作为一种新兴的选择性结构状态空间模型,在长序列建模任务中展现了卓越的表现。通过引入全局感受野和动态加权,Mamba有效地缓解了卷积神经网络的局限性,赋予了其与Transformer相似的高级建模能力。更为重要的是,Mamba在实现这些功能的同时,避免了Transformer常见的二次计算复杂性。鉴于相较于其他主流基础模型的显著优势,Mamba展现出了作为视觉基础模型的巨大潜力。研究人员正在积极探索Mamba在各种计算机视觉任务中的应用,催生了众多新兴研究。为了适应计算机视觉领域的迅速进展,本文旨在全面综述视觉Mamba的方法。我们将首先介绍原始Mamba模型的公式,随后深入分析几个具有代表性的视觉Mamba骨干网络,以揭示其核心见解。接下来,我们将相关文献按照不同模式进行分类,包括图像、视频、点云和多模态等。尤其是对于图像应用,我们将其进一步细分为不同任务,以便于更有条理的讨论。最后,我们将探讨视觉Mamba所面临的挑战和未来的研究方向,为这个快速发展的领域提供新的研究视角。

图片开源链接:https://github.com/Ruixxxx/Awesome-Vision-Mamba-Models

主要贡献总结

  1. Mamba的形成:本文提供了Mamba及其状态空间模型操作原理的初步概述。
  2. 主干网络:我们详细检视了多个具有代表性的视觉Mamba骨干网络,阐明了Visual Mamba框架的核心理念和创新。
  3. 应用领域:根据不同模态对Mamba的其他应用进行分类,包括图像、视频、点云和多模态数据等,深入探讨了每一类别,突出Mamba框架如何适应并受益于每种模态。对于涉及图像的应用,我们将其进一步划分为各种任务,包括分类、检测和分割等。
  4. 面临的挑战:通过分析视觉数据的独特特征、算法的潜在机制以及实际应用中的问题,我们探讨了与计算机视觉相关的挑战。
  5. 未来研究方向:我们展望了视觉Mamba未来的研究方向,特别关注数据利用和算法开发的潜在进展。

Mamba模型的理论基础

Mamba是一个新型的序列模型,旨在通过将其参数作为输入的函数来提升基于上下文的推理能力,特别是在结构化状态空间序列模型(S4)的上下文中。Mamba简化了常见的SSM块,形成了一个简化的SSM架构。下面将详细阐述Mamba的核心概念。

状态空间模型 (SSM)

图片图片图片图片

选择性状态空间模型 (Selective SSM)

图片

Mamba架构

图片Mamba是一种简化的SSM架构,与传统的SSM架构不同,后者通常将类似线性注意力的块和多层感知器(MLP)块叠加为Transformer,而Mamba则将这两个基本模块整合为Mamba块。如图2所示,Mamba块可从两个不同的角度进行分析。首先,它使用激活函数替代线性类注意力或H3块中的乘法门;其次,它将SSM的转化融入MLP的主要路径。Mamba的整体架构由重复的Mamba块构成,这些块与标准的规范化层和残差连接交错排列。

Mamba继承了状态空间模型序列长度的线性可伸缩性,同时兼具Transformer的建模能力。Mamba融合了计算机视觉中两种主要基础模型(CNN和Transformer)的显著优点,使其成为一种极具发展潜力的计算机视觉基础模型。不同于依赖于显式存储全部上下文信息来进行基于上下文推理的Transformer,Mamba则利用了一种选择机制。因此,这种选择机制的1D及因果特性成为研究人员关注Mamba在计算机视觉应用中的一个重点。

表征学习的主干

纯Mamba

1)Vim:Vim是一种基于Mamba的架构,直接在类似于ViT的图像补丁序列上操作。首先将输入图像转换为平坦的2D块,然后使用线性投影层对其进行矢量化,并添加位置嵌入以保留空间信息。随后,类令牌被附加到补丁令牌序列中,并将整个令牌序列馈送到由相同Vim块组成的Vim编码器中。如图6所示,Vim块是一个Mamba块,它将后向SSM路径与前向路径集成在一起。

图片

2)VMamba:VMamba识别了将Mamba应用于2D图像所面临的两个挑战,这些挑战源于Mamba中的选择机制的1D和因果特性。因果处理输入数据使得Mamba无法吸收未扫描数据部分的信息,而1D扫描对于涉及局部和全局尺度的2D空间信息的图像并不是最佳选择。

3)Mamba ND:Mamba ND旨在将Mamba扩展到包括图像和视频在内的多维数据。它将1D Mamba层视为一个黑匣子,并探索如何解构和排序多维数据。Mamba ND主要解决数据缺乏预定义排序的问题,并同时面临固有空间维度带来的挑战。为了实现这一目标,Mamba ND仅通过在其维度轴上前向或后向地平坦化数据,从而进行组合的Mamba ND块应用于序列。研究者进行了广泛的实验,探索不同的排序组合。此外,在输入数据的一维划分为多个排序后,采用不同的Mamba层排列,并将序列分解为更小的序列。结果表明,Mamba层的链条及简单的交替方向排序表现优越。Mamba ND块的最终设计如图4(g)所示。

4)PlainMamba:PlainMamba是一种非层次化结构,旨在实现多个目标:(1)非层次结构有助于多层次特征的融合,增强不同尺度的集成;(2)支持多模态数据的有效融合;(3)其简单的结构往往能提供更好的泛化能力;(4)适用于硬件加速优化。

图片

混合Mamba

1)LocalMamba:LocalMamba解决了在Vim和VMamba模型中观察到的一个显著限制,即在单个扫描过程中空间局部令牌之间的依赖性被破坏。为了克服这一问题,LocalMamba将输入图像划分为多个局部窗口,使得SSM在不同方向上执行,同时保持全局SSM的操作。此外,LocalMamba在补丁合并之前实现了空间和通道注意力模块,以增强方向特征的融合并减少冗余。如图4(f)所示,LocalMamba还采用了一种为每一层选择最有效扫描方向的策略,从而优化计算效率。

2)EfficientVMamba:EfficientVMamba引入了高效二维扫描(ES2D)技术,通过对特征图上斑块的异步采样,以降低计算负担。ES2D用于提取全局特征,而并行卷积分支则用于提取局部特征。之后,通过挤压和激励(SE)块对特征进行处理。ES2D、卷积分支和SE块共同构成了有效视觉状态空间(EVSS)块的核心组件。EVSS块的输出是调制的全局和局部特征的总和,EVSS块的早期阶段构成EfficientVMamba,而EfficientNet块则形成其后期阶段。

3)SiMBA:SiMBA旨在解决Mamba在视觉数据集上扩展到大型网络时的不稳定性问题。它提出了一种新的信道建模技术——EinFFT,并借助Mamba进行序列建模。换言之,SiMBA块由Mamba块和EinFFT块组成,并与层归一化(LN)、丢弃和残差连接交织。

关键提升

1)主干:为处理2D图像,首先通过主干模块将其转化为视觉标记序列,通常包括卷积层和线性投影层。位置嵌入的添加是可选的,因为SSM操作本质上具有因果特性。此外,类标记的包含也是可选的。现有方法将图像序列视为在基于Mamba的块中进行SSM变换和卷积运算的1D或2D结构。鉴于扫描技术在这些过程中的整体作用,我们将在下一节中进行系统分类和更详细的讨论。在此节中,我们将区分基于Mamba的层次结构与非层次结构。

2)扫描:选择性扫描机制是Mamba的关键组成部分。然而,其最初针对1D因果序列的设计在适应2D非因果图像时带来了挑战。为应对这些挑战,进行了大量研究工作。在下一节中,我们将这些研究分为三个主要组:扫描模式、扫描轴和扫描连续性。这一分类基于扫描技术的目标,扫描模式处理视觉数据的非因果特性;扫描轴考虑视觉数据内在的高维特性;扫描连续性关注沿扫描路径的空间连续性;扫描采样则将整个图像划分为子图像。这四组的图示如图5所示。

图片

3)模块:上述提到的扫描技术与选择性SSM变换的不同组合形成了多种模块,这些模块是基于Mamba架构的组成部分。在讨论视觉Mamba骨干网络时,我们对此进行了概述,并在相应图中给出了详细说明。这些数字也验证了我们对扫描技术分类的逻辑。这些模块在应用中被广泛使用,下一节将对此进行详细介绍。为了便于理解,最初的Mamba模块简称为Mamba,代表模块如VSS和Vim等名称,修改模块则用星号(*)表示,诸如+CNN的标签则表示集成了类CNN特征。图3和图4包括了一套视觉Mamba模块,包括Vim模块、VSS模块、EVSS模块、PlainMamba模块、LocalMamba模块和Mamba ND模块,方便进行直接比较,并突出这些模块在视觉领域的演变设计。

在本节中,我们按照标准基准展示了各种视觉Mamba骨干网络的性能。表I、表II和表III展示了在ImageNet-1K上的分类结果、通过Mask R-CNN在MS COCO上的目标检测和实例分割结果,以及利用UperNet在ADE20K上的语义分割结果。

图片图片图片

应用

本节将系统地对Mamba在计算机视觉领域的多种应用进行分类与讨论。分类方案及相关文献的概述如图6所示。

图片

A. 图像

图片

1)分类:除了主干进行图像分类以进行表示学习外,Mamba ND还引入了一种处理多维数据的新方法,通过按照行主顺序交替分解不同维度的输入数据。在自然图像分类方面,与基于Transformer的方法相比,该技术以显著更少的参数展示了更优的性能。此外,Mamba ND可以轻松扩展到涉及多维数据的视频动作识别和3D分割等任务。基于Mamba的架构对更大补丁序列的可扩展性使其适用于高分辨率图像(例如,全幻灯片图像和遥感图像)及高维图像(例如3D医学图像)的分析与识别。

2)分割:分割仍然是计算机视觉领域的一项重要研究,具有广泛的实际应用价值。基于CNN和Transformer的模型在此领域的最新进展取得了显著成果。细胞神经网络方法擅长通过卷积运算捕捉局部特征,而Transformer的自注意机制则展现了理解全局上下文的卓越能力。然而,Transformer方法的一个局限性是,随着输入大小的增加,其自注意的计算复杂度呈二次增长,尤其是在处理高分辨率或高维图像时,表现出对窗口外信息建模能力的不足,从而导致性能下降。

3)生成:Mamba架构在多个生成任务中的应用,尤其是在需要长序列交互的情况下,显示出了令人印象深刻的性能潜力。

4)图像恢复:最近,Mamba架构也被广泛应用于多种低级任务,包括图像去雾、曝光校正、增强锐化、超分辨率、去噪、医学图像重建和水下图像增强等。

B. 视频

视频理解是计算机视觉研究的核心方向之一,主要目标是有效掌握长上下文中的时空表示。Mamba凭借其选择性状态空间模型在此领域表现出色,实现了线性复杂性和有效长期动态建模之间的平衡。这种创新方法促进了其在多种视频分析任务中的广泛应用,包括视频目标分割、视频动作识别、视频生成以及表示学习。

图片

C. 多模态

多模态任务在计算机视觉领域至关重要,它们有助于整合多种信息源,丰富视觉数据的理解与分析。这些任务的目的是聚合多种模态,包括文本与视觉信息、包含附加组件(如深度或热图像)的RGB图像,以及各种形式的医学成像数据。然而,实现多模态目标的一个重大挑战是有效捕捉不同模态间的关联性。近期,已有多种方法将Mamba架构应用于多模态任务,包括多模态大语言模型、多模态配准、参考图像分割、时序视频基础、语义分割、运动生成以及医学应用等。

D. 点云

点云作为一种基本的三维表示,提供具有三维坐标的连续空间位置信息。然而,点云的内在无序性和不规则性一直是三维视觉中的一个挑战。受到Mamba线性复杂性及全局建模能力的启发,在点云处理领域开展了多种基于SSM的通用主干研究。

PointMamba直接利用VSS块作为编码器,并提出了一种重排序策略,通过提供更符合逻辑的几何扫描顺序来增强SSM的全局建模能力。PCM结合几何仿射块与Vim块作为基本模块,并提出了一致遍历串行化(CTS)策略,将点云串行化为1D点序列,同时确保空间的连续性。具体而言,CTS通过排列3D坐标的顺序生成六种变体,以全面观测点云数据。PointMamba采用Vim进行长序列建模,并引入基于八叉树的排序机制来生成输入序列,以建立原始输入点之间的因果关系。3DMamba IPF结合了Mamba架构,以顺序处理来自大型场景的大量点云,并集成了稳健且快速可微分的渲染损失,以约束曲面周围的噪声点。3DMambaC则引入了超点生成模块,生成新的形状表示超点,其中包括用于增强采样点特征和预测超点的Mamba编码器。Mamba3D采用了具有通道翻转的双向SSM,并引入局部范数池(LNP)模块,以提取局部几何特征。

挑战

A. 算法

1)可扩展性和稳定性:目前,Mamba架构在应用于大规模数据集(如ImageNet)时面临着稳定性挑战。Mamba在扩展到更广泛网络配置时的不稳定根源尚不明确,通常会导致梯度消失或爆炸,这对其在大规模视觉任务中的应用构成了障碍。

2)因果关系问题:由于Mamba模型最初为因果序列数据设计,将其选择性扫描技术应用于非因果视觉数据面临重大挑战。现有方法通过采用双向扫描等技术来解决此问题,前向和向后扫描被用来相互补偿感受野中单向扫描的局限性,但这仍然是一个未决问题,持续带来挑战。

3)空间信息:Mamba选择性扫描技术固有的1D特性在应用于2D或更高维度视觉数据时会导致关键空间信息的丢失。为了解决这一限制,当前方法通常从各个方向展开图像块,以便在多个维度中进行空间信息的整合,但该问题依旧是待解的挑战。

4)冗余与计算:如前所述,双向扫描方法和多方向扫描的使用可能导致显著的信息冗余与计算需求的增加,从而削弱模型性能,降低Mamba在复杂度上的优势。研究表明,与Transformer模型相比,Mamba模型的GPU消耗并不一致,这也是亟待进一步研究的重要挑战。

B. 应用

1)可解释性:部分研究提供了实验证据,阐明了Mamba模型在自然语言处理(NLP)中的潜在机制,重点关注其上下文学习能力和事实回忆能力。尽管取得了一定进展,但解释Mamba在视觉任务中有效性的原因依然具有挑战性。然而,视觉Mamba的独特学习特征以及其与其他基础模型(如RNN、CNN和ViTs)的相似性仍需更深入的探讨。

2)泛化与鲁棒性:Mamba中的隐藏状态可能会积累或放大特定领域的信息,进而影响其泛化性能。此外,模型固有的1D扫描策略可能会无意中捕捉到特定领域的偏见,而当前的扫描技术往往无法满足对领域不可知信息处理的需求。研究证明了VMamba在对抗性弹性和整体鲁棒性方面的优势,但在处理这些任务时也指出了可扩展性方面的局限性。该研究包括对VMamba的白盒攻击,以检查其新组件在对抗条件下的行为。结果表明,尽管参数Δ在鲁棒性上表现良好,但参数B和C则容易受到攻击。参数间的这种差异导致了VMamba在保持健壮性过程中的可扩展性挑战。此外,结果还指出,VMamba对其扫描轨迹的连续性和空间信息的完整性特别敏感。增强视觉Mamba的泛化能力和鲁棒性仍然是该领域的未解挑战。

未来方向

A. 数据

1)数据效率:考虑到Mamba与CNN相比的计算成本,即使在不依赖大规模数据集的情况下,Mamba也展示出提供最佳性能的巨大潜力。这一特性使Mamba成为各种下游任务和多任务学习中的有前途候选者,特别是在涉及预训练模型自适应的情境下。

2)高分辨率数据:由于SSM架构在理论上简化了计算复杂性,因此其有效处理高分辨率数据(比如遥感和全切片图像)或长期序列数据(如长时间视频帧)的潜力颇具价值。

3)多模态数据:正如Transformer架构已证明其在统一框架内对自然语言与图像的建模能力一样,Mamba模型在处理扩展序列方面的熟练程度大大拓宽了其在多模态学习中的应用潜力。

4)上下文学习:在深度学习的动态环境中,上下文学习发展出越来越复杂和新颖的方法,以应对NLP、CV和多模态领域的复杂任务。这种方法的进步对于突破现有深度学习框架的极限至关重要。Mamba模型凭借其精湛的上下文建模能力和捕获长程依赖关系的能力,展现出在上下文学习应用中实现更深入语义理解和增强性能的潜力。

B. 算法

1)扫描技术:选择性扫描机制作为Mamba模型的核心组成部分,最初针对1D因果序列数据进行优化。为了解决视觉数据的非因果特性,已有多个现有方法采用双向扫描。此外,为了捕捉2D或高维视觉数据中固有的空间信息,当前方法通常扩展扫描方向。尽管有这些调整,仍需更具创新性的扫描方案,以更有效地利用高维非因果视觉数据的全部潜力。

2)融合技术:使Mamba模型适应视觉任务往往会引入冗余,因此有效融合扫描输出特征是值得深入探索的重要领域。此外,计算机视觉的基础模型各具特殊优势,例如细胞神经网络擅长捕捉归纳偏差,ViT则以其强大的建模能力著称。探索融合这些不同网络架构提取的特征以最大化其优势的方法是一项宝贵的研究机会。

3)计算效率:尽管Mamba在序列长度方面实现了线性可扩展性,但因需在多个路径中进行扫描,导致其在视觉任务中的计算需求增加。因此,开发更高效、更有效的视觉Mamba模型是一个重要的研究机会。此外,Mamba模型在计算效率方面并不总是优于Transformer,这凸显了为视觉任务量身定制的优化和硬件感知的Mamba算法的必要性。这为研究提供了一条有前途的途径,尤其是在开发能够减少计算开销同时保持或提升性能的方法方面,提升视觉Mamba模型的计算效率可显著提升其在现实场景中的适用性。

结论

Mamba迅速崛起为一种变革性的长序列建模架构,以其卓越的性能和高效的计算实现而著称。随着其在计算机视觉领域的不断发展,本文对视觉Mamba方法进行了全面的综述。我们首先对Mamba架构进行了深入概述,然后详细检视了具有代表性的视觉Mamba骨干网络及其在各个视觉应用中的广泛应用。应用按不同模式进行系统分类,包括图像、视频、点云和多模态数据等。最后,我们批判性分析了与视觉Mamba相关的挑战,强调了该架构在推进计算机视觉方面尚未开发的潜力。基于以上分析,我们描绘了视觉Mamba未来的研究方向,提供了有价值的见解,可能将对这一动态发展领域的持续与未来进展产生影响。