Skip to content

Mamba 模型

摘录自 一文读懂:Mamba 模型,transformer的挑战者

Mamba 作为近期的深度学习新架构,被视为为超过transformer的序列建模构架,从文本nlp,语音到视觉处理,相关应用、研究、变体遍地开花。可以说,原本用Transformer做的工作都可以在mamba上重做一遍。

引言

尽管Transformer取得了令人印象深刻的成就,但仍然面临着固有的局限性,尤其是由于注意力计算的二次方计算复杂度,这导致了耗时的推理过程。

2024年,一种名为Mamba的新型架构,从经典的状态空间模型(SSM)中汲取灵感,作为构建基础模型的有前途的替代方案出现,它在保持与Transformer相当的建模能力的同时,对于序列长度具有近线性的可扩展性。

这激发了越来越多的研究人员积极探索Mamba在不同领域实现卓越性能的潜力。鉴于这种快速发展,迫”本文作为一个系统性的回顾,整合现有的Mamba模型资料,介绍这种新兴的模型架构。

image

1 Mamba 的概括

Mamba 是一种基于结构化状态空间序列模型(SSM)的新兴架构,旨在高效捕捉序列数据中的复杂依赖性,成为 Transformer 的强大竞争对手。

Mamba 的高效性和灵活性使其有望在多个研究和应用领域引发革命性变革。

受经典状态空间模型启发,Mamba 融合了循环神经网络(RNN)和卷积神经网络(CNN)的特点,通过递归或卷积操作实现计算成本与序列长度的线性或近线性扩展,显著降低计算复杂度。

具体而言,Mamba 的核心优势包括:

  1. 选择机制:引入简单而有效的选择机制,通过输入参数化 SSM 参数,过滤无关信息,保留必要数据。
  2. 硬件感知算法:采用递归扫描而非卷积计算,优化硬件性能,在 A100 GPU 上实现高达 3 倍的加速。
  3. 建模能力:在保持与 Transformer 相当的建模能力的同时,具备近线性的可扩展性,适用于复杂和长序列数据。

这些特性使 Mamba 成为处理多领域任务的理想基础模型,已在计算机视觉、自然语言处理和医疗保健等领域展现出卓越性能。例如,Vim 模型在高分辨率图像特征提取中比 DeiT 快 2.8 倍,节省 86.8% 的 GPU 内存;而在语言建模任务中,改进的选择性 SSM 架构实现了 2-8 倍的加速。

image

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

由于Mamba强大的长序列建模能力及其高效率,大量的文献出现了,专注于在各种下游任务中使用和改进Mamba。

我们从几个角度对Mamba进行了全面的回顾,为新来者提供了Mamba内部工作原理的基本理解,同时帮助经验丰富的从业者了解其最新发展。

Mamba 的架构设计

image

Mamba 架构的主要创新点包括:对输入信息有选择性处理、硬件感知的算法、更简单的架构。它是第一个真正实现匹配 Transformer 性能的线性时间序列模型,建立在更现代的适用于深度学习的结构化 SSM(S4, Structured SSM)基础上,与经典架构 RNN 有相似之处。

另外,基于 Mamba 架构还出现了一些新的魔改版本,应用于不同领域并带来了更多创新。我们讲在后面进行深入介绍。

例如,u-mamba 用于生物医学图像的分割任务,采用了混合 CNN-SSM 架构,能捕捉局部细粒度特征和长程依赖关系,具有自配置能力、线性扩展能力以及与其他技术集成的潜力;Weak-mamba-unet 是用于医学图像分割的弱监督学习框架,结合了 CNN、ViT 和 vmamba 的优势,采用多视角交叉监督学习方法;graph-mamba 是一种新型的图网络,将选择性状态空间模型与图网络集成,实现了输入相关的节点过滤和自适应上下文选择,在大型图上能减少高达74%的 GPU 内存消耗;swin-umamba 则提出了用于2D 医学图像分割的基于 mamba 的网络及其变体结构,变体具有更少的参数和更低的 FLOPs,适用于高效应用,并且有效整合了基于 ImageNet 的预训练。

image

Mamba 是一种新的选择性结构状态空间模型,它的一些最新创新点包括:

  1. 简化的 SSM 架构:Mamba 将结构化状态空间序列模型(S4)中使用的类似线性注意力的块和多层感知器(MLP)块进行集成,构建了 Mamba 块。总体架构由重复的 Mamba 块与标准规范化层和残差连接交织组成。它继承了状态空间模型序列长度的线性可伸缩性,同时实现了类似于 Transformer 的建模能力。
  2. 选择机制:通过将参数作为输入的函数,提高基于上下文的推理能力,利用选择机制,实现更高效和有效地捕获相关信息,特别是在处理长序列时。这种选择机制可以基于输入内容有选择性地传播或遗忘信息,过滤掉与问题无关的信息,并且可以长期记住与问题相关的信息。
  3. 硬件感知算法:融合了内核和重新计算的硬件感知算法,避免了中间状态的存储,减少内存需求,提高计算效率。例如采用“并行扫描算法”而非“卷积”来进行模型的循环计算,为了减少 GPU 内存层次结构中不同级别之间的 I/O 访问,它没有具体化扩展状态。

image

2 Mamba的基础知识

Mamba与循环框架的循环神经网络(RNNs)、并行计算和Transformer的注意力机制以及状态空间模型(SSMs)的线性属性密切相关。因此,本节旨在介绍这三种突出架构的概述。

2.1 循环神经网络(RNNs)

RNNs在处理序列数据方面表现出色,因为它们能够保留内部记忆。这类网络在涉及分析和预测序列的众多任务中表现出显著的有效性,例如语音识别、机器翻译、自然语言处理和时间序列分析。为了掌握循环模型的基础,本节将提供标准RNN公式的简要概述。

具体来说,在每个离散时间步骤k,标准RNN特别处理一个向量 ∈ RD,连同前一步骤的隐藏状态e ℎ −1 ∈ R ,以产生一个输出向量 ∈ R 并更新隐藏状态到 ℎ ∈ RN

隐藏状态充当网络的内存,并保留有关它所见过去输入的信息。这种动态内存允许RNN处理不同长度的序列。正式地,它可以写成

\[\begin{aligned} h_k &= \tanh(\mathbf W_{hx}x_k + \mathbf W_{hh}h_{k-1} + b_h), \\ o_k &= \mathbf W_{oh}h_k + b_o. \end{aligned}\]

其中Wℎ ∈ R ×D是负责将模型输入处理成隐藏状态的权重矩阵,Wℎℎ ∈ R × 是隐藏状态之间的递归连接,W ℎ ∈ R ×N表示用于从隐藏状态生成输出的权重,, ℎ ∈ R 和 ∈ R 对应偏差,tanh表示引入非线性到RNN模型的双曲正切激活函数。换句话说,RNN是非线性递归模型,通过利用隐藏状态中存储的历史知识有效地捕获时间模式。

然而,RNNs有几个局限性。首先,RNNs在有效提取输入序列中的长距离动态方面能力有限。随着信息通过连续的时间步骤传播,网络中权重的重复乘法可能导致信息的稀释或丢失。

因此,对于RNNs来说,在进行预测时保留和回忆早期时间步骤的信息变得具有挑战性。其次,RNNs以增量方式处理序列数据,限制了它们的计算效率,因为每个时间步骤都依赖于前一个。这使得并行计算对于它们来说很困难。此外,传统的RNNs缺乏内置的注意力机制,这允许网络捕获输入序列中的全局信息。这种注意力机制的缺失限制了网络选择性地建模数据的关键部分的能力。

为了克服这些限制,Transformer和状态空间模型出现了,每种方法都从不同的角度解决了这些挑战。这些两种方法将在后续内容进一步阐述。

2.2 Transformers

Transformer是深度学习领域的开创性模型,彻底改变了各种AI应用。它的引入标志着与传统序列到序列模型的显著偏离,通过采用自我注意力机制,促进了对模型输入中全局依赖性的捕获。例如,在nlp中,这种自我注意力能力允许模型理解序列中不同位置之间的关系。

image

transformer详细介绍参考这个链接←

\[Q = x \cdot \mathbf W^Q,\quad K = x \cdot \mathbf W^K,\quad V = x \cdot \mathbf W^V,\]

这样的程序然后通过Softmax函数传递,以标准化分数 ?并产生注意力权重,定义为:

\[S = \operatorname{Softmax}\left(\frac{QK^T}{\sqrt{d_K}}\right)V,\]

除了执行单个注意力函数外,多头注意力被引入以增强模型捕获不同类型关系的能力,并为输入序列提供多种视角。在多头注意力中,输入序列并行通过多个自注意力模块进行处理。每个头独立操作,执行与标准自注意力机制完全相同的计算。

然后,每个头的注意力权重被结合起来,创建值向量的加权和。这个聚合步骤允许模型利用来自多个头部的信息,并捕获输入序列中的多样化模式和关系。

数学上,多头注意力计算如下:

\[\begin{aligned} \operatorname{MultiHead}(Q,K,V) &= (S_1 \oplus S_2 \oplus \cdots \oplus S_m)\cdot \mathbf W^O, \\ S_i &= \operatorname{Softmax}\left(\frac{Q_iK_i^T}{\sqrt{d_K}}\right)V_i,\quad i\in[1,m]. \end{aligned}\]

其中m是注意力头的数量,⊕是连接操作,WO是将多头注意力分数投影到最终值的线性变换。

2.3 状态空间模型

状态空间模型(SSMs)是一种传统的数学框架,用于描述系统随时间的动态行为。近年来,SSMs在控制理论、机器人技术和经济学等多个领域中发现了广泛的应用。在其核心,SSMs通过一组隐藏变量,即“状态”,来体现系统的行为,使其能够有效地捕获时间数据依赖性。与RNNs不同,SSMs是线性模型,具有关联属性。

具体来说,在经典的状态空间模型中,制定了两个基本方程,即状态方程和观测方程,通过当前时间t的N维隐藏状态ℎ( ) ∈ R 来模拟输入 ( ) ∈ R和输出 ( ) ∈ R之间的关系。该过程可以写成

\[\begin{aligned} h'(t) &= Ah(t) + Bx(t), \\ y(t) &= Ch(t) + Dx(t). \end{aligned}\]

其中ℎ′( ) 是当前状态ℎ( )的导数,A ∈ R × 是描述状态如何随时间变化的状态转移矩阵,B ∈ R ×1是控制输入如何影响状态变化的输入矩阵,C ∈ R1× 表示基于当前状态生成输出的输出矩阵,D ∈ R表示决定输入如何直接影响输出的命令系数。一般来说,大多数SSMs在观测方程中省略了第二项,即,设置D ( ) = 0,这可以被认为是深度学习模型中的跳过连接。

2.3.1 离散化。为了符合机器学习设置对各种现实世界场景的要求,SSMs必须经历一个离散化过程,将连续参数转换为离散参数。离散化方法通常旨在将连续时间划分为 ?个具有相等积分区域的离散间隔。

为了实现这一目标,作为最具代表性的解决方案之一,零阶保持(ZOH)成功地应用于SSMs,它假设函数值在间隔Δ = [ −1, ]内保持恒定。在ZOH离散化之后,SSM方程可以重写为

\[\begin{aligned} h_k &= \bar A h_{k-1} + \bar B x_k, \\ y_k &= C h_k. \end{aligned}\]

其中A = exp(ΔA),B = (ΔA)−1(exp(ΔA) − I) · ΔB,k是离散时间步长。从这些公式中,很明显离散SSM具有类似于循环神经网络的结构,因此离散SSM可以完成与Transformer基础模型相比,具有更高效率的推理过程。

2.3.2 卷积计算。作为线性系统,离散SSM具有关联属性,因此可以无缝集成到卷积计算中。更具体地说,它可以独立地计算每个时间步的输出,如下所示:

\[\begin{aligned} y_0 &= C\bar A^0\bar Bx_0, \\ y_1 &= C\bar A^1\bar Bx_0 + C\bar A^0\bar Bx_1, \\ y_2 &= C\bar A^2\bar Bx_0 + C\bar A^1\bar Bx_1 + C\bar A^0\bar Bx_2, \\ &\cdots \\ y_k &= C\bar A^k\bar Bx_0 + C\bar A^{k-1}\bar Bx_1 + \cdots + C\bar A^1\bar Bx_{k-1} + C\bar A^0\bar Bx_k. \end{aligned}\]

通过创建一组卷积核s K = (CB, ..., CA B, ...),递归计算可以转换为卷积形式:

\[\mathbf y = \mathbf x * \bar{\mathbf K},\]

其中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)的计算算法。

image

图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在某些任务中的性能下降,如选择性复制。

image

表1. 在自回归序列建模任务中,三种主要架构-RNNs、Transformers和SSMs-的优缺点比较。

3 Mamba1

为了解决传统SSM在上下文感知能力方面的不足,Mamba的提出是一种潜在的替代方案,承诺成为一个通用的序列基础模型骨干。最近,Mamba-2提出了结构化状态空间对偶性(SSD),建立了一个强大的理论框架,将结构化SSMs与各种形式的注意力联系起来,允许我们将最初为Transformer开发的算法和系统优化转移到SSMs。在这一节中,我们将简要而清晰地介绍Mamba和Mamba-2。

3.1 硬件感知算法的选择性状态空间模型

传统的SSM在建模文本和其他信息密集型数据方面表现出有限的有效性,阻碍了它们在深度学习中的进展。在追求赋予SSMs与Transformer相当的建模能力的过程中,Gu和Dao基于结构化状态空间模型提出了三种创新技术,即基于高阶多项式投影算子(HiPPO)的内存初始化、选择机制和硬件感知计算,如图3所示。

这些技术旨在增强SSMs在长距离线性时间序列建模方面的能力。特别是,初始化策略建立了一个一致的隐藏状态矩阵,有效地促进了长距离记忆。然后,选择机制使SSM能够获得内容感知表示。最后,Mamba设计了两种硬件感知计算算法,即并行关联扫描和内存重计算,以提高训练效率。

image

3.2 HiPPO-based Memory Initialization。

建模和从序列数据中学习是当代机器学习的基础挑战,构成了各种任务的基础,包括语言建模、语音识别和视频处理。在建模复杂和长期时间依赖性方面,内存是一个基本组成部分,包括存储和整合来自先前时间步骤的信息。与RNNs类似,在SSMs中保留和遗忘历史隐藏状态(即矩阵A)在实现令人满意的性能方面起着关键作用。在以前的结构化状态空间序列模型(SSMs)中,已经提出了特殊的初始化建议,特别是在数据可用性有限的情况下。这些特殊的初始化在各种情况下都证明是有益的,包括在复杂值模型的情况下。同样,Mamba主要关注隐藏状态矩阵A的初始化,以捕获复杂的时间依赖性。这是通过利用HiPPO理论和创新的缩放Legendre度量(LegS)来实现的,确保仔细考虑完整的历史上下文而不是有限的滑动窗口。具体来说,HiPPO-LegS为所有历史数据点分配均匀的权重,可以表示为:

\[A_{nk}^{\mathrm{HiPPO}} = -\begin{cases} (2n+1)^{\frac12}(2k+1)^{\frac12}, & n > k, \\ n+1, & n = k, \\ 0, & n < k. \end{cases}\]

其中n是多项式的数目, k表示特定的离散时间步骤。基于HiPPO理论,Mamba引入了两种简单的初始化方法,即S4D-Lin和S4D-Real,如下所示:

\[A_{dn} = -\begin{cases} \frac12 - ni, & \text{S4D-Lin}, \\ n+1, & \text{S4D-Real}. \end{cases}\]

其中n是A的第n个元素,对于所有输入维度 = 1, 2, ..., D。通过这种初始化,模型可以学习长期依赖性,通过压缩和重建输入信息信号,经历较小的新步骤退化和较大的旧步骤退化。根据公式,HiPPO-LegS具有有利的理论属性:它在输入时间尺度上保持一致,并且提供快速计算。此外,它具有有界的梯度和近似误差,有助于参数学习过程。

3.3 选择机制。

传统的状态空间模型由于时间不变性属性,无法根据特定模型输入(即内容感知建模能力)产生个性化输出。为了为SSM提供类似于注意力机制的能力,Mamba设计了一个时变选择机制,根据模型输入参数化权重矩阵。这种创新使SSM能够无限期地过滤掉无关信息,同时保留相关细节。正式地说,选择机制涉及设置间隔Δ,并将矩阵B、C作为输x ∈ R × ×D的函数,可以公式化为:

\[\begin{aligned} \mathbf B &\to \mathbf S^B = \mathbf W^B\mathbf x, \\ \mathbf C &\to \mathbf S^C = \mathbf W^C\mathbf x, \\ \Delta &\to \mathbf S^\Delta = \tau_\Delta \cdot \operatorname{BroadCast}_D(\mathbf W^\Delta\mathbf x). \end{aligned}\]

其中SB ∈ R × × , SC ∈ R × × ,和SΔ ∈ R × × 是选择性空间矩阵,它们是输入的函数,以实现内容感知建模。 , , , 和 N分别代表批量大小、输入长度、输入特征大小和隐藏通道数。值得注意的是,WB ∈ R × , WC ∈ R × , 和 WΔ ∈ R ×1是相应组成部分的选择权重(即线性参数化投影),而BroadCastD意味着将结果广播到所有维度 = 1, 2, .., 。

随后,选择性SSMs通过使用常见的统计技术,零阶保持(ZOH),进行离散化,如下所示:

\[\begin{aligned} \bar{\mathbf A} &\to \mathbf S^{\bar A} = \exp(\mathbf S^\Delta \mathbf A), \\ \bar{\mathbf B} &\to \mathbf S^{\bar B} = (\mathbf S^\Delta \mathbf A)^{-1}\left(\exp(\mathbf S^\Delta \mathbf A)-\mathbf I\right)\cdot \mathbf S^\Delta\mathbf S^B. \end{aligned}\]

其中SA ∈ R × × × 和 SB ∈ R × × × 分别是选择性状态转移矩阵和输入矩阵,它们成为输入x的函数,即

\[\mathbf y = \operatorname{SSM}(\mathbf A,\mathbf B,\mathbf C)(\mathbf x),\]

它根据输入x生成输出ut y ∈ R × × 。请注意,Mamba中的时变选择机制具有与Transformer中的注意力机制类似的结构,即两者都根据输入及其投影执行操作,这允许Mamba的SSM实现灵活的内容感知建模。然而,它失去了与卷积的等价性,这对其效率产生了负面影响。

3.4 硬件感知计算。

选择机制旨在超越线性时间不变模型的限制。尽管如此,它对有效训练提出了挑战:SSMs的卷积核变得依赖于输入,导致无法进行并行计算。为了解决这个问题,Mamba利用了两种计算技术,即并行关联扫描(也称为并行前缀和)和内存重计算。首先,并行关联扫描利用线性关联计算的属性和现代加速器(GPU和TPU)的并行性,以内存高效的方式执行选择性SSMs的计算。

更具体地说,平行关联扫描围绕构建给定输入的平衡二叉树,并扫掠它从叶子到根。换句话说,平行关联扫描首先通过从叶子到根的遍历(即Sweep-Up),在树的内部节点创建部分和。然后,它逆转遍历,从根开始向上扫掠树,使用部分和构建整个扫描(即Sweep-Down)。

另一方面,Mamba利用了传统的重计算方法,以减少训练选择性SSM层时的整体内存需求。具体来说,Mamba在前向传递的并行关联扫描中避免存储大小为 ( , , , ) 的中间状态,以防止内存扩展。相反,它在后向传递中重新计算这些中间状态以进行梯度计算。通过这样做,重计算避免了在GPU内存单元之间读取 ( ) 元素的必要性。除了优化扫描操作的内存需求外,Mamba-1还将重计算的使用扩展到整个SSM层的效率,包括投影、卷积和激活,这些通常需要大量的内存资源,但可以快速重新计算。

4 Mamba 2

4.1 Mamba-2:状态空间对偶性

Transformers在深度学习的成功中发挥了关键作用,为各种领域带来了革命性的变化。它们启发了各种技术的发展,如参数高效微调、灾难性遗忘缓解和模型量化,旨在从不同角度提高模型性能。为了使状态空间模型能够访问并从最初为Transformer开发的宝贵技术中受益,Mamba-2引入了一个全面的框架,称为结构化状态空间对偶性(SSD),它建立了SSMs和不同形式的注意力之间的理论联系。正式地,

\[\mathbf y = \operatorname{SSD}(\mathbf A,\mathbf B,\mathbf C)(\mathbf x)=\mathbf M\mathbf x,\]

其中M表示SSMs的矩阵形式,使用序列半可分表示,M = CT A : Bi

值得注意的是,C 和 B 分别代表与输入令牌x 和xi相关的选择性空间状态矩阵。A : 表示与输入令牌从j到i的范围相关的选择性隐藏状态矩阵。本质上,SSD表明,Transformer使用的注意力机制和SSM中使用的线性时变系统都可以被视为半可分矩阵变换。此外,Dao和Gu[28]还证明了选择性SSM等同于使用半可分掩码矩阵实现的结构化线性注意力机制。

基于SSD,Mamba-2设计了一种更高效的硬件计算,通过块分解矩阵乘法算法。具体来说,通过将状态空间模型视为半可分矩阵,Mamba-2将计算分解为矩阵块,在这些块中,对角块代表内部块计算。相比之下,非对角块通过SSM的隐藏状态表示内部块计算。这种方法使Mamba-2能够实现比Mamba-1的并行关联扫描快2-8倍的训练过程,同时保持与Transformer的竞争性。

4.2 SSD(State Space Duality)框架

  • 展示了状态空间模型与一类称为半可分矩阵的结构化矩阵族之间的等价性,揭示了状态空间模型的新属性和算法。
  • 显著改进了线性注意力理论,通过张量收缩的语言对其循环形式提供了明确证明,并将其推广到新的结构化掩码注意力(SMA)家族。同时证明了任何具有快速循环形式的核注意方法都是 SSM,且 SSM 和 SMA 有很大的交集,彼此是对偶的,同时具有 SSM 式的线性形式和类似注意力的二次方形式。
  • 提出了一种基于半可分离矩阵块分解的 SSD 算法,该算法利用了 SSM 线性递推和二次对偶形式,在所有主要效率轴上获得了最优的权衡。基于 SSD 的实现比 Mamba 的优化选择性扫描实现快 2 到 8 倍,同时允许使用更大的循环状态大小(是 Mamba 的 8 倍甚至更高,且几乎不影响速度),与优化过的 softmax 注意力实现(FlashAttention-2)具有高度竞争力,在序列长度 2k 时性能相当,在序列长度 16k 时速度快 6 倍。

4.3 Mamba 2 的架构设计

在架构设计方面,Mamba 2 对 Mamba 块进行了修改,引入分组值注意力(GVA,Grouped-Value Attention)头结构以实现张量并行,并将修改后的并行 Mamba 块与作为内部 SSM 层的 SSD 结合,构建了 Mamba-2 架构。该架构的主要改进是从顺序生成转变为并行生成 SSM 参数,更适用于张量并行等扩展方法。采用 Mamba-2 架构的模型在困惑度和实际运行时间上优于 Mamba 和 Transformer++,并在标准下游评估中匹配或超越了 Mamba 和开源 transformers 的表现。

image

Mamba2 有许多创新点,以下是一些主要的方面:

Mamba 2是 Mamba 的进一步发展,它提出了 SSD(State Space Duality)框架,并基于此设计了新的体系架构。

与 Mamba 相比,Mamba 2的核心层是对选择性 SSM 的改进,速度提高了2-8倍,同时在语言建模方面继续与 transformers 竞争。新算法使其能够利用更大的状态维度(从16提升至256),在需要更大状态容量的任务上有显著改进。

研究发现注意力和 SSM 是互补的,将4-6个注意力层与 Mamba-2层混合,其表现优于 Transformer++和纯 Mamba-2。

该研究还展示了状态空间模型与一类称为半可分矩阵的结构化矩阵族之间的等价性,揭示了状态空间模型的新属性和算法;显著改进了线性注意力理论,通过张量收缩的语言对其循环形式提供了明确证明,并推广到新的结构化掩码注意力(SMA)家族;证明了任何具有快速循环形式的核注意方法都是 SSM。

在算法层面,提出了基于半可分离矩阵块分解的 SSD 算法,利用 SSM 线性递推和二次对偶形式,在所有主要效率轴上获得了最优权衡。基于 SSD 的实现比 Mamba 的优化选择性扫描实现快2到8倍,同时允许使用更大的循环状态大小(是 Mamba 的8倍甚至更高,且几乎不影响速度)。

这些创新使得 Mamba 在语言、音频和基因组学等多个领域表现出色,在处理长序列时具有更高的效率和性能。但 Mamba 相关算法仍在不断发展和改进中,新的研究可能会带来更多的创新和优化。

4.4 Mamba-1和Mamba-2的对比

在这一小节中,我们总结了Mamba-1和Mamba-2的块设计。图4展示了这两种架构的比较。Mamba-1是从SSM中心视角出发,其中选择性SSM层负责将输入序列X映射到Y。在这个设计中,(A, B, C)的线性投影应用于初始线性投影之后,以创建X。然后,输入令牌和状态矩阵通过选择性SSM单元传递,使用并行关联扫描,以产生输出Y。之后,Mamba-1使用跳过连接来鼓励特征重用,并缓解在模型训练过程中经常发生的退化问题。最后,通过堆叠这种块与标准归一化和残差连接交错,构建了Mamba模型。

image

至于Mamba-2,它引入了SSD层,旨在从[X, A, B, C]到Y创建一个映射。这是通过同时处理[X, A, B, C]并以与标准注意力架构生成Q、K、V投影相同的方式在块的开始进行单个投影来实现的。换句话说,Mamba-2块通过删除顺序线性投影来简化Mamba-1块,这使得SSD结构的计算速度比Mamba-1中的并行选择性扫描快。此外,在跳过连接后添加了一个归一化层,旨在提高训练稳定性。

image

表2. 利用Mamba架构的代表性开放获取基础模型。

5 Mamba模型的进步

状态空间模型和Mamba最近被探索,并成为基础模型骨干的一个有前途的替代方案。如表2所示,大规模Mamba基础模型不仅在学术研究中蓬勃发展,而且在工业界也取得了显著进展,如Falcon Mamba 7B和Mistral 7B,通过在GPU上成功训练证明了它们的有效性。尽管如此,Mamba架构仍然面临挑战,如内存丢失、对多样化任务的泛化能力不足以及与基于Transformer的语言模型相比在捕获复杂模式方面的劣势。为了克服这些挑战,已经做出了大量努力来改进Mamba架构。现有的研究主要集中在修改块设计、扫描模式和内存管理方面。这一节将介绍这三个方面的几个重要技术,并在表3中提供相关研究的总结。

image

5.1 块设计

Mamba块的设计和结构对Mamba模型的整体性能有重要影响,使其成为一个新的研究焦点。如图5所示,基于构建新Mamba块的不同方法,现有研究可以分为三类:a)集成方法旨在将Mamba块与其他知名模型集成,以便在保持有效性和效率的同时实现平衡;b)替代方法尝试将Mamba块作为高级模型框架中主要层的替代品;c)修改方法专注于修改经典Mamba块内的组件。以下小节将详细介绍这些方法。

image

4.1.1 集成。鉴于Mamba在捕获长期动态方面的卓越能力,它已广泛与其他模型集成,利用其优势为特定场景提供强大的框架。集成特别包括先进的模型,如Transformers、图神经网络(GNNs)、循环神经网络(RNNs)、卷积神经网络(CNNs)和尖峰神经网络(SNNs)。下面描述了一些具体的例子。

• 基于Transformer的模型在许多任务中表现出色,但它们的二次方计算复杂度在推理过程中仍然阻碍了它们。为了实现高效的生成,一些研究人员提出了将Mamba块与基于Transformer的模型结合起来。

例如,Jamba结合了Transformer和Mamba层的块,以应对长期内容的自然语言处理任务,利用了两种模型家族的优势。注意力-Mamba混合模型的性能优于单独的Transformer和Mamba模型,与普通Transformer模型相比,吞吐量更好。Mambaformer利用混合框架预测多个时间序列,包括汇率、小时电力消耗和电力负荷,它在内部结合了Mamba块和Transformer层,分别用于长期和短期依赖性。由于Mamba和Transformer的集成,Mambaformer在长期-短期时间序列预测方面优于基于Transformer的预测器。

• GNN通过消息传递机制在捕获邻近关系方面展现出了有希望的潜力,其中信息通过连接图在堆叠层中传播。然而,这些模型面临一个重大限制,即过度平滑,特别是当试图捕获高阶邻接信号时。为了解决这一挑战,Mamba已被用于图表示学习。例如,Graph Mamba将图结构数据重新格式化为特定顺序的序列令牌,并利用Mamba块中的选择性SSM层构建了一个新的Graph Mamba Network (GMN)架构,它在图表示学习能力方面表现出色,特别是在需要高阶节点依赖性的数据处理方面。

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

• CNN基础方法受到局部接受域的限制,导致在捕获全局和长期语义方面表现不佳。众所周知,状态空间模型在学习长期模式方面具有优越能力,可以利用Mamba块增强CNN基础模型的潜力,特别是在计算机视觉领域。例如,MedMamba和nnMamba展示了如何通过集成视觉Mamba块来提高CNN在图像分析任务中的性能。

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

4.1.2 替代。受到选择性SSM在高效计算和长序列学习方面的卓越能力的启发,采用Mamba模块替代经典建模框架中的关键组件,如U-Net和扩散模型,已经引起了很多关注。通过引入选择性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在残差学习框架中采用了视觉Mamba块,用于食品类别分类。此外,SPMamba采用了TF-GridNet,这是最近开发的一种时频模型,作为其基础架构,然后在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块仍然是一个有前景但尚未充分探索的领域。

5.2 扫描模式

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

image

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

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

• 扫描扫描。如图6所示,扫描扫描技术按特定方向处理模型输入,类似于清洁工仔细清扫地板[189, 216]。例如,Cross-Scan[121]涉及将输入图像划分为补丁,然后将其沿四个不同路径展平,这被视为两个双向扫描的融合。通过采用这些互补的遍历路径,Cross Scan使每个图像补丁能够高效地整合来自不同方向的邻居信息,从而促进了信息丰富、接受域的建立。Omni-Scan[163, 229]结合了从多个方向对图像信息流的建模,例如2(向前和向后)×4(左-右、上-下、右上-左下、左上-右下)。这种策略增强了在各个方向上对上下文信息的全局建模能力,使提取全面的全局空间特征成为可能。

• 连续扫描。为了确保输入序列的连续性,连续扫描技术扫描相邻的列或行之间的令牌[66],如图6所示。例如,为了更好地处理2D空间输入,PlainMamba[203]引入了一种连续扫描方法,称为连续扫描,它扫描列(或行)之间的相邻令牌,而不是在Cross Scan中前往相反的令牌。此外,Hilbert Scan[66]基于Hilbert矩阵沿着曲折的路径行进。根据他们的评估结果,可以推断出,增强输入令牌的语义连续性在各种视觉识别任务中为基于Mamba的模型带来了优越的性能。

• 高效扫描。与上述专注于实现更全面的输入建模的扫描方法不同,高效扫描方法旨在加速训练和推理过程。通常,高效扫描将给定输入分成几个部分,并并行处理它们,从而减少计算需求。例如,Efficient-2D Scan[139]通过跳过补丁来处理图像,从而减少了四倍的计算需求,同时保留了全局特征图。此外,Gao等人[50]在他们的Mamba框架中引入了一种有效的双向子空间扫描方案。该方案旨在为4D光场超分辨率任务高效捕获长期空间-角度对应关系。具体来说,它将补丁序列分解为两部分,并通过两个双向扫描方案进行处理。通过这样做,扫描方法降低了输入长度,并解决了长期记忆问题,而不会牺牲完整的4D全局信息。

4.2.2 立体扫描。通过从额外的角度对模型输入进行建模,立体扫描方法在扫描过程中捕获更广泛知识的能力方面优于扁平扫描方法。这种增强的能力允许更全面地理解模型输入。具体来说,这些方法可以分为三个主要类别:层次扫描、时空扫描和混合扫描。层次扫描从不同层次处理输入,而时空扫描从时间和空间的角度考虑输入模式。此外,混合扫描结合了多种扫描方法,以利用不同扫描技术的益处。

• 层次扫描方法涉及使用不同内核大小的扫描来从全局到局部或从宏观到微观捕获语义知识[26, 63, 162, 181]。例如,Mamba-in-Mamba层次编码器[24]用于红外小目标检测,结合了内部和外部选择性SSM块。内部的一个专门用于捕获视觉补丁之间的相互作用以提取局部模式。相反,外部块旨在描述视觉句子之间的关系以捕获全局特征。HiSS[12]将输入序列划分为块,并连续地对块特征进行层次建模以进行连续序列预测。块首先由低级SSM单元处理,然后由高级SSM块将处理后的特征映射到输出序列。

• 时空扫描。由于现实世界中动态系统的普遍性,对时空扫描方法的兴趣日益增加,以提高Mamba块的性能[207, 209]。例如,VideoMamba[102]将原始的2D扫描扩展为两个3D扫描:空间优先扫描和时间优先扫描。结合这两种扫描方法,VideoMamba在处理长、高分辨率视频方面表现出色。此外,ChangeMamba[21]集成了三种时空扫描机制(顺序建模、交叉建模和并行建模),以实现多时态特征之间的上下文信息交互,用于遥感变化检测。

• 混合扫描。为了全面地建模特征,许多努力专注于结合不同扫描方法的优势[29, 32, 53, 163, 231],即混合扫描。例如,Mambamixer[10]展示了Scan of Scan,它动态地使用一组图像扫描方法,即Cross-Scan、Zigzag Scan和Local Scan,来遍历图像补丁。Mambamixer还引入了双选择机制,以在令牌和通道之间混合信息。通过这样做,他们展示了与其他视觉模型相媲美的性能。

Pan-Mamba[68]引入了两种基于Mamba架构的扫描方法:通道交换扫描和跨模态扫描。通过结合这两种扫描方法,Pan-Mamba增强了其在图像锐化中的高效跨模态信息交换和融合能力。

5.3 内存管理

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

6 Mamba 的应用

Mamba 架构作为选择性状态空间模型(SSM)的扩展,具备处理序列和非序列数据的能力。本章探讨 Mamba 在不同任务中的应用,包括序列数据、非序列数据以及多模态数据,并总结其在实际场景中的显著应用。


6.1 序列数据

序列数据是指以特定顺序组织的数据,Mamba 凭借其高效建模长期依赖性的能力,在多种序列任务中表现出色。

6.1.1 自然语言

Mamba 在自然语言处理(NLP)领域展现出强大的潜力,成为大型语言模型的有力替代品。例如:

  • MambaByte:利用 Mamba 捕获长期依赖性,避免子词标记化的归纳偏差,在长期语言建模任务中超越了最先进的 Transformer。
  • Jamba 和 BlackMamba:结合专家混合(MoE)技术,通过线性复杂度生成与 MoE 的快速推理能力,提升了语言处理性能。

6.1.2 视频

视频理解和生成的核心在于学习时空表示,Mamba 能有效区分短期动作和解释长视频内容。例如:

  • VideoMamba:通过 3D 卷积将输入视频投影到时空补丁中,使用双向 Mamba 块编码向量表示,适用于视频理解和生成任务。
  • Vivim:引入时态 Mamba 块,压缩时空表示,用于医学视频分割。

6.1.3 时间序列

Mamba 在时间序列分析中表现出卓越的效率,适用于股市分析、交通建模和天气预报等任务。例如:

  • TimeMachine:利用 Mamba 捕获多变量时间序列中的持久模式,确保线性复杂度计算和最小内存占用。
  • Mambaformer:结合选择性 SSM 和注意力层,用于长期和短期预测。

6.1.4 语音

Mamba 在语音分离和增强任务中展现了显著优势,降低了计算复杂度。例如:

  • SPMamba 和 DPMamba:利用双向 Mamba 模块捕获更广泛的上下文信息,提升语音分离性能。
  • TRAMBA:结合 Transformer 和 Mamba,显著降低内存消耗,优化移动平台上的语音质量。

6.1.5 运动

Mamba 在人体运动建模中表现优异,能够捕获复杂的时空模式。例如:

  • Motion Mamba:通过层次化 SSM 层捕获时间模式,确保运动一致性。
  • MambaMOS:设计专门的运动感知模型,实现高质量、长序列运动生成。

6.2 非序列数据

非序列数据不遵循特定顺序,Mamba 在图像、图结构数据和点云等领域展现出了出色的成功。

6.2.1 图像

Mamba 在计算机视觉任务中表现出色,显著降低了计算时间和内存占用。例如:

  • Vision Mamba:结合双向 SSM 和位置嵌入,实现全局视觉语义建模,比 Vision Transformers 更高效。
  • VMamba:引入 2D Selective Scan (SS2D),使 Mamba 能有效处理视觉数据。

6.2.2 图结构数据

Mamba 在图建模中展示了强大的长期建模能力和高效率。例如:

  • Graph-Mamba:结合图展平机制和 Mamba 提供的选择机制,促进输入依赖的上下文过滤。
  • GMN:基于选择性 SSMs 的图神经网络格式,超越 GNNs 和基于 Transformer 的模型。

6.2.3 点云

Mamba 在点云分析中显著降低了计算成本,适用于机器人技术和自动驾驶。例如:

  • PointMamba:采用分层扫描策略,利用普通 Mamba 提取序列化点令牌特征。
  • Point Cloud Mamba:显著减少内存使用量,性能优于基于 Transformer 的同类产品。

6.3 多模态数据

Mamba 在多模态任务中整合语言和图像等数据,提供了有价值且互补的信息。例如:

  • VL-Mamba:探索视觉-语言任务的潜力,通过连接器模块对齐视觉补丁与语言令牌。
  • 文本控制的运动 Mamba:根据文本查询动态捕获全局时间信息,增强人体运动理解。
  • Fusion-Mamba 和 Sigma:融合来自不同模态的互补信息,分别改进目标检测和语义分割。

第7章 应用

Mamba 在多个领域中展现出广泛的应用价值,包括自然语言处理、计算机视觉、语音分析、药物发现、推荐系统以及机器人技术和自主系统。

7.1 自然语言处理

Mamba 在问答系统和文本摘要任务中表现出色,解决了传统模型的计算效率问题。

7.1.1 问答系统

  • Mamba-Chat:通过状态空间表示维护对话理解,确保上下文感知。
  • Jamba 和 DenseMamba:结合 MoE 和隐藏状态整合技术,提升问答任务性能。

7.1.2 文本摘要

  • LOCOST:基于状态空间模型,处理长文档摘要,显著降低内存和推理开销。
  • SAMBA:结合滑动窗口注意力和选择性序列压缩,提升吞吐量和性能。

7.2 计算机视觉

Mamba 在疾病诊断和运动识别与生成任务中展现了优越性能。

7.2.1 疾病诊断

  • U-Mamba 和 SegMamba:结合 CNN 和 SSM,提升医学图像分割性能。
  • CMViM 和 ProMamba:通过跨模态对比学习和提示技术,增强疾病诊断和息肉分割能力。

7.2.2 运动识别与生成

  • HARMamba 和 Simba:利用双向 SSM 架构,提升实时人体运动识别和骨架动作识别性能。
  • Motion Mamba 和 InfiniMotion:生成连续、长时间的人体运动,降低计算资源需求。

7.3 语音分析

Mamba 在语音分离、标记和增强任务中表现出显著优势。

7.3.1 语音分离和标记

  • DPMamba 和 SPMamba:捕获动态时间依赖性,提升语音分离性能。
  • DASS 和 MAMCA:结合知识蒸馏和选择性 SSM,优化音频标记和调制分类任务。

7.3.2 语音增强

  • TRAMBA 和 oSpatialNet-Mamba:通过混合架构和长期多通道增强,显著降低内存消耗并提升语音质量。

7.4 药物发现

Mamba 在蛋白质设计、分子设计和基因组分析中降低了复杂性,提升了效率。

  • PTM-Mamba 和 ProtMamba:高效处理长序列,为蛋白质设计提供关键工具。
  • Saturn 和 Caduceus:优化分子设计和基因组建模,提升预测准确性和训练上下文长度。

7.5 推荐系统

Mamba 在个性化推荐任务中展现了高效性和扩展潜力。

  • Mamba4Rec 和 RecMamba:通过选择性 SSM 提升模型性能,显著降低训练时间和内存成本。
  • EchoMamba4Rec 和 Mamba4KT:捕获用户交互复杂模式,提升推荐精度和知识追踪研究效率。

7.6 机器人技术和自主系统

Mamba 在机器人技术和自主系统中展现了优越性能和扩展潜力。

  • RoboMamba 和 MaIL:通过端到端训练和模仿学习,提升视觉常识和机器人推理能力,同时确保高效微调和推理。

挑战与机遇

前面的部分全面回顾了Mamba的最新先进技术和多样化应用。然而,Mamba的研究仍处于起步阶段,存在相当大的挑战和机遇。

基于Mamba的基础模型

通过将模型规模扩大到数十亿级别,并在大规模混合源语料库上进行训练,基础模型(FMs)展现出了令人印象深刻的零样本学习能力,使FMs能够在广泛的通用任务中表现出色[13]。作为代表性的例子,最近几年见证了基于Transformer的大型语言模型的蓬勃发展,特别是ChatGPT,激发了对各个领域基础模型的探索热情。尽管Transformer是成功的主要驱动力,但它们面临着紧迫的计算和内存效率问题[172],这随着基于注意力模型的规模呈指数级增长的训练内存和劳动密集型的自回归解码而增长。作为Transformer的有前途的替代品,Mamba[28, 55]最近出现了。Mamba提供了Transformer的内容感知学习能力,同时将计算与输入长度线性扩展,使其在捕获长期依赖性和提高训练和推理效率方面有效。鉴于这些优势,开发特定领域的基于Mamba的基础模型具有巨大潜力,这为解决基于Transformer的模型所面临的问题提供了机会。

硬件感知计算

基础模型以其庞大的规模和密集的矩阵运算(如矩阵乘法和卷积)而闻名,需要使用GPU和TPU等尖端硬件进行高吞吐量的训练和推理。这些先进的硬件使研究人员能够使用更大的数据集并在各个领域取得最先进的性能。然而,现有的基础模型还没有充分利用硬件的计算能力,导致模型效率有限[172]。作为提高计算效率的有前途的替代品,Mamba-1[55]和Mamba-2[28]提出了硬件感知计算算法,即并行关联扫描和块分解矩阵乘法。这些算法考虑了GPU和TPU的固有特性,包括设备间消息传输等因素,为解决计算效率问题提供了新的视角。受此启发,探索新的硬件高效算法,如FlashButterfly[47],以优化硬件利用,为SSMs以及其他架构如Transformers和RNNs提供了节省资源和加速计算的有前途的途径。

值得信赖的Mamba模型

SSM的发展预计将为电子商务、医疗保健和教育等多个行业带来显著的好处。与此同时,像许多现有的架构一样,Mamba模型可能是数据依赖的,可能对用户和社会构成严重威胁[130]。这些威胁源于几个因素,如不稳定的决策制定、隐私问题等。因此,确保Mamba模型的可信任性在四个关键维度上至关重要[116]:安全性&鲁棒性、公平性、可解释性和隐私。

安全性&鲁棒性。大型基础模型已被证明对输入的小扰动高度敏感,这可能危及这些模型在安全关键应用中的安全性和鲁棒性[44, 135, 191]。同时,基于Mamba的模型也不免于这些漏洞[128]。在追求成为Transformer的可靠替代品的过程中,有必要研究和增强基于Mamba模型的安全性和鲁棒性。具体来说,模型输出应对输入的小扰动保持鲁棒。一个可能的解决方案可能涉及在将提示输入到基于Mamba的模型之前自动预处理它们。此外,作为一种代表性技术,对抗性机器训练[78]可以用来增强基于Mamba模型的安全性和鲁棒性。

公平性。大型基础模型在广泛的数据集上进行训练,往往会无意中暴露出训练语料库中存在的偏见和刻板印象[126],这可能在生成的输出中表现出来。例如,在LLMs领域,偏见可能导致受到用户档案属性如性别和年龄影响的歧视性回应,加强刻板印象并不公平地对待特定用户群体[86]。虽然最近已经做出了努力来解决LLMs中的公平性问题,但在Mamba模型的非歧视性和公平性方面的研究仍然存在差距。因此,需要进一步的探索和研究来弥合这一差距。

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

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

将Transformer的新兴技术应用于Mamba

Transformer作为主导的骨干,已经引领AI社区开发了许多旨在提高基于注意力模型性能的独特工具。幸运的是,通过SSMs和注意力的连接,Mamba-2[28]引入的SSD框架允许我们为Transformer和Mamba开发共享的技术和库。鉴于此,一个重要的未来方向出现了,即探索如何将为基于Transformer的模型设计的新兴技术有效地应用于基于Mamba的模型。

参数高效微调。大型基础模型将其参数规模扩大到数十亿,已经在多个领域取得了突破性进展。然而,它们的庞大规模和计算需求在为特定下游任务定制时提出了重大挑战。为此,已经提出了几种参数高效微调(PEFT)技术,包括LoRA[72]和Adapter家族[49, 91],这些技术涉及在微调期间最小化参数调整或对计算资源的需求。从最近在利用PEFT为基于Transformer层构建的大型语言模型方面取得的成就中汲取灵感,将PEFT应用于Mamba模型已经出现了一个有趣的话题,目标是扩大Mamba在下游任务中的应用范围。例如,LoRA(低秩适应)的部署预计将促进SSD模型的快速微调,从而实现Mamba在各个领域的广泛应用。然而,为Mamba基础模型实施这些PEFT技术的具体细节尚未确定,需要进一步的调查。

缓解灾难性遗忘。灾难性遗忘,也称为灾难性干扰,是指在机器学习模型中观察到的现象,当它们在新任务上进行训练时,会在以前学习的任务上经历显著的性能下降[92]。这个问题对于基础模型来说是一个挑战,因为它们需要保留预训练任务的知识,并在不同的下游领域展示一致的性能。作为基础模型的一个有前途的架构,Mamba需要进行彻底的调查以解决灾难性遗忘问题。最近的研究表明,通过奖励最大化和分布匹配策略[96, 97]来解决这一挑战。此外,还开发了持续学习方法来缓解基于Transformer的语言模型中的灾难性遗忘[90, 187]。这些技术也可以通过连接SSMs和注意力应用于Mamba模型,但尚未探索。

检索增强生成(RAG)。作为AI中最复杂的技术之一,RAG可以提供可靠和最新的外部知识,为多种任务提供重要的实用性[30, 99]。大型语言模型最近展示了突破性的语言理解和生成能力,尽管它们在内部知识过时和幻觉方面的固有局限性。鉴于RAG在提供当前和有价值的补充信息方面的强有力能力,检索增强的LLM已经出现,利用外部知识数据库来增强LLMs的生成质量[22]。同样,RAG可以与Mamba语言模型集成,协助它们产生高质量的输出,这是一个有前途的未来研究方向。

小结

Mamba作为一种新兴的深度学习架构,在多个领域,如语言生成、图像分类、推荐和药物发现等方面展现出了显著的成功,这得益于其强大的建模能力和计算效率。最近,越来越多的努力被投入到开发具有更强大表示学习能力和更低计算复杂度的基于Mamba的深度学习模型中。鉴于Mamba的快速发展,迫切需要一个系统的概述。

Mamba 架构凭借其高效的长期依赖性建模能力和灵活的扩展性,在序列、非序列和多模态数据任务中展现了广泛应用前景。无论是自然语言处理、计算机视觉还是药物发现等领域,Mamba 都为现有方法提供了有力补充和替代方案。

mamba的局限性

Mamba模型当前存在不少局限性。

一,长上下文泛化差,超出训练长度后易出现状态崩溃,方差爆炸致性能骤降,且无法有效遗忘旧信息,新信息难存入,密钥检索等任务精度大幅下滑 。

二,多模态适配受阻,单向序列建模难学好图像、视频的2D空间依赖关系,用于扩散模型时还受因果限制,适配去噪等任务需额外优化 。

三,密集预测能力弱,视觉领域应用中易丢失高层任务特征,且生成的单尺度低分辨率特征,难以满足语义分割等任务对细粒度细节的需求 。

四,训练与架构有短板,输入依赖的卷积核让并行计算困难,层间隐藏信息传输和无损压缩也存在挑战,部分优化方案还会牺牲模型运行速度 。

参考资料:

手把手教你认识学会Mamba-(祝各位蛇来运转)25 赞同 · 0 评论 文章

一文读懂:Mamba 模型,transformer的挑战者189 赞同 · 4 评论 文章

一文读懂: Transformer(详细但无代码)528 赞同 · 32 评论 文章

DeepSeek、Mamba与RWKV:新一代高效AI模型的革新之路6 赞同 · 1 评论 文章

用心记录,持续成长