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

YZRuin

已于 2024-07-24 15:52:26 修改

阅读量20

点赞数

分类专栏： Mamba模型文章标签：机器学习人工智能

于 2024-07-22 19:50:49 首次发布

原文链接：https://goombalab.github.io/blog/2024/mamba2-part1-model/

版权

Mamba模型专栏收录该内容

1 篇文章 0 订阅

订阅专栏

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

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

[ 论文 ] [ 代码 ]

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

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

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

问题 1（理解）

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

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

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

问题 2（效率）

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

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

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

SSD 模型

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

SSD 模型是指可以纳入深度神经网络的特定独立层，如注意力机制或 SSM
SSD 框架是推理此模型（以及更多理论联系）的通用框架
SSD 算法是一种比以前的 SSM 更有效地计算 SSD 层的算法

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

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

线性（SSM）模式

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

$\begin{aligned} h_{t} &= A_t h_{t-1} + B_t x_t \\ y_t &= C_t^{\top} h_t \end{aligned} \cdots (1)$

回顾一下， 结构化状态空间模型 (SSM)定义一个映射 $x \in \mathbb{R}^\mathtt{T} \to y \in \mathbb{R}^\mathtt{T}$ 。考虑到 $x_t$ 和 $y_t$ 作为标量，隐藏状态 $h_t$ 作为 $\mathtt{N}$ 维向量，其中 $\mathtt{N}$ 是一个独立的超参数，称为状态大小、状态维度或状态扩展因子 。

选择性状态空间模型（Mamba）允许 $(A, B, C)$ SSM 参数随时间变化。我们分别将它们视为具有形状的张量 $A \in \mathbb{R}^\mathtt{(T, N, N)}$ ， $B \in \mathbb{R}^\mathtt{(T, N)}$ ，和 $C \in \mathbb{R}^\mathtt{(T, N)}$ 。与 Mamba-1 一样，我们采用实数 $\mathbb{R}$ ，尽管与 S4 谱系等其他结构化 SSM 一样，存在复杂的变体也是可能的。

结构化 SSM 需要 $A$ 具有可有效计算的结构，例如最常用的对角线结构。在这种情况下 $A$ 有形状 $\mathtt{(T, N)}$ 其中只有对角线元素 $\mathtt{N} \times \mathtt{N}$ 矩阵被存储。

SSD：标量结构化SSM

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

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

多头 SSM

方程 (1) 仅针对单维输入进行定义 $x \in \mathbb{R}^\mathtt{T}$ 。如果 $X \in \mathbb{R}^\mathtt{(T, P)}$ 有 $\mathtt{P}$ 不同的通道，我们可以使用相同的动态（即相同的 SSM $(A, B, C)$ ) 独立地针对每个通道。这可以解释为 SSM 模型的 单个头。

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

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

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

$Y^ \mathtt{(T,P)} = \mathsf{SSM} (A^ \mathtt{(T, \cdots)}, B^\mathtt{(T,N)}, C^ \mathtt{(T,N)})(X^ \mathtt{(T,P)})\cdots (2)$

一些变化轴包括

结构 $A$ ，这会影响其参数形状：
1. ... = (N,N) 对于一般（非结构化）SSM
2. ... = (N) 对于对角线 SSM（或其他结构，例如对角线加低秩）
3. ... = () 对于标量 SSM（即 SSD）
状态维度 $\mathtt{N}$ （i.e. d_state ）
头部维度 $\mathtt{P}$ （i.e. d_head ）

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

二次（注意）模式

但首先，让我们换个话题，暂时忘掉状态空间模型。给定上面具有相同形状的相同张量 $(A^ \mathtt{(T)}, B^ \mathtt{(T, N)}, C^ \mathtt{(T, N)})$ ，让我们定义一个不同的对象。

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

$L = \begin{bmatrix} 1 & \\ a_1 & 1 & \\ a_2a_1 & a_2 & 1 \\ \vdots & \vdots & \ddots & \ddots \\ a_{\mathtt{T}-1}\dots a_1 & a_{\mathtt{T}-1}\dots a_2 & \dots & a_{\mathtt{T}-1} & 1 \\ \end{bmatrix}$

然后，我们定义以下矩阵

$M = L \circ CB^\top \in \mathbb{R}^{\mathtt{(T,T)}}\cdots (3)$

最后， $M$ 编码一个序列变换 $x \in \mathbb{R}^\mathtt{T} \to y \in \mathbb{R}^\mathtt{T}$ 通过基本矩阵乘法 $y = Mx$ 将一维输入映射到一维输出（如方程 (1) 所示）。

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

$Y = (L \circ QK^\top) V$

如果我们将其重命名，则这与方程 (3) 完全相同 $(C,B,X)\mapsto(Q,K,V)$ ！

状态空间对偶

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

$(A^\mathtt{(T)}, B^\mathtt{(T, N)}, C^\mathtt{(T, N)}, X^\mathtt{(T, P)}) \mapsto Y^\mathtt{(T, P)}$

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

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

SSD vs. 状态空间模型

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

Mamba-1（S6）采用对角结构 $A$ ，而 Mamba-2 (SSD) 使用标量乘以恒等结构 $A$ 。
Mamba-1 的头部尺寸为 $\mathtt{P}=1$ （即所有通道都由单独的 SSM 完全独立控制），而 Mamba-2 使用的头部尺寸为 $\mathtt{P}>1$ （就像是 $\mathtt{P}=64$ 默认情况下）。

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

通过限制对角线结构 $A$ 对于标量乘以恒等式，递归动力学在所有 $\mathtt{N}$ 状态空间的元素。
这些动态也适用于所有 $\mathtt{P}$ 给定头部的通道。

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

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

底线：Mamba-1 vs. Mamba-2

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

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

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

SSD vs. 注意力机制

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

softmax 标准化被删除。
单独的元素掩码矩阵以乘法方式应用。

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

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

$a_{i:j}^\times = a_i \cdots a_{j+1}$

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

两全其美

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

效率：SSM 和注意模式

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

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

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

效率：SSD 模式

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

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

定义 SSD“令牌混合”序列转换的特定结构化矩阵的块分解。
一种“分块”算法，将序列分成几段，计算每段上的二次注意形式，并通过在段之间传递 SSM 状态来调整结果。

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

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

Mamba-2 架构

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

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

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

语言建模

就实证结果而言，我们没有像 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 层所声称的“二元性”及其强大的概括。

脚注

与 Mamba-1 一样，我们采用实数𝑅R，尽管与S4谱系等其他构造SSM一样，存在复杂的变体[17]也有人这么说过。 [↩]
通常还有额外的批次维度B在实现这些模型时，我们将在整个演示过程中忽略这些模型。 [↩]

参考

Vision Mamba: Efficient Visual Representation Learning with Bidirectional State Space Model
Zhu, L., Liao, B., Zhang, Q., Wang, X., Liu, W. and Wang, X., 2024. The International Conference on Machine Learning (ICML).
U-Mamba: Enhancing Long-Range Dependency for Biomedical Image Segmentation
Ma, J., Li, F. and Wang, B., 2024. arXiv preprint arXiv:2401.04722.
VMamba: Visual State Space Model
Liu, Y., Tian, Y., Zhao, Y., Yu, H., Xie, L., Wang, Y., Ye, Q. and Liu, Y., 2024. arXiv preprint arXiv:2401.10166.
Caduceus: Bi-directional Equivariant Long-Range DNA Sequence Modeling
Schiff, Y., Kao, C., Gokaslan, A., Dao, T., Gu, A. and Kuleshov, V., 2024. The International Conference on Machine Learning (ICML).
Graph-Mamba: Towards Long-Range Graph Sequence Modeling with Selective State Spaces
Wang, C., Tsepa, O., Ma, J. and Wang, B., 2024. arXiv preprint arXiv:2402.00789.
Graph Mamba: Towards Learning on Graphs with State Space Models
Behrouz, A. and Hashemi, F., 2024. arXiv preprint arXiv:2402.08678.
Repeat After Me: Transformers Are Better Than State Space Models at Copying
Jelassi, S., Brandfonbrener, D., Kakade, S.M. and Malach, E., 2024. The International Conference on Machine Learning (ICML).
In-Context Language Learning: Architectures and Algorithms
Akyurek, E., Wang, B., Kim, Y. and Andreas, J., 2024. The International Conference on Machine Learning (ICML).
Is Mamba Capable of In-Context Learning?
Grazzi, R., Siems, J., Schrodi, S., Brox, T. and Hutter, F., 2024. arXiv preprint arXiv:2402.03170.
Can Mamba Learn How to Learn? A Comparative Study on In-Context Learning Tasks
Park, J., Park, J., Xiong, Z., Lee, N., Cho, J., Oymak, S., Lee, K. and Papailiopoulos, D., 2024. The International Conference on Machine Learning (ICML).
The Illusion of State in State-Space Models
Merrill, W., Petty, J. and Sabharwal, A., 2024. The International Conference on Machine Learning (ICML).
The Expressive Capacity of State Space Models: A Formal Language Perspective
Sarrof, Y., Veitsman, Y. and Hahn, M., 2024. arXiv preprint arXiv:2405.17394.
Combining Recurrent, Convolutional, and Continuous-time Models with the Linear State Space Layer
Gu, A., Johnson, I., Goel, K., Saab, K., Dao, T., Rudra, A. and R\'{e}, C., 2021. Advances in Neural Information Processing Systems (NeurIPS).
Modeling Sequences with Structured State Spaces
Gu, A., 2023.
Neural Machine Translation by Jointly Learning to Align and Translate
Bahdanau, D., Cho, K. and Bengio, Y., 2015. The International Conference on Learning Representations ({ICLR}).
Attention Is All You Need
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, L. and Polosukhin, I., 2017. Advances in Neural Information Processing Systems (NeurIPS).
Efficiently Modeling Long Sequences with Structured State Spaces
Gu, A., Goel, K. and Re, C., 2022. The International Conference on Learning Representations (ICLR).
Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Gu, A. and Dao, T., 2023. arXiv preprint arXiv:2312.00752.
On the Parameterization and Initialization of Diagonal State Space Models
Gu, A., Gupta, A., Goel, K. and R\'e, C., 2022. Advances in Neural Information Processing Systems (NeurIPS).
Diagonal State Spaces are as Effective as Structured State Spaces
Gupta, A., Gu, A. and Berant, J., 2022. Advances in Neural Information Processing Systems, Vol 35, pp. 22982--22994.
Simplified State Space Layers for Sequence Modeling
Smith, J.T., Warrington, A. and Linderman, S.W., 2023. The International Conference on Learning Representations (ICLR).
Simplifying and Understanding State Space Models with Diagonal Linear {RNN}s
Gupta, A., Mehta, H. and Berant, J., 2022. arXiv preprint arXiv:2212.00768.
Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
Katharopoulos, A., Vyas, A., Pappas, N. and Fleuret, F., 2020. International Conference on Machine Learning, pp. 5156--5165.
GateLoop: Fully Data-Controlled Linear Recurrence for Sequence Modeling
Katsch, T., 2023. arXiv preprint arXiv:2311.01927.
An Empirical Analysis of Compute-Optimal Large Language Model Training
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Rutherford, E., de Las Casas, D., Hendricks, L.A., Welbl, J., Clark, A. and others,, 2022. Advances in Neural Information Processing Systems (NeurIPS), Vol 35, pp. 30016--30030.
The {P}ile: An 800{G}{B} Dataset of Diverse Text for Language Modeling
Gao, L., Biderman, S., Black, S., Golding, L., Hoppe, T., Foster, C., Phang, J., He, H., Thite, A., Nabeshima, N., Presser, S. and Leahy, C., 2020. arXiv preprint arXiv:2101.00027.
Zoology: Measuring and Improving Recall in Efficient Language Models
Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A. and Re, C., 2024. The International Conference on Learning Representations ({ICLR}).
Simple Linear Attention Language Models Balance the Recall-Throughput Tradeoff
Arora, S., Eyuboglu, S., Zhang, M., Timalsina, A., Alberti, S., Zinsley, D., Zou, J., Rudra, A. and Re, C., 2024. The International Conference on Machine Learning (ICML).