什么是 Mamba 模型？

Mamba 模型或是首个在语言建模任务中（Transformer 成名领域）媲美 Transformer 效能的深度学习架构。最显著的是，在主流 LLM 基准测试中，Mamba 架构展现出与同等规模 Transformer 相当的性能，同时延迟与内存需求往往显著降低。

Mamba 架构由 Tri Dao 与 Albert Gu 在 2023 年论文《Mamba：基于选择性状态空间的线性时间序列建模》中首次提出。一年后，二人发表续作论文，既探索了 SSM 与 Transformer 的关联性，又提出性能显著增强的 Mamba 架构升级版——Mamba-2。

尽管 Transformer 在 Mamba 问世后的两年内仍是 LLM 主导架构，但 Mamba 已被日益增多的开源模型集成。其中如 Mistral AI 的 Codestral Mamba 采用纯 Mamba 架构。而 AI2I 的 Jamba 系列及 IBM Granite 4.0 等更多模型则融合了注意力 (Transformer) 层与 SSM (Mamba) 层的混合架构。除性能优势外，基于 Mamba 的模型浪潮有望通过低成本硬件流畅运行，显著降低 AI 使用门槛。

SSM 最初设计用于预测连续序列（如电信号、天气模式或运动轨迹）在给定输入下的下一状态。从概念与数学层面，它们与 2017 年 Transformer 出现前主导自然语言处理 (NLP) 的循环神经网络 (RNN) 相关，也与卷积神经网络 (CNN) 和隐马尔可夫模型 (HMM) 等算法存在关联。

顾名思义，SSM 通过对状态空间的建模预测动态系统的下一状态：该数学空间表征描述系统状态的所有变量及其相互关联的可能取值范围。

SSM 接收输入序列 x(t)，将其映射至隐藏状态表示 h(t)——类比 RNN 的隐藏状态——以预测输出序列 y(t)。SSM 核心由两个方程构成：

• 状态方程， $h^{\prime }\left(t\right)=A*h\left(t\right)+B*x\left(t\right)$ 
• 输出方程， $y\left(t\right)=C*h\left(t\right)+D*x\left(t\right)$ 

模型关键参数 A、B、C、D 通常为权重矩阵形式。在控制理论等传统 SSM 应用领域，这些矩阵常被视为固定值：表征既定系统的动态特性，SSM 用于寻找能产生理想输出 y 的输入 x。现代 SSM 概念中，这些矩阵本身成为可通过机器学习优化的参数。深度学习模型中，它们由神经网络的可学习权重表示。

状态方程描述状态演化机制。矩阵 A 值决定各状态变量的自主演化规律。矩阵 B 值控制输入（如文本序列的下一词元）对各状态变量的影响。

在语言建模中，当前状态代表文本序列的上下文信息，每个词元处理后更新。其作用等同于 Transformer 模型中的 KV 缓存。

输出方程描述了当前状态如何通过矩阵 C 影响输出，以及输入如何通过矩阵 D 直接影响输出。由于矩阵 D 本质上独立于 h (t) 自身的建模过程，在 SSM 的图示和讨论中通常被省略，以聚焦于核心矩阵 A、B 和 C。

在 Mamba LLM 中，输出方程用于生成下一个词元。

传统 SSM 旨在建模连续输入，但文本序列（及现代深度学习模型处理的大多数数据模态）属于离散输入。利用 SSM 建模离散序列时，需将其离散的特定时间步表示为连续信号的一部分。

从概念上，离散化本质是将连续函数在特定时刻采样取值。这需要引入新参数——步长（记作 ∆），它决定每个离散时间步 t 采样值的持续时长（即“保持”时间）。调整 ∆ 类似于改变数据属性，例如时间序列数据的分辨率或视频数据的帧率。存在多种“离散化”方法，但现代 SSM 变体（包括 Mamba）普遍采用简单的零阶保持器 (ZOH) 方法。

离散化 SSM 使其能像 RNN 一样用于序列到序列任务。离散化 SSM 的参数和方程通常经过改写，采用 RNN 常用的下标表示法以区别于连续时间模型。该表示法中，h_t 代表模型将生成的新状态空间，h_t-1 则代表前一时刻的状态——即当前的状态空间。

 $h_t=\overline{A}{h}_{t-1}+\overline{B}{x}_t$ 
 $y_t=\overline{C}{h}_t$

由于标准离散化 SSM 存在与 RNN 相同的若干缺陷，直接用于文本数据建模并不现实。Albert Gu 等人于 2021 年提出的结构化状态空间序列模型（简称“S4 模型”）解决了其中两大缺陷：训练效率低下以及长序列建模能力不足。

尽管 S4 模型及其众多衍生版本（如对角 SSM (DSS)、对角 S4 (S4D) 和 H3 模型）的成功直接为 Mamba 的诞生铺平了道路。

离散化 SSM 等效于 RNN 特定实例的优势在于：RNN 在推理时速度极快。但缺点在于 RNN 训练过程极其缓慢。

所幸离散化 SSM 具备一项区别于其他 RNN 的关键特性：专有线性关系建模能力。换句话说，仅使用基础乘法与加法运算。正如 S4 论文所示，这些简单、重复且互依的线性递归展开为一维卷积核 ，实现单步输入 x 到输出 y 的映射： 。该过程可通过快速傅里叶变换高效计算。

唯一局限在于：此方法需要预知完整输入序列的所有步骤。这种操作在推理时无法实现，但在训练阶段可行。因此结构化 SSM 兼具双重优势：训练时可高效作为 CNN 运行；推理时可高效作为 RNN 运行。

与多数 RNN 类似，标准 SSM 天生不擅长建模长距离依赖关系。换言之，它们难以理解序列中远距离步骤（如段落首尾词语）的关联性，导致整体长序列建模能力薄弱。

为解决此问题，Gu 及其合著者（其中包括 Tri Dao）采用名为 HiPPO（高阶多项式投影算子）的技术，通过正交多项式推导的公式结构化初始化 A、B 矩阵的行为方式。这有别于标准机器学习实践——模型权重在训练初始阶段采用随机初始化。对于 S4 模型，Dao 和 Gu 提出了基于 Legendre 多项式的初始化方案。他们在后续论文《如何训练你的 HiPPO》中探索了更多公式变体。¹

S4 论文指出：“仅将 SSM 的随机矩阵 A 替换为 HiPPO 矩阵，其在序列 MNIST 基准测试中的性能即从 60% 提升至 98%”，彻底解决 SSM 长期记忆问题。结构化 SSM 的后续变体（如 DSS、S5 和 Mamba）为矩阵 A 和 B 采用了不同（通常更简化）的初始化方案，但仍遵循 HiPPO 核心原则：通过实施对角矩阵结构确保更新稳定性，并维持矩阵各值间的部分独立性。

Mamba 架构的核心包含两大创新。首先是选择性状态空间模型，它赋予 Mamba 此前仅 Transformer 模型具备的关键能力——能根据当前相关性，选择性地聚焦或忽略历史输入中的特定部分。其次是硬件感知并行扫描算法，该算法优化了图形处理器 (GPU) 在内存层级中处理模型计算的方式，从而实现速度与计算效率的最大化。

在 Transformer 中，这种能力由注意力机制提供：该机制通过调整注意力权重，基于历史词元与当前输入词元的相关性来增强或削弱其影响力。传统 SSM 显式依赖完整输入历史实现输入-输出映射。虽适用于部分序列建模任务，却成为高级语言建模的显著缺陷。

为解决无法动态忽略或强调历史输入的问题，Dao 和 Gu 提出了一类具备“选择性扫描”机制的新型状态空间模型。Mamba 论文中作者指出：“我们有时将选择性 SSM 简称为 S6 模型，因其本质是带选择机制的 S4 模型，并通过扫描计算实现”。他们将基于 S6 的架构命名为“Mamba”，部分原因在于多个 S 字母连读形似蛇的嘶鸣声。

Mamba 可以理解为以选择性状态空间模型为核心的神经网络架构。类比而言：Mamba 之于选择性 SSM，如同 Transformer 模型之于注意力机制。

传统 SSM 具有固定动态特性——控制隐藏状态逐步演化的规则（即模型参数）对所有输入及序列每一步均相同。该特性称为线性时不变性 (LTI)。为使 SSM 能基于当前语境选择性筛选历史信息，Dao 和 Gu 重构了 SSM 框架，使关键模型参数值随输入动态变化。

具体而言，选择性 SSM 使步长 ∆_t 及矩阵 B_t、C_t 成为当前输入词元 x_t 的直接函数。实现方式是将 xt 的_向量嵌入通过三个并行线性投影层（即标准前馈神经网络层或 MLP 层）处理。这类似于 Transformer 模型中并行查询头、键头和值头生成输入对应的 Q、K、V 向量。

将 x_t 向量嵌入乘以线性投影网络的权重和偏置项，即可得到 ∆_t、B_t 和 C_t 的值。这些线性投影层的权重与偏置项本身，需在模型预训练阶段通过海量文本样本学习获得，后续（可选）通过微调进一步优化。

• ∆_t 的值决定了 x_t 对模型当前上下文记忆的影响强度：即控制从隐藏状态 h_t-1 到 h_t 的更新幅度。较大步长 ∆_t：导致状态内旧信息发生更剧烈变化并加速衰减（即“遗忘”）。较小步长 ∆t：产生更细微的更新。步长足够小时：当前输入对隐藏状态完全不产生影响。
  
  
• 矩阵 B_k 的变化决定当前输入词元如何更新隐藏状态。例如，若 x_t是“yesterday”一词的词元，则 B_t 的调整可能使状态更新为：后续上下文很可能涉及过去时态。
  
  
• 矩阵 C_t 的变化决定该上下文信息如何影响模型输出 y_t。延续以“yesterday”为词元 x_k 的示例：矩阵 C_k 的调整将导致模型后续输出的所有动词均以过去时态变位形式呈现。

值得注意的是，矩阵 A 不会基于输入进行此类调整。其作用与 S4 模型一致：高效记忆全部历史输入。在特定时刻决定使用哪些历史信息的职责由矩阵 B、C 承担。

但当模型失去时不变性后，训练阶段无法再使用卷积捷径——因为转移核不再恒定：选择性机制的核心在于，从 h_t-1 到 h_t 的转移过程现在依赖于上下文。

作为替代方案，Mamba 采用巧妙变通实现类似并行化优势。由于 SSM 仅涉及乘法与加法运算，其计算遵循数学结合律：可通过不同方式分组计算而不改变最终结果。这使得大量顺序计算可拆解为独立小块，由 GPU 通过并行前缀和扫描算法并行处理。

此外，计算结果通过特定分层机制进行整合，该机制借鉴 Tri Dao 开发的 FlashAttention 技术原理，充分利用 GPU 不同类型硬件内存实现最优效率（该技术现已广泛用于现代大语言模型）。 

在 Mamba 架构中，S6 模型作为更大“模块”的组成部分运行，类似于注意力机制作为“注意力模块”中的子模块。该模块将 S6 模型与门控神经网络架构相结合。Mamba 模型通常在最终输出预测层之前包含多个连续堆叠的模块（即神经网络中的连续 Mamba 层）。

在进入 Mamba 模块前，输入的副本作为残差连接直接传输至末端。Mamba 模块的核心功能不仅在于识别全局上下文中与当前输入相关的部分，更在于确定上下文信息对原始输入含义的修正强度。

在 Mamba 模块内部，原始输入向量按以下流程处理：

• 首先，输入通过宽度倍增的线性层（即输出维度为输入向量两倍）投影至高维空间。例如，若模型原将每个输入词元 x 表示为 512 维向量嵌入，则经线性投影层权重扩展后，x 将转换为 1024 维向量。
  
  
• 随后，该扩展向量被均分为两部分。其中一半（称为 x_proj）输入至 SSM 处理路径，另一半（称为 z_proj）输入至独立门控机制处理路径。为清晰起见，前述扩展步骤通常图示为两个并行线性层实现。
  
  
• 在 x_proj 进入 SSM 前，需经一维卷积层处理。该卷积层专门提取局部特征模式（如相邻词元间的依存关系：典型如简单主谓搭配）。此处理使 SSM 能专注于建模长距离全局依赖的上下文理解。
  
  
• 卷积层的输出经非线性激活函数处理。神经网络引入非线性是捕获复杂模式的标志。Mamba 论文采用 Sigmoid 线性单元 (SiLU)。所得向量记为 x_act。
  
  
• 同时，在独立门控机制路径中，z_proj 也经非线性激活函数处理生成 z_act。
  
  
• 在 SSM 路径中，x_act 输入三个并行线性投影层，分别生成 ∆_x、B_x 和 C_x 的值。
  
  
• SSM 利用这些输入相关参数（及矩阵 A、D）计算状态空间更新值及 SSM 输出 y。
  
  
• 随后，SSM 的输出向量 y 现在与门控路径的输出向量 z_act 按元素相乘。本质上，z_act 中的每个元素就像音频混音台的音量旋钮：如果 z_act 的某个元素接近零，那么与 y 对应部分相乘后，结果会更接近零，其影响也会随之减弱。相反，若 zact 某元素值较大，则将放大 y 中上下文信息的影响力。
  
  
• 最终所得向量被投影压缩至原始尺寸。可理解为对原始输入向量各元素的加权上下文更新（或非更新）结果集。
  
  
• 最后，该更新向量与作为残差连接直达模块末端的原始输入副本相加。
  
  
• 此时原始输入向量已完成上下文感知更新（由选择性 SSM 提供）。现在可输送至下一 Mamba 层；若处于模型最终层，则作为 softmax 函数的输入，输出经全更新向量对应词表中各单词的概率分布。

Mamba 论文发表一年后，Dao 与 Gu 发表续作《Transformer 即状态空间模型：基于结构化状态空间对偶性的通用模型与高效算法》。该论文做出三大贡献：

• 理论探索：揭示 Mamba 与 Transformer 的架构关联性并建立统一术语体系
• 系统阐释 Mamba 模型不同设计选项并进行深度探索
• 架构升级：基于设计研究推出改进版 Mamba-2 架构

Mamba-2 算法较初版显著提速且更易实现——作者提供约 25 行代码的“极简 SSD”代码库实现选择性 SSM。²其高效性支持 Mamba-2 在模型速度不变前提下大幅扩展隐藏状态维度，从而构建规模更大、能力更强、表达更丰富的模型。经测试，Mamba-2 模型在多项下游任务中确定持平或超越同规模 Mamba 及 Transformer 模型。

如论文引言所述，Dao 与 Gu 的“核心目标在于建立结构化 SSM 与注意力变体间的理论关联体系”。这催生了统一两者的新概念框架——“状态空间对偶性” (SSD)^。3此举为 Mamba 架构开启了吸收 Transformer 多年优化成果的通道。

显著成果之一是开发出多头注意力 (MHA) 的 Mamba 等效机制：Mamba 模块可拆分为多个“处理头”，类似于 Transformer 中的“注意力头”。其中被类比分组查询注意力的变体方案，通过 GPU 张量并行实现更高效率。

在 Mamba-2 模块（作者称为并行 Mamba 模块，区别于初版“串行”模块）中：输入相关参数 ∆、B、C 在初始投影层并行生成。具体而言，矩阵 B 和 C 的生成方式简化为：直接复制 x_proj 的片段数据，而非通过专用线性层对 x_proj 执行矩阵乘法运算。此设计在简化模型结构、减少参数总量的同时，实现了显著更高效的大规模训练。⁴

Mamba 与 Transformer 架构各有优势，但在内存使用和速度相关领域，基于 Mamba 的模型普遍更优：据 Mamba 论文显示，其吞吐量可达同等 Transformer 的 5 倍。

Transformer 虽具备极高精度与通用性，但对计算资源需求极其严苛。在预训练（及微调）阶段，自注意力机制的内存需求随序列长度呈二次方增长：若序列上下文长度加倍，注意力机制的资源消耗量将增至四倍。随着上下文窗口扩展，这种“二次方瓶颈”对速度和内存可用性的制约日益加剧。而在推理阶段，其内存需求仅呈线性增长。

在训练阶段，Mamba 模型的内存消耗仅随序列长度线性增长。更重要的是，其在推理阶段的内存使用保持恒定：无论模型处理过多少词元，SSM 始终以固定大小表征历史输入。这使得理论上支持无限上下文长度，实际仅受硬件限制。

尽管如此，Transformer 虽内存密集且计算冗余，却具备独特优势。例如，研究表明，在需要上下文学习（如少样本提示）、文本复制或长上下文推理任务中，Transformer 的性能仍优于 Mamba 及 Mamba-2。

值得庆幸的是，Transformer 与 Mamba 的优势并非互斥。Mamba-2 论文指出，混合模型可能超越纯 Transformer 或 SSM 架构——该观点于 2024 年获英伟达研究正式验证。⁵广义而言，混合模型兼具 Mamba 的高效性与 Transformer 注意力机制（尽管资源消耗更大）带来的精细建模及上下文学习能力。

为深入探索，IBM Research 联合 Dao、Gu 及伊利诺伊大学香槟分校 (UIUC) Minjia Zhang 团队开发 Bamba 与 Bamba V2。Bamba 反过来为 IBM Granite 4.0 大模型提供了核心架构设计基础。

混合模型研究仍是活跃领域，开源社区尤其聚焦于此方向。