回顾Transformer，并深入讲解替代方案Mamba原理（图解）

一种语言建模中 Transformer 的替代方案

Transformer 架构是大语言模型（LLMs）成功的关键组成部分。几乎所有今天使用的大语言模型都采用了该架构，从开源模型如 Mistral 到闭源模型如 ChatGPT。

为了进一步改进大语言模型，新的架构被开发出来，这些架构可能甚至会超越 Transformer 架构。其中一种方法是 Mamba，一种 状态空间模型。

Mamba 在论文 Mamba: Linear-Time Sequence Modeling with Selective State Spaces1 中被提出。你可以在其 仓库 中找到其官方实现和模型检查点。

在这篇文章中，我将介绍在语言建模背景下的状态空间模型领域，并逐步探索各个概念，以帮助理解这一领域。然后，我们将讨论 Mamba 如何可能挑战 Transformer 架构。

作为一个视觉指南，本文将通过许多可视化内容来帮助理解 Mamba 和状态空间模型！

第 1 部分：Transformers 的问题

为了说明 Mamba 是多么有趣的架构，我们首先简要回顾一下 Transformers，并探索其一个缺点。

Transformer 将任何文本输入视为由 tokens 组成的 sequence。

 

Transformers 的一个主要优点是，无论它接收到什么输入，它都可以回溯到序列中任何早期的 tokens 来推导其表示。

 

 

Transformers 的核心组件

记住，Transformer 由两个结构组成，一个用于表示文本的编码器块集和一个用于生成文本的解码器块集。结合起来，这些结构可用于多个任务，包括翻译。

 

我们可以采用这种结构，通过仅使用解码器来创建生成模型。这个基于 Transformer 的模型，生成式预训练 Transformer（GPT），使用解码器块来完成一些输入文本。

 

让我们看看这是如何工作的！

训练中的祝福…

单个解码器块由两个主要组件组成，遮蔽的自注意力机制和前馈神经网络。

 

自注意力机制是这些模型工作得如此出色的一个重要原因。它使整个序列的未压缩视图得以实现，并且训练速度快。

那么它是如何工作的呢？

它创建了一个矩阵，将每个 Token 与之前的每个 Token 进行比较。矩阵中的权重取决于 Token 对彼此的相关性。

 

在训练过程中，这个矩阵一次性创建。“My”和“name”之间的注意力不需要先计算，才能计算“name”和“is”之间的注意力。

它实现了并行化，这极大地加快了训练速度！

推理中的问题！

然而存在一个缺陷。当生成下一个 Token 时，我们需要重新计算整个序列的注意力，即使我们已经生成了一些 Token。

 

生成长度为L的序列需要大约L²次计算，如果序列长度增加，计算代价会很高。

 

这种需要重新计算整个序列的情况是 Transformer 架构的主要瓶颈之一。

让我们看看一种“经典”技术——循环神经网络（Recurrent Neural Networks, RNN）是如何解决这个推理速度慢的问题。

RNN 是解决方案吗？

循环神经网络（RNN）是一种基于序列的网络。它在每个时间步中接收两个输入，即时间步t的输入和上一个时间步t-1的隐藏状态，用以生成下一个隐藏状态并预测输出。

RNN 有一个循环机制，允许它将信息从前一步传递到下一步。我们可以将这个可视化过程“展开”，使其更加明确。

 

在生成输出时，RNN 只需要考虑前一个隐藏状态和当前输入。它避免了 Transformer 所需的重新计算所有之前隐藏状态的问题。

换句话说，RNN 能够快速进行推理，因为它随着序列长度线性扩展！理论上，它甚至可以拥有无限的上下文长度。

为了说明这一点，让我们将 RNN 应用于我们之前使用的输入文本。

 

每个隐藏状态都是所有之前隐藏状态的聚合，通常是一个压缩的视图。

然而，这里有一个问题……

注意，当生成名字“Maarten”时，最后一个隐藏状态已经不再包含关于单词“Hello”的信息了。RNN 随时间推移往往会忘记信息，因为它们只考虑上一个状态。

虽然 RNN 在训练和推理方面速度较快，但它们缺乏 Transformer 模型所能提供的精度。

因此，我们研究状态空间模型 (State Space Models) 来高效地使用 RNN（有时还会使用卷积）。

第 2 部分：状态空间模型 (SSM)

状态空间模型 (SSM) 像 Transformer 和 RNN 一样处理信息序列，比如文本和信号。在本节中，我们将介绍 SSM 的基本概念以及它们与文本数据的关系。

什么是状态空间？

状态空间包含完全描述一个系统所需的最少变量数。它是一种通过定义系统可能状态来数学表示问题的方法。

让我们简化一下。想象我们正在穿越迷宫。“状态空间”就是所有可能位置（状态）的地图。每个点代表迷宫中的一个独特位置，并带有特定的细节，比如你距离出口有多远。

“状态空间表示”是对这张地图的简化描述。它显示了你当前的位置（当前状态）、你可以前往的下一个位置（可能的未来状态）、以及你如何到达下一个位置（向右或向左移动）。

 

虽然状态空间模型使用方程和矩阵来跟踪这种行为，但它实际上只是跟踪你在哪里、你可以去哪里、以及如何到达那里的一种方式。

变量描述了一个状态，在我们的例子中，X 和 Y 坐标，以及到出口的距离，可以表示为“状态向量”。

 

听起来很熟悉吗？这是因为在语言模型中，嵌入或向量也经常用于描述输入序列的“状态”。例如，描述你当前位置的向量（状态向量）可能类似于如下所示：

 

在神经网络中，“状态”通常指的是系统的隐藏状态，在大语言模型中，这是生成新 token 时最重要的方面之一。

什么是状态空间模型？

SSM（状态空间模型）是一类用于描述这些状态表示并预测其下一个状态的模型，预测基于某些输入。

传统上，在时间 t，SSM：

• 将输入序列 x(t)（例如，在迷宫中向左和向下移动）映射到潜在状态表示 h(t)（例如，到出口的距离和 x/y 坐标）
• 并推导出预测的输出序列 y(t)（例如，再次向左移动以更快到达出口）

然而，SSM 并不是使用离散序列（如向左移动一次），而是接受连续序列作为输入并预测输出序列。

 

SSM 假设动态系统（如在 3D 空间中移动的物体）可以通过时间 t 的状态和两个方程来预测。

 

通过求解这些方程，我们假设可以根据观测到的数据（输入序列和先前状态）揭示预测系统状态的统计原则。

其目标是找到这个状态表示 h(t)，使我们能够从输入到输出序列。

 

这两个方程是状态空间模型的核心。

这两个方程将在整个指南中引用。为了让它们更直观，我们使用颜色编码，这样你可以快速引用它们。

状态方程 描述了状态如何根据输入（通过 矩阵 B）影响状态（通过 矩阵 A）而变化。

 

如前所述，h(t) 指的是任意给定时间 t 的潜在状态表示，x(t) 指的是某个输入。

输出方程 描述了状态如何通过 矩阵 C 转换为输出，以及输入如何通过 矩阵 D 影响输出。

 

注意：矩阵 A、B、C 和 D 也常被称为 参数，因为它们是可学习的。

将这两个方程可视化，我们得到以下架构：

 

让我们一步步了解这些矩阵如何影响学习过程。

假设我们有一些输入信号 x(t)，该信号首先与 矩阵 B 相乘，矩阵 B 描述了输入如何影响系统。

 

更新后的状态（类似于神经网络的隐藏状态）是一个潜在空间，其中包含了环境的核心“知识”。我们将状态与 矩阵 A 相乘，该矩阵描述了所有内部状态之间的关联，表示系统的基本动态。

 

你可能已经注意到，矩阵 A 在创建状态表示之前应用，并在状态表示更新后再次应用。

然后，我们使用 矩阵 C 来描述状态如何转换为输出。

 

最后，我们可以使用 矩阵 D 提供从输入到输出的直接信号。这通常也被称为 跳跃连接。

 

由于 矩阵 D 类似于跳跃连接，SSM 通常被认为是不包含跳跃连接的以下形式。

 

回到我们的简化视角，我们现在可以将重点放在矩阵 A、B 和 C 上，作为 SSM 的核心。

 

我们可以更新原始方程（并添加一些漂亮的颜色）来标明每个矩阵的用途，就像我们之前做的那样。

 

这两组方程一起旨在从观察数据中预测系统的状态。由于输入是连续的，SSM 的主要表示形式是 连续时间表示。

从连续信号到离散信号

如果你有一个连续信号，找到状态表示 h(t) 在解析上具有挑战性。此外，由于我们通常有离散输入（例如文本序列），我们希望对模型进行离散化。

为此，我们使用了 零阶保持技术。其工作原理如下：首先，每当我们接收到离散信号时，我们保持其值，直到我们接收到新的离散信号为止。这个过程会生成一个连续信号，供 SSM 使用：

 

我们保持数值的时间由一个可学习的新参数 步长 ∆ 表示。它表示输入的分辨率。

现在我们为输入生成了一个连续信号，接下来我们可以生成连续输出，并根据输入的时间步长对这些值进行采样。

 

这些采样的值就是我们离散化的输出！

从数学上讲，我们可以如下应用零阶保持：

 

它们共同使我们可以从一个连续的 SSM 过渡到一个离散的 SSM。此时模型不再是 函数对函数 x(t) → y(t)，而是 序列对序列 xₖ → y_ₖ：

 

这里，矩阵 A 和 B 现在表示模型的离散化参数。

我们用 k 来代替 t，以区分我们何时在谈论连续的 SSM 与离散的 SSM。

注意： 在训练期间，我们仍然保留 矩阵 A 的连续形式，而不是其离散化版本。在训练过程中，连续表示会被离散化。

现在我们已经有了一个离散化表示的公式，让我们探讨如何实际 计算 该模型。

递归表示

我们的离散化 SSM 使我们能够在特定的时间步中，而不是连续信号中，构建问题。正如我们之前在 RNN 中看到的，递归方法在这里非常有用。

如果我们考虑离散的时间步而不是连续信号，我们可以使用时间步重新表述这个问题：

 

在每个时间步中，我们计算当前输入 (Bxₖ) 如何影响前一个状态 (Ahₖ₋₁)，然后计算预测输出 (Chₖ)。

 

这个表示可能已经让你感到熟悉！我们可以像之前对 RNN 的处理方法一样进行分析。

 

我们可以将其展开（或展开成一系列时间步）如下：

 

请注意，我们可以使用 RNN 的基本方法来使用这种离散化的版本。

卷积表示

我们可以使用卷积来表示 SSM。记住在经典图像识别任务中，我们应用滤波器（kernels）来提取聚合特征：

 

由于我们处理的是文本而不是图像，我们需要使用一维视角：

 

我们用来表示这个“滤波器”的内核源自 SSM 公式：

 

让我们来看看这个内核在实践中的作用。像卷积一样，我们可以使用 SSM 内核遍历每组 tokens 并计算输出：

 

这也展示了填充可能对输出产生的影响。我更改了填充的顺序以改善可视化，但我们通常在句子末尾应用填充。

在下一步中，内核会移动一次以执行计算的下一步：

 

在最后一步中，我们可以看到内核的全部效果：

 

将 SSM 表示为卷积的一个主要好处是它可以像卷积神经网络（CNNs）一样并行训练。然而，由于内核大小固定，它们的推理速度不如 RNNs 快且不受限制。

三种表示方式

这三种表示方式——连续、递归 和 卷积 各有不同的优缺点：

 

有趣的是，我们现在可以利用递归 SSM 进行高效推理，同时利用卷积 SSM 进行并行训练。

利用这些表示方式，我们可以使用一个巧妙的技巧，即根据任务选择表示方式。在训练过程中，我们使用可以并行化的卷积表示，而在推理过程中，我们使用高效的递归表示：

 

这个模型被称为 线性状态空间层（LSSL）。 2

这些表示方式共享一个重要属性，即线性时间不变性（LTI）。LTI 表示 SSM 的参数 A、B 和 C 在所有时间步中都是固定的。这意味着矩阵 A、B 和 C 对每个生成的 Token 都是相同的。

换句话说，无论你给 SSM 任何序列，A、B 和 C 的值都保持不变。我们有一个不关心内容的静态表示方式。

在探讨 Mamba 如何解决这个问题之前，让我们探讨最后一个拼图碎片——矩阵 A。

矩阵 A 的重要性

可以说，SSM 公式中最重要的方面之一是 矩阵 A。正如我们之前在递归表示中看到的，它捕捉了有关 上一个 状态的信息，以构建 新的 状态。

 

本质上，矩阵 A 生成隐藏状态：

 

因此，创建 矩阵 A 可能会决定我们能否记住仅仅几个之前的 token 或捕捉到我们迄今为止看到的每一个 token。特别是在递归表示的背景下，因为它只 回顾 上一个状态。

如何以保留大量记忆（上下文大小）的方式创建 matrix A？

我们使用 Hungry Hungry Hippo！或者 HiPPO3 来实现高阶 多项式 投影 运算符。HiPPO 试图将其至今看到的所有输入信号压缩为一个系数向量。

 

它使用 matrix A 来构建一个状态表示，该表示能够较好地捕捉最近的 token，并衰减较旧的 token。其公式可以表示如下：

 

假设我们有一个方阵 matrix A，这给我们提供了：

 

使用 HiPPO 构建 matrix A 被证明比将其初始化为随机矩阵要好得多。因此，它在重建 更新 信号（最近的 token）方面比 较旧 信号（初始 token）更为准确。

HiPPO 矩阵的核心思想是生成一个记忆其历史的隐藏状态。

从数学上讲，它通过跟踪 Legendre 多项式 的系数来实现这一点，这使得它能够近似所有的历史记录。4

然后 HiPPO 被应用于我们之前看到的递归和卷积表示，以处理长距离依赖关系。结果是 Structured State Space for Sequences (S4)，一种可以高效处理长序列的 SSM 类。5

它由三部分组成：

• 状态空间模型
• HiPPO 用于处理 长距离依赖
• 离散化用于创建 递归 和 卷积 表示

 

这类 SSM 具有多个优点，具体取决于你选择的表示方式（递归与卷积）。它还可以处理长文本序列，并通过基于 HiPPO 矩阵来高效地存储记忆。

注意：如果你想深入了解如何计算 HiPPO 矩阵并自己构建 S4 模型，我强烈建议你阅读 Annotated S4。

 

第 3 部分：Mamba – 一种选择性状态空间模型

我们终于覆盖了理解 Mamba 特殊之处所需的所有基础知识。状态空间模型可以用于建模文本序列，但仍然有一系列我们希望避免的缺点。

在这一部分，我们将讨论 Mamba 的两个主要贡献：

1. 一种 选择性扫描算法，允许模型筛选（不）相关信息
2. 一种 硬件感知算法，通过 并行扫描、内核融合 和 重新计算 来高效存储（中间）结果。

这两者共同创建了 选择性状态空间模型 或 S6 模型，可以像自注意力一样用于创建 Mamba 块。

在探索这两个主要贡献之前，让我们首先探讨一下它们为何必要。

试图解决什么问题？

状态空间模型，甚至 S4（结构化状态空间模型），在语言建模和生成中某些关键任务上表现不佳，即 关注或忽略特定输入的能力。

我们可以用两个合成任务来说明这一点，即 选择性复制 和 诱导头。

在 选择性复制 任务中，SSM 的目标是复制输入的部分并按顺序输出：

 

然而，由于 SSM 是 线性时间不变的，它在这个任务上表现不佳。正如我们之前看到的，矩阵 A、B 和 C 对于 SSM 生成的每个 token 都是相同的。

因此，SSM 无法执行 内容感知推理，因为它由于固定的 A、B 和 C 矩阵而平等对待每个 token。这是一个问题，因为我们希望 SSM 对输入（提示）进行推理。

SSM 在另一个任务上表现不佳，即 诱导头，其目标是重现输入中找到的模式：

 

在上面的示例中，我们本质上是在进行一次性提示，其中我们尝试“教”模型在每个 “Q:” 后提供一个 “A:” 响应。然而，由于 SSM 是时间不变的，它无法选择从历史中召回哪些先前的 tokens。

让我们通过关注 矩阵 B 来说明这一点。无论输入 x 是什么，矩阵 B 始终保持不变，因此与 x 无关：

 

同样，A 和 C 也始终保持固定，与输入无关。这表明我们迄今看到的 SSM 的 静态 特性。

 

相比之下，这些任务对于 Transformer 来说相对简单，因为它们根据输入序列 动态 改变注意力。它们可以选择性地“观察”或“关注”序列的不同部分。

SSM 在这些任务上的较差表现说明了时间不变 SSM 的潜在问题，矩阵 A、B 和 C 的静态特性导致了 内容感知 的问题。

选择性保留信息

SSM 的递归表示创建了一个较小的状态，这种状态非常高效，因为它压缩了整个历史。然而，与 Transformer 模型相比，Transformer 模型不会对历史进行压缩（通过注意力矩阵），因此它的能力更强。

Mamba 旨在兼具两者的优势。一个与 Transformer 状态一样强大的小状态：

 

如上所述，它通过选择性地将数据压缩到状态中来实现。当你有一个输入句子时，通常会有一些信息，例如停用词，没有太多意义。

为了选择性地压缩信息，我们需要参数依赖于输入。为此，让我们首先探索 SSM 在训练过程中输入和输出的维度：

 

在结构化状态空间模型（S4）中，矩阵 A、B 和 C 与输入无关，因为它们的维度 N 和 D 是静态的，不会改变。

 

相反，Mamba 通过结合输入的序列长度和批量大小，使矩阵 B 和 C，甚至 步长 ∆_，依赖于输入：

 

这意味着对于每个输入 token，我们现在有不同的 B 和 C 矩阵，这解决了内容感知的问题！

注意：矩阵 A 保持不变，因为我们希望状态本身保持静态，但它的影响方式（通过 B 和 C）是动态的。

它们共同 选择性地 选择保留在隐藏状态中的内容和忽略的内容，因为它们现在依赖于输入。

较小的 步长 ∆ 导致忽略特定的词汇，而更多地使用之前的上下文，而较大的 步长 ∆ 则更加关注输入词汇，而不是上下文：

 

扫描操作

由于这些矩阵现在是 动态 的，它们不能使用卷积表示进行计算，因为卷积表示假设一个 固定 的卷积核。我们只能使用递归表示，这样就失去了卷积提供的并行化优势。

为了实现并行化，让我们探索如何使用递归计算输出：

 

每个状态都是前一个状态（乘以 A）与当前输入（乘以 B）的和。这被称为 扫描操作，可以通过 for 循环轻松计算。

相比之下，并行化似乎是不可能的，因为每个状态只能在有了前一个状态之后才可以计算。然而，Mamba 通过 并行扫描 算法使这一点成为可能。

它假设我们进行操作的顺序不重要，利用了结合律的属性。因此，我们可以将序列分成若干部分进行计算，然后迭代地将它们组合起来：

 

动态矩阵 B 和 C 以及并行扫描算法共同创造了 选择性扫描算法，以表示使用递归表示的动态和快速特性。

硬件感知算法

最近的 GPU 的一个缺点是它们的小型但高效的 SRAM 与大型但稍微不那么高效的 DRAM 之间的传输（IO）速度有限。频繁地在 SRAM 和 DRAM 之间复制信息会成为瓶颈。

 

Mamba 类似于 Flash Attention，试图限制从 DRAM 到 SRAM 及反向的次数。它通过 kernel fusion 实现这一点，允许模型防止写入中间结果，并持续执行计算直到完成。

 

我们可以通过可视化 Mamba 的基本架构来查看 DRAM 和 SRAM 分配的具体实例：

 

在这里，以下内容被融合成一个内核：

• 离散化步骤与 step size ∆
• 选择性扫描算法
• 与 C 的乘法

硬件感知算法的最后一部分是 recomputation。

中间状态不会被保存，但在反向传递中计算梯度时是必要的。相反，作者在反向传递期间重新计算这些中间状态。

尽管这看起来可能效率低下，但比从相对较慢的 DRAM 中读取所有这些中间状态要便宜得多。

我们现在已经涵盖了其架构的所有组件，该架构在其文章中的图像如下所示：

 

选择性 SSM。 取自：Gu, Albert 和 Tri Dao. “Mamba: 具有选择性状态空间的线性时间序列建模。” arXiv 预印本 arXiv:2312.00752 (2023)。

这个架构通常被称为 选择性 SSM 或 S6 模型，因为它本质上是使用选择性扫描算法计算的 S4 模型。

Mamba 模块

我们迄今为止探索的 选择性 SSM 可以作为一个模块进行实现，就像我们可以在解码器模块中表示自注意力一样。

 

像解码器一样，我们可以堆叠多个 Mamba 模块，并将它们的输出作为下一个 Mamba 模块的输入：

 

它以线性投影开始，以扩展输入嵌入。然后，在 选择性 SSM 之前应用卷积，以防止独立的 token 计算。

选择性 SSM 具有以下特性：

• 通过 离散化 创建的 递归 SSM
• 在矩阵 A 上进行 HiPPO 初始化，以捕捉 长程依赖性
• 选择性扫描算法 以选择性地压缩信息
• 硬件感知算法 以加快计算速度

当我们查看代码实现时，我们可以对这个架构进行更多扩展，并探索一个端到端的示例会是什么样的：

 

注意一些变化，例如增加了归一化层和用于选择输出 token 的 softmax。

当我们将所有内容整合在一起时，我们得到了快速的推理和训练，甚至是无限上下文。使用这种架构，作者发现其性能与同等大小的 Transformer 模型相匹配，有时甚至超过！

结论

这就结束了我们对状态空间模型和令人难以置信的 Mamba 架构使用选择性状态空间模型的探索。希望这篇文章能让你更好地理解状态空间模型，特别是 Mamba。谁知道这是否会取代 Transformers，但现在，看到如此不同的架构获得应有的关注真是令人惊叹！

要查看更多与大语言模型相关的可视化内容并支持这份通讯，请查看我与 Jay Alammar 合著的书籍。