Mamba架构的综述

24年8月来自香港理工的论文“A survey of Mamba”。

作为最具代表性的架构，Transformers 赋能众多高级模型，尤其是包含数十亿个参数的大语言模型 (LLM)，成为深度学习的基石。尽管取得了令人瞩目的成就，但 Transformers 仍然存在固有的局限性，尤其是注意机制的二次计算复杂度导致推理过于耗时。最近，一种名为 Mamba 的新架构从经典状态空间模型中汲取灵感，成为构建基础模型有前途的替代方案，它提供与 Transformers 相当的建模能力，同时保留序列长度的近线性可扩展性。这引发了越来越多积极探索 Mamba 在不同领域实现出色性能的研究。

本综述对近期与 Mamba 相关的研究进行深入调查，主要涵盖三个方面：基于 Mamba 模型的进展、使 Mamba 适应各种数据的技术以及 Mamba 可以擅长的应用。全面回顾相关研究，重点关注 Mamba 模型的架构设计、数据适应性和应用。

结构化状态空间序列模型 (S4M) [58] ，有效地捕捉序列数据中的复杂依赖关系，成为 Transformer 的强大竞争对手。这些模型受到经典状态空间模型 [89] 的启发，可以看作是循环神经网络和卷积神经网络的融合。它们可以使用递归或卷积运算进行高效计算，实现序列长度的线性或近线性缩放，从而显著降低计算成本。更具体地说，作为最成功的 SSM 变体之一，Mamba 实现了与 Transformer 相当的建模能力，同时保持了序列长度的线性可扩展性 [55]。Mamba 首先引入一种简单而有效的选择机制，基于输入参数化 SSM 参数，使模型能够过滤掉不相关的信息，同时无限期地保留必要和相关的数据。然后，Mamba 提出一种硬件-觉察算法，通过一个扫描而不是卷积来循环地计算模型，在 A100 GPU 上实现高达 3 倍的计算速度。

如图给出基于 Mamba 的模型在不同下游任务中的应用示例：

状态空间模型 (SSM) 是一种传统的数学框架，用于描述系统随时间变化的动态行为 [89]。近年来，SSM 在控制理论、机器人技术和经济学等不同领域得到了广泛应用 [58, 59]。从本质上讲，SSM 通过一组称为“状态”的隐变量来体现系统的行为，使其能够有效地捕获时间数据依赖关系。与 RNN 不同，SSM 是具有关联属性的线性模型。具体而言，在经典的状态空间模型中，通过 N 维隐藏状态 h(𝑡) 建立两个关键方程，即状态方程和观测方程，以模拟当前时间 𝑡 的输入 𝑥(𝑡) 和输出 𝑦(𝑡) 之间的关系。该过程可以写为

一般来说，大多数 SSM 排除了观测方程中的第二项，即设置 D𝑥 (𝑡 ) = 0，这可以认为是深度学习模型中的跳连接。

离散化方法通常旨在将连续时间划分为具有尽可能相等积分面积的𝐾个离散区间。为了实现这一目标，作为最具代表性的解决方案之一，零阶保持（ZOH）[138, 223] 被成功应用于 SSM，它假设函数值在区间 Δ = [𝑡𝑘−1, 𝑡𝑘] 上保持不变。ZOH 离散化后，SSM 方程可以重写为

从这些公式中可以看出，离散 SSM 具有与循环神经网络相似的结构，因此，与基于 Transformer 的模型（在每个自回归解码迭代中计算所有输入的注意机制）相比，离散 SSM 可以更高效地完成推理过程。

更具体地说，它可以独立计算每个时间步骤的输出，如下所示：

通过创建一组卷积核 K = (CB, …, CA^kB, …)，可以将循环计算转换为卷积形式，如下所示：

这种卷积计算允许 SSM 充分利用现代矩阵计算硬件（例如 GPU），从而在训练过程中实现并行计算，而使用非线性激活函数的 RNN 则无法做到这一点。值得注意的是，给定一个输入 𝑥 (𝑘 )，SSM 计算将针对每个维度分别计算产生输出 𝑦(𝑡)。在这种情况下，输入矩阵 B、输出矩阵 C 和命令矩阵 D 保持不变，而状态转换矩阵保持不变，即 A。

如图是RNN、Transformer和SSM的算法代表性模型架构。(a) RNN 在非线性循环框架内运行，有助于在自回归推理期间快速输出。(b) Transformer 在众多Q-K对上同时执行矩阵乘法，从而促进并行训练。© 由于 SSM 具有线性特性，因此可以同时容纳循环和卷积计算，从而表现出多功能性。这种融合充分利用了 RNN 和 Transformer 的优势，使 SSM 能够进行循环推理和并行训练。尽管如此，传统的时不变 SSM 在上下文-觉察建模方面仍存在不足，导致特定任务的性能下降。

为了解决传统 SSM 在上下文-觉察能力较差方面的上述缺点，[55] 提出 Mamba 作为潜替代方案，有望成为通用序列基础模型主干。最近，Mamba-2 [29] 提出了结构化空间状态对偶 (SSD)，建立一个将结构化 SSM 与各种形式的注意机制连接起来的稳健理论框架，能够将最初为 Transformers 开发的算法和系统优化转移到 SSM。

Mamba-1：具有硬件-觉察算法的选择性状态空间模型

传统的 SSM 在对文本和其他信息密集型数据进行建模时效果有限 [55]，阻碍了它们在深度学习中的进展。为了使 SSM 具备 Transformers 的建模能力，[55] 提出三种基于结构化状态空间模型的创新技术，即基于高阶多项式投影算子 (HiPPO) 的内存初始化、选择机制和硬件感知计算，如图所示。这些技术旨在增强 SSM 在长距离线性时间序列建模中的能力。具体而言，初始化策略建立一个连贯的隐状态矩阵，有效地促进长距离记忆。然后，选择机制使 SSM 能够获得内容-觉察表示。最后，Mamba 设计两种硬件-觉察计算算法，即并行联想扫描和内存重计算，以提高训练效率。

基于HiPPO的内存初始化

从序列数据建模和学习，代表了当代机器学习的基础挑战，构成了各种任务的基石，包括语言建模、语音识别和视频处理。对复杂和长时间依赖性进行建模的基本组成部分是记忆，包括存储和整合前一时间步骤信息的能力 [73]。与 RNN 类似，保留和遗忘历史隐藏状态（即矩阵 A）在 SSM 中起着关键作用，以实现令人满意的性能。

在以前的结构化状态空间序列模型 (S4M) 中，有人建议进行特殊初始化，尤其是在复数值模型的情况下。这些特殊初始化已被证明在各种情况下都很有用，包括数据可用性有限的情况。同样，Mamba 主要关注隐状态矩阵 A 的初始化，以捕捉复杂的时间依赖性。这是通过利用采用一个尺度 Legendre 度量 (LegS) 的HiPPO 理论[56] 实现的，确保全面考虑完整的历史上下文，而不是有限的滑动窗口。具体来说，HiPPO-LegS 为所有历史数据点分配统一的权重，可以表示为：

基于 HiPPO 理论，Mamba 为复杂和真实情况引入了两种简单的初始化方法，即 S4D-Lin 和 S4D-Real [57]，如

给定这样的初始化，该模型可以通过压缩和重建输入信息信号来学习长时依赖的记忆，这种记忆经历较新阶段的较小退化和较老阶段的较大退化。根据公式，HiPPO-LegS 具备有利的理论特性：它在输入时间尺度上保持一致，并提供快速计算 [56]。此外，它具备有界的梯度和近似误差，有助于参数学习过程。

选择机制

由于时间不变性的特性，传统的状态空间模型无法根据特定的模型输入（即内容感知建模能力）产生个性化输出。为了使 SSM 具有类似于注意机制的能力，Mamba 设计一种随时间变化的选择机制，可根据模型输入参数化权重矩阵。这种创新使 SSM 能够过滤掉无关信息，同时无限期地保留相关细节。正式地说，选择机制涉及将间隔 Δ 和矩阵 B、C 设置为输入 x 的函数，可以表述为：

其中BroadCast𝐷（） 表示将结果广播到所有维度 𝑑 = 1，2，…，𝐷。随后，选择性 SSM 使用常见的统计技术零阶保持 (ZOH) [138] 进行离散化，如下所示

这样做，离散 SSM 从时不变变为时变（即内容-觉察）

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

硬件-觉察计算

选择机制旨在超越线性时不变模型的限制。然而，它在高效训练面临挑战：SSM 的卷积核变得依赖于输入，导致无法执行并行计算。为了解决这个问题，Mamba 采用两种计算技术，即并行联想扫描（也称为并行前缀和）[64] 和内存重计算。首先，并行联想扫描利用线性联想计算的属性和现代加速器（GPU 和 TPU）的并行性，以内存高效的方式执行选择性 SSM 计算。更具体地说，并行联想扫描将模型训练的计算复杂度从 O(𝑁²𝑑) 降低到 O(𝑁/𝑡)。扫描的核心是围绕在给定输入上构建平衡二叉树并将其扫描到根和从根开始。换句话说，并行联想扫描首先从叶子节点遍历到根节点（即向上扫描），在树的内部节点处创建部分和。然后，它反向遍历，从根节点沿树向上移动，使用部分和构建整个扫描（即向下扫描）。

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

Mamba-2：状态空间对偶

在深度学习各个领域的成功中，Transformers 发挥了至关重要的作用，它激发了各种技术的开发，例如参数高效微调 [95]、灾难性遗忘缓解 [96] 和模型量化 [195]，旨在从不同角度提高模型性能。为了使状态空间模型能够访问并受益于最初为 Transformers 开发的宝贵技术，Mamba-2 [29] 引入一个称为结构化状态空间对偶 (SSD) 的综合框架，它建立了 SSM 与不同形式的注意之间的理论联系。正式来说，

其中 M 表示采用顺序半-可分离表示的 SSM 矩阵形式，Mji = Cj^T Aj:i Bi。值得注意的是，Cj 和 Bi 分别表示与输入 token xj 和 x𝑖 相关的选择性空间状态矩阵。A𝑗:𝑖 表示与从 𝑗 到 𝑖 的输入 token 相对应隐状态的选择性矩阵。本质上，SSD 证明 Transformers 使用的注意机制和 SSM 中采用的线性时变系统，都可以看作是半-可分离矩阵变换。此外， [29] 还证明选择性 SSM 等同于用半-可分离掩码矩阵实现的结构化线性注意机制。

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

Mamba 块

总结 Mamba-1 和 Mamba-2 的块设计，如图所示是这两种架构的比较。Mamba-1 以 SSM 为中心，其中选择性 SSM 层负责从输入序列 X 到 Y 进行映射。在此设计中，在创建 X 的初始线性投影之后应用 (A，B，C) 的线性投影。然后，输入tokens和状态矩阵通过选择性 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 结构。此外，在跳连接之后添加了一个规范化层，旨在提高训练稳定性。

虽然 Mamba 在自然语言处理方面表现出色，但它仍然面临挑战，例如记忆丧失、泛化到不同任务以及捕捉复杂模式的能力，其不如基于 Transformer 的语言模型。为了克服这些挑战，人们付出了很多努力来改进 Mamba 架构。现有的研究主要集中在修改块设计、扫描模式和内存管理方面。相关研究的摘要如表所示：

块设计

Mamba 模块的设计和结构对 Mamba 模型的整体性能有重大影响，使其成为一个新兴的研究重点。如图所示，根据构建新 Mamba 模块的不同方法，现有研究可分为三类：a) 集成方法，旨在将 Mamba 模块与其他知名模型集成，以在有效性和效率之间取得平衡；b) 替代方法，试图利用 Mamba 模块替代高级模型框架中的主要层；c) 修改方法，侧重于修改经典 Mamba 模块内的组件。

扫描模式

并行联想扫描操作，是 Mamba 模型中的关键组件，旨在解决选择机制导致的计算问题，加速训练过程并减少内存需求。它通过利用时变 SSM 的线性特性在硬件级别设计内核融合和重计算来实现这一点。然而，Mamba 的单向序列建模范式阻碍了对各种数据（如图像和视频）的全面学习过程。为了缓解这个问题，一些研究专注于设计有效的扫描方法，以提高模型性能并促进 Mamba 模型的训练过程。如图所示，现有的专注于开发扫描模式技术的研究可以分为两类：1）扁平扫描方法，从扁平的 token 序列角度处理模型输入（a-c）；2）立体扫描方法，跨维度、通道或尺度扫描模型输入（d-f）。

内存管理

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

Mamba 架构代表了选择性状态空间模型的扩展，它具有循环模型的基本属性，非常适合作为对文本、时间序列、语音等序列进行操作的通用基础模型。同时，最近的开创性研究将 Mamba 架构的用途扩展到序列数据之外，涵盖了图像和图等领域，如图所示。（a-e）Mamba 架构具有循环模型的基本特征，是处理语言、时间序列和语音等序列的理想基础模型。（f-h）为了处理非序列数据，一种常见的方法是将数据分割或采样为离散tokens，并按照定义的规则将它们组织成序列。此外，Mamba 还具有通过连接各自的序列或投影来处理多模态数据的能力。这些研究旨在利用 Mamba 在捕获长距离依赖关系方面的卓越能力，同时利用其在学习和推理过程中的效率。

下表是各种数据下的Mamba研究概要：

基于Mamba架构模型的应用包括：自然语言处理、计算机视觉、语音分析、药物发现、推荐系统、以及机器人和自主系统。

目前存在的挑战和机会包括以下领域：

1. 基于Mamba的基础模型
2. 硬件觉察计算
3. 可信任的Mamba模型
4. 从Transformer到Mamba的现有技术应用