Mamba In-1.Mamba介绍

Mamba In

RNN、LSTM的训练是很漫长的过程，因为它们：

• 训练过程通常需要通过时间展开（时间步骤）进行，因为它们具有记忆功能，每个时间步骤的计算依赖于前一个时间步骤的输出。
• 这种依赖关系导致了训练过程中的序列化，即必须按照时间步骤一个接一个地处理数据。

但它们的推理是很快速的

而2017年Transformer的提出，其训练速度是很快的，因为其加入了并行式训练，而推理是慢的。

Part 1: Transformers的缺点

为了说明 Mamba 为什么是一个如此有趣的架构，让我们先对Transformer做一个简短的回顾，并探讨它的一个缺点。

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

Transformer的一个主要优点是，无论它接收到多长的输入，它都使用序列中的任何令牌信息（无论序列有多长)来对输入数据进行处理。这就是Transformer中注意力机制的作用，但是为了获得全局信息，注意力机制在长序列上非常耗费显存。

Transformers的核心部分

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

我们可以采用这种结构来创建仅使用解码器的生成模型。比如基于Transformer的GPT，使用解码器块来完成一些输入文本。

让我们来看看它是如何工作的！

高效的训练过程…

单个解码器块由两个主要部分组成，一个是自注意力模块，另一个是前馈神经网络。

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

自我关注是这些模型如此有效的主要原因。它可以通过快速训练，对整个序列进行无压缩浏览。

那么，它是如何工作的呢？

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

在训练过程中，这个矩阵是一次性创建的。在计算 "name "和 "is "之间的注意力之前，无需先计算 "my "和 "name "之间的注意力。

因此它支持并行化，所以可以极大地加快训练速度!

推理的诅咒

不过，这也有缺陷。但是当生成下一个token时，我们需要重新计算整个序列的注意力，即使我们已经生成了一些新的token。

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

并且在这里需要计算所有令牌的注意力，所以如果序列很长，那么内存占用也会很大。所以需要重新计算整个序列是Transformer体系结构的主要瓶颈。这里暂时不提，只看最经典的原始论文。

让我们看看 "经典 "技术循环神经网络是如何解决推理速度慢的问题的。

RNN是一个解决方案？

下面我们介绍更早的序列模型RNN。循环神经网络(RNN)是一种基于序列的网络。它在序列的每个时间步长取两个输入，即时间步长t的输入和前一个时间步长t-1的隐藏状态，以生成下一个隐藏状态并预测输出。

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

在生成输出时，RNN只需要考虑之前的隐藏状态和当前的输入。这样不会重新计算以前的隐藏状态，这正Transformer没有的。

这种流程可以让RNN进行快速推理，因为的时间与序列长度线性关系！并且可以有无限的上下文长度（理论上），因为每次推理他只取一个隐藏状态和当前输入，内存的占用是非常稳定的。

我们将RNN应用于之前使用过的输入文本。

每个隐藏状态都是以前所有隐藏状态的聚合。但是这里就出现了问题，在生成名称”Maarten“时，最后一个隐藏状态不再包含关于单词“Hello”的信息（或者说最早的信息会被新加入的信息覆盖）。这会导致随着时间的推移，rnn会忘记信息，因为它们只考虑前一个状态。

并且rnn的这种顺序性产生了另一个问题。训练不能并行进行，因为它需要按顺序完成每一步。

与Transformer相比，rnn的问题完全相反！它的推理速度非常快，但不能并行化导致训练很慢。

人们一直在寻找一种既能像Transformer那样并行化训练，能够记住先前的信息，并且在推理时间还是随序列长度线性增长的模型。

有，那就是Mamba，在其后也有人提出了 Vision Mamba

在介绍Mamba之前，让我们还需要介绍以下状态空间模型

Part2：The State Space Model (SSM)

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

什么是状态空间？

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

比如说我们正在通过一个迷宫。 “状态空间” 就是所有可能位置（状态）的地图。每个点代表迷宫中的一个独特位置，具有特定的细节，比如你离出口有多远。

“状态空间表示” 是对这个地图的简化描述。它展示了你当前所处的位置（当前状态）、下一步可能去的地方（未来可能的状态），以及哪些变化会让你进入下一个状态（向右或向左）。

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

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

听起来熟悉吗?这不就是强化学习中的状态吗，我个人认为是可以这么理解的，那么怎么和序列有关呢？

因为语言模型中的嵌入或向量也经常用于描述输入序列的“状态”。例如，你当前位置的向量(状态向量)可能看起来像这样:

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

什么是状态空间模型？

状态空间模型（SSMs）是用于描述这些状态表示并根据某些输入进行下一个状态预测的模型。

在时间t，状态空间模型(SSMs)：

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

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

外链图片转存失败,源站可能有防盗链机制,建议将图片保存下来直接上传

SSMs假设动态系统，例如在三维空间中移动的物体，可以通过两个方程从时间t的状态预测。

通过求解这些方程，假设可以揭示基于观测数据(输入序列和先前状态)预测系统状态的统计原理。

它的目标是找到这个状态表示h(t)，这样我们就可以从输入序列进入输出序列。

这两个方程就是状态空间模型的核心。这两个方程将在本指南中贯穿始终。为了使它们更加直观，我们用颜色对它们进行了编码，以便您能快速参考。状态方程描述了状态如何根据输入对状态的影响（通过矩阵 B）发生变化（通过矩阵 A）。

正如我们之前看到的，**h(t)**指的是我们在任何给定时间t的潜在状态表示，**x(t)**指的是某种输入。

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

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

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

下面我们看看这些矩阵如何影响学习过程。

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

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

这里可以看到，在创建状态表示之前应用矩阵A，并在状态表示更新之后更新矩阵A。

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

最后利用矩阵D提供从输入到输出的直接信号。这通常也被称为跳过（残差）连接。

由于矩阵D类似于跳过连接，所以SSM通常被视为不进行跳过连接的部分

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

更新原始方程并添加一些颜色来表示每个矩阵的目的

这两个方程的共同目的是根据观测数据预测系统状态。由于输入预期是连续的，因此 SSM 的主要表示方法是连续时间表示法。

从连续变为离散信号

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

这里就要使用Zero-order hold技术，其工作原理如下：

每次我们接收到一个离散信号，都会保证他的值不变，直到接收到一个新的离散信号再改变。这个过程创建了一个SSM可以使用的连续信号：

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

现在我们有了连续的输入信号，就可以生成连续的输出，只需根据输入的时间步长对数值进行采样即可。

这些采样值就是我们的离散输出！在数学上，我们可以应用Zero-order hold如下:

它们共同使我们从连续的 SSM 变为离散的 SSM。因为我们SSM处理的是离散信号，所以这里不是一个函数到函数，x(t)→y(t)，而是一个序列到序列，xₖ→yₖ，我们用公式表示如下：

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

我们用 k 代替 t 来表示离散化的时间步长，这样可以更清楚地说明连续 SSM 和离散 SSM 的区别。

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

既然我们已经有了离散表示法的表述，那么让我们来探讨一下如何实际计算模型。

循环表示

离散化的SSM允许在特定的时间步中处理信息。就像我们之前在循环神经网络（RNNs）中看到的那样，循环方法在这里也非常有用。

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

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

这种表示法可能看起来已经有点熟悉了！我们可以用之前处理 RNN 的方法来处理它。

也可以这样展开:

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

这种技术同时具备 RNN 的优点和缺点，即推理速度快，训练速度慢。

卷积表示

我们可以用于 SSM 的另一种表示方法是卷积。还记得在经典的图像识别任务中，我们应用滤波器（内核）来获得聚合特征：

由于我们处理的是文本而不是图像，因此我们需要用一维视角来代替：

我们用来表示这个“过滤器”的核是由SSM公式推导出来的:

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

上图也说明了padding 可能对输出产生的影响，所以我们一般都会在末尾padding而不是在前面。在下一步中，内核被移动一次，以执行下一步计算：

在最后一步，我们可以看到核的完整效果：

卷积的一个主要好处是它可以并行训练。但是由于核大小是固定，它们的推理不如rnn快速并且对序列长度有限制。

三种表示

上面的三种SSM 连续、递归和卷积这三种表达方式各有利弊：

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

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

听起来有点奇幻，但是有人就是实现出来了，这个模型被称为 Linear State-Space Layer (LSSL).

它结合了线性动态系统理论和神经网络的概念，可以有效地捕获数据中的时序信息和动态特征。LSSL基于线性动态系统理论，这种系统可以用状态空间模型表示。在这个模型中，系统的行为由状态变量的演化和外部控制信号的影响决定。状态变量是系统的内部表示，可以捕获系统的动态特性。

这些表示有一个共同的特性，即线性时间不变性（Linear Time Invariance，LTI）。LTI 说明 SSM 的参数 A、B 和 C 在所有时间步中都是固定的。这意味着 SSM 生成的每个标记的矩阵 A、B 和 C 都是相同的。

换句话说，无论你给 SSM 什么样的序列，A、B 和 C 的值都是一样的。我们得到了一个不感知内容的静态表示，所以Mamba解决的就是这个问题。

在探讨 Mamba 如何解决这个问题之前，我们先来看看图的最后一块–矩阵 A。

矩阵A的重要性

矩阵A可以说是 SSM 表述中最重要的一个方面。正如我们之前在循环表示中看到的那样，它捕捉了前一个状态的信息，用于构建新的状态。如果矩阵A如果跟RNN一样会遗忘掉非常靠前的信息那么SSM将没有任何的意义。

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

因此，创建矩阵 A 可能是只记住之前的几个标记与捕捉到迄今为止我们看到的每一个标记之间的区别。特别是在递归表示法的情况下，因为它只能回顾之前的状态。

那么，如何创建矩阵 A 才能保留较大的记忆（上下文大小）呢？

我们使用 Hungry Hungry Hippo！或 HiPPO（High-order Polynomial Projection Operators，高阶多项式投影运算符）。HiPPO 尝试将迄今为止看到的所有输入信号压缩成一个系数向量。HiPPO 的模型结合了递归记忆（Recurrent Memory）和最优多项式投影（Optimal Polynomial Projections）的概念，这种投影技术可以显著改善递归记忆的性能，特别是在处理长序列和长期依赖关系时。

在线函数逼近——随着时间序列，我们能不能记住前面状态的同时来很好的拟合后面的状态，他们得到了HiPPO函数，并发现这样的微分方程和SSM是一样的，或者说与SSM的形式一样，所以他们就围绕SSM来做了

它使用矩阵A来建立一种状态表示法，能很好地捕捉最近的标记并衰减较早的标记。其计算公式如下

假定我们有一个正方形矩阵 A，就可以得出

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

HiPPO 矩阵背后的理念是，它产生的隐藏状态会记住自己的历史。

在数学上，它是通过跟踪 Legendre 多项式的系数来实现这一功能的，这使得它能够近似于之前的所有历史。

然后，HiPPO被应用于我们之前看到的递归和卷积表示，以处理长程依赖关系。其结果就是序列结构化状态空间（S4），这是一类能有效处理长序列的 SSM。

它包括三个部分：

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

这类 SSM 具有多种优势，具体取决于您选择的表示方式（递归表示与卷积表示）。它还可以处理长文本序列，并通过建立在 HiPPO 矩阵基础上的有效的存储记忆。

注：如果您想深入了解如何计算 HiPPO 矩阵和自己建立 S4 模型的更多技术细节，我强烈建议您阅读 S4。

Part 3: Mamba — 一个选择性的SSM

我们终于了解了所有必要的基础知识，从而理解了 Mamba 的特别之处。状态空间模型可用于文本序列建模，但仍有一系列我们希望避免的缺点。

在本节中，我们将介绍 Mamba 的两大贡献：

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

它们共同创建了选择性 SSM 或 S6 模型，可以像自注意力一样用于创建Mamba块。在探讨这两项主要贡献之前，让我们先来探讨一下为什么需要它们。

它试图解决什么问题？

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

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

在选择性复制任务中，SSM 的目标是复制部分输入内容并按顺序输出（下面例子的任务是提取所有的名词）：

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

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

SSM 表现不佳的第二项任务是感应头，其目标是再现输入中发现的模式：

在上面的例子中，我们基本上是在进行单次提示，我们试图 "教会 "模型在每次 "Q: "之后做出 "A: "的回答。然而，由于 SSM 是时间不变的，因此它无法从其历史记录中选择要调用的先前的token。

让我们以矩阵B为例加以说明。无论输入 x 是什么，矩阵B都保持不变，因此与 x 无关：

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

相比之下，这些任务对Transformers 来说相对容易，因为它们会根据输入序列动态地改变注意力。它们可以有选择地 "看 "或 "注意 "序列的不同部分。再加上位置编码，这使得Transformers对于这种任务非常的简单。

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

有选择地保留信息

SSM的循环表示创建了一个非常有效的小状态，因为它压缩了整个历史信息，所以与不压缩历史(注意力矩阵)的Transformer模型相比，它的功能要弱得多。

Mamba 的目标是获得Transformer一样强大的“小”状态

通过有选择地将数据压缩到状态，当输入一个句子时，通常会有一些信息，比如停顿词，这些信息没有太多的意义。

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

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

相反，Mamba 通过将输入的序列长度和批量大小纳入其中，使矩阵 B 和 C，甚至步长 ∆*，*依赖于输入：

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

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

它们一起选择将哪些内容保留在隐藏状态，将哪些内容忽略不计，因为它们现在依赖于输入。

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

扫描操作

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

为了实现并行化，让我们来探讨一下如何利用递归计算输出：

每个状态都是前一个状态（乘以 A）加上当前输入（乘以 B）的总和。这就是所谓的扫描运算，可以用 for 循环轻松计算。相比之下，并行化似乎是不可能的，因为只有在我们拥有前一个状态的情况下，才能计算出每个状态。

然而，Mamba 通过并行扫描 算法实现了这一点。

通过关联属性，它假定我们进行操作的顺序并不重要。因此，我们可以分段计算序列，然后迭代合并：

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

这样还有一个好处是因为顺序不重要，也可以省略掉Transformer的位置编码。

硬件感知算法

最近推出的 GPU 的一个缺点是，在体积小但效率高的SRAM和体积大但效率稍差的 DRAM 之间的传输（IO）速度有限。经常在 SRAM 和 DRAM 之间复制信息会成为瓶颈。

与Flash Attention一样，Mamba 试图限制从 DRAM 到 SRAM 以及从 SRAM 到 DRAM 的次数。它通过内核融合来实现这一目标，使模型能够防止写入中间结果，并持续执行计算，直到完成为止。

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

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

• 步长为 ∆ 的离散化步骤
• 选择性扫描算法
• 与 C 相乘

硬件感知算法的最后一个环节是重新计算。

中间状态不会被保存，但却是计算梯度的反向传递所必需的。相反，作者在反向过程中重新计算了这些中间状态。

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

我们现在已经介绍了其架构的所有组件，下面是文章中的图片：

选择性 SSM. 检索自：Gu、Albert 和 Tri Dao。"Mamba: Linear-time sequence modeling with selective state spaces.” arXiv preprint arXiv:2312.00752 (2023).

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

The Mamba Block

我们迄今为止所探索的选择性 SSM 可以作为一个block来实现，就像我们可以在解码器区块中表示self-attention一样。

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

它首先通过线性投影来扩展输入嵌入。然后，在应用选择性SSM之前进行卷积，以防止独立token计算。

选择性 SSM 具有以下特性：

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

我们可以在查看代码实现时进一步扩展这一架构，并探讨端到端的示例是怎样的：

请注意一些变化，比如加入了归一化层和用于选择输出标记的 softmax。

当我们把所有东西放在一起时，我们就能获得快速的推理和训练，甚至是无限制的语境！

作者发现，使用这种结构，它的性能可以与相同尺寸的Transformer models相媲美，有时甚至超过它们！