深度学习(DL)作为一种重要的技术,已经在人工智能(AI)中引发了一场显著的革命,极大地改变了人类的生活方式。作为最具代表性的深度学习技术之一,Transformer架构已经赋能了众多先进的模型,尤其是包含数十亿参数的大型语言模型(LLMs),成为深度学习的基石。
尽管取得了令人印象深刻的成就,但Transformer仍然面临着固有的局限性,特别是由于注意力计算的二次方计算复杂度导致的耗时推理。最近,一种名为Mamba的新型架构,从经典的状态空间模型(SSMs)中汲取灵感,作为构建基础模型的有前途的替代方案出现,它在保持与Transformer相当的建模能力的同时,对于序列长度具有近线性的可扩展性。这激发了越来越多的研究积极探索Mamba在不同领域实现卓越性能的潜力。鉴于这种快速发展,迫切需要一个系统性的回顾,整合现有的Mamba赋能模型,为这种新兴的模型架构提供全面的理解。因此,在这项调查中,我们对最近的与Mamba相关的研究进行了深入调查,涵盖了三个主要方面:基于Mamba的模型的进步、适应Mamba的多样化数据的技术,以及Mamba可以大放异彩的应用领域。具体来说,我们首先回顾了各种代表性深度学习模型的基础知诿,以及Mamba-1&2的详细信息作为初步了解。然后,为了展示Mamba对AI的重要性,我们全面回顾了关注Mamba模型架构设计、数据适应性和应用的相关研究。最后,我们提出了当前局限性的讨论,并探索了各种有希望的研究方向,为未来的调查提供更深入的见解。
我们翻译解读最新论文:Mamba综述,文末有论文链接。
作者:
张长旺,图源:
旺知识
1 引言
在过去的二十年中,深度学习(DL)作为最突出的人工智能(AI)技术,已经在医疗保健[88]、自主系统[36, 60]、推荐系统[104, 230]和金融服务[144, 218]等多个领域引发了一场革命。这一时期见证了众多深度神经网络(DNNs)的出现,这些网络显著改变了人类的生活方式,为个人提供了极大的便利。一个值得注意的例子是U-Net[151, 164],这是视觉领域内的一个强大的深度学习模型,它在医学成像中被广泛用于检查放射线扫描,如MRI和CT扫描。其应用协助识别和诊断疾病,展示了其在这一关键医疗保健领域的有效性[112, 192]。此外,图神经网络(GNNs)被用于处理图结构数据,以支持智能服务,例如推荐系统,这些系统向用户推荐个性化的内容、产品或服务[40, 41, 195]。此外,循环神经网络(RNNs)由于其能够捕捉对准确翻译至关重要的序列和上下文信息,而被广泛采用于机器翻译[119, 167],使来自不同语言背景的个人能够有效地沟通和理解彼此的想法、观点和信息。
在各种深度学习架构中,Transformer最近脱颖而出,并在广泛的应用领域确立了其主导地位[33, 176]。例如,作为最具代表性的大型基础模型,像ChatGPT和GPT4这样的大型语言模型(LLMs)基本上是基于Transformer架构构建的[2, 146, 230]。通过将模型规模扩大到数十亿,并在多样化的数据源上进行训练,这些基于Transformer的模型展示了人类级别的智能,以其在语言理解、常识推理和内容学习方面令人印象深刻的能力[43, 219]。这一显著的成功得益于注意力机制[175],它使基于Transformer的模型能够集中关注输入序列的相关部分,并促进更好的上下文理解。然而,注意力机制也引入了显著的计算开销,随着输入大小的增加而呈二次方增加[124, 235],这在处理长输入时提出了挑战。例如,计算成本的快速增长使得Transformer在处理大量序列时变得不切实际或不可行,从而限制了它们在文档级机器翻译[131]或长文档摘要[94]等任务中的适用性。
最近,一种有前途的架构,结构化状态空间序列模型(SSMs)[58],已经出现,以有效地捕捉序列数据中的复杂依赖性,成为Transformer的强大竞争对手。这些模型,受到经典状态空间模型[89]的启发,可以被认为是循环神经网络和卷积神经网络的融合。它们可以通过递归或卷积操作高效计算,实现与序列长度的线性或近线性扩展,从而显著降低计算成本。更具体地说,作为最成功的SSM变体之一,Mamba实现了与Transformer相当的建模能力,同时保持了与序列长度的线性可扩展性[55],这推动了它成为焦点话题。Mamba首先引入了一个简单但有效的选择机制,使模型能够通过基于输入的参数化SSM参数,无限期地过滤掉无关信息,同时保留必要和相关的数据。然后,Mamba提出了一种硬件感知算法,通过扫描而不是卷积,以递归方式计算模型,实现了在A100 GPU上高达3倍的计算速度。如图1所示,Mamba强大的建模能力,用于复杂和长序列数据,以及近线性的可扩展性,定位Mamba作为新兴的基础模型,准备在各种研究和应用领域引发革命,如计算机视觉[200, 236]、自然语言处理[111, 228]、医疗保健[152, 180, 199]等。例如,朱等人[236]提出了Vim,其速度比DeiT[174]快2.8倍,并且在提取高分辨率图像的特征时节省了86.8%的GPU内存。Dao和Gu[28]展示了SSMs与注意力变体之间的联系,并提出了一种新架构,完善了选择性SSM,实现了在语言建模方面的2-8倍加速。
图1. Mamba模型在不同下游任务中的应用示例。
由于Mamba强大的长序列建模能力及其高效率,大量的文献出现了,专注于在各种下游任务中使用和改进Mamba。鉴于与Mamba相关的研究的显著增长,进行一个全面的文献回顾,讨论未来研究的潜在方向是至关重要的。因此,在这项调查中,我们从几个角度对Mamba进行了全面的回顾,为新来者提供了Mamba内部工作原理的基本理解,同时帮助经验丰富的从业者了解其最新发展。具体来说,剩余的调查组织如下:第2节,我们回顾了各种代表性深度神经网络的背景知识,包括RNNs、Transformers和状态空间模型,而Mamba的细节在第3节中介绍。随后,我们在第4节总结了Mamba基础研究的最新进展,从块设计、扫描模式和内存管理的角度。然后,第5节介绍了将Mamba适应多样化数据的技术,包括序列和非序列数据。此外,第6节介绍了Mamba模型的代表性应用,而挑战和未来方向在第7节中呈现。最后,我们在第8节总结了整个调查。
与我们的调查同时,一些相关的调查已经发布,纯粹关注状态空间模型[137, 184]和视觉Mamba[120, 200, 220]。与这些调查不同,本文以Mamba为中心。它从一个新的角度系统地分析了现有文献,探索了Mamba架构的演变和Mamba模型中使用的数据适应方法。
2 基础知识
Mamba与循环框架的循环神经网络(RNNs)、并行计算和Transformer的注意力机制以及状态空间模型(SSMs)的线性属性密切相关。因此,本节旨在介绍这三种突出架构的概述。
2.1 循环神经网络(RNNs)
RNNs在处理序列数据方面表现出色,因为它们能够保留内部记忆[54]。这类网络在涉及分析和预测序列的众多任务中表现出显著的有效性,例如语音识别、机器翻译、自然语言处理和时间序列分析[69, 170]。为了掌握循环模型的基础,本节将提供标准RNN公式的简要概述。
具体来说,在每个离散时间步骤k,标准RNN特别处理一个向量 𝑥𝑘 ∈ R
D
,连同前一步骤的隐藏状态e ℎ𝑘−1 ∈ R𝑁 ,以产生一个输出向量𝑜𝑘 ∈ R𝑂 并更新隐藏状态到 ℎ𝑘 ∈ R
N
。
隐藏状态充当网络的内存,并保留有关它所见过去输入的信息。这种动态内存允许RNN处理不同长度的序列。正式地,它可以写成
其中Wℎ𝑥 ∈ R𝑁 ×
D
是负责将模型输入处理成隐藏状态的权重矩阵,Wℎℎ ∈ R𝑁 ×
𝑁
是隐藏状态之间的递归连接,W𝑜ℎ ∈ R𝑂×
N
表示用于从隐藏状态生成输出的权重,, 𝑏ℎ ∈ R𝑁 和𝑏𝑜 ∈ R𝑂 对应偏差,tanh表示引入非线性到RNN模型的双曲正切激活函数。换句话说,RNN是非线性递归模型,通过利用隐藏状态中存储的历史知识有效地捕获时间模式。
然而,RNNs有几个局限性。首先,RNNs在有效提取输入序列中的长距离动态方面能力有限。随着信息通过连续的时间步骤传播,网络中权重的重复乘法可能导致信息的稀释或丢失。因此,对于RNNs来说,在进行预测时保留和回忆早期时间步骤的信息变得具有挑战性。其次,RNNs以增量方式处理序列数据,限制了它们的计算效率,因为每个时间步骤都依赖于前一个。这使得并行计算对于它们来说很困难。此外,传统的RNNs缺乏内置的注意力机制,这允许网络捕获输入序列中的全局信息。这种注意力机制的缺失限制了网络选择性地建模数据的关键部分的能力。为了克服这些限制,Transformer和状态空间模型出现了,每种方法都从不同的角度解决了这些挑战。这些两种方法将在后续的小节中进一步阐述。
2.2 Transformers
Transformer[175]是深度学习领域的开创性模型,彻底改变了各种AI应用。它的引入标志着与传统序列到序列模型的显著偏离,通过采用自我注意力机制,促进了对模型输入中全局依赖性的捕获。例如,在自然语言处理中,这种自我注意力能力允许模型理解序列中不同位置之间的关系。
这样的程序然后通过Softmax函数传递,以标准化分数𝐴?并产生注意力权重,定义为:
除了执行单个注意力函数外,多头注意力被引入以增强模型捕获不同类型关系的能力,并为输入序列提供多种视角。在多头注意力中,输入序列并行通过多个自注意力模块进行处理。每个头独立操作,执行与标准自注意力机制完全相同的计算。然后,每个头的注意力权重被结合起来,创建值向量的加权和。这个聚合步骤允许模型利用来自多个头部的信息,并捕获输入序列中的多样化模式和关系。
数学上,多头注意力计算如下:
其中m是注意力头的数量,⊕是连接操作,W
O
是将多头注意力分数投影到最终值的线性变换。
2.3 状态空间模型
状态空间模型(SSMs)是一种传统的数学框架,用于描述系统随时间的动态行为[89]。近年来,SSMs在控制理论、机器人技术和经济学等多个领域中发现了广泛的应用[58 59]。在其核心,SSMs通过一组隐藏变量,即“状态”,来体现系统的行为,使其能够有效地捕获时间数据依赖性。与RNNs不同,SSMs是线性模型,具有关联属性。具体来说,在经典的状态空间模型中,制定了两个基本方程,即状态方程和观测方程,通过当前时间t的N维隐藏状态ℎ(𝑡) ∈ R𝑁 来模拟输入𝑥 (𝑡) ∈ R和输出 𝑦(𝑡) ∈ R之间的关系。该过程可以写成
其中ℎ′(𝑡) 是当前状态ℎ(𝑡)的导数,A ∈ R𝑁 ×
𝑁
是描述状态如何随时间变化的状态转移矩阵,B ∈ R𝑁 ×1是控制输入如何影响状态变化的输入矩阵,C ∈ R1×
𝑁
表示基于当前状态生成输出的输出矩阵,D ∈ R表示决定输入如何直接影响输出的命令系数。一般来说,大多数SSMs在观测方程中省略了第二项,即,设置D𝑥 (𝑡) = 0,这可以被认为是深度学习模型中的跳过连接。
2.3.1 离散化。
为了符合机器学习设置对各种现实世界场景的要求,SSMs必须经历一个离散化过程,将连续参数转换为离散参数。离散化方法通常旨在将连续时间划分为𝐴?个具有相等积分区域的离散间隔。为了实现这一目标,作为最具代表性的解决方案之一,零阶保持(ZOH)[138, 225]成功地应用于SSMs,它假设函数值在间隔Δ = [𝑡𝑘−1, 𝑡𝑘]内保持恒定。在ZOH离散化之后,SSM方程可以重写为
其中A = exp(ΔA),B = (ΔA)−1(exp(ΔA) − I) · ΔB,k是离散时间步长。从这些公式中,很明显离散SSM具有类似于循环神经网络的结构,因此离散SSM可以完成与Transformer基础模型相比,具有更高效率的推理过程。
2.3.2 卷积计算。
作为线性系统,离散SSM具有关联属性,因此可以无缝集成到卷积计算中。更具体地说,它可以独立地计算每个时间步的输出,如下所示:
通过创建一组卷积核s K = (CB, ..., CA𝑘B, ...),递归计算可以转换为卷积形式:
其中x = [𝑥0, 𝑥1, ...]和y = [𝑦0, 𝑦1, ...] ∈ R L分别表示输入和输出序列,而𝐿 是序列长度。在这个情况下,输入矩阵B ∈ R𝑁 ×
𝑁
,输出矩阵C ∈ R𝐷×
𝑁
,和命令矩阵D ∈ R𝐷×
𝐷
,而状态转移矩阵保持不变,即A ∈ R𝑁 ×
𝑁
。
2.3.3 RNN、Transformer和SSM之间的关系。
图2描述了循环神经网络(RNN)、Transformer和状态空间模型(SSM)的计算算法。
图2. 代表性模型架构的图示,即循环神经网络(RNN)、Transformer和状态空间模型(SSM)。
(a) RNNs在非线性递归框架内运作,便于在自回归推理过程中快速输出。
(b) Transformers在多个查询-键对上并行执行矩阵乘法,便于并行训练。
(c) SSMs通过线性属性适应递归和卷积计算,融合了RNNs和Transformers的优势,允许SSMs进行递归推理和并行训练。
尽管如此,传统的时间不变SSM在上下文感知建模方面不足,导致在特定任务中的性能下降。
一方面,传统的RNN在非线性递归框架内运作,每个计算仅依赖于前一个隐藏状态和当前输入。虽然这种格式允许RNN在自回归推理过程中快速生成输出,但它阻碍了它们充分利用GPU并行计算,导致模型训练速度变慢。另一方面,Transformer架构在多个查询-键对上并行执行矩阵乘法,可以有效地跨硬件资源分布,从而加快了基于注意力模型的训练。然而,当从基于Transformer的模型生成响应或预测时,推理过程可能会很耗时。例如,语言模型的自回归设计需要顺序生成输出序列中的每个标记,这要求在每一步重复计算注意力分数,导致推理时间变慢。如表1所示,与RNNs和Transformers不同,它们仅限于支持一种类型的计算,离散SSMs具有支持递归和卷积计算的灵活性,这得益于它们的线性属性。这种独特的能力允许SSMs不仅实现有效的推理,还实现并行训练。然而,应该注意的是,大多数传统的SSM是时间不变的,这意味着它们的A、B、C和Δ与模型输入x无关。这将限制上下文感知建模,导致SSMs在某些任务中的性能下降,如选择性复制[55]。
表1. 在自回归序列建模任务中,三种主要架构-RNNs、Transformers和SSMs-的优缺点比较。
3 Mamba
为了解决传统SSM在上下文感知能力方面的不足,Mamba被[55]提出作为一种潜在的替代方案,承诺成为一个通用的序列基础模型骨干。最近,Mamba-2[28]提出了结构化状态空间对偶性(SSD),建立了一个强大的理论框架,将结构化SSMs与各种形式的注意力联系起来,允许我们将最初为Transformer开发的算法和系统优化转移到SSMs。在这一节中,我们将简要而清晰地介绍Mamba和Mamba-2。
3.1 Mamba-1:具有硬件感知算法的选择性状态空间模型
传统的SSM在建模文本和其他信息密集型数据方面表现出有限的有效性[55],阻碍了它们在深度学习中的进展。在追求赋予SSMs与Transformer相当的建模能力的过程中,Gu和Dao[55]基于结构化状态空间模型提出了三种创新技术,即基于高阶多项式投影算子(HiPPO)的内存初始化、选择机制和硬件感知计算,如图3所示。这些技术旨在增强SSMs在长距离线性时间序列建模方面的能力。特别是,初始化策略建立了一个一致的隐藏状态矩阵,有效地促进了长距离记忆。然后,选择机制使SSM能够获得内容感知表示。最后,Mamba设计了两种硬件感知计算算法,即并行关联扫描和内存重计算,以提高训练效率。
3.1.1 HiPPO-based Memory Initialization。
建模和从序列数据中学习是当代机器学习的基础挑战,构成了各种任务的基础,包括语言建模、语音识别和视频处理。在建模复杂和长期时间依赖性方面,内存是一个基本组成部分,包括存储和整合来自先前时间步骤的信息[73]。与RNNs类似,在SSMs中保留和遗忘历史隐藏状态(即矩阵A)在实现令人满意的性能方面起着关键作用。在以前的结构化状态空间序列模型(SSMs)中,已经提出了特殊的初始化建议,特别是在数据可用性有限的情况下。这些特殊的初始化在各种情况下都证明是有益的,包括在复杂值模型的情况下。同样,Mamba主要关注隐藏状态矩阵A的初始化,以捕获复杂的时间依赖性。这是通过利用HiPPO理论[56]和创新的缩放Legendre度量(LegS)来实现的,确保仔细考虑完整的历史上下文而不是有限的滑动窗口。具体来说,HiPPO-LegS为所有历史数据点分配均匀的权重,可以表示为:
其中
n
是多项式的数目, k表示特定的离散时间步骤。基于HiPPO理论,Mamba引入了两种简单的初始化方法,即S4D-Lin和S4D-Real[57],如下所示:
其中n是A的第n个元素,对于所有输入维度𝑑 = 1, 2, ..., D。通过这种初始化,模型可以学习长期依赖性,通过压缩和重建输入信息信号,经历较小的新步骤退化和较大的旧步骤退化。根据公式,HiPPO-LegS具有有利的理论属性:它在输入时间尺度上保持一致,并且提供快速计算[56]。此外,它具有有界的梯度和近似误差,有助于参数学习过程。
3.1.2 选择机制。
传统的状态空间模型由于时间不变性属性,无法根据特定模型输入(即内容感知建模能力)产生个性化输出。为了为SSM提供类似于注意力机制的能力,Mamba设计了一个时变选择机制,根据模型输入参数化权重矩阵。这种创新使SSM能够无限期地过滤掉无关信息,同时保留相关细节。正式地说,选择机制涉及设置间隔Δ,并将矩阵B、C作为输x ∈ R𝐵×𝐿×
D
的函数,可以公式化为:
其中SB ∈ R𝐵×𝐿×𝑁 , SC ∈ R𝐵×𝐿×
𝑁
,和SΔ ∈ R𝐵×𝐿×𝐷 是选择性空间矩阵,它们是输入的函数,以实现内容感知建模。𝐵, 𝐿, 𝐷, 和
N
分别代表批量大小、输入长度、输入特征大小和隐藏通道数。值得注意的是,WB ∈ R𝑁 ×𝐷 , WC ∈ R𝑁 ×𝐷 , 和 WΔ ∈ R𝐷×1是相应组成部分的选择权重(即线性参数化投影),而BroadCast
D
意味着将结果广播到所有维度𝑑 = 1, 2, .., 𝐷。随后,选择性SSMs通过使用常见的统计技术,零阶保持(ZOH)[138],进行离散化,如下所示:
其中SA ∈ R𝐵×𝐿×𝐷×𝑁 和 SB ∈ R𝐵×𝐿×𝐷×𝑁 分别是选择性状态转移矩阵和输入矩阵,它们成为输入x的函数,即
它根据输入x生成输出ut y ∈ R𝐵×𝐿×𝐷。请注意,Mamba中的时变选择机制具有与Transformer中的注意力机制类似的结构,即两者都根据输入及其投影执行操作,这允许Mamba的SSM实现灵活的内容感知建模。然而,它失去了与卷积的等价性,这对其效率产生了负面影响。
3.1.3 硬件感知计算。
选择机制旨在超越线性时间不变模型的限制。尽管如此,它对有效训练提出了挑战:SSMs的卷积核变得依赖于输入,导致无法进行并行计算。为了解决这个问题,Mamba利用了两种计算技术,即并行关联扫描(也称为并行前缀和)[64]和内存重计算。首先,并行关联扫描利用线性关联计算的属性和现代加速器(GPU和TPU)的并行性,以内存高效的方式执行选择性SSMs的计算。更具体地说,平行关联扫描围绕构建给定输入的平衡二叉树,并扫掠它从叶子到根。换句话说,平行关联扫描首先通过从叶子到根的遍历(即Sweep-Up),在树的内部节点创建部分和。然后,它逆转遍历,从根开始向上扫掠树,使用部分和构建整个扫描(即Sweep-Down)。
另一方面,Mamba利用了传统的重计算方法,以减少训练选择性SSM层时的整体内存需求。具体来说,Mamba在前向传递的并行关联扫描中避免存储大小为 (𝐵, 𝐿,𝐷, 𝑁) 的中间状态,以防止内存扩展。相反,它在后向传递中重新计算这些中间状态以进行梯度计算。通过这样做,重计算避免了在GPU内存单元之间读取𝑂(𝐵𝐿𝑁𝐷) 元素的必要性。除了优化扫描操作的内存需求外,Mamba-1还将重计算的使用扩展到整个SSM层的效率,包括投影、卷积和激活,这些通常需要大量的内存资源,但可以快速重新计算。
3.2 Mamba-2:状态空间对偶性
Transformers在深度学习的成功中发挥了关键作用,为各种领域带来了革命性的变化。它们启发了各种技术的发展,如参数高效微调[95]、灾难性遗忘缓解[96]和模型量化[196],旨在从不同角度提高模型性能。为了使状态空间模型能够访问并从最初为Transformer开发的宝贵技术中受益,Mamba-2[28]引入了一个全面的框架,称为结构化状态空间对偶性(SSD),它建立了SSMs和不同形式的注意力之间的理论联系。正式地,
其中M表示SSMs的矩阵形式,使用序列半可分表示,M𝑗𝑖 = CT𝑗A𝑗:𝑖B
i
。
值得注意的是,C𝑗 和 B𝑖 分别代表与输入令牌x𝑗和x
i
相关的选择性空间状态矩阵。A𝑗:𝑖 表示与输入令牌从j到i的范围相关的选择性隐藏状态矩阵。本质上,SSD表明,Transformer使用的注意力机制和SSM中使用的线性时变系统都可以被视为半可分矩阵变换。此外,Dao和Gu[28]还证明了选择性SSM等同于使用半可分掩码矩阵实现的结构化线性注意力机制。
基于SSD,Mamba-2设计了一种更高效的硬件计算,通过块分解矩阵乘法算法。具体来说,通过将状态空间模型视为半可分矩阵,Mamba-2将计算分解为矩阵块,在这些块中,对角块代表内部块计算。相比之下,非对角块通过SSM的隐藏状态表示内部块计算。这种方法使Mamba-2能够实现比Mamba-1的并行关联扫描快2-8倍的训练过程,同时保持与Transformer的竞争性。
3.3 Mamba块
在这一小节中,我们总结了Mamba-1和Mamba-2的块设计。图4展示了这两种架构的比较。Mamba-1是从SSM中心视角出发,其中选择性SSM层负责将输入序列X映射到Y。在这个设计中,(A, B, C)的线性投影应用于初始线性投影之后,以创建X。然后,输入令牌和状态矩阵通过选择性SSM单元传递,使用并行关联扫描,以产生输出Y。之后,Mamba-1使用跳过连接来鼓励特征重用,并缓解在模型训练过程中经常发生的退化问题。最后,通过堆叠这种块与标准归一化和残差连接交错,构建了Mamba模型。
至于Mamba-2,它引入了SSD层,旨在从[X, A, B, C]到Y创建一个映射。这是通过同时处理[X, A, B, C]并以与标准注意力架构生成Q、K、V投影相同的方式在块的开始进行单个投影来实现的。换句话说,Mamba-2块通过删除顺序线性投影来简化Mamba-1块,这使得SSD结构的计算速度比Mamba-1中的并行选择性扫描快。此外,在跳过连接后添加了一个归一化层,旨在提高训练稳定性。
表2. 利用Mamba架构的代表性开放获取基础模型。
4 Mamba模型的进步
状态空间模型和Mamba最近被探索,并成为基础模型骨干的一个有前途的替代方案。如表2所示,大规模Mamba基础模型不仅在学术研究中蓬勃发展,而且在工业界也取得了显著进展,如Falcon Mamba 7B和Mistral 7B,通过在GPU上成功训练证明了它们的有效性。尽管如此,Mamba架构仍然面临挑战,如内存丢失、对多样化任务的泛化能力不足以及与基于Transformer的语言模型相比在捕获复杂模式方面的劣势。为了克服这些挑战,已经做出了大量努力来改进Mamba架构。现有的研究主要集中在修改块设计、扫描模式和内存管理方面。这一节将介绍这三个方面的几个重要技术,并在表3中提供相关研究的总结。
4.1 块设计
Mamba块的设计和结构对Mamba模型的整体性能有重要影响,使其成为一个新的研究焦点。如图5所示,基于构建新Mamba块的不同方法,现有研究可以分为三类:a)集成方法旨在将Mamba块与其他知名模型集成,以便在保持有效性和效率的同时实现平衡;b)替代方法尝试将Mamba块作为高级模型框架中主要层的替代品;c)修改方法专注于修改经典Mamba块内的组件。以下小节将详细介绍这些方法。
4.1.1 集成。
鉴于Mamba在捕获长期动态方面的卓越能力,它已广泛与其他模型集成,利用其优势为特定场景提供强大的框架。集成特别包括先进的模型,如Transformers、图神经网络(GNNs)、循环神经网络(RNNs)、卷积神经网络(CNNs)和尖峰神经网络(SNNs)。下面描述了一些具体的例子。
• 基于Transformer的模型在许多任务中表现出色,但它们的二次方计算复杂度在推理过程中仍然阻碍了它们[58]。为了实现高效的生成,一些研究人员提出了将Mamba块与基于Transformer的模型结合起来。例如,Jamba[111]结合了Transformer和Mamba层的块,以应对长期内容的自然语言处理任务,利用了两种模型家族的优势。注意力-Mamba混合模型的性能优于单独的Transformer和Mamba模型,与普通Transformer模型相比,吞吐量更好。Mambaformer[201]利用混合框架预测多个时间序列,包括汇率、小时电力消耗和电力负荷,它在内部结合了Mamba块和Transformer层,分别用于长期和短期依赖性。由于Mamba和Transformer的集成,Mambaformer在长期-短期时间序列预测方面优于基于Transformer的预测器。
• GNN通过消息传递机制在捕获邻近关系方面展现出了有希望的潜力,其中信息通过连接图在堆叠层中传播。然而,这些模型面临一个重大限制,即过度平滑[20],特别是当试图捕获高阶邻接信号时。为了解决这一挑战,Mamba已被用于图表示学习[103, 115, 179, 205]。例如,Graph Mamba[9]将图结构数据重新格式化为特定顺序的序列令牌,并利用Mamba块中的选择性SSM层构建了一个新的Graph Mamba Network (GMN)架构,它在图表示学习能力方面表现出色,特别是在需要高阶节点依赖性的数据处理方面。
• RNN基础模型在捕获时间动态方面取得了显著成果。然而,RNNs仍然面临重大挑战,包括耗时的递归训练和隐藏状态的内存容量限制。受到最近Mamba基础架构出现的启发,一些研究人员开发了Mamba块与RNNs的融合。例如,VMRNN[171]在时空预测方面取得了最先进的性能,与基于递归和无递归的方法相比,减少了浮点运算(FLOPs)。它通过引入一种新的递归单元来实现这一点,该单元将Mamba块与长短期记忆(LSTM)结合起来。
• CNN基础方法受到局部接受域的限制,导致在捕获全局和长期语义方面表现不佳[55]。众所周知,状态空间模型在学习长期模式方面具有优越能力,一些研究[107, 188, 206]探索了利用Mamba块增强CNN基础模型的潜力,特别是在计算机视觉领域。例如,MedMamba[216]和nnMamba[53]展示了如何通过集成视觉Mamba块来提高CNN在图像分析任务中的性能。
• SNN最近被提出作为一种有前景的网络架构,灵感来自大脑中生物神经元的行为:通过离散尖峰在神经元之间传递知识。SNN的一个关键优势在于其潜在的低功耗实现,因为它们可以利用神经活动的稀疏和事件驱动特性。受到SNN和SSM在低功耗实现方面的启发,以及SSM在长期学习能力方面的优越性,一些开创性的研究深入研究了将这两种方法结合起来。例如,SpikeMba[106]将它们结合起来处理对突出对象的信心偏差,并在视频序列中捕获持久依赖性。通过广泛的评估,作者声称将这两种模型结合起来提高了视频定位任务的有效性,精确地时刻检索和高光检测。
4.1.2 替代。
受到选择性SSM在高效计算和长序列学习方面的卓越能力的启发,采用Mamba模块替代经典建模框架中的关键组件,如U-Net[151]和扩散模型[70],已经引起了很多关注。通过引入选择性SSM层,这些方法实现了长距离学习和高效计算,以满足特定任务的需求。下面,我们展示了使用Mamba模块替代的一些实例,特别是高级框架,如U-Net和扩散模型。
• U-Net。许多努力[110, 163, 180, 181]已经做出,以将U-Net与Mamba在捕获复杂和广泛语义方面的能力结合起来,以推进计算机视觉任务中的模型性能。例如,Mamba UNet[180]专门使用视觉Mamba块构建了一个类似U-Net的模型(即,一个编码器-解码器模型,融合了跳跃连接),用于医学图像分割。他们的评估表明,Mamba-UNet超过了几种U-Net变体,这可以归因于Mamba块在处理长距离补丁序列方面的有效性和效率。
• 扩散模型。一些努力[46, 48, 136]已经尝试构建一种新型的扩散模型,扩散状态空间模型(DiS),它用状态空间骨干替代了典型的骨干(例如,CNNs、注意力、U-Nets)。鉴于Mamba块在适应长距离依赖性方面的显著效率和有效性,DiS被区别于使用扩散模型生成更长序列[46]。例如,Oshima等人[136]提出了一个基于Mamba的扩散模型,显著降低了长视频序列的内存消耗,同时与基于Transformer的模型相比,在性能指标上仍然保持了竞争力。此外,MD-Dose[48]和P-Mamba[211]在扩散模型的反向过程中构建了使用Mamba块的噪声预测器,最终为医学图像处理生成了特定目标。
• 其他。除了U-Net和扩散模型,还有一些替代品。例如,Res-VMamba[19]在残差学习框架中采用了视觉Mamba块,用于食品类别分类。此外,SPMamba[101]采用了TF-GridNet[190],这是最近开发的一种时频模型,作为其基础架构,然后在Transformer组件之后使用双向Mamba块。这种调整使模型能够高效地包含更广泛的上下文信息,用于语音分离任务。
4.1.3 修改。
除了直接使用Mamba块的集成和替代方法外,还进行了一些努力来修改Mamba块,目的是在不同场景中增强其性能。例如,Jamba[111]借鉴了专家混合(MoE)[45, 82]的概念,使他们的混合(Transformer-Mamba)仅解码器模型能够以更少的计算进行预训练,并允许灵活的目标特定配置。值得注意的是,Jamba模型(56B可用参数,12B活动参数,4GB KV缓存)需要比代表性的基于Transformer的语言模型LLaMA-2-7B(6.7B可用参数,12B活动参数,128GB KV缓存)小32倍的KV缓存,同时提供了更多的可用和活动参数。这使得Jamba能够在单个A100 GPU(80GB)上吞吐140K的上下文长度,这是LLaMA-2-70B支持的长度的七倍。除了MoE,一些研究提出了将SSM层修改为K路结构,涉及使用并行SSM单元处理模型输入,允许从多个角度捕获信息和知识。例如,Sigma[178]开发了一种新的基于Mamba的视觉编码器,它通过使用并行SSM层来处理多模态输入。UltraLight VM-UNet[194]提出了一种具有并行SSM单元的视觉Mamba层,该层在不同通道中处理深度特征。总之,通过实施这些修改(即,K路,MoE),这些基于Mamba的模型获得了增强的学习能力,特别是在处理多模态输入和快速适应多尺度任务方面。此外,一项开创性研究,Mamba®,提出了一种新的方法,建议在将输入通过SSM层之前,将寄存器均匀地整合到视觉输入令牌中。这种修改旨在增强图像补丁序列方向的表示,从而使Mamba块的单向推理范式适用于视觉任务。尽管取得了这些成功,但探索修改Mamba块仍然是一个有前景但尚未充分探索的领域。
4.2 扫描模式
并行关联扫描操作是Mamba模型中的一个关键组件,旨在解决选择机制引起的计算问题,加速训练过程,并减少内存需求。它通过利用时间变化SSMs的线性属性来设计硬件级别的内核融合和重计算。然而,Mamba的单向序列建模范式阻碍了对各种数据(如图像和视频)的全面学习过程。为了缓解这个问题,一些研究专注于设计有效的扫描方法,以增强模型性能并促进Mamba模型的训练过程。如图6所示,现有研究集中在开发扫描模式技术,可以分为两类:1) Flatten Scan方法从平铺的角度处理模型输入;2) Stereo Scan方法跨维度、通道或尺度扫描模型输入。