Mamba方法精读

背景：Mamba的研究目标

1. 设计一个新的结构，性能不弱于Transformer，并具备更高的计算效率

大型语言模型（Large Language Models，LLMs）的成功离不开Transformer结构，然而，Transformer的推理效率较低。以文本生成为例，文本序列的生成是逐步进行的，Transformer每生成一个token，就需要和已有的tokens计算一次注意力。注意力的时间复杂度是$O(n^2)$，这意味着随着序列长度的增加，Transformer的计算量将呈平方级递增，导致难以处理很长的序列。

不过，Transformer的训练过程却很高效。与RNN不同，Transformer在处理输入序列时并不依赖于前一个时间步的输出。这意味着在训练过程中，整个序列的所有位置可以同时输入模型进行处理，无需按时间步顺序逐步计算。这使得Transformer可以在一个批次中并行处理整个输入序列。

Mamba的研究目标即保留Transformer的训练效率和生成精度的同时，提高模型的推理效率。具体来说，训练过程中，Mamba的时间和空间复杂度都是$O(n)$；而推理时可以采用自回归的方式，使得每一步的生成只需要常量的时间复杂度。

研究思路

1. 概述

作者发现状态空间模型（State Space Model，SSM）具有良好的特性，使其可以进行高效的训练和推理。然而，现有SSM性能还不够，因此设计了Selective State Space Model；并将其和MLP整合成一个block，多个blocks构成了Mamba结构。

2. 什么是SSM？

备注： 论文¹中的SSM实际代指S4（Structured State Space Sequence models），本文也使用SSM代指S4。

结论： SSM本质上是RNN，区别是当有序列输入（多个时间步的输入）时，它能够以卷积的形式快速生成序列输出（每个时间步对应的输出），而非RNN中串行的方式。

简单回顾RNN： 下图展示了RNN的工作流程。在每个时间步t，网络接收当前输入$x_t$和上一步的状态$h_{t-1}$，得到当前步的输出$y_t$。

当生成当前时间t的输出时，RNN只用考虑当前输入$x_t$和上一步的状态$h_{t-1}$，不需要和Transformer一样重新计算之前的状态（即注意力），因此可以高效推理，并且理论上拥有无限长的序列窗口。

然而，RNN在训练时只能串行处理，降低了训练效率。此外，在实际应用中，RNN往往倾向于记住和利用短期依赖信息，容易遗忘长程信息。

连续状态空间模型： SSM启发于连续状态空间模型，该模型被用于描述动态系统，被广泛应用于控制系统和信号处理等领域。和RNN相似，它也通过一个潜在状态$h(t)\in {\mathbb{R}}^N$将输入序列$x(t)\in \mathbb{R}$映射为输出序列$y(t)\in \mathbb{R}$。连续状态空间模型的公式为：
$$\begin{array}{rl}{h}^{\prime }(t)& =\mathbf{A}h(t)+\mathbf{B}x(t)\\ y(t)& =\mathbf{C}h(t)\end{array}\text{\hspace{1em}(1a)(1b)}$$
其中公式（1a）是状态转移方程，公式（1b）是输出方程；$\mathbf{A}$，$\mathbf{B}$和$\mathbf{C}$是参数矩阵。由于参数矩阵的值不随时间步变化，且无非线性操作，因此上述模型是线性时不变的（linear time invariance，LTI）。

离散化： 上述公式无法直接应用到离散的时间步上。为了获得$h_t$和$h_{t-1}$之间的关系，作者采用零阶保持法（Zero-Order Hold，ZOH）离散化上述公式，得到以下结果：
$$\begin{array}{rl}{h}_t& =\bar{\mathbf{A}}{h}_{t-1}+\bar{\mathbf{B}}{x}_t\\ y_t& =\mathbf{C}{h}_t\\ \bar{\mathbf{A}}& =exp(\mathbf{\Delta A})\bar{\mathbf{B}}={\mathbf{\Delta A}}^{-1}(exp(\mathbf{\Delta A})-\mathbf{I})\cdot \mathbf{\Delta B}\end{array}\text{\hspace{1em}(2a)(2b)(3)}$$
其中$\mathbf{\Delta }$也是可学习参数。这部分不理解不要紧。

通过卷积快速得到每一步的输出： 因为SSM是LTI，所以可以得到每一步$y_t$的closed-form计算公式。举个例子，假设我想要计算$y_3$：
$$\begin{array}{rl}{y}_3& =(\bar{\mathbf{A}}{h}_2+\bar{\mathbf{B}}{x}_3)\cdot \mathbf{C}\\ y_3& =(\bar{\mathbf{A}}(\bar{\mathbf{A}}{h}_1+\bar{\mathbf{B}}{x}_2)+\bar{\mathbf{B}}{x}_3)\cdot \mathbf{C}\\ y_3& =\mathbf{C}\bar{\mathbf{B}}{x}_3+\mathbf{C}\bar{\mathbf{B}}\bar{\mathbf{A}}{x}_2+\mathbf{C}\bar{\mathbf{B}}{\bar{\mathbf{A}}}^2{x}_1\\ y_k& =\sum _{i=1}^k\mathbf{C}\bar{\mathbf{B}}{\bar{\mathbf{A}}}^{k-i}{x}_i\end{array}\text{\hspace{1em}(4)(5)(6)(7)}$$
可以发现，公式（7）和卷积的形式相同。假设输出序列$y=[y_k,...,y_1]$，输入序列$x=[x_k,...,x_1]$，公式（7）又可写为：
$$\begin{array}{rl}\bar{\mathbf{K}}& =(\mathbf{C}\bar{\mathbf{B}},...,\mathbf{C}\bar{\mathbf{B}}{\bar{\mathbf{A}}}^{k-1})\\ y& =x\ast \bar{\mathbf{K}}.\end{array}\text{\hspace{1em}(8)}$$
因此每个forward就是一次卷积操作。

3. Mamba中的Selective SSM对SSM做了什么改进？

$B,C\to B_t,C_t$: SSM是时不变模型，即模型参数$A,B,C$在每个时间步相同，与输入无关，因此难以进行content-aware reasoning。在论文¹中，作者列举了两个的任务：Copying & Selective Copying来证明content-aware learning的重要性，如下图所示：

时不变模型可以完成Copying任务，但无法完成Selective Copying，因为它无法区分当前步是有效还是无效输入。因此，作者认为模型参数应该是input-dependent，即时变的。这也是Selective SSM的改进动机。算法示意图如下所示：

此外，作者认为权衡序列模型的效率和精度的关键在于how well they compress their state：

In summary, the efficiency vs. effectiveness tradeoff of sequence models is characterized by how well they compress their state: efficient models must have a small state, while effective models must have a state that contains all necessary information from the context. In turn, we propose that a fundamental principle for building sequence models is selectivity: or the context-aware ability to focus on or filter out inputs into a sequential state.

例如，attention is both effective and inefficient because it explicitly does not compress context at all.

4. 如何高效实现Selective SSM？

将SSM变成时变模型后，将遇到两个问题：

1. 无法直接利用卷积，不能享受卷积带来的高效性。因为此时“卷积核”是时变的；
2. 参数$B,C$为了变成$B_t,C_t$，会多一个L维度（时间/序列维度），参数量的变大增加了对memory IOs带宽和显存的占用，增加了计算和显存成本。

对此，作者设计了三个改进：

1. 相比在GPU HBM（high-bandwidth memory）中准备形状为$(B,L,D,N)$的$(\bar{\mathbf{A}},\bar{\mathbf{B}})$，先将SSM参数$(\mathbf{\Delta },\mathbf{A},\mathbf{B},\mathbf{C})$加载到更快的SRAM，在SRAM中执行discretization和recurrence，最后将结果$(B,L,D)$写回HBM。这避免了在HBM和SRAM之间的带宽消耗，因为搬运$(B,L,D,N)$d的参数显然比$(B,L,D)$的更占用带宽。
2. 为了避免串行recurrence减慢训练速度，采用work-efficient parallel scan algorithm并行化。
3. 为了减少显存占用，前向传播时不保存intermediate states，而在反向传播时重新计算。

参考资料

• Mamba: Linear-Time Sequence Modeling with Selective State Spaces (Paper Explained)
• A Visual Guide to Mamba and State Space Models

---

1. Mamba: Linear-Time Sequence Modeling with Selective State Spaces ↩︎ ↩︎