状态空间对偶 (Mamba-2) 第一部分 - 模型_mamba2论文原文

 本文是对Albert Gu和Tri Dao两位大佬的Mamba2博客的翻译

 原文链接：State Space Duality (Mamba-2) Part I - The Model | Goomba Lab

[ 论文 ] [ 代码 ]

本系列文章已在tridao.me 上发布

1. 第一部分-模型

2. 第二部分-理论

3. 第三部分-算法

4. 第四部分-系统

自 6 个月前Mamba 发布以来 ，我们惊喜地发现 社区反响 热烈。我们多年来一直追求的高效序列模型研究方向真正引起了机器学习社区的共鸣，并且取得了超出我们预期的进展，这让我们感到非常欣慰。我们已经看到大量令人兴奋的后续工作，从直接应用（例如视觉、基因组学、图表等）对理解的影响（例如回忆能力、情境学习以及形式语言表达能力)，以及大量的 在线 博客 、 教程 和 视频 。我们对这项研究的方向感到无比兴奋 ！ 

尽管 Mamba 目前具有巨大潜力，但我们对它的第一个版本并不完全满意……

问题 1（理解）

从概念的角度来看，我们发现 SSM 如此迷人的原因之一是它们给人的感觉就像 基础 一样。其中一个例子就是它们与许多主要的序列模型范式有着密切的联系。正如我们之前在结构化 SSM 方面的工作中所阐述的那样，它们似乎抓住了连续、卷积和循环序列模型的本质——所有这些都包含在一个简单而优雅的模型中。

当然，除了这些之外，还有另一个主要的序列模型范式：无处不在的 注意力 机制的变体。SSM 总是让人感觉与注意力有些脱节，我们尝试了一段时间来更好地理解它们之间的关系。

问题 1： 状态空间模型和注意力之间的概念联系是什么？ 我们可以将它们结合起来吗？

问题 2（效率）

从计算的角度来看，尽管 Mamba 为提高速度付出了很多努力（特别是其硬件感知选择性扫描实现），但它的硬件效率仍然远低于注意力机制。缺少的一点是，现代加速器（如 GPU 和 TPU） 高度 专门用于矩阵乘法。虽然这对于推理来说不是问题，因为推理的瓶颈在于不同的考虑因素，但这在训练期间可能是一个大问题。

问题 2： 我们可以通过将其重铸为矩阵乘法来加快 Mamba 模型的训练速度吗？

这些是 Mamba-2（特别是其新的状态空间模型变体）试图解决的主要问题。

SSD 模型

Mamba-2 论文的重点是我们所说的 结构化状态空间对偶性 （SSD），它指的是几件事：

1. SSD 模型 是 指可以纳入深度神经网络的特定独立层，如注意力机制或 SSM

2. SSD 框架 是 推理此模型（以及更多理论联系）的通用框架

3. SSD 算法 是 一种比以前的 SSM 更有效地计算 SSD 层的算法

主 SSD 模型或“状态空间对偶模型”本身其实并不复杂！在本系列博客文章的第一部分中，我们将单独提供 SSD 层（和 Mamba-2）的独立描述，以及它与相关模型（特别是 Mamba-1）的比较。

在本系列的下一部分中，我们将描述总体框架和理论联系，这些对于实际使用 Mamba-2 并不是必需的。

线性（SSM）模式

SSD 与 Mamba 从同一组方程式开始：

回顾一下， 结构化状态空间模型 (SSM)定义一个映射。 考虑到和作为标量，隐藏状态作为维向量，其中是一个独立的超参数，称为状态大小、状态维度或状态扩展因子 。

选择性 状态 空间模型（Mamba）允许SSM 参数随时间变化。我们分别将它们视为具有形状的张量，，和 。与 Mamba-1 一样，我们采用实数，尽管与 S4 谱系等其他结构化 SSM 一样，存在复杂的变体也是可能的。

结构化 SSM 需要具有可有效计算的结构，例如最常用的对角线结构。 在这种情况下有形状其中只有对角线元素矩阵被存储。

SSD：标量结构化SSM

原始的Mamba（或者更准确地说是其核心的“S6”层）正是具有对角结构的选择性SSM。 

Mamba-2 的 SSD 层只做了一个小修改 ：它限制了对角线甚至进一步变成 标量乘以单位阵 的结构；换句话说必须全部为相同值。在这种情况下可以用形状来表示还可以识别只是一个标量（所以我们有时会把它表示为）。

多头 SSM

方程 (1) 仅针对单维输入进行定义。 如果有不同的通道，我们可以使用相同的动态（即相同的 SSM) 独立地针对每个通道。这可以解释为 SSM 模型的 单个头。

在这里，我们想到作为形状张量，其中是序列（时间）维度，是“头部维度”。通常会有一个额外的批次维度在实现这些模型时，我们将在整个演示过程中忽略这些模型。

多个头部可以完全独立构建；在本文的其余部分，我们假设我们使用单个头部。请注意，这些头部与多头注意力模型中的头部的工作方式完全类似，在 Mamba-2 中，我们也选择了与现代 Transformer 类似的维度，例如 = 64或者=128. （要缩放到更大的模型宽度 = ，我们保持这个固定并增加独立头的数量。）

我们可以将一般（选择性）状态空间模型表示为

一些变化轴包括

1. 结构，这会影响其参数形状：

   1. ... = (N,N) 对于一般（非结构化）SSM

   2. ... = (N) 对于对角线 SSM（或其他结构，例如对角线加低秩）

   3. ... = () 对于标量 SSM（即 SSD）

2. 状态维度（i.e. d_state ）

3. 头部维度（i.e. d_head ）

结构化 SSM 还有其他变化轴（例如时间不变性与选择性、SISO 与 MIMO）、真实与复杂等），但我们突出显示这些内容，以便我们可以在短短一秒钟内对比 Mamba-2 与 Mamba-1……

二次（注意）模式

但首先，让我们换个话题，暂时忘掉状态空间模型。给定上面具有相同形状的相同张量，让我们定义一个不同的对象。

首先，我们将定义以下矩阵（不用担心，我们将在本系列的第二部分中进行更多解释并为其命名！）

然后，我们定义以下矩阵

最后，编码一个序列变换 通过基本矩阵乘法将一维输入映射到一维输出（如方程 (1) 所示）。

这有什么特别之处？好吧，你可能会注意到它看起来非常类似于注意力计算。事实上，如果所有， 然后只是下三角 因果掩码 ，而 方程(3) 相当于 因果线性注意力：

如果我们将其重命名，则这与方程 (3) 完全相同  ！

状态空间对偶

所谓“对偶性”，是指方程 (1)（对于标量恒等结构的情况) 和 方程(3) 实际上是 完全相同的模型 ，我们可以将其视为一个特定的函数

在通用的 SSD 框架 （本系列的第二部分）中，我们将以两种完全不同的方式展示这种等价性，这两种方式实际上都更为通用并且都很有启发性。

不过，如果您相信我们的话，那么 SSD 与 SSM 或注意力相比相对简单。

SSD vs. 状态空间模型

与之前的 SSM 相比，SSD 与 Mamba 的核心层基本相同，但在循环结构上具有更多结构矩阵。

1. Mamba-1（S6）采用对角结构，而 Mamba-2 (SSD) 使用标量乘以恒等结构。

2. Mamba-1 的头部尺寸为（即所有通道都由单独的 SSM 完全独立控制），而 Mamba-2 使用的头部尺寸为（就像是默认情况下）。

具体而言，这可以从两个方面视为与权重相关：

• 通过限制对角线结构对于标量乘以恒等式，递归动力学在所有状态空间的元素。

• 这些动态也适用于所有给定头部的通道。

换句话说，单个 SSM 头拥有整个状态大小，在 Mamba-1 中它们各自由单独的标量递归控制，但在 Mamba-2 中它们由单个共享递归控制。

为什么要做出这些限制？主要动机是效率：这些变化对于能够以 [ 双重注意形式 ] 查看模型是必要的，这允许使用矩阵乘法。

底线：Mamba-1 vs. Mamba-2

与 Mamba-1 相比，Mamba-2 允许 更大的状态维度 （从 Mamba-1 中的 N=16到Mamba-2 中的 N=64 N=256 甚至更高），同时 训练速度更快 。

但这会伤害我们吗？有些人直觉地认为不会。选择性的一个主要原因是（例如这取决于输入) 是为了让 SSM 能够控制是否记住或忽略特定的信息；例如，如果在文本记录中遇到填充词“um”。但如果应该忽略此类信息，则整个状态可以一起忽略它，因此如果状态的动态在所有特征之间共享，那就没问题了。

从经验上看，我们还没有发现证据表明 Mamba-2 的受限表达能力可能会带来不利影响，但目前尚无定论！从一个角度来看，Mamba-2 并不完全 优于 Mamba-1：虽然从 训练角度来看，Mamba-2 有了显著的改进，但从纯 推理 角度来看，Mamba-1 可能更好 。由于 SSM 的推理速度完全由状态维度决定，因此如果想要最大化目标推理效率（即特定状态大小）的性能\mathtt{N})，那么 Mamba-1 的增强表现力可能会更好。我们尚未完全分析这里的（理论或经验）权衡，并认为这将是社区进一步深入研究的一个很酷的方向！

SSD vs. 注意力机制

和标准（自我）注意力机制相比，SSD 也只有两点不同：

1. softmax 标准化被删除。

2. 单独的元素掩码矩阵以乘法方式应用。

一阶差分可以解释为将模型的有效状态大小从线性降低到常数，并将其效率从二次提高到线性。

第二个区别是 SSD 与标准线性注意力的区别。一种思考掩码的方式是将其视为 依赖于输入的相对位置编码 。由于掩码在 方程(3) 中，标准注意力得分被权重所减弱

这可以解释为一个“折扣因素”，基于位置和距离的远近。（这一解释同时得到了Tobias Katsch的GateLoop论文的支持[24]）在其注意力形式中，这种依赖于输入的位置掩码可以被解释为编码 Mamba“选择性”的关键因素！

两全其美

那么我们为什么要关心这个模型的两种观点呢？首先，它在数学上非常有趣，我们将在 第二部分 中介绍，我们希望它能启发未来的发展方向。但它也有直接的实际好处！

效率：SSM 和注意模式

SSM 方程(1) 和注意力 方程(3) 模式代表计算同一函数的两种不同方式，因此让我们对它们进行对比。

首先，请记住，SSM 之所以有趣的一个主要原因是，计算 方程(1) 作为递归需要维护一个 恒定大小的状态 （大小每通道）并 在序列长度上线性缩放缺点是，由于硬件方面的考虑，原始 FLOP 不能反映实际速度……

另一方面，计算这个序列变换通过方程 (3) 需要序列长度的二次时间，因为我们正在实现这一点矩阵。但在实践中它可以很快，因为它只使用矩阵乘法，这在 GPU 和 TPU 上得到了极大的优化。

效率：SSD 模式

那么，如果有两种等效的方法来计算同一个模型，我们什么时候应该使用其中一种模式？在推理过程中，没有权衡：SSM 模式专为快速自回归推理而设计。但训练呢？这里 FLOP 和硬件效率之间存在矛盾，注意模式使用更多 FLOP，但通过矩阵乘法更有效地使用它们。

事实证明，通过结合这些算法，我们可以同时实现两全其美！这种“状态空间对偶”算法有两种等效解释，要么是

1. 定义 SSD“令牌混合”序列转换的特定结构化矩阵的块分解。

2. 一种“分块”算法，将序列分成几段，计算每段上的二次注意形式，并通过在段之间传递 SSM 状态来调整结果。

我们将把该算法的细节留到 第三部分（或  全文 的第 6 节 ），因为它需要从理论中推导一些机制。但我们确实强调，该算法的实现并不太复杂——我们提供的最小实现只有大约 30 行 PyTorch 代码！

SSD 算法的优势在于，它保留了与 SSM 相同的高效 FLOP 计数（与二次注意力相比），并且通过利用 matmuls，与一般状态空间模型相比，还显著加快了训练速度。

Mamba-2 架构

虽然 Mamba-2 的核心贡献是新的 SSD 层和理论，但我们也对 Mamba 的神经网络架构做了一些小的改动。

主要的变化是并行的产生SSM 参数与输入，而不是顺序输入。这部分是出于与注意力的联系；但更务实的是，它更简单，更适合诸如张量并行之类的扩展技术，这将在本系列的第四部分中讨论！

还有一些其他细微的差异，本文将对此进行更详细的介绍。然而，我们确实想强调的是，这些架构变化并不是该模型的真正重点。

语言建模

就实证结果而言，我们没有像 Mamba-1 那样广泛地测试 Mamba-2，但相信它应该在各方面都表现相当甚至更好。我们的完整语言模型结果使用与 Mamba 相同的协议，并发现在 Chinchilla 法则下的扩展性略好。

在 Pile 数据集上进行全面训练的模型并且标准零样本下游评估也显示出类似的趋势。我们强调，即使性能相当，Mamba-2 的 训练速度也比 Mamba-1 快得多 ！

合成语言建模：MQAR

更有趣的是，我们重点介绍了我们尝试的一项合成任务。自最初的 Mamba 论文研究了合成 复制和感应头等合成任务以来，许多后续工作已经开始研究更难的联想回忆任务。 由动物学和基于工作方式已成为事实上的标准。

我们运行了该任务的一个版本，该版本比文献中通常报道的版本要难得多，发现 Mamba-2 比 Mamba-1 好得多。性能改进的一个原因是更大的状态大小（在此远比 Mamba-1 大的多），这也是 Mamba-2 最初设计的主要动机之一。

有趣的是，即使控制状态大小，Mamba-2 在这个特定任务上的表现也明显优于 Mamba-1。老实说，我们不太清楚为什么，如果能研究一下模型的其他方面就好了……例如，[ SSD 的受限结构 ] 在这里真的 有用 吗？

下一步

在 本系列的下一篇文章 中，我们将更深入地介绍完整的 SSD 框架，包括如何证明 SSD 层所声称的“二元性”及其强大的概括。

脚注

1. 与 Mamba-1 一样，我们采用实数

   声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/木道寻08/article/detail/911947