Mamba最新综述《A Survey of Mamba》-2024.8.5

1、摘要

深度学习作为一项重要技术，引发了人工智能（AI）领域的一场重大革命，导致人类生活方式发生了巨大变化。作为最具代表性的架构，Transformers为众多高级模型提供了支持，尤其是包含数十亿参数的大型语言模型（LLM），成为深度学习的基石。尽管取得了令人印象深刻的成就，变形金刚仍然面临着固有的局限性，特别是由于注意力计算的二次计算复杂性而导致的耗时推理。最近，一种名为Mamba的新架构，从经典的状态空间模型中汲取灵感，已经成为构建基础模型的一种有前途的替代方案，提供与Transformer相当的建模能力，同时保留关于序列长度的近线性可扩展性。这引发了越来越多的研究积极探索Mamba在不同领域取得令人印象深刻的性能的潜力。鉴于这种快速的演变，有一个系统的审查，巩固现有的Mamba授权的模型，提供了一个全面的了解这个新兴的模型架构的迫切需要。因此，在本次调查中，我们对最近与Mamba相关的研究进行了深入调查，涵盖了三个主要方面：基于Mamba的模型的进展，Mamba适应不同数据的技术，以及Mamba可以脱颖而出的应用。具体来说，我们首先回顾了各种代表性深度学习模型的基础知识以及Mamba-1&2的细节。然后，为了展示Mamba对人工智能的意义，我们全面回顾了Mamba模型的架构设计，数据适应性和应用程序的相关研究。最后，我们提出了目前的局限性进行了讨论，并探讨各种有前途的研究方向，为未来的调查提供更深入的见解。

Additional Key Words and Phrases: State Space Model, Mamba, Sequence Modeling, Foundation Models, Language Models

2、引用

在过去的二十年里，深度学习作为最突出的人工智能（AI）技术，已经在各个领域带来了革命，如医疗保健[88]，自治系统[37，60]，推荐系统[104，228]和金融服务[144，216]。在这一时期，出现了许多深度神经网络（DNN），这些网络显著改变了人类的生活方式，为个人提供了巨大的便利。一个值得注意的例子是U-Net [151，164]，这是一种在视觉领域内强大的深度学习模型，广泛用于医学成像，用于检查MRI和CT扫描等放射学扫描。其应用有助于疾病的识别和诊断，展示了其在这一关键医疗领域的有效性[112，191]。此外，图神经网络（GNN）用于处理图结构数据以支持智能服务，例如向用户推荐个性化内容，产品或服务的推荐系统[41，42，194]。此外，递归神经网络（RNN）被广泛用于机器翻译，因为它们能够捕获准确翻译所必需的顺序和上下文信息[119，167]，使来自不同语言背景的个人能够有效地沟通和理解彼此的想法，意见和信息。

在各种深度学习架构中，Transformers最近脱颖而出，并在广泛的应用中确立了其主导地位[34，175]。例如，作为最具代表性的大型基础模型，像ChatGPT这样的大型语言模型（LLM）从根本上构建在Transformer架构上[2，146，228]。通过将其模型大小扩展到数十亿并在不同数据源的混合上进行训练，这些基于Transformer的模型已经展示了人类水平的智能，其在语言理解，常识推理和内容学习方面的能力令人印象深刻[44，217]。这种显著的成功得到了注意力机制的支持[174]，它使基于transformer的模型能够专注于输入序列的相关部分，并促进更好的上下文理解。然而，注意力机制也引入了显著的计算开销，其随着输入大小的平方增加[124，233]，这在处理冗长的输入时提出了挑战。例如，计算成本的快速增长使得Transformers处理大量序列不切实际或不可行，从而限制了它们在文档级机器翻译[131]或长文档摘要[94]等任务中的适用性。

最近，出现了一种很有前途的架构，结构化状态空间序列模型（SSM）[58]，可以有效地捕获序列数据中的复杂依赖关系，成为Transformer的强大竞争对手。这些模型受到经典状态空间模型的启发[89]，可以被认为是递归神经网络和卷积神经网络的融合。它们可以使用递归或卷积运算来有效地计算，实现与序列长度的线性或近线性缩放，从而显著降低计算成本。更具体地说，作为最成功的SSM变体之一，Mamba实现了与Transformers相当的建模能力，同时保持了序列长度的线性可扩展性[55]，将其推向了焦点话题的领域。Mamba首先引入了一种简单而有效的选择机制，该机制使模型能够过滤掉不相关的信息，同时通过根据输入参数化SSM参数来无限期地保留必要和相关的数据。然后，Mamba提出了一种硬件感知算法，通过扫描而不是卷积来循环计算模型，在A100 GPU上实现了高达3倍的计算速度。如图1所示，复杂和冗长的序列数据的强大建模能力，沿着接近线性的可扩展性，将Mamba定位为新兴的基础模型，有望彻底改变各种研究和应用领域，如计算机视觉[199，234]，自然语言处理[111，226]，医疗保健[152，179，198]等。Zhu等人。[234]提出了Vim，比DeiT [173]快2.8倍，在提取高分辨率图像的特征时节省86.8%的GPU内存。Dao和Gu [29]展示了SSM和注意力变体之间的联系，并提出了一种新的架构，可以改进选择性SSM，实现2-8倍的语言建模速度。

3、文章安排

第2节，我们回顾了各种代表性深度神经网络的背景知识，包括RNN，Transformer和状态空间模型。

第3节，介绍Mamba的细节。

第4节，从块设计、扫描模式和内存管理的角度总结了基于Mamba的研究的最新进展。

第5节，介绍了使Mamba适应不同数据的技术，包括顺序和非顺序数据。

第6节，介绍了Mamba模型的代表性应用。

第7节，提出了挑战和未来的发展方向。

最后，我们在第8节中总结了整个调查。

4、回顾知识点

4.1、RNN

递归神经网络（RNN）在处理连续数据方面表现出色，因为它们能够保留内部存储器[54]。这样的网络已经在涉及分析和预测序列的大量任务中表现出显著的有效性，例如，语音识别、机器翻译、自然语言处理和时间序列分析[69，169]。

具体地说，在每个离散时间步长k，标准RNN具体处理一个向量𝑥，沿着前一步的隐藏状态 hk−1 ，以产生一个输出向量，并将隐藏状态更新为hk 。隐藏状态充当网络的记忆，并保留有关它所看到的过去输入的信息。这种动态记忆允许RNN处理不同长度的序列。形式上，它可以写成

NN是非线性递归模型，通过利用存储在隐藏状态中的历史知识来有效地捕获时间模式。

然而，RNN存在一些与RNN相关的限制。首先，RNN在有效提取输入序列中的长程动态方面的能力有限。当信息遍历连续的时间步时，网络中权重的重复相乘会导致信息的稀释或丢失。因此，RNN在进行预测时保留和回忆早期时间步的信息变得具有挑战性。其次，RNN递增地处理顺序数据，限制了它们的计算效率，因为每个时间步都依赖于前一个。这使得并行计算对他们来说具有挑战性。此外，传统的RNN缺乏内置的注意力机制，这使得网络能够在输入序列中捕获全局信息。这种注意力机制的缺乏阻碍了网络选择性地对数据的关键部分进行建模的能力。为了克服这些限制，出现了Transformers和State Space Models，每个都从不同的角度应对这些挑战。这两种方法将在随后的小节中进一步阐述。

4.2、Transformer

Transformer 是深度学习领域的一个突破性模型，彻底改变了各种人工智能应用。它的引入标志着与传统的序列到序列模型的重大偏离，采用了自注意机制，便于捕获模型输入中的全局依赖关系。例如，在自然语言处理中，这种自注意能力允许模型理解序列的不同位置之间的关系。它通过根据每个位置相对于其他位置的重要性为每个位置分配权重来实现这一点。更具体地说，首先通过利用原始输入的线性变换将输入向量序列x变换为三种类型的向量：Query、Key和Value，定义为：

注意力权重

除了执行单个注意力功能外，还引入了多头注意力，以增强模型捕捉不同类型关系的能力，并提供对输入序列的多个视角。在多头注意中，输入序列由多组自注意模块并行处理。每个头独立操作，执行标准自我注意机制中的精确计算。然后将每个头部的注意力权重组合以创建值向量的加权和。这个聚合步骤允许模型利用来自多个头部的信息，并捕获输入序列中存在的不同模式和关系。在数学上，多头注意力计算如下：

4.3、SSM

状态空间模型（SSM）是一种传统的数学框架，用于描述系统随时间的动态行为[89]。近年来，SSM在控制理论、机器人技术和经济学等不同领域得到了广泛应用[58，59]。在其核心，SSM通过称为“状态”的隐藏变量的集合来体现系统的行为，使其能够有效地捕获时态数据依赖性。与RNN不同的是，SSM是一种线性模型，其特征在于它们的关联性。具体地说，在经典的状态空间模型中，建立了两个关键方程，状态方程和观测方程，通过一个N维的隐状态，来模拟当前时刻的输入和输出之间的关系。该过程可由以下人员编写：

离散化。为了符合机器学习设置的要求，SSM必须经历离散化过程，将连续参数转换为离散参数。离散化方法的目的通常是将连续时间分割成离散区间，并尽可能地使积分面积相等。为了实现这一目标，作为最具代表性的解决方案之一，零阶保持（ZOH）[138，223]被成功地应用于SSM中，它假设函数值在区间Δ = [tk−1，tk]上保持恒定。在ZOH离散化之后，SSM方程可改写为:

从这些公式中可以清楚地看出，离散SSM具有与递归神经网络相似的结构，因此，与在每个自回归解码迭代中计算对所有输入的注意力的基于变换器的模型相比，离散SSM可以以更高的效率完成推理过程。

卷积计算等特性...

4.4、RNN、Transformer和SSM之间的关系

RNN、Transformer和SSM的计算算法如图所示。一方面，传统的RNN在非线性递归框架，其中每个计算仅取决于先前的隐藏状态和当前输入。虽然这种格式允许RNN在自回归推理期间快速生成输出，但它阻碍了它们充分利用GPU并行计算的能力，导致模型训练速度较慢。另一方面，Transformer架构跨多个查询密钥对并行执行矩阵乘法，这些查询密钥对可以有效地分布在硬件资源上，这使得能够更快地训练基于注意力的模型。但是，当涉及到从基于Transformer的模型生成响应或预测时，推理过程可能非常耗时。例如，语言模型的自回归设计需要按顺序生成输出序列中的每个标记，这需要在每一步重复计算注意力分数，导致推理时间变慢。RNN和Transformer，它们仅限于支持一种类型的计算，与RNN和Transformer不同的离散SSM具有灵活性，可以支持递归和卷积计算，因为它具有线性特性。这种独特的能力使SSM不仅可以实现高效的推理，还可以实现并行训练。然而，应该注意的是，最传统的SSM是时不变的，这意味着它们的A、B、C和Δ与模型输入无关。𝑥这将限制上下文感知建模，从而导致SSM在某些任务（如选择性复制）中的性能较差。

5、Mamba

为了解决传统SSM在其较差的上下文感知能力方面的上述缺点，提出了Mamba作为一种潜在的替代方案，有望成为通用序列基础模型骨干。最近，Mamba-2 提出了结构化空间-状态对偶（SSD），它建立了一个强大的理论框架，将结构化空间-状态对偶和各种形式的注意力连接起来，使我们能够将最初为Transformer开发的算法和系统优化转移到空间-状态对偶。在本节中，我们将对Mamba和Mamba-2进行简洁明了的介绍。

传统的SSM在对文本和其他信息密集型数据建模方面表现出有限的有效性[55]，阻碍了它们在深度学习方面的进展。为了使SSM具有Transformer的建模能力，Gu和Dao [55]介绍了三种基于结构化状态空间模型的创新技术，即，基于高阶多项式投影算子（HiPPO）的内存优化、选择机制和硬件感知计算，如图3所示。这些技术旨在提高SSM在长距离线性时间序列建模中的能力。特别是，初始化策略建立了一个连贯的隐藏状态矩阵，有效地促进远程记忆。然后，选择机制授权SSM获取内容感知表示。最后，Mamba提出了两种硬件感知的计算算法，并行关联扫描和内存重新计算，以提高训练效率。

基于HiPPO的内存初始化。对序列数据进行建模和学习是当代机器学习面临的基本挑战，为语言建模、语音识别和视频处理等各种任务奠定了基础。复杂和长期时间依赖性建模的一个基本组成部分在于存储器，包括存储和整合先前时间步长信息的能力[73]。类似于RNN，保留并遗忘历史隐藏状态（即，基体A）在SSM中起着关键作用以实现令人满意的性能。在先前的结构化状态空间序列模型（SSM）中，已经建议了特殊的初始化，特别是在复值模型的情况下。这些特殊的初始化已被证明在各种情况下都是有益的，包括数据可用性有限的情况。类似地，Mamba主要关注隐藏状态矩阵A的初始化，以捕获复杂的时间依赖性。这是通过利用HiPPO理论[56]和创新的缩放勒让德测度（LegS）来实现的，确保了对完整历史背景的全面考虑，而不是有限的滑动窗口。具体地说，HiPPO-LegS为所有历史数据点分配统一权重，可表示为：

其中，n是多项式的数量，并且n表示特定的离散时间步长。基于HiPPO理论，Mamba为复杂和真实的情况引入了两种简单的初始化方法，即，S4 D-Lin和S4 D-Real [57]，如

其中，n是A矩阵的第n个元素，A矩阵的所有输入维度= 1，2，...，𝑛𝐷。给定这样的初始化，该模型可以通过压缩和重构输入信息信号来学习经历较新步骤的较小退化和较旧步骤的较大退化的长相关记忆。根据公式，HiPPO-LegS具有有利的理论性质：它在输入时间尺度上保持一致，并提供快速计算。此外，它具有有界的梯度和逼近误差，便于参数学习过程。

选择机制。传统的状态空间模型不能基于特定的模型输入（即，内容感知建模能力）。为了给SSM提供类似于注意力机制的能力，Mamba设计了一种时变选择机制，该机制根据模型输入来参数化权重矩阵。这种创新使SSM能够过滤掉无关信息，同时无限期地保留相关细节。形式上，选择机制涉及将区间Δ和矩阵B、C设置为输入x 的函数，其可以公式化为：

离散化后的参数

这样，离散SSM从时不变变为时变，依赖于输入x。

输出可转化为

Mamba中的时变选择机制与Transformer中的注意力机制具有类似的结构，即，两者都基于输入及其投影执行操作，这使得Mamba的SSM能够实现灵活的内容感知建模。然而，它失去了与卷积的等价性，这对其效率产生了负面影响。

硬件感知型计算。该选择机制被精心设计以超越线性时不变模型的限制。尽管如此，它仍然挑战有效的训练：SSM的卷积核变得依赖于输入，导致无法执行并行计算。为了解决这个问题，曼巴使用了两种计算技术，即，并行联想扫描（也称为并行前缀和）[64]和记忆重算。首先，并行联想扫描利用线性联想计算的特性和现代加速器（GPU和TPU）的并行性，以节省内存的方式执行选择性SSM的计算。更具体地说，并行关联扫描将模型训练的计算复杂度从O（2）降低到O（/）。𝑁𝑑𝑁在其核心，扫描围绕着在给定的输入上构造一个平衡的二叉树，并从根开始扫描。换句话说，并行关联扫描开始于从叶子遍历到根（即，Sweep-Up），在树的内部节点创建部分和。然后，它反转遍历，从根向上移动回到树，以使用部分和来构造整个扫描（即，下扫）。

另一方面，Mamba利用传统的重新计算方法来减少训练选择性SSM层的总体内存需求。特别地，Mamba避免在并行关联扫描的前向传递期间存储大小为（B,L,D,N）的中间状态，以防止内存扩展。相反，它在梯度计算的反向过程中重新计算这些中间状态。通过这样做，重新计算回避了在GPU存储器单元之间阅读缓冲（缓冲）元素的必要性。除了优化扫描操作的内存需求外，Mamba-1还扩展了重新计算的使用，以提高整个SSM层的效率。这种优化包括投影、卷积和激活，它们通常需要大量的内存资源，但可以快速重新计算。

6、Mamba2-SSD

曼巴-2：状态空间对偶变换器在各个领域的深度学习成功中发挥了至关重要的作用，激发了各种技术的发展，例如参数有效微调[95]，灾难性遗忘缓解[96]和模型量化[195]，旨在从不同角度提高模型性能。为了使状态空间模型能够访问并受益于最初为变形金刚开发的有价值的技术，Mamba-2 [29]引入了一个称为结构化状态空间对偶（SSD）的综合框架，该框架建立了SSM和不同形式的注意力之间的理论联系。从形式如下：

其中M表示使用顺序半可分表示的SSM的矩阵形式，并且M= CT *A*B。值得注意的是，C和B分别表示与输入令牌相关联的选择性空间状态矩阵。A表示对应于范围从A到B的输入标记的隐藏状态的选择矩阵。从本质上讲，SSD表明，变压器使用的注意力机制和SSM的线性时变系统可以看作是半可分离的矩阵变换。此外，Dao和Gu [29]还证明了选择性SSM等价于用半可分掩蔽矩阵实现的结构化线性注意机制。

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

在本小节中，我们总结了Mamba-1和Mamba-2的区组设计。图4显示了这两种体系结构的比较。Mamba-1的动机是以SSM为中心的观点，其中选择性SSM层的任务是执行从输入序列X到Y的映射。在该设计中，在创建X的初始线性投影之后应用（A，B，C）的线性投影。然后，利用并行关联扫描，使输入令牌和状态矩阵通过选择性SSM单元，以产生输出Y。在此基础上，Mamba-1采用了跳跃连接的方法，以促进特征的重用，并缓解模型训练过程中经常出现的退化问题。最后，将该块与标准归一化连接和残差连接交错叠加，构建了Mamba模型。

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

这篇内容太多了，剩下内容见下一篇博客！！！