【大模型架构】Mamba 最新综述

Mamba 综述

论文链接:https://arxiv.org/pdf/2408.01129

深度学习作为一项重要技术,已在人工智能(AI)领域引发显著的革命,导致人类生活方式的巨大变化。作为最具代表性的架构,变换器(Transformers)赋能了众多先进模型,特别是包含数十亿参数的大型语言模型(LLMs),并成为深度学习的基石。尽管取得了令人印象深刻的成就,变换器仍然面临固有的局限性,尤其是由于注意力计算的二次计算复杂度而导致的耗时推断。最近,一种名为 Mamba 的新架构应运而生,受到了经典状态空间模型的启发,成为构建基础模型的有前景的替代方案,提供了与变换器相当的建模能力,同时在序列长度方面保持近线性的可扩展性。这引发了越来越多的研究积极探索 Mamba 在不同领域取得优异表现的潜力。鉴于这种快速的发展,迫切需要对现有的 Mamba 赋能的模型进行系统的审查,以提供对这一新兴模型架构的全面理解。因此,在这项调查中,我们深入研究了最近与 Mamba 相关的研究,重点关注三个主要方面:基于 Mamba 的模型的进展、将 Mamba 适应于不同数据的技术,以及 Mamba 能够出色表现的应用。具体而言,我们首先回顾了各种代表性深度学习模型的基础知识以及 Mamba-1&2 的细节作为前提。然后,为了展示 Mamba 对 AI 的重要性,我们全面回顾了聚焦于 Mamba 模型架构设计、数据适应性和应用的相关研究。最后,我们讨论了当前的局限性,并探索了各种有希望的研究方向,以提供对未来研究的更深入见解。

附加关键词和短语:状态空间模型,Mamba,序列建模,基础模型,语言模型


1 引言

在过去的二十年中,深度学习作为最杰出的人工智能(AI)技术,已在医疗保健 [88]、自主系统 [37, 60]、推荐系统 [104, 228] 和金融服务 [ 144 , 216 ] \left\lbrack {{144},{216}}\right\rbrack [144,216] 等多个领域引发了革命。这一时期见证了许多深度神经网络(DNNs)的出现,这些网络显著改变了人类的生活方式,为个人提供了巨大的便利。其中一个显著的例子是 U-Net [ 151 , 164 ] \left\lbrack {{151},{164}}\right\rbrack [151,164] ,这是一个在视觉领域内强大的深度学习模型,广泛应用于医疗成像,用于检查放射学扫描,如 MRI 和 CT 扫描。它的应用有助于识别和诊断疾病,展示了其在这一重要医疗领域的有效性 [112, 191]。此外,图神经网络(GNNs)被用来处理图结构数据,以支持智能服务,比如向用户建议个性化内容、产品或服务的推荐系统 [41, 42, 194]。此外,由于其捕捉准确翻译所需的顺序和上下文信息的能力,递归神经网络(RNNs)在机器翻译中被广泛采用 [ 119 , 167 ] \left\lbrack {{119},{167}}\right\rbrack [119,167] ,使来自不同语言背景的个人能够有效地交流和理解彼此的思想、观点和信息。

在各种深度学习架构中,变压器最近脱颖而出,并在广泛的应用中建立了其主导地位 [34, 175]。例如,作为最具代表性的基础大型模型,大型语言模型(LLMs)如 ChatGPT 基本上是基于变压器架构构建的 [2, 146, 228]。通过将模型规模扩展到数十亿规模,并在各种数据源上进行训练,这些基于变压器的模型展示了人类级别的智能,其在语言理解、常识推理和内容内学习方面表现出色 [44, 217]。这一显著成功得益于注意力机制 [174],该机制使基于变压器的模型能够集中关注输入序列的相关部分,促进更好的上下文理解。然而,注意力机制也引入了显著的计算开销,这种开销随着输入大小呈平方增长 [ 124 , 233 ] \left\lbrack {{124},{233}}\right\rbrack [124,233] ,这在处理较长输入时带来了挑战。例如,计算成本的快速增长使得变压器在处理较大序列时变得不切实际或不可行,从而限制了其在文档级机器翻译 [131] 或长文档摘要 [94] 等任务中的适用性。

最近,一种前景广阔的架构,结构状态空间序列模型(SSMs) [58],已出现以有效捕获序列数据中的复杂依赖关系,成为对变压器的强大竞争者。这些模型受到经典状态空间模型的启发 [89],可以被视为递归神经网络和卷积神经网络的融合。它们可以通过递归或卷积操作高效计算,随着序列长度实现线性或近线性缩放,从而显著降低计算成本。更具体地说,作为最成功的SSM变体之一,Mamba 在建模能力上与变压器相当,同时在序列长度上保持线性可扩展性 [55],使其成为焦点话题。Mamba 首先引入了一种简单而有效的选择机制,使模型能够过滤掉无关的信息,同时无限期地保留必要且相关的数据,通过根据输入对SSM参数进行参数化。然后,Mamba 提出了一个硬件感知算法,通过扫描而不是卷积以递归方式计算模型,在 A100 GPU 上实现了高达 3 × 3 \times 3× 的计算速度。正如图 1 所示,对于复杂和长序列数据的强大建模能力以及近线性可扩展性,使 Mamba 成为一个新兴的基础模型,准备在计算机视觉 [199, 234]、自然语言处理 [ 111 , 226 ] \left\lbrack {{111},{226}}\right\rbrack [111,226] 、医疗保健 [ 152 , 179 , 198 ] \left\lbrack {{152},{179},{198}}\right\rbrack [152,179,198] 等多个研究和应用领域实现革命性变化。例如,Zhu 等 [234] 提出 Vim,其提取高分辨率图像特征的速度比 DeiT [173] 快 2.8 × {2.8} \times 2.8× ,并节省 86.8 % {86.8}\% 86.8% 的 GPU 内存。Dao 和 Gu [29] 显示 SSM 与各种注意力变体之间的联系,并提出了一种新架构,改进选择性 SSM,实现语言建模上快 2 − 8 × 2 - 8 \times 28× 的速度。

图 1. Mamba 基础模型在不同下游任务中的应用示例。

受到 Mamba 强大长序列建模能力和高效性的启发,涌现出大量文献,专注于在各种下游任务中应用和改进 Mamba。鉴于与 Mamba 相关的研究显著增加,对现有文献进行全面审查并讨论未来研究的潜在方向至关重要。在本次调查中,我们因此从多个角度对 Mamba 进行全面审查,以为新手提供 Mamba 内部工作原理的基本理解,同时帮助经验丰富的从业者了解其最新进展。具体而言,余下的调查组织如下:在第二节中,我们回顾各种代表性深度神经网络的背景知识,包括 RNN、变压器和状态空间模型,而 Mamba 的详细信息将在第三节介绍。随后,我们在第四节总结了基于 Mamba 的研究在块设计、扫描模式和内存管理方面的最新进展。然后,第五节介绍将 Mamba 适配于多样数据的技术,包括序列数据和非序列数据。此外,第六节介绍 Mamba 模型的代表性应用,第七节则呈现挑战和未来方向。最后,我们在第八节总结整个调查。

与我们的调查同时发布了几篇相关的调查研究,纯粹聚焦于状态空间模型 [137, 183] 和视觉 Mamba [120, 199, 218]。与这些调查不同,本文集中于与 Mamba 相关的研究。它从一个新颖的角度系统分析现有文献,以探讨 Mamba 架构的演变和在 Mamba 基础模型中利用的数据适配方法。

2 初步

Mamba 深入地与递归神经网络(RNN)的递归框架、变压器的并行计算和注意力机制,以及状态空间模型(SSMs)的线性属性紧密相连。因此,本节旨在概述这三种突出的架构,为全面理解 Mamba 提供前提。

2.1 递归神经网络

递归神经网络(RNNs)由于其保持内部记忆的能力,在处理顺序数据方面表现卓越 [54]。这样的网络已在涉及分析和预测序列的众多任务中证明了其显著的有效性,例如,语音识别、机器翻译、自然语言处理和时间序列分析 [ 69 , 169 ] \left\lbrack {{69},{169}}\right\rbrack [69,169] 。为了掌握递归模型的基础,本节将简要概述标准 RNN 公式。

具体而言,在每个离散时间步 k k k ,标准 RNN 特别处理一个向量 x k ∈ R D {x}_{k} \in {\mathbb{R}}^{D} xkRD 和前一步的隐藏状态 h k − 1 ∈ R N {h}_{k - 1} \in {\mathbb{R}}^{N} hk1RN ,生成一个输出向量 o k ∈ R O {o}_{k} \in {\mathbb{R}}^{O} okRO ,并将隐藏状态更新为 h k ∈ R N {h}_{k} \in {\mathbb{R}}^{N} hkRN 。隐藏状态作为网络的记忆,保持对其见到的过去输入的信息。这种动态记忆允许 RNN 处理不同长度的序列。形式上,可以写为 h k = tanh ⁡ ( W h x x k + W h h h k − 1 + b h ) (1) {h}_{k} = \tanh \left( {{\mathbf{W}}_{hx}{x}_{k} + {\mathbf{W}}_{hh}{h}_{k - 1} + {b}_{h}}\right) \tag{1} hk=tanh(Whxxk+Whhhk1+bh)(1) o k = W o h h k + b o (2) {o}_{k} = {\mathrm{W}}_{oh}{h}_{k} + {b}_{o} \tag{2} ok=Wohhk+bo(2) 其中 W h x ∈ R N × D {\mathbf{W}}_{hx} \in {\mathbb{R}}^{N \times D} WhxRN×D 是负责将模型输入处理为隐藏状态的权重矩阵, W h h ∈ R N × N {\mathbf{W}}_{hh} \in {\mathbb{R}}^{N \times N} WhhRN×N 是隐藏状态之间的递归连接, W o h ∈ R O × N {\mathbf{W}}_{oh} \in {\mathbb{R}}^{O \times N} WohRO×N 表示用于生成来自隐藏状态的输出的权重, b h ∈ R N {b}_{h} \in {\mathbb{R}}^{N} bhRN b o ∈ R O {b}_{o} \in {\mathbb{R}}^{O} boRO 分别对应偏差,而 tanh 表示引入非线性的双曲正切激活函数,应用于 RNN 模型。换言之,RNN 是一种非线性递归模型,通过利用保存在隐藏状态中的历史知识有效捕获时间模式。然而,与 RNN 相关联的几个限制。首先,RNN 有效提取输入序列中的长程动态的能力受到限制。当信息通过连续的时间步传递时,网络中权重的重复乘法可能导致信息的稀释或丢失。因此,RNN 在进行预测时,很难保留和回忆早期时间步的的信息。其次,RNN 以增量的方式处理顺序数据,限制了它们的计算效率,因为每个时间步都依赖于前一个时间步。这使得它们的并行计算变得具有挑战性。此外,传统的 RNN 缺乏内置的注意力机制,这使得网络无法捕捉输入序列中的全局信息。缺少注意力机制妨碍了网络选择性建模数据重要片段的能力。为克服这些限制,Transformers 和状态空间模型应运而生,它们从不同的角度解决这些挑战。这两种方法将在后续小节中进一步阐述。

2.2 Transformers

Transformer [174] 是深度学习领域的一个突破性模型,彻底改变了各种人工智能应用。它的引入标志着对传统序列到序列模型的重大变革,通过采用自注意机制,促进了模型输入中全局依赖关系的捕捉。例如,在自然语言处理领域,这种自注意能力使模型能够理解序列中不同位置之间的关系。它通过根据每个位置相对于其他位置的重要性分配权重来实现这一点。更具体地说,一系列输入向量 x \mathbf{x} x 首先通过利用原始输入的线性变换转化为三种类型的向量:查询 Q Q Q 、键 K K K 和值 V V V ,定义为: Q = x ⋅ W Q , K = x ⋅ W K , V = x ⋅ W V , (3) Q = \mathbf{x} \cdot {\mathbf{W}}^{Q}, K = \mathbf{x} \cdot {\mathbf{W}}^{K}, V = \mathbf{x} \cdot {\mathbf{W}}^{V}, \tag{3} Q=xWQ,K=xWK,V=xWV,(3) 其中 W Q , W K {\mathbf{W}}^{Q},{\mathbf{W}}^{K} WQ,WK W V {\mathbf{W}}^{V} WV 是可训练的参数。注意力得分通过计算 Q Q Q K K K 的点积来计算,然后将结果缩放为 d K \sqrt{{d}_{K}} dK ,其中 d K {d}_{K} dK 是键向量的维度。这些过程然后通过 Softmax 函数进行规范化,以生成注意力权重,定义为: S = Softmax ⁡ ( Q K T d K ) V (4) S = \operatorname{Softmax}\left( \frac{Q{K}^{T}}{\sqrt{{d}_{K}}}\right) V \tag{4} S=Softmax(dK QKT)V(4) 除了执行单一的注意力功能外,还引入了多头注意力,以增强模型捕捉不同类型关系的能力,并提供多种对输入序列的视角。在多头注意力中,输入序列通过多组自注意力模块并行处理。每个头独立运作,执行与标准自注意机制相同的计算。然后将每个头的注意力权重结合起来,创建值向量的加权和。这个聚合步骤使模型能够利用来自多个头的信息,捕捉输入序列中存在的多样化模式和关系。在数学上,多头注意力的计算如下: MultiHead ⁡ ( Q , K , V ) = ( S 1 ⊕ S 2 … ⊕ S m ) ⋅ W O , \operatorname{MultiHead}\left( {Q, K, V}\right) = \left( {{S}_{1} \oplus {S}_{2}\ldots \oplus {S}_{m}}\right) \cdot {\mathbf{W}}^{O}, MultiHead(Q,K,V)=(S1S2Sm)WO, where S i = Softmax ⁡ ( Q i K i T d K ) V i , i ∈ [ 1 , m ] , (5) \text{where}{S}_{i} = \operatorname{Softmax}\left( \frac{{Q}_{i}{K}_{i}^{T}}{\sqrt{{d}_{K}}}\right) {V}_{i}, i \in \left\lbrack {1, m}\right\rbrack \text{,} \tag{5} whereSi=Softmax(dK QiKiT)Vi,i[1,m],(5) 其中 m m m 是注意力头的数量。 ⊕ \oplus 是串联操作, W O {\mathbf{W}}^{O} WO 是线性变换,用于将多头注意力得分投影到最终值。

2.3 状态空间模型

状态空间模型(SSMs)是一种传统的数学框架,用于描述系统随时间的动态行为 [89]。近年来,SSMs 在控制理论、机器人技术和经济学等多个领域得到了广泛应用 [58, 59]。SSMs 的核心通过一组称为“状态”的隐藏变量体现系统的行为,使其能够有效捕捉时间数据依赖关系。与 RNN 不同,SSMs 是线性模型,其特征在于它们的关联属性。具体来说,在经典状态空间模型中,制定了两个关键方程,即状态方程和观测方程,以建模在当前时间 t t t 处输入 x ( t ) ∈ R x\left( t\right) \in \mathbb{R} x(t)R 和输出 y ( t ) ∈ R y\left( t\right) \in \mathbb{R} y(t)R 之间的关系,通过一个 N \mathrm{N} N 维隐藏状态 h ( t ) ∈ R N h\left( t\right) \in {\mathbb{R}}^{N} h(t)RN 。该过程可以写成 h ′ ( t ) = A h ( t ) + B x ( t ) (6) {h}^{\prime }\left( t\right) = \mathrm{A}h\left( t\right) + \mathrm{B}x\left( t\right) \tag{6} h(t)=Ah(t)+Bx(t)(6) y ( t ) = C h ( t ) + D x ( t ) (7) y\left( t\right) = \mathrm{C}h\left( t\right) + \mathrm{D}x\left( t\right) \tag{7} y(t)=Ch(t)+Dx(t)(7) 其中 h ′ ( t ) {h}^{\prime }\left( t\right) h(t) 是当前状态 h ( t ) h\left( t\right) h(t) 的导数, A ∈ R N × N \mathbf{A} \in {\mathbb{R}}^{N \times N} ARN×N 是状态转移矩阵,描述了状态如何随时间变化, B ∈ R N × 1 \mathbf{B} \in {\mathbb{R}}^{N \times 1} BRN×1 是输入矩阵,控制输入如何影响状态变化, C ∈ R 1 × N \mathrm{C} \in {\mathbb{R}}^{1 \times N} CR1×N 表示输出矩阵,指示如何根据当前状态生成输出, D ∈ R \mathrm{D} \in \mathbb{R} DR 代表命令系数,决定输入如何直接影响输出。一般来说,大多数状态空间模型(SSMs)在观测方程中排除了第二项,即设置 D x ( t ) = 0 \mathrm{D}x\left( t\right) = 0 Dx(t)=0 ,这一点可以视为深度学习模型中的跳跃连接。

2.3.1 离散化。为了符合机器学习环境的要求,状态空间模型必须经历离散化过程,将连续参数转化为离散参数。离散化方法通常旨在将连续时间划分为 K K K 个尽可能具有相等积分面积的离散区间。为了达到目标,作为最具代表性的解决方案之一,零阶保持(ZOH)[138, 223] 被成功应用于状态空间模型,该方法假设函数值在区间 Δ = [ t k − 1 , t k ] \Delta = \left\lbrack {{t}_{k - 1},{t}_{k}}\right\rbrack Δ=[tk1,tk] 内保持不变。经过 ZOH 离散化后,状态空间模型方程可以被重写为 h k = A ‾ h k − 1 + B ‾ x k (8) {h}_{k} = \overline{\mathbf{A}}{h}_{k - 1} + \overline{\mathbf{B}}{x}_{k} \tag{8} hk=Ahk1+Bxk(8) y k = C h k (9) {y}_{k} = \mathrm{C}{h}_{k} \tag{9} yk=Chk(9) 其中 A ‾ = exp ⁡ ( Δ A ) \overline{\mathbf{A}} = \exp \left( {\Delta \mathbf{A}}\right) A=exp(ΔA) B ‾ = ( Δ A ) − 1 ( exp ⁡ ( Δ A ) − I ) ⋅ Δ B , k \overline{\mathbf{B}} = {\left( \Delta \mathbf{A}\right) }^{-1}\left( {\exp \left( {\Delta \mathbf{A}}\right) - \mathbf{I}}\right) \cdot \Delta \mathbf{B},k B=(ΔA)1(exp(ΔA)I)ΔBk 是离散时间步。从这些公式中可以明显看出,离散的状态空间模型具有与递归神经网络相似的结构,因此,离散的状态空间模型可以比基于 Transformer 的模型在每次自回归解码迭代中计算所有输入的注意力时更高效地完成推理过程。

2.3.2 卷积计算。离散状态空间模型作为线性系统,具有相关属性,因此与卷积计算无缝集成。更具体地,它可以独立地在每个时间步计算输出,如下所示: y 0 = C A 0 B ‾ x 0 (10) {y}_{0} = {\mathbf{{CA}}}^{0}\overline{\mathbf{B}}{x}_{0} \tag{10} y0=CA0Bx0(10) y 1 = C A ‾ 1 B ‾ x 0 + C A ‾ 0 B ‾ x 1 (11) {y}_{1} = {\mathbf{C\overline{A}}}^{1}\overline{\mathbf{B}}{x}_{0} + {\mathbf{C\overline{A}}}^{0}\overline{\mathbf{B}}{x}_{1} \tag{11} y1=CA1Bx0+CA0Bx1(11) y 2 = C A ‾ 2 B ‾ x 0 + C A ‾ 1 B ‾ x 1 + C A ‾ 0 B ‾ x 2 (12) {y}_{2} = {\mathbf{C\overline{A}}}^{2}\overline{\mathbf{B}}{x}_{0} + {\mathbf{C\overline{A}}}^{1}\overline{\mathbf{B}}{x}_{1} + {\mathbf{C\overline{A}}}^{0}\overline{\mathbf{B}}{x}_{2} \tag{12} y2=CA2Bx0+CA1Bx1+CA0Bx2(12) … … \ldots \ldots …… (13) y k = C A ‾ k B ‾ x 0 + C A ‾ k − 1 B ‾ x 1 + … + C A ‾ 1 B ‾ x k − 1 + C A ‾ 0 B ‾ x k . (14) {y}_{k} = \mathrm{C}{\overline{\mathbf{A}}}^{k}\overline{\mathbf{B}}{x}_{0} + \mathrm{C}{\overline{\mathbf{A}}}^{k - 1}\overline{\mathbf{B}}{x}_{1} + \ldots + \mathrm{C}{\overline{\mathbf{A}}}^{1}\overline{\mathbf{B}}{x}_{k - 1} + \mathrm{C}{\overline{\mathbf{A}}}^{0}\overline{\mathbf{B}}{x}_{k}. \tag{14} yk=CAkBx0+CAk1Bx1++CA1Bxk1+CA0Bxk.(14) 通过创建一组卷积核 K ‾ = ( C B ‾ , … , C A ‾ k B ‾ , … ) \overline{\mathbf{K}} = \left( {\mathrm{C}\overline{\mathbf{B}},\ldots ,\mathrm{C}{\overline{\mathbf{A}}}^{k}\overline{\mathbf{B}},\ldots }\right) K=(CB,,CAkB,) ,递归计算可以转换为卷积形式如下: y = x ∗ K ‾ (15) \mathbf{y} = \mathbf{x} * \overline{\mathbf{K}} \tag{15} y=xK(15) 其中 x = [ x 0 , x 1 , … ] \mathbf{x} = \left\lbrack {{x}_{0},{x}_{1},\ldots }\right\rbrack x=[x0,x1,] y = [ y 0 , y 1 , … ] ∈ R L \mathbf{y} = \left\lbrack {{y}_{0},{y}_{1},\ldots }\right\rbrack \in {\mathbb{R}}^{L} y=[y0,y1,]RL 分别表示输入和输出序列,而 L L L 是序列长度。此卷积计算使得状态空间模型(SSMs)充分利用现代矩阵计算硬件(如GPU),以实现训练过程中的并行计算,这在使用非线性激活函数的递归神经网络(RNNs)中是不可行的。值得注意的是,给定一个具有 D D D 维的输入 x ( k ) x\left( k\right) x(k) ,SSM计算将针对每个维度单独计算,以产生 D D D 维的输出 y ( t ) y\left( t\right) y(t) 。在这种情况下,输入矩阵 B ∈ R N × D \mathbf{B} \in {\mathbb{R}}^{N \times D} BRN×D ,输出矩阵 C ∈ R D × N \mathbf{C} \in {\mathbb{R}}^{D \times N} CRD×N ,命令矩阵 D ∈ R D × D \mathbf{D} \in {\mathbb{R}}^{D \times D} DRD×D ,而状态转移矩阵保持不变,即 A ∈ R N × N \mathbf{A} \in {\mathbb{R}}^{N \times N} ARN×N

2.3.3 RNN、Transformer 和 SSM 之间的关系。RNN、Transformer 和 SSM 的计算算法如图 2 所示。一方面,传统的 RNN 在一个非线性的递归框架内进行操作,其中每次计算仅依赖于前一个隐藏状态和当前输入。虽然这种格式允许 RNN 在自回归推理期间快速生成输出,但它限制了它们充分利用 GPU 并行计算的能力,导致模型训练速度较慢。另一方面,Transformer 架构在多个查询-键对之间并行执行矩阵乘法,这可以高效分布到硬件资源,进而加快基于注意力的模型的训练。然而,当涉及到从基于 Transformer 的模型生成响应或预测时,推理过程可能会消耗较长时间。例如,语言模型的自回归设计要求顺序生成输出序列中的每个标记,这需要在每个步骤重复计算注意力分数,导致推理时间变慢。与仅支持单一计算类型的 RNN 和 Transformer 不同,离散 SSM 具有支持递归和卷积计算的灵活性,因为其线性属性。这一独特能力使得 SSM 不仅实现高效推理而且实现并行训练。然而,应该注意的是,最传统的 SSM 是时间不变的,这意味着它们的 A , B , C \mathbf{A},\mathbf{B},\mathbf{C} A,B,C Δ \Delta Δ 与模型输入 x x x 无关。这会限制上下文感知建模,从而导致 SSM 在某些任务(例如选择性复制)中的性能较差 [55]。

图 2. 代表性模型架构的示意图,即递归神经网络(RNN)、Transformer 和状态空间模型(SSM)。(a) RNN 在非线性递归框架内工作,在自回归推理期间能快速生成输出。(b) Transformer 在多个查询-键对之间并行执行矩阵乘法,加快了训练速度。© SSM 由于其线性特性展现了灵活性,能够同时支持递归和卷积计算。这种融合结合了 RNN 和 Transformer 的优势,使得 SSM 能够进行递归推理和并行训练。尽管如此,传统的时间不变 SSM 在上下文感知建模方面仍显不足,导致特定任务的性能下降。

3 Mamba

为了解决传统 SSM 在上下文感知能力不足的缺点,[55] 提出了 Mamba,作为一个潜在的替代方案,承诺成为通用序列基础模型支架。最近,Mamba-2 [29] 提出了结构化状态-空间二元性(SSD),建立了一个坚实的理论框架,将结构化 SSM 和各种形式的注意力联系在一起,使我们能够将最初为 Transformer 开发的算法和系统优化转移到 SSM。在本节中,我们将简明扼要地介绍 Mamba 和 Mamba-2。

3.1 Mamba-1:具有硬件感知算法的选择性状态空间模型

图 3. 具有硬件感知状态扩展的选择性状态空间模型概览。选择性机制引入输入相关的动态,而硬件计算算法优化了 GPU 内存层次中更高效级别的内存使用。传统的结构化状态空间模型(SSMs)在建模文本和其他信息密集型数据方面的有效性有限 [55],阻碍了其在深度学习中的发展。为提升SSMs的建模能力,Gu和Dao [55]基于结构化状态空间模型引入了三种创新技术,即基于高阶多项式投影算子(HiPPO)的记忆初始化、选择机制和硬件感知计算,如图3所示。这些技术旨在增强SSMs在长程线性时间序列建模中的能力。特别是,初始化策略建立了一致的隐状态矩阵,有效促进长程记忆。接着,选择机制使得SSMs能够获取内容感知的表示。最后,Mamba设计了两种硬件感知计算算法,平行关联扫描和内存重计算,以提高训练效率。

3.1.1 基于HiPPO的记忆初始化。对顺序数据的建模和学习代表了当代机器学习中的基础挑战,形成了包括语言建模、语音识别和视频处理在内的各种任务的基础。建模复杂和长期时间依赖关系的一个基本组成部分是记忆,包含存储和整合来自前面时间步的信息的能力 [73]。与RNNs类似,保留和遗忘历史隐状态(即矩阵A)在SSMs中发挥着关键作用,以实现令人满意的性能。在之前的结构化状态空间序列模型(SSMs)中,曾提出针对特别初始化的建议,特别是在复杂值模型的情况下。这些特殊初始化在多种场景中被证明是有利的,包括数据可用性有限的情况。类似地,Mamba主要关注隐状态矩阵 A \mathbf{A} A 的初始化,以捕捉复杂的时间依赖关系。这是通过利用HiPPO理论 [56] 和一种创新的缩放勒让德测量(LegS)来完成的,确保全面考虑完整的历史上下文,而不是有限的滑动窗口。具体而言,HiPPO-LegS

提交给ACM的手稿将所有历史数据点赋予均匀权重,可以表示为: A n k HiPPO  = − { ( 2 n + 1 ) 1 2 ( 2 k + 1 ) 1 2  如果  n > k n + 1  如果  n = k 0  如果  n < k (16) {\mathrm{A}}_{nk}^{\text{HiPPO }} = - \left\{ \begin{array}{ll} {\left( 2n + 1\right) }^{\frac{1}{2}}{\left( 2k + 1\right) }^{\frac{1}{2}} & \text{ 如果 }n > k \\ n + 1 & \text{ 如果 }n = k \\ 0 & \text{ 如果 }n < k \end{array}\right. \tag{16} AnkHiPPO = (2n+1)21(2k+1)21n+10 如果 n>k 如果 n=k 如果 n<k(16) 其中 n n n 是多项式的数量, k k k 表示特定的离散时间步。基于HiPPO理论,Mamba提出两种针对复杂和实数情况的简单初始化方法,即S4D-Lin和S4D-Real [57],如以下所示: A d n = − { 1 2 − n i  S4D-Lin  n + 1  S4D-Real  , (17) {\mathbf{A}}_{dn} = - \left\{ {\begin{array}{ll} \frac{1}{2} - {ni} & \text{ S4D-Lin } \\ n + 1 & \text{ S4D-Real } \end{array},}\right. \tag{17} Adn={21nin+1 S4D-Lin  S4D-Real ,(17) 其中 n n n 是所有输入维度 d = 1 , 2 , … , D d = 1,2,\ldots, D d=1,2,,D 的 A 的第 n n n 个元素。通过这样的初始化,模型可以学习长期依赖的记忆,从而通过压缩和重构输入信息信号,经历较小的新步骤降解和较大的旧步骤降解。根据公式,HiPPO-LegS 具有有利的理论属性:它在输入时间尺度上保持一致,并提供快速计算 [56]。此外,它具有有界的梯度和近似误差,从而促进参数学习过程。

3.1.2 选择机制。传统的状态空间模型由于时间不变性的特性,无法根据特定模型输入(即内容感知建模能力)生成个性化的输出。为了为SSMs提供类似于注意力机制的这种能力,Mamba设计了一种时间变化的选择机制,根据模型输入参数化权重矩阵。这一创新使得SSMs能够过滤掉多余的信息,同时无限期地保留相关细节。形式上,选择机制涉及设置间隔 Δ \Delta Δ ,并将矩阵 B , C \mathbf{B},\mathbf{C} B,C 作为输入 x ∈ R B × L × D \mathbf{x} \in {\mathbb{R}}^{B \times L \times D} xRB×L×D 的函数,可以表示为: B → S B = W B x (18) \mathrm{B} \rightarrow {\mathrm{S}}^{\mathrm{B}} = {\mathrm{W}}^{\mathrm{B}}\mathrm{x} \tag{18} BSB=WBx(18) C → S C = W C x (19) \mathrm{C} \rightarrow {\mathrm{S}}^{\mathrm{C}} = {\mathrm{W}}^{\mathrm{C}}\mathrm{x} \tag{19} CSC=WCx(19) Δ → S Δ = τ Δ ⋅ BroadCast ⁡ D ( W Δ x ) (20) \Delta \rightarrow {\mathrm{S}}^{\Delta } = {\tau }_{\Delta } \cdot {\operatorname{BroadCast}}_{D}\left( {{\mathbf{W}}^{\Delta }\mathbf{x}}\right) \tag{20} ΔSΔ=τΔBroadCastD(WΔx)(20) 其中 S B ∈ R B × L × N ,    S C ∈ R B × L × N {\mathrm{S}}^{\mathrm{B}} \in {\mathbb{R}}^{B \times L \times N},{\mathrm{\;S}}^{\mathrm{C}} \in {\mathbb{R}}^{B \times L \times N} SBRB×L×N,SCRB×L×N , 和 S Δ ∈ R B × L × D {\mathrm{S}}^{\Delta } \in {\mathbb{R}}^{B \times L \times D} SΔRB×L×D 是针对输入的选择性空间矩阵,用于实现内容感知建模。 B , L , D B, L, D B,L,D , 和 N N N 分别表示批处理大小、输入长度、输入特征大小和隐含通道数。值得注意的是, W B ∈ R N × D , W C ∈ R N × D {\mathbf{W}}^{\mathbf{B}} \in {\mathbb{R}}^{N \times D},{\mathbf{W}}^{\mathbf{C}} \in {\mathbb{R}}^{N \times D} WBRN×D,WCRN×D ,和 W Δ ∈ R D × 1 {\mathbf{W}}^{\Delta } \in {\mathbb{R}}^{D \times 1} WΔRD×1 是对应组件的选择权重(即线性参数化投影),而 BroadCast D {}_{D} D 意味着将结果广播到所有维度 d = 1 , 2 , … , D d = 1,2,\ldots, D d=1,2,,D 。随后,选择性 SSM 经过使用常见统计技术零阶保持 (ZOH) [138] 的离散化,如下所示: A ‾ → S A ‾ = exp ⁡ ( S Δ A ) (21) \overline{\mathrm{A}} \rightarrow {\mathrm{S}}^{\overline{\mathrm{A}}} = \exp \left( {{\mathrm{S}}^{\Delta }\mathrm{A}}\right) \tag{21} ASA=exp(SΔA)(21) B ‾ → S B ‾ = ( S Δ A ) − 1 ( exp ⁡ ( S Δ A ) − I ) ⋅ S Δ S B , (22) \overline{\mathbf{B}} \rightarrow {\mathbf{S}}^{\overline{\mathbf{B}}} = {\left( {\mathbf{S}}^{\Delta }\mathbf{A}\right) }^{-1}\left( {\exp \left( {{\mathbf{S}}^{\Delta }\mathbf{A}}\right) - \mathbf{I}}\right) \cdot {\mathbf{S}}^{\Delta }{\mathbf{S}}^{B}, \tag{22} BSB=(SΔA)1(exp(SΔA)I)SΔSB,(22) 其中 S A ‾ ∈ R B × L × D × N {\mathrm{S}}^{\overline{\mathbf{A}}} \in {\mathbb{R}}^{B \times L \times D \times N} SARB×L×D×N S B ‾ ∈ R B × L × D × N {\mathrm{S}}^{\overline{\mathbf{B}}} \in {\mathbb{R}}^{B \times L \times D \times N} SBRB×L×D×N 分别是选择性状态转移矩阵和输入矩阵,它们成为输入 x \mathbf{x} x 的函数。通过这样做,离散 SSM 从时间不变变为时间可变(即,内容感知) 如下所示: y = SSM ⁡ ( A , B , C ) ( x ) (23) \mathbf{y} = \operatorname{SSM}\left( {\mathbf{A},\mathbf{B},\mathbf{C}}\right) \left( \mathbf{x}\right) \tag{23} y=SSM(A,B,C)(x)(23) 它生成的输出 y ∈ R B × L × D \mathbf{y} \in {\mathbf{R}}^{B \times L \times D} yRB×L×D 依赖于输入 x \mathbf{x} x 。注意,Mamba 中的时间可变选择机制与 Transformer 中的注意机制具有相似的结构,即二者都基于输入及其投影进行操作,这使得 Mamba 的 SSM 能够实现灵活的内容感知建模。然而,它失去了与卷积的等价性,这对其效率产生负面影响。

3.1.3 硬件感知计算。选择机制旨在超越线性时间不变模型的局限性。然而,它对高效训练提出了挑战:SSM 的卷积核变得依赖于输入,从而导致无法进行并行计算。为了解决这个问题,Mamba 利用两种计算技术,即并行关联扫描(也称为并行前缀和) [64] 和内存重计算。首先,并行关联扫描利用线性关联计算的性质和现代加速器(GPU 和 TPU)的并行性,以一种内存高效的方式进行选择性 SSM 的计算。更具体地说,并行关联扫描将模型训练的计算复杂度从 O ( N 2 d ) \mathrm{O}\left( {{N}^{2}d}\right) O(N2d) 降低到 O ( N / t ) \mathrm{O}\left( {N/t}\right) O(N/t) 。其核心在于在给定输入上构建一个平衡的二叉树,并从树的根部进行扫过。换句话说,并行关联扫描首先从叶子节点遍历到根节点(即,向上扫描),在树的内部节点创建部分和。然后,它反向遍历,沿着树从根节点回到上面,利用部分和构建整个扫描(即,向下扫描)。

另一方面,Mamba 利用传统的重计算方法来降低选择性 SSM 层的整体内存需求。特别地,Mamba 在执行并行关联扫描的前向传递过程中避免存储大小为 ( B , L (B, L (B,L , D , N D, N D,N ) 的中间状态,以防止内存扩展。相反,它在反向传递中重新计算这些中间状态以用于梯度计算。通过这样做,重计算避免了读取 O ( B L N D ) O\left( {BLND}\right) O(BLND) 个元素在 GPU 内存单元之间的必要性。除了优化扫描操作的内存需求外,Mamba-1 还扩展了重计算的使用,以提高整个 SSM 层的效率。此优化包括投影、卷积和激活,这通常需要大量内存资源,但可以快速重新计算。## 3.2 Mamba-2: 状态空间对偶性

变换器在各种领域的深度学习成功中发挥了关键作用,激发了许多技术的发展,例如参数高效微调 [95]、灾难性遗忘缓解 [96] 和模型量化 [195],旨在从不同角度提高模型性能。为了使状态空间模型能够访问和受益于最初为变换器开发的有价值技术,Mamba-2 [29] 引入了一个名为结构化状态空间对偶性(SSD)的综合框架,该框架建立了状态空间模型(SSMs)与不同形式注意力之间的理论联系。形式上, y = SSD ⁡ ( A , B , C ) ( x ) = M x (24) \mathbf{y} = \operatorname{SSD}\left( {\mathbf{A},\mathbf{B},\mathbf{C}}\right) \left( \mathbf{x}\right) = \mathbf{M}\mathbf{x} \tag{24} y=SSD(A,B,C)(x)=Mx(24) 其中 M \mathbf{M} M 表示使用顺序半可分离表示的状态空间模型的矩阵形式,而 M j i = C j T A j : i B i {\mathbf{M}}_{ji} = {\mathbf{C}}_{j}^{\mathrm{T}}{\mathbf{A}}_{j : i}{\mathbf{B}}_{i} Mji=CjTAj:iBi 。值得注意的是, C j {\mathrm{C}}_{j} Cj B i {\mathrm{B}}_{i} Bi 分别表示与输入标记 x j {\mathbf{x}}_{j} xj x i {\mathrm{x}}_{i} xi 相关的选择性空间状态矩阵。 A j : i {\mathbf{A}}_{j : i} Aj:i 表示对应于输入标记范围从 j j j i i i 的选择状态矩阵。从本质上讲,SSD 证明了变换器使用的注意力机制和 SSM 中采用的线性时变系统可以看作是半可分离矩阵变换。此外,Dao 和 Gu [29] 还证明了选择性 SSM 等同于使用半可分离掩蔽矩阵实现的结构化线性注意力机制。

图 4. Mamba-1 和 Mamba-2 的块架构比较。

基于 SSD,Mamba-2 通过块分解矩阵乘法算法设计了一种更硬件高效的计算。具体来说,通过将状态空间模型视为半可分离矩阵,Mamba-2 将计算分解为矩阵块,其中对角块代表内部计算,而非对角块代表通过 SSM 的隐藏状态因子的块间计算。此方法使 Mamba-2 的训练过程比 Mamba-1 的并行关联扫描快 2-8 倍,同时仍与变换器具有竞争力。

3.3 Mamba 块

在本小节中,我们总结了 Mamba-1 和 Mamba-2 的块设计。图 4 展示了这两种架构的比较。Mamba-1 的设计受到 SSM 中心观点的激励,其中选择性 SSM 层的任务是将输入序列 X \mathbf{X} X 映射到 Y \mathbf{Y} Y 。在此设计中, ( A , B , C ) (\mathbf{A},\mathbf{B}, \mathbf{C}) (A,B,C) 的线性投影在创建 X \mathrm{X} X 的初始线性投影之后应用。输入标记和状态矩阵然后通过选择性 SSM 单元,利用并行关联扫描,生成输出 Y。之后,Mamba-1 采用跳过连接以鼓励特征重用并缓解模型训练过程中常出现的退化问题。最后,Mamba 模型通过堆叠这个块并与标准归一化和残差连接交织构建。

至于 Mamba-2,它引入了 SSD 层,旨在从 [ X , A , B , C ] \left\lbrack {\mathbf{X},\mathbf{A},\mathbf{B},\mathbf{C}}\right\rbrack [X,A,B,C] 创建到 Y \mathbf{Y} Y 的映射。这是通过在块的开始同时处理 [ X , A , B , C ] \left\lbrack {\mathrm{X},\mathrm{A},\mathrm{B},\mathrm{C}}\right\rbrack [X,A,B,C] 实现的,类似于标准注意力架构如何并行生成 Q , K , V \mathbf{Q},\mathbf{K},\mathbf{V} Q,K,V 投影。换句话说,Mamba-2 块通过去除顺序线性投影简化了 Mamba-1 块。这使 SSD 结构的计算比 Mamba-1 中的并行选择扫描更快。此外,在跳过连接之后添加了规范化层,旨在提高训练稳定性。

4 Mamba 模型的进展

状态空间模型和 Mamba 最近被研究,成为基础模型骨干的一个有前景的替代方案。尽管 Mamba 在自然语言处理方面表现出色,但仍面临一些挑战,如记忆损失、对各种任务的泛化以及捕捉复杂模式的能力不及基于变换器的语言模型。为了克服这些挑战,进行了大量努力以改善 Mamba 架构。现有研究主要集中在修改块设计、扫描模式和内存管理方面。本节将介绍这三个方面的一些重要技术,并在表 1 中呈现相关研究的总结。

表1. 改进Mamba模型现有研究的总结。

4.1 块设计

Mamba块的设计和结构对Mamba模型的整体性能产生显著影响,使其成为新兴的研究焦点。如图5所示,基于构建新Mamba块的不同方法,现有研究可分为三类:a)集成方法旨在将Mamba块与其他知名模型集成,以在有效性和效率之间取得平衡;b)替换方法试图将Mamba块作为高阶模型框架中主要层的替代;c)修改方法则专注于修改经典Mamba块中的组件。因此,我们将在以下小节中详细回顾这些方法。

提交给ACM的手稿

图5. 基于块设计视角的改进Mamba模型的代表性示例:(a)集成方法结合了正交的架构设计(如变压器)与Mamba,从而提高了模型性能和吞吐量,同时保持可控的内存占用 [ 141 , 200 ] \left\lbrack {{141},{200}}\right\rbrack [141,200] ; (b)替换方法通过将Mamba块集成作为它们主要层的替代来增强标准学习框架的学习能力(例如,U-Net)[110, 214];(c)修改方法采用尖端技术,如专家混合(MoE),来优化Mamba块 [111]。

4.1.1 集成。鉴于Mamba在捕获长期动态方面的卓越能力,它与其他模型的集成受到广泛关注,利用它们自身的优势来提供一个针对特定场景的强大框架。集成特别包括变压器、卷积神经网络 (CNN)、图神经网络 (GNN)、循环神经网络 (RNN) 和脉冲神经网络 (SNN) 等先进模型。具体示例如下。

  • 基于变压器的模型在许多任务中表现出色,但它们的平方计算复杂性仍在推理过程中限制了它们的性能 [58]。为了追求高效生成,一些研究者提出将Mamba块与基于变压器的模型相结合。例如,Jamba [111] 将变压器和Mamba层的块结合起来,以应对长内容自然语言处理任务,借助于两种模型家族的优势。Attention-Mamba混合模型的性能优于单独的变压器和Mamba模型,吞吐量也优于普通的变压器模型。Mambaformer [200] 运用混合框架来预测多个时间序列,包括汇率、每小时电力消费和电力负荷,内部结合了Mamba块和变压器层,以分别处理长短期依赖性。由于Mamba和变压器的结合,Mambaformer在长短期时间序列预测中优于基于变压器的预测模型。

  • 基于CNN的方法受限于局部感受野,导致在捕获全球和长范围语义时的表现不佳 [55]。由于状态空间模型在学习长范围模式方面的优越能力,一些研究 [ 107 , 187 , 204 ] \left\lbrack {{107},{187},{204}}\right\rbrack [107,187,204] 探索利用Mamba块增强基于CNN的模型的潜力,尤其是在计算机视觉领域。例如,MedMamba [214] 和nnMamba [53] 展示了视觉Mamba块的集成如何提高CNN在图像分析任务中的表现。

  • GNN在通过消息传递机制捕获邻接关系方面显示出良好的潜力,其中信息通过堆叠层在连接图上进行传播。然而,这些模型面临一个被称为过平滑的显著限制 [20],尤其是在试图捕获高阶邻接信号时。为了解决这一挑战,Mamba已被用于图表示学习 [ 103 , 115 , 178 , 203 ] \left\lbrack {{103},{115},{178},{203}}\right\rbrack [103,115,178,203] 。例如,Graph Mamba [9] 将图结构数据重新整理为特定顺序的序列标记,并利用Mamba块内的选择性SSM层构建了一种新的图Mamba网络 (GMN) 架构,在图表示学习能力方面优于特别需要节点之间高阶依赖性的数据集。

  • 基于RNN的模型在捕获时间动态方面取得了卓越成果。然而,RNN仍面临重大挑战,包括耗时的递归训练和对隐藏状态内存容量的限制。受到最近基于Mamba的架构出现的启发,一些研究者开发了Mamba块与RNN的融合。例如,VMRNN [170] 在时空预测中取得了最先进的性能,同时相比基于递归和非递归方法减少了浮点运算 (FLOPs)。它通过引入一种新的递归单元,将Mamba块与长短期记忆 (LSTM) 结合。

  • SNN最近被提出作为一种受生物神经元行为启发的有前景的网络架构:通过离散脉冲在神经元之间传递知识。SNN的一个关键优势在于其低功耗实现的潜力,因为它们能够利用神经活动的稀疏和事件驱动特性。受到SNN的能效实现和SSM的优越长程学习能力的启发,开创性的研究深入探讨了将这两种方法结合起来。例如,SpikeMba [106] 将两者结合以处理对显著物体的置信偏差,并捕获视频序列中的持久依赖性。通过广泛评估,作者声称整合这两种模型提高了时间视频定位任务的有效性,特别是时刻检索和高亮检测。

4.1.2 替换。受到选择性SSM在高效计算和长序列学习中的卓越能力的启发,采用Mamba模块替换U-Net [151] 和扩散模型 [70] 等经典建模框架中的关键组件受到了广泛关注。通过引入选择性SSM层,这些方法在它们特定任务中实现了长程学习和高效计算。下面,我们将展示使用Mamba模块替换的实例,特别是针对U-Net和扩散模型等先进框架。- U-Net。一方面,许多努力 [ 110 , 163 , 179 , 180 ] \left\lbrack {{110},{163},{179},{180}}\right\rbrack [110,163,179,180] 被用于将 U-Net 与 Mamba 在捕捉复杂和广泛语义方面的能力协同,以提高计算机视觉任务中的模型性能。例如,Mamba-UNet [179] 独家利用视觉 Mamba 块构建了一种类似 U-Net 的模型(即,注入跳跃连接的编码器-解码器模型)用于医学图像分割。他们的评估表明,Mamba-UNet 超越了几种 U-Net 变体,这可以归因于 Mamba 块在处理长距离补丁序列方面的有效性和效率。

  • 扩散模型。另一方面,一些工作 [47, 49, 136] 已经致力于构建一种新型的扩散模型,即扩散状态空间模型 (DiS),它用状态空间骨干替代了典型的骨干(例如 CNN、注意力、U-Nets)。鉴于 Mamba 块在适应长距离依赖方面的显著效率和效能,DiS 的特点在于使用扩散模型生成更长的序列 [47]。例如,Oshima 等 [136] 提出了一种基于 Mamba 的扩散模型,显著减少了长视频序列的内存消耗,同时仍保持与基于 Transformer 的模型相比的竞争性能指标。此外,MD-Dose [49] 和 P-Mamba [209] 在扩散模型的反向过程中使用 Mamba 块构建噪声预测器,最终生成用于医学图像处理的特定目标。

  • 其他。除了 U-Net 和扩散模型外,还有一些替代方法。例如,Res-VMamba [19] 在残差学习框架中采用了视觉 Mamba 块用于食品类别分类。此外,SPMamba [101] 采用 TF-GridNet [189],一种最近开发的时频模型,作为其基础架构,然后用双向 Mamba 块后续 Transformer 组件。这种适配使模型能够有效地涵盖更广泛的上下文信息,用于语音分离任务。

4.1.3 修改。除直接使用 Mamba 块的整合和替代方法外,还进行了其他一些努力,以修改 Mamba 块,以期在不同场景下提高其性能。例如,Jamba [111] 借鉴了专家混合 (MoE) [46, 82] 的概念,使他们的混合 (Transformer-Mamba) 解码器模型能够以更少的计算进行预训练,并允许灵活的目标特定配置。值得注意的是,Jamba 模型(56B 可用参数,12B 活跃参数,4GB KV 缓存)相比于代表性的基于 Transformer 的语言模型 LLaMA-2-7B(6.7B 可用参数, 12    B {12}\mathrm{\;B} 12B 活跃参数, 128 G B K V {128}\mathrm{{GB}}\mathrm{{KV}} 128GBKV 缓存)需要 32 倍更小的 KV 缓存,同时提供更广泛的可用和活跃参数。这使得 Jamba 能够在单个 A100 GPU(80GB)上应对 140    K {140}\mathrm{\;K} 140K 的上下文长度,是 LLaMA-2-70B 支持的长度的七倍。除了 MoE,一些研究提出将 SSM 层修改为 K-way 结构,这涉及使用并行 SSM 单元处理模型输入,从多个角度捕捉信息和知识。例如,Sigma [177] 开发了一种新型基于 Mamba 的视觉编码器,通过利用并行 SSM 层处理多模态输入。UltraLight VM-UNet [193] 提出了一种具有并行 SSM 单元的视觉 Mamba 层,可以在不同通道中处理深度特征。综上所述,通过实施这些修改(即 K-way,MoE),这些基于 Mamba 的模型获得了增强的学习能力,特别是在处理多模态输入和快速适应多尺度任务方面。此外,一项开创性的研究 Mamba ® {}^{\circledR } ® 引入了一种新方法,建议在将输入通过 SSM 层之前将寄存器均匀地整合到视觉输入标记中。此修改旨在增强图像补丁的序列方向表示,从而使 Mamba 块的单向推断范式适用于视觉任务。尽管取得了这些成功,对修改 Mamba 块的探索仍然是一个有前景但尚未充分探索的领域。

4.2 扫描模式

并行关联扫描操作是 Mamba 模型中的一个关键组成部分,旨在解决选择机制造成的计算问题,加速训练过程,并减少内存需求。它通过利用时间变化 SSM 的线性特性在硬件级别设计内核融合和重新计算来实现这一目标。然而,Mamba 的单向序列建模范式阻碍了对各种数据(如图像和视频)进行综合学习的过程。为了缓解这一问题,一些研究集中在设计高效的扫描方法,以增强模型性能并促进 Mamba 模型的训练过程。如图 6 所示,现有研究集中于开发扫描模式技术,可以分为两类:1)扁平扫描方法从标记序列的平面视角处理模型输入;2)立体扫描方法跨维度、通道或尺度扫描模型输入。

4.2.1 扁平扫描。扁平扫描是指将模型输入展平为标记序列并相应地从不同方向扫描的过程。这种扫描类型通常用于一维(例如时间序列)和二维(例如图像)数据。在本节中,我们进一步将其细分为四类,即双向扫描、扫掠扫描、连续扫描和高效扫描。

图 6. 最近开发的基于 Mamba 的模型的扫描方法:扁平扫描 (a-c) 涉及将模型输入展平为标记序列,并从不同方向相应地扫描它们,而立体扫描 (d-e) 从额外的角度处理输入,与扁平扫描方法相比,捕捉更广泛的知识谱。

  • 双向扫描。借用双向递归神经网络 (Bi-RNNs) [155] 的概念,视觉 Mamba [234] 引入了一种视觉数据的扫描方法,称为双向扫描 (Bi-Scan),该方法通过同时前向和后向 SSM 处理输入标记,从而增强模型在空间感知处理方面的能力。最近,一些研究利用 Bi-Scan 方法来促进其基于 Mamba 的模型的学习能力 [105]。例如,DPMamba [87] 和 SPMamba [101] 都利用了一对双路径(前向和后向)选择性 SSM 来建模语音信号的依赖性,从而实现语音分离的双向知识处理。这些显著的成功可以归因于 Bi-Scan 的有效性及其易于部署性。

  • 扫掠扫描。如图 6 所示,扫掠扫描技术以特定方向处理模型输入,类似于清洁工仔细扫地 [188, 214]。例如,Cross-Scan [121] 涉及将输入图像分成补丁,然后沿四条不同路径展平,这被视为两次双向扫描的融合。通过采用这些互补的遍历路径,Cross Scan 使图像中的每个补丁能够有效地从不同方向整合邻域信息,从而促进形成信息丰富的感受野。Omni-Scan [163, 227] 综合了从多个方向建模图像信息流,例如 2 (前向和后向) × 4 \times 4 ×4 (左-右,上-下,右上-左下,左上-右下)。这种策略增强了各种方向上下文信息的全局建模能力,从而能够提取全面的全局空间特征。

  • 连续扫描。为了确保输入序列的连续性,连续扫描技术在列或行之间扫描相邻的标记 [66],如图 6 所示。例如,为了更好地应对 2D 空间输入,PlainMamba [201] 引入了一种名为连续扫描的新方法,该方法在列(或行)之间扫描相邻的标记,而不是在 Cross Scan 中旅行到对面的标记。此外,Hilbert Scan [66] 基于 Hilbert 矩阵沿着弯曲的路径移动。根据他们的评估结果,可以推断出增强输入标记的语义连续性会导致基于 Mamba 的模型在各种视觉识别任务中表现更佳。

  • 高效扫描。与上述关注实现更全面输入建模的扫描方法相比,高效扫描方法旨在加速训练和推断过程。一般来说,高效扫描将给定输入分成若干部分并并行处理,从而减少计算时间。例如,高效 2D 扫描 [139] 通过跳过补丁来处理图像,从而在保留全局特征图的同时减少四倍的计算需求。此外,Gao 等 [51] 在他们的 Mamba 框架中引入了一种有效的双向子空间扫描方案。这一方案旨在有效捕捉 4D 光场超分辨率任务中的长期空间-角对应关系。具体而言,它将补丁序列分解为两部分,并通过两个双向扫描方案进行处理。通过这样做,扫描方法降低了输入长度,并解决了长期记忆问题,而不会牺牲完整的 4D 全局信息。4.2.2 立体扫描。通过从额外的视角建模输入,立体扫描方法在扫描过程中能够捕捉到更广泛的知识谱系,相较于扁平扫描方法。这种增强的能力使得对模型输入有更彻底的理解。具体来说,这些方法可以分为三类:层次扫描、时空扫描和混合扫描。层次扫描从不同层次处理输入,而时空扫描则从时间和空间两个角度考虑输入模式。此外,混合扫描将多种扫描方法结合,以充分发挥不同扫描技术的优势。

  • 层次扫描方法涉及运用不同的扫描核大小,从全局到局部或从宏观到微观的视角捕捉语义知识 [ 26 , 63 , 162 , 180 ] \left\lbrack {{26},{63},{162},{180}}\right\rbrack [26,63,162,180] 。例如,[24]提出了一种针对红外小目标检测的Mamba-in-Mamba层次编码器,结合了内外选择性SSM模块。内模块特别设计用于捕捉视觉补丁之间的相互作用,以进行局部模式提取。而外模块则旨在表征视觉句子之间的关系,以捕捉全局特征。HiSS [12]将输入序列分成多个块,并对块特征进行层次建模,以实现连续序列预测。块首先由低级SSM单元处理,处理后的特征由高级SSM模块映射到输出序列中。

  • 时空扫描。受到现实世界中动态系统普遍存在的推动,时空扫描方法的关注度日益增长,以增强Mamba模块的性能 [205, 207]。例如,VideoMamba [102]将图像的原始二维扫描扩展为两个三维扫描:空间优先扫描和时间优先扫描。通过结合这两种扫描方法,VideoMamba在处理长时间高分辨率视频方面表现出卓越的效率。此外,ChangeMamba [21]集成了三种时空扫描机制(顺序建模、交叉建模和并行建模),以实现多时态特征之间的上下文信息交互,用于遥感变化检测。

  • 混合扫描。在追求全面特征建模的过程中,许多研究努力结合不同扫描方法的优点 [ 30 , 33 , 53 , 163 , 229 ] \left\lbrack {{30},{33},{53},{163},{229}}\right\rbrack [30,33,53,163,229] ,即所谓的混合扫描。例如,Mambamixer [ 10 ] \left\lbrack {10}\right\rbrack [10] 提出了动态使用一套图像扫描方法的扫描切换,这些方法包括交叉扫描、之字形扫描和局部扫描,以遍历图像补丁。Mambamixer还引入了一种双重选择机制,以混合跨通道和跨标记的信息。通过这种方式,他们在与其他视觉模型的竞争中表现出色。Pan-Mamba [68]引入了两种建立在Mamba架构之上的扫描方法:通道交换扫描和跨模态扫描。通过结合这两种扫描方法,Pan-Mamba提升了其在图像增强中的高效跨模态信息交换和融合能力。

4.3 内存管理

与RNN一样,状态空间模型中的隐藏状态的内存有效地存储来自以前步骤的信息,因此在SSM的整体功能中扮演着关键角色。尽管Mamba引入了基于HiPPO的内存初始化方法 [55],但在SSM单元的内存管理中仍然存在挑战,包括在层之间转移隐藏信息和实现无损内存压缩。为此,一些开创性的研究提出了不同的解决方案,包括内存初始化、压缩和连接。例如,Ezoe和Sato [36]尝试通过在模型重新训练过程中使用平衡截断方法来优化选择性SSM的初始化过程。此外,DGMamba [123]引入了一种隐藏状态抑制方法,以增强状态空间模型中隐藏状态的领域泛化能力。该方法旨在减轻这些隐藏状态带来的负面影响,从而缩小不同领域之间隐藏状态的差距。同样,DenseMamba [67]提出了一种密集连接方法,以增强SSM中层之间隐藏信息的传播。这一策略旨在缓解内存退化,并通过选择性地将较浅层的隐藏状态整合到较深层中来保留输出生成的详细信息。

5 将Mamba适应多样化数据

Mamba架构代表了选择性状态空间模型的扩展,具有循环模型的基本属性,使其非常适合作为在文本、时间序列、语音等序列上操作的一般基础模型。同时,近期一些开创性研究已经将Mamba架构的应用扩展到超越序列数据的领域,包括图像和图表,如图7所示。这些研究旨在利用Mamba在捕获长距离依赖性方面的卓越能力,同时发挥其在学习和推理过程中的高效性。因此,在本节中,我们旨在研究将Mamba适应各种数据类型的新兴技术。

图7. 适应Mamba以应对多类型数据的代表性策略。(a-e) Mamba架构,融入循环模型的基本特征,是处理语言、时间序列和语音等序列的理想基础模型。(f-h) 为了解决非序列数据的处理,常见的方法是将数据分段或采样成离散的标记,按照定义的规则将它们组织成序列。此外,Mamba具备通过连接其各自的序列或投影来处理多模态数据的能力。

表2. 不同类型数据中与Mamba相关的研究汇总。

5.1 序列数据

序列数据是指以特定顺序收集或组织的数据,其中数据点的顺序具有重要意义。为了探索利用Mamba作为处理序列数据任务的基础模型的潜力,我们提供了在后续部分呈现的全面回顾,涵盖了自然语言、视频、时间序列、语音和人类运动等各种序列数据。5.1.1 自然语言。作为最具代表性的架构之一,Mamba 在保证效率的同时进行基于内容的推理,被认为是大型语言模型主干的一种有前景的替代方案,以解决 Transformers 在长序列上的计算低效性。在此基础上,许多研究探索了 Mamba 在自然语言处理 (NLP) 中各种下游任务的潜力 [10, 12, 67, 134]。例如,MambaByte [181] 被提出用于在字节序列上利用 Mamba,旨在利用 Mamba 在捕捉长程依赖关系方面的优势以适应无 token 的语言模型。他们的评估显示,MambaByte 避免了子词标记化的归纳偏差,并在长期语言建模任务中超越了最先进的子词 Transformers。此外,Jamba [111] 和 BlackMamba [5] 融入了专家混合 (MoE) 的概念,通过将 SSMs 提供的线性复杂度生成与 MoE 提供的快速而经济的推理能力结合,提升 Mamba 在语言处理上的性能。

5.1.2 视频。视频理解和生成的核心目标在于学习时空表示,这固有地面临两个巨大的挑战:短视频片段内的大量时空冗余和长上下文之间复杂的时空依赖关系 [6, 122]。在同时解决这两个挑战的追求中,Mamba 以其区分短期动作和解释长视频的能力脱颖而出 [ 106 , 136 , 205 , 235 ] \left\lbrack {{106},{136},{205},{235}}\right\rbrack [106,136,205,235] 。例如,VideoMamba [102] 首先通过 3D 卷积将输入视频投影为一组不重叠的时空补丁,然后利用堆叠的双向 Mamba 块将这些补丁编码为向量化表示,以用于视频理解和生成等下游任务。

5.1.3 时间序列。作为典型的序列数据,时间序列数据在我们生活的各个方面随处可见,包括股市分析、交通建模和天气预报 [147, 230]。受到 Mamba 在建模长程序列方面的最新进展的启发,许多努力致力于研究其在时间序列数据中的潜力 [ 3 , 4 , 108 , 200 ] \left\lbrack {3,4,{108},{200}}\right\rbrack [3,4,108,200] 。例如,TimeMachine [3] 利用 Mamba 捕捉多变量时间序列数据中的持久模式,确保线性复杂度计算和最小内存占用,以简化时间序列处理。此外,Mambaformer [200] 结合选择性 SSM 和注意力层,用于天气、交通流等的长期和短期预测。

5.1.4 语音。语音特指人类沟通的声音表达形式,涉及使用特定的音素、单词、语法和语调模式进行的表达 [1]。最近,在与语音相关的任务领域,研究者们 [27] 在开发基于 Mamba 的模型以解决现有模型架构(如 RNN 和 Transformers)所遇到的挑战方面取得了显著进展。例如,SPMamba [101] 和 DPMamba [87] 利用双向 Mamba 模块捕捉更广泛的上下文信息,用于语音分离,展示了在处理语音分离任务时,模型性能提高了 13 % {13}\% 13% ,计算复杂度相比基于 Transformer 的基线降低了 566 % {566}\% 566%

5.1.5 动作。人类动作理解和生成在计算机动画、游戏开发和机器人操作等广泛的实际应用中是一个重要的追求。然而,长时间运动序列中发生的不频繁的语义动作使得长程动作建模变得困难。为了解决这个问题,一些研究提出利用 Mamba 捕捉运动序列中的时空模式 [182]。例如,Motion Mamba [225] 提出了一个混合 Mamba 模型,利用分层 SSM 层捕捉时间模式,并引入双向 SSM 层以学习空间知识,保持帧之间的运动一致性。根据综合实验,基于 Mamba 的模型在人体动作生成任务中超越了代表性的扩散方法,实现了 50 % {50}\% 50% 的 FID 改进和四倍更快的性能。此外,MambaMOS [215] 设计了一种运动感知状态空间模型,明确地关注于捕捉连续时间步之间运动的变化,进一步强调了 Mamba 在实现高质量、长序列运动建模方面的卓越能力。

5.2 非序列数据

非序列数据不同于序列数据,不依赖于特定的顺序。其数据点可以以任何顺序组织或访问,而不会显著影响数据的含义或解释 [80]。这种缺乏固有顺序的特性给专门设计以捕捉数据时序依赖关系的递归模型(如 RNN 和 SSM)带来了困难。令人惊讶的是,作为 SSM 的代表,Mamba 在处理非序列数据方面显示出了卓越的成功。在这一部分,我们将回顾有关 Mamba 如何有效处理非序列数据的相关研究,包括图像、图形和点云。

5.2.1 图像。作为最常见的模态之一,图像数据构成了各种计算机视觉应用的基础,例如,人脸识别、医疗视觉 [214] 和遥感 [21, 23]。受到 Mamba 在序列建模成功的启发,存在一个有趣的机会将这一成就从文本处理转移到图像分析。这涉及将图像视为一系列补丁,可能为计算机视觉领域的新探索开辟了新途径。因此,最近开发了许多基于 Mamba 的视觉模型,以缓解重大的计算资源和内存压力,同时展示出竞争性的建模能力 [114, 152, 179, 179]。例如,Vision Mamba [234] 融入了双向 SSM,以促进全局视觉语义建模,并引入位置嵌入以实现位置感知的视觉理解。Vision Mamba 不再需要注意力机制,其建模能力与 Vision Transformers 相匹配,同时显著降低了计算时间至次平方级别,并保持线性内存复杂度。具体而言,其在速度上超越了最先进的基线 DeiT [173],快了 2.8 × {2.8} \times 2.8× ,并在高分辨率图像 ( 1248 × 1248 ) \left( {{1248} \times {1248}}\right) (1248×1248) 的特征提取过程中,GPU 内存使用量减少了 86.8 % {86.8}\% 86.8% 。此外,VMamba [121] 引入了 2D 选择性扫描 (SS2D),作为 1D 数组扫描与 2D 平面遍历之间的桥梁,使 Mamba 能够有效处理视觉数据。

5.2.2 图结构数据。图建模在管理复杂结构和关系方面得到了广泛应用,包括社交网络 [41] 和分子相互作用 [77] 等领域的应用。受到 ChatGPT 等语言基础模型的显著发展的启发,许多研究者开始构建图基础模型,以为下游预测任务提供节点和边嵌入 [28]。由于 Mamba 在长程建模和高效率方面的强大能力,一些开创性的研究采用了选择性状态空间模型(SSM)处理非序列图数据 [115]。这些研究利用状态空间模型通过隐藏状态在递归扫描时编码上下文,允许输入流控制,这类似于图上的注意力稀疏化过程,展示了在图建模上下文中数据依赖的节点选择过程 [208]。此外,预计 Mamba 将在大图训练任务中提高模型效率。

例如,Graph-Mamba [178] 引入了一种新型的基于 Mamba 的模块,作为图建模的基础组件。该模块结合了图扁平化机制与 Mamba 提供的选择机制,分别将子图转换为节点序列并促进输入依赖的上下文过滤。在最近的工作中,Behrouz 和 Hashemi [9] 提出了图 Mamba 网络 (GMN),一种基于选择性 SSM 的新图神经网络格式。作者将选择性 SSM 重新格式化为图学习格式,并为所提出网络的强大能力提供了理论依据。通过解决图消息传递关键步骤中出现的挑战,GMN 在多个基准数据集上取得了优异的表现,超过了 GNN 和基于 Transformer 的模型,在具有多样化图规模的多个基准数据集上皆表现优异。此外,Huang 等 [81] 引入了图状态空间卷积 (GSSC),作为针对图结构数据的 SSM 的系统扩展。具体而言,GSSC 将基于距离的图卷积核纳入 SSM 单元,旨在增强表达能力并捕捉长程依赖性。通过在十个基准数据集上进行的评估,该研究 [81] 强调了 GSSC 作为图机器学习的强大且可扩展模型的潜力。5.2.3 点云。点云是计算机视觉中的一种关键模态,具有广泛的实际应用,涉及机器人技术、自动驾驶和增强现实等领域 [62]。与图像处理和图学习不同,点云分析面临着独特的挑战,这些挑战源于点云固有的不规则性和稀疏性,属于三维非结构化数据。为了应对这些挑战,基于深度学习的方法取得了显著的进展,特别强调基于Transformer的模型 [211]。然而,注意力机制的复杂性是二次的,这带来了显著的计算成本,对低资源设备并不友好。最近,状态空间模型(SSMs)在处理一维序列(例如语言和语音)和二维数据(例如图像和图形)方面取得了进展,因此有人尝试将Mamba的应用扩展到三维点云 [210]。一般来说,这些基于Mamba的点云分析方法采用两步过程 [ 63 , 232 ] \left\lbrack {{63},{232}}\right\rbrack [63,232] 。首先,使用特定的扫描方法将点云数据分割成离散的标记。然后,利用Mamba捕捉这些标记中的潜在模式。例如,PointMamba [109] 提出了一个层次扫描策略,以编码三维点云的局部和全局信息,然后利用普通的Mamba作为骨干提取从序列化点标记中提取特征,而不引入额外的复杂技术。Point Cloud Mamba [220] 将Mamba作为基础模型骨干,显著降低内存使用,表现出与基于Transformer的对应物相当(或更优)的性能。

5.3 多模态数据

集成多种模态,如语言(序列数据)和图像(非序列数据),为人工智能感知和场景理解提供了有价值的互补信息。最近,针对继承了大型语言模型 [192] 的先进能力的多模态大型语言模型(MLLMs)进行了显著的研究关注,包括强大的语言表达和逻辑推理。尽管Transformer在该领域一直是主导方法,但Mamba作为一个强有力的竞争者,通过在混合源数据对齐和序列长度线性复杂度扩展方面表现出色,使得Mamba成为多模态学习中对Transformer的有希望的替代品 [117, 206]。例如,Qiao等人 [145] 提出了VL-Mamba,旨在探索Mamba高效架构在解决视觉-语言任务中的应用,利用预训练的Mamba模型进行语言理解,并结合连接模块将视觉补丁与手稿提交给ACM语言标记对齐。Wang等人 [182] 提出了文本控制运动Mamba [182],利用Mamba动态捕捉基于文本查询的全球时间信息,以增强人类运动理解。此外,Fusion-Mamba [33] 和 Sigma [177] 尝试融合热成像、深度和RGB等不同模态的互补信息。Fusion-Mamba专注于改善物体检测,而Sigma旨在增强语义分割。

6 应用

在本节中,我们介绍一些基于Mamba模型的显著应用。为了提供全面的概述,我们将这些应用分类为:自然语言处理、计算机视觉、语音分析、药物发现、推荐系统以及机器人和自主系统。

6.1 自然语言处理

在自然语言处理领域,最近出现了一些基于Mamba的模型作为针对语言建模的Transformer模型的替代方案 [ 5 , 14 , 67 , 111 , 176 , 226 ] \left\lbrack {5,{14},{67},{111},{176},{226}}\right\rbrack [5,14,67,111,176,226] ,尤其在涉及广泛上下文的应用中,如问答系统和文本摘要。

6.1.1 问答系统。问答(QA)涉及 AI 模型使用广泛知识库进行理解、推理和响应,从而实现连贯且具有丰富上下文的对话,广泛应用于聊天机器人和虚拟助手中。纳入先前交互的上下文对于准确处理多轮对话中的后续问题至关重要。然而,现有模型在推理速度和计算效率方面面临挑战,特别是在复杂推理任务中。这导致显著的内存使用和计算开销,限制了可扩展性和实时应用效率。为了解决这些局限性,最近的研究探索了基于Mamba的模型,以改善问答系统中的长期对话管理 [111, 111, 132]。例如,Mamba-Chat [132] 是第一个利用状态空间框架的聊天语言模型。该模型通过使用状态空间表示来维持和更新对话的理解,从而确保上下文意识。Jamba [111] 在Transformer和Mamba层之间战略性地交替,结合MoE以提升模型容量,同时优化参数使用。在常识推理和阅读理解任务中,Jamba的表现可与更大Llama-2模型相媲美,但拥有更少的参数,展示了高效和有效性。类似地,DenseMamba [67] 提出了通过选择性地将浅层的隐藏状态融入深层,丰富SSMs中隐藏信息传播的全新方法。与传统基于Transformer的模型相比,保持了重要的细粒度信息,从而在问答任务中表现更优。总体而言,整合基于Mamba的模型显示出推动问答系统发展的良好潜力,通过改善对话管理和增强复杂推理任务的性能。

6.1.2 文本摘要。文本摘要旨在通过保留关键信息来压缩长文本。在这个任务中,保持一致性和相关性至关重要。基于Transformer的模型通常在处理长序列依赖方面面临困难,这可能会影响一致性和相关性。相比之下,基于Mamba的模型利用强大的长序列处理能力,使其非常适合处理连贯且丰富上下文的文本。其强大的架构使其在摘要任务中表现出色,能够精准捕捉和浓缩大量文档的核心内容。例如,LOCOST [14] 基于状态空间模型,处理的序列长度显著超过稀疏注意力模型。在长文档的抽象摘要任务中,LOCOST的性能可媲美等效维度的稀疏Transformer的最高表现,同时在训练中内存使用减少了多达 50 % {50}\% 50% ,在推理中减少了 87 % {87}\% 87% 。此外,SAMBA [149] 将Mamba与滑动窗口注意力集成,使得能够将选择的序列压缩到递归隐藏状态中,同时通过注意力机制保留准确的记忆回调。SAMBA在处理输入长度为 128    K {128}\mathrm{\;K} 128K 的情况下,其吞吐量比Transformer高出 3.73 倍,展示了在需要长上下文摘要的任务中的优越性能。

6.2 计算机视觉

除了NLP应用,基于Mamba的模型在计算机视觉领域也显示出潜力,代表性应用包括疾病诊断和运动识别与生成。

6.2.1 疾病诊断。在临床实践中,医学图像和视频提供了关于器官或组织形态的重要见解。高效分析生物医学对象,如大型二维/三维图像或视频中的病变,能够显著提升疾病诊断和临床治疗。然而,像UNet这样的基于CNN的模型在处理长距离依赖时面临挑战,因为其感受野有限。由于医学图像通常比自然图像的大小更大、分辨率更高,这一挑战进一步加剧。同时,基于Transformer的算法计算量大,限制了它们在资源有限的临床环境中的实用性。为克服这些局限性,许多研究在真实医学环境中采用了基于Mamba的模型 [110, 127, 152, 187]。例如,U-Mamba [127] 和 SegMamba [198] 都集成了一个混合CNN-SSM模块,将卷积层的局部特征提取能力与SSMs提供的长距离依赖建模相结合。这种混合方法在如 CT 和 MR 图像中腹部器官的3D分割、内窥镜图像中的仪器分割以及显微镜图像中的细胞分割等任务中优于现有模型。类似地,CMViM [202] 通过利用掩蔽的Vim自编码器和跨模态对比学习,解决了阿尔茨海默病(AD)诊断成像中的挑战,在AD 诊断成像分类中取得了最佳性能。此外,ProMamba [196] 专门用于息肉分割。通过集成 Vision-Mamba 结构和提示技术,该模型比先前的方法取得了更高的准确性和更好的泛化能力。对于视频中的动态医学物体分割,Vivim [205] 通过使用时间Mamba模块有效地将长期的时空表示在不同尺度之间压缩为序列。这一方法在疾病诊断中表现出增强的性能和计算效率,例如超声乳腺病变分割和结肠镜视频中的息肉分割。6.2.2 动作识别与生成。 动作识别和生成在动作监测 [52]、计算机动画 [165]、游戏开发 [133] 和电影制作 [184] 中至关重要。然而,基于变压器的模型在计算和内存需求方面遇到挑战,限制了它们在资源受限环境中的适用性。此外,基于变压器和GCN的模型在有效捕获视频和4D点云中的长动作序列和复杂时空模式方面也面临困难。最近的研究探索了使用Mamba来应对这些挑战,利用其强大的性能和较低的计算需求 [18, 105, 224, 225]。例如,HARMamba [105] 利用双向SSM架构处理来自穿戴式传感器的数据,显著降低了计算负载和内存使用,同时保持实时人类动作识别的高准确性。类似地,Simba [18] 在U-ShiftGCN框架中整合了Mamba,有效处理较长序列和复杂的时空交互,在视频中的骨骼动作识别中取得了最佳结果。此外,Motion Mamba [225] 和InfiniMotion [224] 均用于动作生成。具体而言,Motion Mamba [225] 利用分层时间Mamba块处理时间数据,双向空间Mamba块处理潜在姿态,确保帧间动作一致性,并提高在时间帧内的生成精度。InfiniMotion [224] 引入了带有双向Mamba内存的运动记忆变换器,提高了变换器的记忆能力,从而有效生成持续的长时间人类动作(长达一个小时和80,000帧),而不至于使计算资源超负荷。

6.3 语音分析

语音信号本质上由数千个样本组成。虽然这种广泛的时间上下文提供了丰富的声学特征,但它也带来了显著的计算需求。为了有效处理语音信号,几种基于Mamba的模型已成功应用于多种语音应用,特别是在语音分离和标记及语音增强中。

6.3.1 语音分离与标记。 语音分离涉及从多说话者环境中孤立个体语音信号。它对提高音频通信的可懂性和质量至关重要。同时,音频标记或分类涉及将音频样本映射到其对应类别。两个任务都依赖于捕捉短期和长期的音频顺序模式。尽管基于变压器的模型在这些应用中一直是主导架构,但由于其自注意机制导致的平方级计算和内存成本,它们面临显著的挑战。最近,语音分离 [ 87 , 101 ] \left\lbrack {{87},{101}}\right\rbrack [87,101] 和音频标记 [ 11 , 222 ] \left\lbrack {{11},{222}}\right\rbrack [11,222] 的状态空间模型得到了越来越多的应用。具体而言,DPMamba [ 87 ] \left\lbrack {87}\right\rbrack [87] 利用选择性状态空间捕获语音信号中的动态时间依赖关系,涵盖了短期和长期的前向与后向依赖。SPMamba [101] 整合了TF-GridNet模型,将其变压器组件替换为双向Mamba模块。DASS [11] 将知识蒸馏与状态空间模型结合,允许对持续时间最长达2.5小时的音频文件中的声音事件进行标记。同时,MAMCA [222] 通过引入选择性状态空间模型作为其骨干,专注于自动调制分类(AMC),有效解决了与长序列AMC相关的准确性和效率挑战。采用状态空间模型后,这些模型展示了定性改进,捕获了更广泛的上下文信息,增强了整体效果,从而证明了SSM在处理长时间时的优越可扩展性。

6.3.2 语音增强。 语音增强(SE)旨在从失真信号中提取清晰的语音成分,产生具有改进声学特征的增强信号。作为前端处理器,SE在多种语音应用中至关重要,包括助听技术 [98]、说话人识别 [7] 和自动语音识别 [129]。移动音频设备面临由于资源有限而带来的挑战。最近的研究探索了Mamba的应用,利用其强大性能和减少的计算需求来进行SE任务 [17, 148, 159, 168, 221]。例如,TRAMBA [168] 利用结合变压器和Mamba的混合架构改善移动和可穿戴平台的语音质量,特别针对声学和骨导。它实现了与当前领先模型相比,内存消耗减少十倍的显著效果。此外,oSpatialNet-Mamba [148] 利用Mamba进行长期多通道语音增强,对于静态和移动的说话者取得了卓越的结果。

6.4 药物发现

蛋白质设计、分子设计和基因组分析在推进药物发现和生物技术中至关重要 [100, 156]。借助基于MAMBA的模型显著减少了这些领域中建模长序列的复杂性 [61, 140, 140, 154, 154]。具体而言,PTM-Mamba [140] 和ProtMamba [158] 是基于Mamba架构的蛋白质语言模型。PTM-Mamba 利用双向门控Mamba块和结构化状态空间模型,高效处理长序列,同时降低计算需求。ProtMamba 旨在具备同源性意识而不依赖比对,能够处理数百个蛋白质序列中的广泛上下文。这两种模型即使在处理大数据集时也能保持高效率和准确性,为蛋白质设计提供了关键工具。与此同时,生成性分子设计旨在从特定分布中模拟具有定制属性特征的分子。然而,当前模型缺乏优化高保真预言所需的效率,导致成功率低下。Saturn [61] 应用Mamba架构,利用其线性复杂性和计算效率,超过22个竞争模型用于药物发现。此外,理解基因组对于获得细胞生物学的见解至关重要。基因组建模中的挑战包括捕捉远距离符号之间的相互作用,考虑上下游区域的影响,以及确保DNA序列的互补性。Caduceus [154] 和MSAMamba [172] 流行,均利用Mamba模型,在解决这些挑战方面表现出色。Caduceus作为一个DNA基础模型,通过增强Mamba架构中的BiMamba和MambaDNA组件以进行双向建模,确保反向互补等变性,在长距离基因组任务中大大优于现有模型。同样,MSAMamba [172] 通过在序列维度上实现选择性扫描操作,解决了基于变压器的模型在DNA多重序列比对中的局限性。该设计将之前方法的训练上下文长度扩展至八倍,使得对广泛DNA序列的综合分析变得可行。

6.5 推荐系统

推荐系统广泛应用于电子商务 [ 25 , 219 , 231 ] \left\lbrack {{25},{219},{231}}\right\rbrack [25,219,231] 和社交网络 [ 38 , 39 , 43 ] \left\lbrack {{38},{39},{43}}\right\rbrack [38,39,43] ,旨在捕捉用户不断变化的偏好及其过去行为之间的相互依赖关系 [40, 228]。尽管基于变压器的模型在推荐系统中展示了有效性,但由于注意机制的平方复杂性,尤其是在处理较长行为序列时,它们面临着计算效率的挑战。最近,多个基于Mamba的模型已应用于分析长期用户行为以进行个性化推荐 [ 16 , 115 , 166 , 185 , 203 ] \left\lbrack {{16},{115},{166},{185},{203}}\right\rbrack [16,115,166,185,203] 。例如,Mamba4Rec [115] 首创使用选择性状态空间模型进行高效的序列推荐,提升模型性能的同时保持推理效率。类似地,RecMamba [203] 探索了Mamba在终身序列推荐场景中的有效性(即序列长度 ≥ 2 k \geq 2\mathrm{k} 2k ),在压缩训练时间70%和减少内存成本80%的情况下,取得了与基准模型相当的性能。此外,EchoMamba4Rec [185] 将双向Mamba模块与频域滤波相结合,以准确捕捉用户交互数据中的复杂模式和相互依赖关系。它在现有模型中显示出了优越的性能,提供了更精准和个性化的推荐。此外,Mamba4KT [16] 明确设计用于智能教育中的知识追踪,利用Mamba模型捕捉练习与学生知识水平之间的持久相关性。随着教育数据集的扩展,这种方法为提升知识追踪研究中的预测准确性、模型效率和资源利用率提供了一个有前景的途径。

6.6 机器人技术与自主系统

机器人和自主系统的主要目标是开发能够理解视觉环境并执行复杂动作的模型。目前在机器人领域使用的多模态大型语言模型(MLLMs)面临两方面重大挑战:1)处理复杂任务所需的高级推理能力有限,2)微调和推理任务带来了巨大的计算开销。由于在推理速度、内存利用和整体效率方面的优势,基于Mamba的模型正在成为自主和智能系统的有希望基础 [ 15 , 85 , 117 ] \left\lbrack {{15},{85},{117}}\right\rbrack [15,85,117] ,承诺提供卓越的性能和显著的可扩展性。例如,RoboMamba [117] 将视觉编码器与Mamba结合,创建了一个端到端的机器人MLLM。这种方法通过共同训练将视觉数据与语言嵌入对齐,增强模型的视觉常识和特定于机器人的推理能力,同时确保高效的微调和推理能力。类似地,Jia等人 [85] 提出了MaIL,一种使用Mamba作为骨干的模仿学习(IL)策略架构。MaIL在处理观察序列时填补了效率与性能之间的差距。真实机器人实验的广泛评估表明,MaIL提供了一种有竞争力的替代方案,优于传统的、庞大且复杂的基于变压器的IL策略。

7 挑战与机会

前面的章节已彻底考察了与Mamba相关的最新先进技术和多样化应用。然而,对于Mamba的研究仍处于初期阶段,未来还有大量挑战和机遇。## 7.1 基于 Mamba 的基础模型

通过将模型规模扩大到十亿级,并在大规模混合源语料库上进行训练,基础模型(FMs)展现出了令人印象深刻的零-shot 学习能力,使得 FMs 在各种一般任务中表现出色 [13]。作为一个代表性例子,近年来 Transformer 基础的大型语言模型,尤其是 ChatGPT 的蓬勃成功,激发了人们对在各种领域探索基础模型的日益热情。尽管 Transformers 是成功的主要推动力,但它们面临着迫切的计算和内存效率问题 [171],这种问题伴随着与基于注意力的模型规模成指数增长的训练内存和推理过程中繁琐的自回归解码。为应对这些问题,一个有前景的替代基础,亦即 Mamba [29, 55],最近已出现。Mamba 提供了 Transformers 的内容感知学习能力,同时使计算与输入长度线性扩展,有效捕捉长距离依赖并提高训练和推理的效率。鉴于这些优点,针对特定领域开发基于 Mamba 的基础模型具有巨大的潜力,为解决基于 Transformer 的模型面临的问题提供了机会。

7.2 硬件感知计算

基础模型的特点是其庞大的规模和密集的矩阵运算,例如矩阵乘法和卷积,这需要尖端硬件如 GPU 和 TPU 来进行高通量的训练和推理。这些先进的硬件使研究人员能够处理更大的数据集,并在各个领域实现最先进的性能。然而,现有的基础模型仍未能充分利用硬件的计算能力,导致模型效率有限 [171]。作为提高计算效率的有前景的替代方案,Mamba-1 [55] 和 Mamba-2 [29] 提出了硬件感知计算算法,即并行关联扫描和块分解矩阵乘法。这些算法考虑了 GPU 和 TPU 的固有特性,包括设备之间的消息传输等因素,为解决计算效率问题提供了新的视角。受到此启发,探索新颖的硬件高效算法,例如 FlashButterfly [48],以优化硬件利用率,为节约资源和加速计算提供了有希望的途径,不仅有利于 SSM,还可用于其他架构,如 Transformers 和 RNNs。

7.3 可信赖的 Mamba 模型

SSM 的发展预计将为电子商务、医疗保健和教育等各个行业带来显著利益。同时,作为数据驱动的模型,Mamba 模型可能对用户和社会构成严重威胁 [130]。这些威胁源于多个因素,如不可预测的决策、隐私问题等。因此,确保 Mamba 模型在四个关键维度上的可信性至关重要 [116]:安全性与鲁棒性、公平性、可解释性和隐私。

7.3.1 安全性与鲁棒性。大型基础模型已被证明对对抗扰动高度脆弱,这可能危及这些模型在安全关键应用中的安全性和鲁棒性 [ 45 , 135 , 190 ] \left\lbrack {{45},{135},{190}}\right\rbrack [45,135,190] 。与此同时,基于 Mamba 的模型也未能免受这些脆弱性的影响 [128]。为了成为 Transformer 的可靠替代品,必须研究和提高基于 Mamba 的模型的安全性和鲁棒性。具体而言,模型输出应对输入中的小扰动保持鲁棒性。一种潜在的解决方案可能涉及在将提示输入 Mamba 模型之前自动进行预处理。此外,作为一种代表性技术,对抗机器训练 [78] 可用于增强 Mamba 模型的安全性和鲁棒性。

7.3.2 公平性。经过大量数据集训练的大型基础模型,往往会在无意中暴露于广泛训练语料库中存在的偏见和刻板印象 [126],这些偏见可能会在生成的输出中显现。例如,在 LLMs 领域中,偏见可能导致受用户资料属性(如性别和年龄)影响的歧视性响应,从而强化刻板印象并不公平地对待特定用户群体 [86]。尽管近期已经有努力解决 LLM 中的公平性问题,但有关 Mamba 模型的非歧视性和公平性的研究仍然有待加强。因此,有必要进一步探索和研究以弥补这一空白。

7.3.3 可解释性。深度学习模型常常因其“黑箱”性质受到批评,深度学习模型的可解释性已成为研究界的热门话题,这表明理解和解释模型生成的决策或预测的能力 [35]。通过解释模型的预测,用户能够基于模型的输出做出更明智的决策。为此,已经提出了几种技术,以提供基于注意力机制的神经架构的合理内在解释 [74]。此外,研究人员还考察了基于 Transformer 的语言模型生成自然语言描述以解释其答案的能力 [212]。尽管越来越多的研究试图充分利用 Mamba,但对理解 Mamba 模型功能的研究仍处于早期阶段,仍需进一步调查。

7.3.4 隐私。隐私保护在用户与基于 Mamba 的模型之间建立了信任。当用户相信他们的隐私得到了尊重时,他们更可能与 AI 系统互动、分享相关信息,并在不担心数据被滥用的情况下寻求帮助。因此,这种信任对 Mamba 模型的广泛采用和接受至关重要。减轻隐私风险的一种有效策略是交叉验证 Mamba 模型的输出并筛选敏感内容 [93]。此外,联邦学习有望在 Mamba 模型的训练过程中增强隐私,其中模型在众多分散的边缘设备或服务器上进行训练,侧重于本地数据样本,而无需数据交换。这种方法有助于保持数据的本地化和隐私。此外,在训练过程中整合隐私意识的正则化技术,如差分隐私约束,表明在防止对敏感数据过拟合方面具有良好的前景。

7.4 将 Transformer 的新兴技术应用于 Mamba

Transformer 作为主导基础,促使 AI 社区开发了众多旨在增强基于注意力模型性能的独特工具。幸运的是,通过将 SSM 和注意力连接起来,Mamba-2 [29] 提出的 SSD 框架使我们能够为 Transformer 和 Mamba 开发共享的词汇和技术库。有鉴于此,一个重要的未来方向出现了,即探索为基于 Transformer 的模型设计的新兴技术如何有效地应用于基于 Mamba 的模型。

7.4.1 参数高效微调。大型基础模型将其参数扩大到十亿级,在多个领域见证了突破性进展。然而,它们庞大的规模和计算需求在为特定下游任务定制时面临重大挑战。为此,提出了几种参数高效微调(PEFT)技术,包括 LoRA [72] 和适配器系列 [50, 91],其中涉及在微调过程中最小化对参数的调整或对广泛计算资源的需求。受到最近在基于 Transformer 层构建的大型语言模型中采用 PEFT 的成就的启发,Mamba 模型的 PEFT 采用成为一个引人注目的话题,旨在扩大它们在下游任务中的应用范围。例如,预计 LoRA(低秩适应)的部署将为 SSD 模型提供快速微调的便利,从而使 Mamba 在各个领域得到广泛应用。然而,针对基于 Mamba 的模型实施这些 PEFT 技术的具体细节尚待确定,并需要进一步研究。

7.4.2 灾难性遗忘的缓解。灾难性遗忘,也称为灾难性干扰,是指机器学习模型在训练新任务时,其在先前学习任务上的性能显著下降的现象 [92]。这个问题对基础模型构成了挑战,因为它们需要保留来自预训练任务的知识,并在不同下游领域表现出一致的性能。作为基础模型的有前景架构,Mamba 需要深入调查以解决灾难性遗忘的问题。最近的研究已经提出通过奖励最大化和分布匹配策略来封装任务特定需求以解决这一挑战 [96, 97]。此外,还开发了持续学习方法以缓解基于 Transformer 的语言模型中的灾难性遗忘 [90, 186]。这些技术也可以通过将 SSM 和注意力连接应用于 Mamba 模型,但仍待进一步探索。

7.4.3 检索增强生成(RAG)。作为 AI 中最复杂的技术之一,RAG 能够提供可靠和最新的外部知识,为多种任务提供显著的效用 [31, 99]。尽管大型语言模型最近展示了突破性的语言理解和生成能力,但仍然存在固有限制,如幻觉和过时的内部知识。鉴于 RAG 提供最新和有价值的补充信息的强大能力,检索增强的 LLMs 已应运而生,它们利用外部知识数据库提高 LLMs 的生成质量 [22]。同样,RAG 也可以与 Mamba 语言模型结合,帮助它们生成高质量的输出,这也是一个有前景的未来研究方向。

8 结论

Mamba,一种新兴的深度学习架构,由于其强大的建模能力和计算效率,在语言生成、图像分类、推荐和药物发现等多个领域展现了显著的成功。最近,越来越多的努力集中在开发基于Mamba的深度学习模型,以实现更强大的表征学习能力和更低的计算复杂性。考虑到Mamba的快速发展,迫切需要一个系统的概述。为了弥补这一空白,本文提供了对Mamba的全面回顾,重点关注其架构进展、数据适应性和应用领域,为研究人员提供深入的理解和对Mamba最新发展的概述。此外,鉴于Mamba研究仍处于初期阶段,我们还讨论了当前的局限性,并提出了未来研究的有希望的方向。

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

数智笔记

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值