另辟蹊径--Mamba介绍

当前流行的大模型架构几乎都是基于Transformer，为了改进优化Transformer，研究人员也提出了新的架构，如 Mamba（a State Space Model）

• https://newsletter.maartengrootendorst.com/p/a-visual-guide-to-mamba-and-state

• Paper: Mamba： Linear-Time Sequence Modeling with Selective State Spaces: https://arxiv.org/abs/2312.00752

• implementation & model checkpoints: https://github.com/state-spaces/mamba

Part 1: The Problem with Transformers

Transformers作为目前主流的架构，有很多优点，但也存在一些缺点。

take-aways

• Transformer的优势在于 self attention，由于对attention的计算支持并行计算，结合GPU硬件支持，可以加速模型训练，尤其在处理大规模数据集时具有更高的效率。

• 同样的原因，self attention在推理的时候需要计算整个序列长度，导致推理速度较慢，成为 Transformer 架构的一个主要瓶颈

Transformer 将任何文本 input 视为由tokens组成的序列（a sequence that consists of tokens）。

这样做的好处是无论它接收到什么输入，它都可以回溯查看序列中的任何早期 Token 以得出其表示。

The Core Components of Transformers

Transformer 主要包括：

• 用于表示文本的 encoder blocks

• 用于生成文本的 decoder blocks。

这个encoder-decoder结构可用于多项任务，如翻译；其中encoder-only的代表如 BERT 系列模型。

decoder-only 往往用于生成模型。比如基于 Transformer 的模型 Generative Pre-trained Transformers （GPT） 使用解码器模块来complete input text

A Blessing with Training

single decoder block由两个主要组件组成：

• 掩码自注意力（masked self-attention）一般也用MMHA（masked multi head attention）

• 前馈神经网络（FNN，feed-forward neural network）。

Self-attention 是 Transformer "work well"的主要原因， 它可以通过快速训练实现整个序列的未压缩视图.

how does Self-attention work?

它创建一个attention矩阵，将每个token与之前的每个token进行比较。矩阵中的权重由token对彼此之间的相关性决定。

• 在模型训练期间，此矩阵是一次性创建的。在计算 “name” 和 “is” 之间的 attention 之前，不需要先计算 “my” 和 “name” 之间的关注度。

★

attention计算支持并行化，从而极大地加快了训练速度！

Transformer其他优点

• Transformer中的自注意力机制可以直接捕捉任意两个位置之间的依赖关系，不受序列长度的影响，从而更好地处理长文本序列，方便对全部上下文建模。

• Transformer架构可以处理任意长度的序列数据，无需进行填充或截断操作，从而避免了由此产生的信息损失。

• 模块化设计使得Transformer更易于扩展和调整，可以根据具体任务需求进行灵活的配置。

• Transformer架构的设计非常灵活，可以通过增加层数、增加头数、调整嵌入层维度等方式来扩展模型的能力。此外，Transformer还可以与其他深度学习模型进行结合，形成更强大的复合模型。

And the Curse with Inference!

这种架构在推理的时候存在一个缺陷。在生成下一个 Token 时，我们需要重新计算整个序列的 attention ，即使我们已经生成了一些 token。

为长度为 L 的序列生成标记大约需要 L² 计算，如果序列长度增加，这可能会很昂贵。

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

如何解决slow inference的问题？

Are RNNs a Solution?

递归神经网络 （RNN） 是一种基于序列的网络。序列中的每个时间步都需要两个输入，即时间步 t 的输入和前一个时间步 t-1 的隐藏状态，以生成下一个隐藏状态并预测输出。

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

在生成 output 时，RNN 只需要考虑之前的 hidden state 和当前 input。它可以防止重新计算所有以前的隐藏状态。

也就是说，RNN 可以快速进行推理，因为它与序列长度呈线性关系缩放！理论上，它甚至可以具有无限的上下文长度。

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

比如在计算最后一个输出“Maarten”的时候，只需要有上一个time setp的hidden state和当前输入“is”即可进行计算。

但是RNN有个问题，“信息遗忘”，比如上图中，上一个time step的hidden state中对于第一个input “Hello”的信息块（红色）已经丢失，随着时间的推移，RNN 往往会忘记以前信息，因为它们只考虑一个以前的状态。

尽管 RNN 对于训练和推理来说都很快，但它们缺乏 Transformer 模型所能提供的准确性。

我们着眼于状态空间模型（State Space Models），以有效地使用 RNN（有时使用卷积）。

Part 2: The State Space Model (SSM)

状态空间模型 （SSM） 与 Transformer 和 RNN 一样，处理信息序列，如文本和信号。在本节中，我们将介绍 SSM 的基础知识以及它们与文本数据的关系。

What is a State Space?

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

想象一下，我们正在穿过一个迷宫。“状态空间” 是所有可能的位置（状态）的地图。每个点代表迷宫中的一个独特位置，具有特定的详细信息，例如您离出口有多远。

“state space representation” 是此映射的简化描述。它显示你在哪里 （当前状态），你可以去哪里 （可能的未来状态），以及哪些变化会把你带到下一个状态 （向右或向左）。

尽管状态空间模型使用方程和矩阵来跟踪此行为，但它只是一种跟踪您所在位置、可以到达何处以及如何到达那里的方法。

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

语言模型中的嵌入或向量也经常用于描述输入序列的 “状态”。

例如，您当前位置的 vector （state vector） 可能看起来有点像这样

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

What is a State Space Model?

SSM 是用于描述这些状态表示（state representations）的模型，并根据某些输入预测它们的下一个状态可能是什么。

• at time t, SSMs

• 映射输入序列 x（t） — （例如，在迷宫中向左和向下移动）

• 到潜在状态表示 h（t） — （例如，到 Exit 的距离和 x/y 坐标）

• 并推导出预测的输出序列 y（t） — （例如，再次向左移动以更快地到达出口）

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

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

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

它的目标是找到这个状态表示 h（t），这样我们就可以从 input 转到 output 序列。

这两个方程是 State Space Model 的核心。

state equation

describes how the state changes (through matrix A) based on how the input influences the state (through matrix B)（状态方程描述了 state如何变化（通过矩阵 A）基于input 通过矩阵 B 影响state）。

正如我们之前看到的，h（t） 指的是我们在任何给定时间 t 的潜伏状态表示，而 x（t） 指的是一些输入。

output equation

describes how the state is translated to the output (through matrix C) and how the input influences the output (through matrix D).（输出方程描述了状态如何转换为输出（通过矩阵 C）以及输入如何影响输出（通过矩阵 D）。）

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

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

具体步骤

1. 假设我们有一些输入信号 x（t），这个信号首先乘以矩阵 B，它描述了输入如何影响系统。

2. 更新状态（类似于神经网络的隐藏状态）是一个包含环境核心“知识”的潜在空间。我们将状态与矩阵 A 相乘，矩阵 A 描述了所有内部状态是如何连接的，因为它们代表了系统的基本动态。

• 注意，矩阵 A 在创建 state representations之前应用，并在 state representations更新之后更新。

3. 使用矩阵 C 来描述如何将状态转换为输出。

4. 可以使用矩阵 D 提供从输入到输出的直接信号。这通常也称为 skip-connection。

• 由于矩阵 D 类似于 skip-connection，因此 SSM 通常被视为the following without the skip-connection.

回到我们的简化视角，我们现在可以关注矩阵 A、B 和 C 作为 SSM 的核心。

我们可以像以前一样更新原始方程（并添加一些漂亮的颜色）来表示每个矩阵的用途。

这两个方程总之旨在从观察到的数据中预测系统的状态。由于输入预期是连续的，因此 SSM 的主要表示形式是连续时间表示形式。

From a Continuous to a Discrete Signal（从连续信号到离散信号）

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

为此，我们使用 Zero-order hold 技术。它的工作原理为通过Zero-order hold将离散输入转化为连续信号，通过SSM对连续信号进行处理后，对于输出的连续信号基于同步长进行离散采样，将连续信号转换为离散信号。

具体如下：

1. 首先，每次我们接收到一个离散信号时，我们都会保持它的值，直到我们接收到一个新的离散信号。此过程会创建一个 SSM 可以使用的连续信号：

2. 我们保持该值的时间由一个新的可学习参数表示，称为步长∆（step size）。它表示输入的分辨率（the resolution of the input）。

3. 现在我们有一个 continuous 信号作为我们的输入，我们可以生成一个 continuous output，并且只根据 input 的时间步长对值进行再次采样。

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

Zero-order hold实现

在数学上，我们可以按如下方式应用 Zero-order hold：

它们共同使我们能够从连续的 SSM 转变为离散的 SSM，该公式不是函数到函数（function-to-function） x（t） → y（t），现在是序列到序列（equence-to-sequence） xk → yk：

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

我们使用 k 而不是 t 来表示离散化的时间步长，并在我们提到连续 SSM 与离散 SSM 时更加清晰。

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

现在我们有了离散表示的公式，让我们探索一下如何实际计算模型。

The Recurrent Representation

我们的离散 SSM 允许我们以特定的时间步长而不是连续信号来表述问题。正如我们之前看到的 RNN 一样，recurrent approach在这里非常有用。

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

在每个时间步长，我们计算输入 （Bxk） 如何影响先前的状态 （Ahk₋₁），然后计算预测输出 （Chk）。

我们可以像以前看到的一样处理 RNN。

• 使用RNN结构来替换SSM模块

我们可以如下展开：

注意我们如何使用 RNN 的底层方法来使用这个离散化版本。

The Convolution Representation

用于 SSM 的另一种表示形式是卷积的表示形式。

在经典的图像识别任务中，我们应用了（filters (kernels) to derive aggregate features）过滤器（内核）来派生聚合特征：

由于我们处理的是文本而不是图像，因此我们需要一维透视：

我们用来表示此 “filter” 的内核源自 SSM 公式：

让我们来探索一下这个内核在实践中是如何工作的。与卷积一样，我们可以使用 SSM 内核来检查每组 tokens 并计算输出：

这也说明了 padding 可能对输出产生的影响。我更改了填充的顺序以改善可视化效果，但我们经常将其应用于句子末尾。

在下一步中，将 kernel 移动一次以执行计算的下一步：

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

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

The Three Representations

这三种表示形式（连续、循环和卷积）都有不同的优点和缺点：

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

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

此模型称为线性状态空间层 （LSSL，Linear State-Space Layer）：https://proceedings.neurips.cc/paper_files/paper/2021/hash/05546b0e38ab9175cd905eebcc6ebb76-Abstract.html

这些表示都有一个重要的特性，即线性时间不变性 （LTI）。LTI 指出 SSM 参数 A、B 和 C 对于所有时间步长都是固定的。这意味着矩阵 A、B 和 C 对于 SSM 生成的每个标记都是same的。

换句话说，无论您为 SSM 提供什么序列，A、B 和 C 的值都保持不变。我们有一个不感知内容的静态表示。

在我们探索 Mamba 如何解决这个问题之前，让我们先探索拼图的最后一块，矩阵 A

The Importance of Matrix A

SSM 公式最重要的方面之一是矩阵 A。正如我们之前看到的 recurrent 表示，它捕获有关先前 state 的信息以构建新 state。

实质上，矩阵 A 产生隐藏状态：

因此，创建矩阵 A 可能是只记住几个先前的标记和捕获我们目前看到的每个标记之间的区别。尤其是在 Recurrent 表示的上下文中，因为它只回顾之前的状态。

那么我们如何以保留大内存（上下文大小）的方式创建矩阵 A呢？

我们使用 Hungry Hungry Hippo！或 HiPPO for High-order Polynomial Projection Operators：https://proceedings.neurips.cc/paper/2020/hash/102f0bb6efb3a6128a3c750dd16729be-Abstract.html

HiPPO 尝试将迄今为止看到的所有输入信号压缩为系数向量。

它使用矩阵 A 构建一个状态表示形式，该表示形式可以很好地捕获最近的token并衰减较旧的token。其公式可以表示如下：

假设我们有一个方阵 A，这得到：

使用 HiPPO 构建矩阵 A 比将其初始化为随机矩阵要好得多。因此，与旧信号 （初始标记） 相比，它更准确地重建了较新的信号 （最近的标记）

HiPPO 矩阵背后的想法是，它会产生一个隐藏状态来记住其历史记录。

在数学上，它通过跟踪Legendre polynomial（勒让德多项式：https://proceedings.neurips.cc/paper/2019/hash/952285b9b7e7a1be5aa7849f32ffff05-Abstract.html）的系数来实现这一点，这使得它能够近似所有以前的历史

然后将 HiPPO 应用于我们之前看到的递归和卷积表示，以处理长距离依赖性。结果是 Structured State Space for Sequences (S4)（序列的结构化状态空间 ：https://arxiv.org/abs/2111.00396），这是一类可以有效处理长序列的 SSM

它由三个部分组成：

• State Space Models 状态空间模型

• HiPPO for handling long-range dependencies 用于处理长距离依赖的 HiPPO

• Discretization for creating recurrent and convolution representations 用于创建循环和卷积表示的离散化

此类 SSM 具有多种优势，具体取决于您选择的表示形式（递归与卷积）。它还可以通过构建 HiPPO 矩阵来处理长文本序列并有效地存储内存。

注意：如果您想深入了解有关如何计算 HiPPO 矩阵并自己构建 S4 模型的更多技术细节，https://srush.github.io/annotated-s4/

Part 3: Mamba - A Selective SSM

我们终于涵盖了了解 Mamba 的特别之处所需的所有基础知识。状态空间模型 可以用来对文本序列进行建模，但仍然有一系列我们想要防止的缺点。

本节中，我们将介绍 Mamba 的两个主要贡献：

1. A selective scan algorithm, which allows the model to filter (ir)relevant information 一种选择性扫描算法，允许模型过滤（非）相关信息

2. A hardware-aware algorithm that allows for efficient storage of (intermediate) results through parallel scan, kernel fusion, and recomputation.一种硬件感知算法，允许通过并行扫描、内核融合和重新计算来高效存储（中间）结果。

它们一起创建了 selective SSM or S6 models，这些模型可以像self attention一样用于创建 Mamba 块。

在探讨这两个主要贡献之前，让我们先探讨一下为什么它们是必要的。

What Problem does it attempt to Solve?

State Space Models, and even the S4 (Structured State Space Model)在某些对语言建模和生成至关重要的任务上表现不佳，即关注或忽略特定输入的能力（ability to focus on or ignore particular inputs）。

我们可以用两个合成任务来说明这一点，即选择性复制和归纳头（elective copying and induction heads）。

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

但是，（递归/卷积）SSM 在此任务中表现不佳，因为它是线性时间不变的。 正如我们之前看到的，SSM 生成的每个标记的矩阵 A、B 和 C 都是相同的。

因此，SSM 无法执行内容感知推理，因为它将每个标记视为固定 A、B 和 C 矩阵的结果。这是一个问题，因为我们希望 SSM 对输入 （提示） 进行推理。

SSM 执行不佳的第二个任务是归纳头，其目标是重现输入中的模式：

在上面的例子中，我们基本上是在执行一次性提示，我们试图 “教” 模型在每个 “Q：” 之后提供 “A：” 响应。但是，由于 SSM 是时间不变（ time-invariant）的，因此它无法从其历史记录中选择要调用哪些以前的标记。

让我们通过关注矩阵 B 来说明这一点。无论输入 x 是什么，矩阵 B 都保持完全相同，因此独立于 x：

同样，无论输入如何，A 和 C 也保持固定。这证明了我们迄今为止看到的 SSM 的静态性质。

相比之下，这些任务对 Transformer 来说相对容易，因为它们会根据输入序列动态地改变注意力。他们可以有选择地 “查看” 或 “参与” 序列的不同部分。

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

Selectively Retain Information 选择性地保留信息

SSM 的递归表示创建了一个非常有效的小状态，因为它压缩了整个历史记录。然而，与不压缩历史记录（通过注意力矩阵）的 Transformer 模型相比，它的功能要弱得多。

Mamba 的目标是两全其美。一个与 Transformer 的状态一样强大的小状态：

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

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

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

相反，Mamba 通过合并输入的序列长度和批量大小，使矩阵 B 和 C，甚至步长 ∆，都取决于输入：

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

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

它们一起有选择地选择要保留在隐藏状态中的内容和要忽略的内容，因为它们现在依赖于 Input。

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

The Scan Operation 扫描操作

由于这些矩阵现在是动态的，因此无法使用卷积表示来计算它们，因为它假定一个固定的内核。我们只能使用递归表示，而失去了卷积提供的并行化。

为了实现并行化，我们来探索如何使用重新并发来计算输出：

每个状态是前一个状态（乘以 A）加上当前输入（乘以 B）的总和。这称为 scan 操作，可以使用 for 循环轻松计算。

相比之下，并行化似乎是不可能的，因为只有当我们拥有前一个状态时，才能计算每个状态。但是，Mamba 通过并行扫描算法实现了这一点。

• parallel scan algorithm：https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda

它假设我们执行操作的顺序与 associate 属性无关紧要。因此，我们可以分部分计算序列并迭代组合它们：

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

Hardware-aware Algorithm 硬件感知算法

最近 GPU 的一个缺点是它们在小型但高效的 SRAM 和大型但效率稍低的 DRAM 之间的传输 （IO） 速度有限。频繁地在 SRAM 和 DRAM 之间复制信息成为瓶颈。–内存墙问题

Mamba 和 Flash Attention 一样，试图限制我们需要从 DRAM 到 SRAM 的次数，反之亦然。它通过内核融合来实现这一点，这允许模型防止写入中间结果并持续执行计算，直到完成。

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

在这里，以下内容融合到一个内核中：

• Discretization step with step size ∆步长为 ∆ 的离散化步骤

• Selective scan algorithm 选择性扫描算法

• Multiplication with C 使用 C 的乘法

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

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

虽然这看起来效率低下，但比从相对较慢的 DRAM 中读取所有这些中间状态的成本要低得多

我们现在已经介绍了其架构的所有组件，使用其文章中的下图进行了描述：

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

The Mamba Block

到目前为止，我们探索的selective SSM 可以实现为一个块，就像我们可以在解码器块中表示自我注意一样。

与解码器一样，我们可以堆叠多个 Mamba 块，并将其输出用作下一个 Mamba 块的输入：

它从线性投影开始，以扩展输入嵌入。然后，在 Selective SSM 之前应用卷积以防止独立标记计算。

The Selective SSM has the following properties:

• Recurrent SSM created through discretization 通过离散化创建的递归 SSM

• HiPPO initialization on matrix A to capture long-range dependencies 在矩阵 A 上进行 HiPPO 初始化以捕获长距离依赖关系

• Selective scan algorithm to selectively compress information扫描算法，选择性压缩信息

• Hardware-aware algorithm to speed up computation 硬件感知算法可加快计算速度

在查看代码实现时，我们可以进一步扩展此体系结构，并探索端到端示例的外观：

请注意一些更改，例如包含normalization layers and softmax for choosing the output token.

当我们把所有东西放在一起时，我们既可以得到快速推理和训练，也可以得到无限的上下文。使用这种架构，作者发现它与相同大小的 Transformer 模型的性能相当，有时甚至超过了它们！