Mamba 与状态空间模型：面试问题与推荐回答

SSM 与 Mamba 基础

Q1 [基础] 什么是状态空间模型？

Q： 什么是状态空间模型（SSM）？它如何应用于序列建模？

A： 状态空间模型（State Space Model，SSM）是一种通过潜在隐藏状态将输入序列映射为输出序列的数学框架。其连续形式由两个方程定义：状态方程 $\frac{dx}{dt} = Ax + Bu$（描述隐藏状态如何随时间演化）和输出方程 $y = Cx + Du$（描述如何从状态生成输出）。其中，$A$ 是状态转移矩阵，$B$ 将输入映射到状态更新，$C$ 将状态映射到输出，$D$ 是可选的跳跃连接。

对于离散数据（文本、音频采样）上的序列建模，连续系统被离散化为递推形式：$x_t = \bar{A}x_{t-1} + \bar{B}u_t$，$y_t = Cx_t$。这种递推使 SSM 在计算上与 RNN 类似：推理时每步只需固定大小的状态更新，无论序列长度如何，每步均需 $O(1)$ 内存。

一个关键数学性质是：递推可等价地表示为卷积 $y = K * u$，其中 $K$ 是由 SSM 参数导出的卷积核。这种对偶性使 SSM 可以通过基于 FFT 的卷积高效训练（在序列上可并行），同时以递推形式高效部署（推理时内存恒定）。这种训练-推理对偶性是 SSM 相比 Transformer（长序列推理慢）和 RNN（训练慢）的核心优势之一。

Q2 [基础] 什么是 HiPPO 矩阵？为何重要？

Q： 什么是 HiPPO 矩阵？它在现代 SSM 中扮演什么角色？

A： HiPPO（High-order Polynomial Projection Operators，Gu et al., 2020）是一种有原则的方法，用于初始化 SSM 中的状态转移矩阵 $A$，使隐藏状态能够最优地压缩输入信号的历史。具体而言，HiPPO 框架构造了一个 $A$ 矩阵，使得每个时间步的隐藏状态包含到目前为止输入的最优多项式近似（Legendre 基）的系数，直至 $N$ 阶。

其实际意义在于：经 HiPPO 初始化的 $A$ 矩阵为 SSM 提供了一种结构化的方式来同时表示近期和远期历史，每个多项式基函数在不同的时间尺度上捕获信息。若不使用 HiPPO 初始化，以随机 $A$ 矩阵初始化的 SSM 会像普通 RNN 一样遭受梯度消失或爆炸问题，迅速丢失遥远过去的信息。

HiPPO 是奠定了 S4 和 Mamba 等后续所有 SSM 工作的基础性记忆机制。虽然 Mamba 并未使用 HiPPO 的完整连续时间构造，但结构化初始化的设计理念——确保 $A$ 编码了对长程记忆有用的归纳偏置——在 Mamba 的设计中得以延续。

Q3 [基础] S4 的核心贡献是什么？

Q： S4 解决了什么问题？其关键创新是什么？

A： S4（Structured State Space Sequence Model，Gu et al., 2021）解决了将 SSM 应用于长序列时的计算瓶颈。此前的 SSM 每序列计算量为 $O(n \cdot N^2)$（$N$ 为状态维度），因为计算 SSM 卷积核需要涉及 $A$ 的密集矩阵运算。S4 的核心洞见是将 $A$ 参数化为对角加低秩（DPLR）矩阵。在这种结构下，SSM 卷积核可使用 Cauchy 核高效计算，通过 FFT 将复杂度降至 $O(n \log n)$。

S4 在 Long Range Arena 基准测试上取得了最优结果——该测试需要建模跨越数千个时间步的依赖关系，远超 Transformer 和此前 SSM 的表现。它证明了具有原则性结构的 SSM 可以在长程任务上与基于注意力的模型竞争。

S4 的局限在于其参数（$A$、$B$、$C$）在所有输入位置上是固定的——它是一个线性、时不变系统。这意味着 S4 无论内容如何都以相同方式处理每个输入，无法选择性地关注相关词元或忽略无关词元，而这是 Transformer 注意力的基本能力。这一局限推动了 Mamba 选择性机制的诞生。

Q4 [基础] Mamba 的选择性机制是什么？

Q： Mamba 使用"选择性"状态空间模型意味着什么？为何重要？

A： 在 S4 和早期 SSM 中，矩阵 $A$、$B$、$C$ 是时不变的：它们是不依赖输入的固定参数。这意味着模型对每个输入词元应用相同的变换，无论其内容如何——无法决定对某些词元"更加关注"。

Mamba（Gu & Dao, 2023）通过使参数 $B$、$C$ 和 $\Delta$（离散化步长）成为输入的函数来引入选择性。具体而言，对于每个词元 $x_t$，模型计算 $B_t = \text{Linear}(x_t)$、$C_t = \text{Linear}(x_t)$，以及 $\Delta_t = \text{softplus}(\text{Linear}(x_t))$。$A$ 矩阵保持固定（以结构化对角形式初始化），但其他所有参数现在随词元变化。

其直觉与注意力机制类似：模型可以学习在遇到希望记住的词元时"打开门"（大 $\Delta$，大状态更新），在希望保留现有记忆时"关闭门"（小 $\Delta$，状态变化缓慢）。这种基于内容的选择性使 Mamba 能够选择性地过滤信息，并推理上下文中哪些部分与当前预测相关——而这是 S4 完全不具备的能力。

Q5 [基础] Mamba 中的 B、C、Δ 各自控制什么？

Q： 参数 B、C 和 Δ 在 Mamba 的选择性状态空间模型中各扮演什么角色？

A： 在 Mamba 的 SSM 中，$B$、$C$ 和 $\Delta$ 是实现选择性的三个输入依赖参数，各自控制信息流经隐藏状态的不同方面。

$B$ 是输入投影矩阵，决定当前输入 $x_t$ 对状态更新的影响程度。某维度上较大的 $B_t$ 值意味着输入会强烈写入状态的该维度。通过使 $B$ 依赖输入，Mamba 可以学习根据内容选择性地将特定类型的信息编码到状态中。

$C$ 是输出投影矩阵，决定当前状态如何被读出以产生输出 $y_t$。通过使 $C$ 依赖输入，模型可以根据当前词元的需求选择性地检索累积状态的不同方面。

$\Delta$（delta）是离散化步长，控制状态更新的时间尺度。较大的 $\Delta$ 将连续 SSM 折叠为一个大的离散步，使新状态受当前输入强烈影响（关注当下）；较小的 $\Delta$ 对应小的离散步，意味着状态变化缓慢，更多保留其先前值（关注历史）。$\Delta$ 是最关键的选择性参数：它充当一个可学习的门，调节记忆过去与响应当前输入之间的权衡，其作用类似于 LSTM 中的遗忘门。

Q6 [进阶] 并行扫描如何实现 Mamba 的高效训练？

Q： Mamba 的核心是递推结构。它如何避免训练时 RNN 的串行瓶颈？

A： 训练递推模型的根本挑战在于朴素递推本质上是串行的：$x_t$ 依赖 $x_{t-1}$，后者又依赖 $x_{t-2}$，形成 $n$ 个依赖链，无法在序列维度上并行化。这正是 RNN 相比 Transformer 训练慢的原因——Transformer 能同时计算所有位置。

Mamba 通过并行扫描算法（也称并行前缀和）解决这一问题。核心洞见在于：Mamba 的状态递推 $x_t = \bar{A}_t x_{t-1} + \bar{B}_t u_t$ 是线性递推，而线性递推是满足结合律的运算。任何满足结合律的运算都可以使用分治法在所有前缀子序列上同时计算，深度为 $O(\log n)$，总工作量为 $O(n)$，而串行方法需要 $O(n)$ 深度。

具体而言，并行扫描首先合并相邻元素对，然后合并元素对的对，依此类推，在 $O(\log n)$ 轮并行计算中构建所有前缀归约结果。对于长度为 $n$ 的序列，总操作数为 $O(n)$，但串行步骤仅需 $O(\log n)$——从而实现完全的 GPU 并行性。Mamba 将其实现为优化的 CUDA 内核，完全在快速片上 SRAM 中运行并行扫描，避免了从慢速 HBM 的昂贵读写。最终结果是训练吞吐量与优化的 Transformer 实现相当，尽管底层仍是递推结构。

Q7 [进阶] Mamba 的硬件感知算法与 FlashAttention 有何相似之处？

Q： 什么是 Mamba 的硬件感知算法？与 FlashAttention 的方法相比如何？

A： Mamba 和 FlashAttention 解决的是同一类问题：一种数学上直观的计算，其朴素实现在现代 GPU 上受内存带宽瓶颈制约，而非算术吞吐量。两者的解决方案都是重构计算，通过将中间结果保留在快速片上 SRAM 中，最小化从慢速 HBM（GPU 全局内存）的读取次数。

在标准 SSM 计算中，所有时间步 $t$ 的中间状态 $x_t$ 都会在 HBM 中实体化——需要 $O(n \cdot d_\text{state})$ 内存以及 $O(n)$ 次 SRAM 与 HBM 之间的往返。Mamba 的硬件感知算法将整个前向传播融合为单个 CUDA 内核：输入 $u_t$、离散化参数 $\bar{A}_t$ 和 $\bar{B}_t$，以及输出 $y_t$ 全部在 SRAM 分块内处理，状态永远不写回 HBM。在反向传播中，状态从存储的输入中重新计算，而非从 HBM 加载——以算术运算（廉价）换取内存带宽（昂贵），与 FlashAttention 的做法完全一致。

两者的实际效果相同：实际内存使用远低于朴素实现，由于减少了 I/O，实际运行速度显著更快。FlashAttention 将 Transformer 注意力的 $O(n^2)$ 内存降至 $O(n)$；Mamba 的硬件感知内核类似地降低了 SSM 的内存开销。两者都已成为各自模型家族中事实上的标准实现。

Q8 [进阶] Mamba 如何实现线性时间复杂度？有哪些注意事项？

Q： Mamba 被描述为具有线性时间复杂度。这究竟意味着什么？有哪些权衡？

A： Mamba 的线性复杂度主张主要适用于推理（自回归生成）。在每个推理步骤中，Mamba 使用当前输入更新维度为 $d_\text{model} \times d_\text{state}$ 的固定大小隐藏状态——这是与序列长度 $n$ 无关的 $O(d)$ 操作。因此生成 $n$ 个词元总共需要 $O(n \cdot d)$ 工作量和 $O(d)$ 内存（状态在每步被覆盖，不累积）。相比之下，Transformer 解码器每步需要 $O(n)$ 工作量（关注 KV 缓存）和 $O(n \cdot d)$ 内存用于 KV 缓存——两者都随序列长度增长。

训练时，Mamba 使用并行扫描方法，工作量为 $O(n \log n)$——并非严格线性，但对 Transformer 的 $O(n^2)$ 有显著改善。训练期间用于反向传播存储中间状态的内存为 $O(n \cdot d_\text{state})$，但硬件感知内核通过重新计算状态而非存储，在实践中降低了这一开销。

一个重要注意事项是，Mamba 的固定大小状态是一把双刃剑。$O(1)$ 推理内存之所以可能，是因为模型将整个历史压缩为固定维度 $d_\text{state}$ 的状态。这种压缩是有损的：与 Transformer 的 KV 缓存（精确存储所有过去词元，可以完整精度检索任意词元）不同，Mamba 的状态是一个可能无法忠实表示所有过去信息的摘要。这是根本性的权衡——以完美召回换取效率——也正是为何 Mamba 在需要精确检索特定过去词元的任务上往往不如 Transformer 的原因。

Q9 [基础] Mamba2 的核心改进是什么？

Q： Mamba2 与原始 Mamba 有何不同？它解决了什么问题？

A： Mamba2（Dao & Gu, 2024，“Transformers are SSMs”）引入了一个关键的架构简化：将 $A$ 矩阵限制为单位矩阵的标量倍（$A = -\exp(a) \cdot I$，其中 $a$ 是每通道学习的标量）。在原始 Mamba 中，$A$ 是全对角矩阵，每个状态维度有一个可学习值。Mamba2 中这一限制看似削减了表达能力，但它开启了一个关键的理论联系：SSM 计算可以被重写为结构化矩阵乘法，并被证明等价于一种线性注意力形式。

这种等价关系——状态空间对偶性（State Space Duality，SSD）——有两个实践意义。首先，Mamba2 的 SSD 层更适合跨多 GPU 的张量并行，因为它可以表示为块对角矩阵运算，而非自定义扫描内核，从而提高大规模训练吞吐量。其次，Mamba2 既可实现为 SSM（对长序列高效），也可实现为一种注意力形式（对短序列使用 FlashAttention 高效），根据给定的上下文长度选择速度更快的方式。

实验上，Mamba2 在等效参数量下达到与 Mamba 相当或更好的语言建模质量，并在多 GPU 训练中具有更高的吞吐量。它代表了 SSM 与注意力文献的融合，提供了一个统一框架，而非将两者视为完全独立的范式。

Q10 [进阶] 什么是状态空间对偶性？它如何统一 SSM 与注意力？

Q： Mamba2 中引入的状态空间对偶性（SSD）框架是什么？它揭示了 SSM 与注意力之间怎样的关系？

A： 状态空间对偶性（SSD）是 Mamba2 论文引入的理论框架，表明在特定结构约束下，SSM 和一种线性注意力形式是计算同一函数的两种不同方式。具体而言，当 $A$ 被限制为标量乘以单位矩阵（如 Mamba2 中）时，SSM 计算 $x_t = A_t x_{t-1} + B_t u_t$、$y_t = C_t^\top x_t$ 可以被重组为矩阵形式 $Y = M \cdot (B^\top U)$，其中 $M$ 是由 $A$ 标量决定的结构化下三角矩阵（因果掩码），$B$、$C$、$U$ 是输入投影。这正是具有特定结构化掩码的线性注意力形式。

这种对偶性具有深刻含义。标准 softmax 注意力使用 $n \times n$ 注意力矩阵 $A_{ij} = \exp(q_i^\top k_j) / Z$；线性注意力通过去掉 softmax 并将注意力表示为矩阵乘积来近似，这可以重组为递推形式。SSD 框架表明，Mamba2 等 SSM 是这个线性注意力家族的一个具体实例，具有特定的结构化（因果、衰减）注意力模式。

实践意义是架构灵活性：SSD 层可以实现为递推（用于自回归推理，$O(1)$ 内存）、分块算法（用于训练，在 SRAM 中处理词元块）或原则上实现为注意力（用于极短序列）。这使 Mamba2 比 Mamba 的自定义选择性扫描更加硬件友好，并为根据不同上下文长度选用最高效算法的混合实现打开了大门。

Mamba 与 Transformer 的比较

Q11 [基础] Mamba 与 Transformer 在效率上如何比较？

Q： Mamba 与 Transformer 架构之间的关键效率差异是什么？

A： Mamba 与 Transformer 的效率比较取决于是比较训练还是推理，以及在何种序列长度下。

训练时，Transformer 注意力相对于序列长度 $n$ 具有 $O(n^2)$ 的时间和内存复杂度——注意力矩阵为 $n \times n$。Mamba 的并行扫描具有 $O(n \log n)$ 的时间和 $O(n \cdot d_\text{state})$ 的内存，对于长序列有显著优势。实践中，对于最多约 2,000 个词元的序列，高度优化的 Transformer 内核（FlashAttention）可以与 Mamba 媲美甚至超越；超过 8,000 个词元后，Mamba 的优势变得显著。

推理时（自回归生成），差距更为明显。Transformer 解码器必须维护以 $O(n \cdot d_\text{model})$ 增长的 KV 缓存——存储所有过去的键值向量。每步关注该缓存的代价为 $O(n)$，内存随上下文长度无限增长。Mamba 在推理时维护大小为 $O(d_\text{state})$ 的固定隐藏状态，无论处理了多少词元，每步生成代价为 $O(d)$——真正与序列长度无关。这使 Mamba 在长上下文推理和流式应用中特别具有吸引力，因为 Transformer 的 KV 缓存内存在这些场景下会变得难以承受。

Q12 [基础] Mamba 与 Transformer 在归纳偏置上有何不同？

Q： Mamba 与 Transformer 在根本性归纳偏置上有何差异？

A： Transformer 注意力的序列归纳偏置极少：它对称地对待所有位置，使用显式位置编码注入顺序信息。注意力原则上可以为任意两个位置赋予相等的权重，无论其距离远近。这使 Transformer 灵活，但需要从数据中学习位置关系。

Mamba 通过其递推结构具有强烈的序列归纳偏置：信息从左到右流动（在标准因果模型中），过去词元对当前输出的影响随学习的 $\Delta$ 参数和中间词元的作用而衰减。除非模型显式学习保留，否则近期词元对隐藏状态的影响自然强于远期词元。这种近期偏置隐含在架构中，而 Transformer 的近期偏置必须通过位置编码和注意力模式来学习。

此外，Mamba 通过固定大小的瓶颈（隐藏状态）处理信息，迫使其压缩上下文——这对需要摘要的任务是有用的归纳偏置，但对需要精确召回特定过去词元的任务则是局限。Transformer 通过缓存所有过去的键值对，没有这种压缩瓶颈，可以以完整精度检索任意过去词元。这种记忆结构的根本差异是两种架构观察到的性能权衡的主要驱动因素。

Q13 [进阶] Mamba 在长序列任务上有哪些优势和局限？

Q： 在长序列任务上，Mamba 相比 Transformer 在哪些方面表现出色，在哪些方面存在不足？

A： Mamba 在长序列上的主要优势是效率：训练内存（线性 vs 二次方）和推理内存（常数 vs 线性）的扩展性远优于 Transformer。这使 Mamba 在 Transformer 代价过高的 10 万词元以上序列长度下可行。相关领域包括基因组学（超长 DNA 序列）、音频处理（长波形）和长文档语言建模。

对于受益于长上下文压缩的任务——如摘要、全局模式提取或时间序列预测——Mamba 的递推结构非常适合。隐藏状态自然地积累整个历史的压缩表示，选择性机制允许保留相关信息同时过滤无关信号。

然而，Mamba 在需要从遥远过去精确检索特定词元的任务上存在有据可查的弱点。针对"大海捞针"基准测试（在长文档中找到特定事实）和联想回忆任务（给定键，检索之前见过的关联值）的研究表明，Mamba 的固定大小状态难以忠实保留所有过去信息。Transformer 注意力将所有过去词元存储在 KV 缓存中，可以精确检索任何过去词元。Mamba 只能寄希望于相关信息已被编码进其有限维状态，而未被后续输入覆盖。这是使 Mamba 不太适合检索增强生成和长上下文多跳推理的核心局限。

Q14 [进阶] 何时应选择 Mamba 而非 Transformer？

Q： 实践中，如何在基于 Mamba 和基于 Transformer 的架构之间做出选择？

A： 决策取决于三个因素：序列长度、任务性质和部署约束。

对于超长序列（超过 8K–16K 词元），Mamba 线性推理内存和次二次方训练复杂度使其在 Transformer 代价过高的场景中切实可行。基因组学、音频、长上下文文档理解和流式推理是天然的适用场景。若用例需要处理 10 万以上词元且无法获得超大 GPU 集群，Mamba（或基于 Mamba 的混合模型）可能是唯一可行的选择。

对于需要精确词元级检索的任务——如少样本上下文学习、长上下文多跳问答或检索增强生成——Transformer 注意力的精确 KV 缓存具有结构优势。Mamba 的压缩状态可能丢失检索所需的特定信息，导致精度下降。类似地，若需要利用大型预训练 Transformer 检查点（数量远超可用的 Mamba 检查点），微调 Transformer 是务实之选。

对于中等上下文长度（2K–32K）的通用语言建模，两种架构更具竞争性，混合设计（交错 Mamba 和注意力层）通常提供最佳权衡：Mamba 层处理序列大部分的高效压缩，注意力层在关键处提供精确检索能力。若不确定，小比例注意力层（如 1/8）的混合模型可以保留 Mamba 的大部分效率增益，同时恢复 Transformer 的大部分检索能力。

Q15 [进阶] Mamba 能取代 Transformer 吗？

Q： Mamba 是 Transformer 的替代品，还是两者互补？

A： 截至 2024 年，Mamba 最好被理解为 Transformer 的互补而非替代。在等效参数量和训练算力下，Mamba 在语言建模困惑度上与 Transformer 相当，并在吞吐量敏感的长上下文基准测试上超越 Transformer。然而，Transformer 在指令遵循、少样本上下文学习和精确检索任务上保持优势——这些能力是当今大语言模型最具实用价值的应用的核心。

生态系统差距同样显著：Transformer 已有多年的预训练规模、基础设施工具和微调研究积累。Mamba 最大的可用预训练模型（截至 2024 年初）在 30 亿参数量级，而领先的 Transformer 模型已达到数千亿参数。弥合这一差距需要社区刚刚开始投入的大量资源。

最有前景的近期方向是混合架构。Jamba（AI21 Labs）和 Zamba 等模型将 Mamba 层与少量 Transformer 注意力层交错，在恢复 Transformer 检索和上下文学习能力的同时，保留了 Mamba 的大部分效率。这些混合模型证明两种架构不是相互竞争的范式，而是可以组合的互补工具。更长远的问题——纯 SSM 架构在适当训练下能否在规模上与 Transformer 媲美——仍然开放，但混合方案已经在今天带来了切实的实践价值。

应用与扩展

Q16 [基础] Mamba 在语言建模中表现如何？

Q： Mamba 与基于 Transformer 的语言模型在实践中如何比较？

A： Mamba 在语言建模中表现强劲，尤其随序列长度增加优势更加明显。在等效训练词元和 10 亿–30 亿参数量级下，Mamba 达到与 GPT-NeoX 等 Transformer 基线相当的困惑度，且由于线性推理复杂度，在更长上下文时实现更高的吞吐量（更高词元/秒）。原始 Mamba 论文证明 Mamba-3B 在标准语言建模基准上与 GPT-NeoX-3B 持平或超越。

Mamba 在语言建模中的优势在受益于长上下文的任务上最为突出：长文档摘要、以大型代码库为上下文的代码补全，以及长输入上的结构化预测。其弱点出现在少样本提示任务（上下文学习）上，在等效规模下，Transformer 模型持续优于 Mamba。这一差距归因于 Mamba 的固定大小状态无法以相同保真度维护所有少样本示例。

在实际部署方面，Mamba 为需要在有限 GPU 内存下服务长上下文的应用提供了有意义的优势——常数 KV 缓存等价物（固定大小状态）意味着上下文长度不会增加推理时的内存需求，从而为长上下文请求实现更大的批处理大小和更低的延迟。

Q17 [基础] Mamba 如何应用于视觉任务（Vision Mamba）？

Q： Mamba 如何处理二维图像数据？有哪些将其应用于视觉的方法？

A： 将 Mamba 应用于图像需要解决一个根本性的不匹配：Mamba 的 SSM 设计用于一维序列，而图像是二维空间结构，局部性和二维位置关系至关重要。最直接的方法——Vision Mamba（Zhu et al., 2024）所采用——是将图像分块为图像块（如 ViT），展平为一维序列，然后应用双向 Mamba 模型。

双向性对视觉很重要，因为图像块在所有方向上都有空间邻居，而非仅有标准 Mamba 使用的因果（从左到右）方向。Vision Mamba 通过在两个相反扫描方向上运行两个 Mamba 模型并合并输出来引入双向 SSM，使每个图像块能够在扫描顺序中聚合前后图像块的信息。其他扫描策略也已被提出：VMamba 在四个方向（水平、垂直和两个对角线）扫描；LocalMamba 使用基于窗口的局部扫描，以降低高分辨率输入的复杂度。

性能方面，Vision Mamba 在 ImageNet 分类上达到与参数量相当的 ViT 模型竞争性的精度，在 ViT 二次方注意力代价高昂的高分辨率图像上内存使用更低。对于密集预测任务（检测、分割），产生多尺度特征图的分层变体（类似 Swin Transformer）显示出最大的实践前景。

Q18 [进阶] 什么是 Mamba + Transformer 混合架构？为何有前景？

Q： Jamba 等混合架构的动机是什么？它们如何工作？

A： 混合架构将 Mamba（SSM）层与 Transformer 注意力层在同一模型中交错，旨在融合两者的互补优势：Mamba 在大多数层的效率和压缩能力，以及 Transformer 注意力在少数层的精确检索能力。

Jamba（AI21 Labs，2024）是最具代表性的例子。它每七个 Mamba 层交错一个 Transformer 注意力层，并配合混合专家（MoE）前馈层。结果是一个支持 256K 词元上下文窗口、在等效质量下吞吐量比纯 Transformer 显著更高的模型——Mamba 层廉价地处理序列处理的主体，而注意力层提供 Mamba 状态无法保证的精确全局检索。Jamba 在长上下文长度下实现了约 3 倍于可比 Transformer 的吞吐量。

关键设计洞见是：并非每层都需要完整的全局注意力——注意力的大部分优势（精确检索、上下文学习）可以通过少量注意力层捕获，其余层由 Mamba 层更高效地处理。这类似于 Swin Transformer 在大多数层使用局部注意力，仅偶尔计算跨窗口交互：并非每层都需要最昂贵的机制。混合方法也使基于 Mamba 的模型在上下文学习基准测试上更具竞争力——目前纯 Mamba 模型在这方面落后于纯 Transformer。

Q19 [进阶] Mamba 的上下文学习能力如何？与 Transformer 相比怎样？

Q： Mamba 能否进行上下文学习（ICL）？其 ICL 能力与 Transformer 相比如何？

A： 上下文学习（In-Context Learning，ICL）是指模型仅凭提示中的少量输入-输出示例就能完成新任务，而无需任何权重更新。这一能力被广泛认为是大型 Transformer 语言模型的标志性特征，对实际部署至关重要。

Mamba 确实表现出上下文学习能力——它可以利用提示中的示例改善在新任务上的输出。然而，研究一致表明，随着上下文示例数量增加，Mamba 的 ICL 能力退化速度比 Transformer 更快。提出的解释具有机制性：Transformer 注意力可以从 KV 缓存中以完整精度检索任何过去示例，使其直接将早期示例的模式应用于当前查询。Mamba 必须将所有示例压缩进固定大小的状态，示例越多，状态越成为有损摘要——早期示例更可能被后来的示例覆盖或稀释。

这一局限具有实践意义：基于 Mamba 的模型可能需要更大的参数量或专门的训练才能达到 Transformer 的 ICL 性能。含少量注意力层的混合架构可以弥补大部分 ICL 差距，因为注意力层维护了上下文示例的精确记录，而 Mamba 层处理周围上下文的高效压缩。因此，Jamba 等混合模型在等效参数量下表现出比纯 Mamba 模型强得多的 ICL 性能。

Q20 [进阶] Mamba 目前有哪些局限性？未来研究方向是什么？

Q： Mamba 的主要开放问题和局限性是什么？该领域走向何方？

A： Mamba 最根本的局限是其固定大小状态作为历史有损压缩器的本质。这表现为在联想回忆、大海捞针检索和长上下文多跳推理任务上性能较弱——这些任务中 Transformer 的精确 KV 缓存具有结构优势。这不是会随更多数据自动解决的训练或规模问题，而是有限状态递推的架构属性。

第二个局限是生态系统的相对不成熟。截至 2024 年初，最大的可用 Mamba 预训练模型在 30 亿参数量级，而最先进的 Transformer 模型已达 700 亿至 7,000 亿参数。Mamba 在极大模型规模和训练预算下的缩放行为尚未得到充分表征。目前在 30 亿参数下观察到的 Mamba 与 Transformer 困惑度差距在 700 亿以上参数时是扩大还是缩小，尚不明确。

在研究前沿，若干方向正在活跃探索。混合架构（结合 Mamba 和注意力）正显示出最直接的实践成果。将 Mamba 扩展到二维及更高维数据（视觉、视频、图）需要对一维扫描结构进行非平凡的适应。Mamba2 引入的 SSD 框架与线性注意力建立了理论联系，可能带来进一步的架构改进。Mamba 在非 NVIDIA 硬件（TPU、AMD GPU）上的高效内核仍不成熟。最后，SSM 表达能力的理论研究——作为状态大小和深度的函数，它们能表示哪些函数、不能表示哪些——是一个活跃领域，将为更好的架构设计提供指导。

快速参考

参考文献

• Gu & Dao, Mamba: Linear-Time Sequence Modeling with Selective State Spaces (2023)
• Dao & Gu, Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Mamba2, 2024)
• Gu et al., Efficiently Modeling Long Sequences with Structured State Spaces (S4, 2021)
• Gu et al., HiPPO: Recurrent Memory with Optimal Polynomial Projections (2020)
• Zhu et al., Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model (2024)
• Lieber et al., Jamba: A Hybrid Transformer-Mamba Language Model (2024)