• 文献阅读：Mamba: Linear-Time Sequence Modeling with Selective State Spaces
  • 1. 文章简介
  • 2. 方法介绍
    • 1. State Space Models
    • 2. Selective State Space Models
  • 3. 实验考察 & 结论
    • 1. 简单问题上的验证
    • 2. 实际场景效果
      • 1. 语言模型
      • 2. DNA模型
      • 3. 语音模型
    • 3. 细节考察
      • 1. 速度和内存考察
      • 2. 消融实验
  • 4. 结论 & 思考

• 文献链接：https://arxiv.org/pdf/2312.00752.pdf
• GitHub链接：https://github.com/state-spaces/mamba

1. 文章简介

这篇文章23年12月普林斯顿的一篇文章，文中提出了一个Mamba的模型结构，尝试挑战了一下Transformer的霸权地位。

众所周知，自从BERT和GPT分别在NLP任务以及生成任务上展现出绝对的统治力之后，transformer框架的模型同时开始进军CV和ASR领域，不断在各个领域当中出圈，大有一统天下的趋势。

不过，针对Transformer框架的挑战也是一直存在，从简单的对于Transformer框架中self attention结构的计算量优化到尝试复兴RNN的RetNet等等，整体的思路前者基本就是希望减少self attention的计算量从而使得可以容纳的context窗口长度，而后者干脆就回到RNN的框架来完全舍弃掉窗口的设置，通过设置并行训练的方式来修改掉RNN只能串行训练的问题。

这里，Mamba走的也是后者这个路线，完全舍弃掉了self-attention的框架，使用文中提到的state space model的框架来进行实现。

文中宣称：

• Mamba模型不但可以无视掉context长度限制进行任意长文本的生成，还可以并行高速地训练，甚至有着很好的可扩展性，可以容纳大量参数，和transformer一样，在大数据预训练的框架下依然没有看到效果的瓶颈。

更牛逼的是，这篇文章的作者几乎是一己之力推着这个模型框架往前走，最早是一个S4的模型框架，然后优化成了H3的结果，到现在的Mamba，都是同一个团队沿着同一条路子走下来的，也是牛逼的厉害。不过可惜的是S4和H3那两篇文章我还没看过，所以这里对于Mamba的结构理解多少还有一点难度，后面会找时间去把剩下那两篇文章也看一下，或许对这个文章会有更好的一个理解。

2. 方法介绍

下面，我们首先来看一下Mamba的具体模型结构。

整体来说，Mamba的模型结构是在这篇文章的前作中提出的State Space Model（SSM）的基础上进行优化得到的，加上了选择机制并使之适应GPU的并行加速机制。

因此，我们下面就会遵循稳重的思路首先来看一下State Space Model，然后来看一下文中优化得到的Selective State Space Model，也就是文中的Mamba模型框架。

1. State Space Models

我们首先来考察一下State Space Model（SSM）。

本质上来说，包含输入$x_t$输出$y_t$以及态函数$h_t$，且训练过程可以并行，然后infer过程可以做到迭代串行的模型结构都可以称之为SSM。

下面就是一个SSM的典型case，它参数主要包括4个部分$(\Delta ,A,B,C)$，而整体的操作则是包括两个部分：

1. Discretization

   $$\begin{array}{rl}\bar{A}& =exp(\Delta A)\\ \bar{B}& =(\Delta A{)}^{-1}(exp(\Delta A)-I)(\Delta B)\end{array}$$

2. Computation

   Computation部分则可以有以下两种等价地表述：

   1. 线性recurrence实现

      $$\begin{array}{rl}{h}_t& =\bar{A}{h}_{t-1}+\bar{B}{x}_t\\ y_t& =C{h}_t\end{array}$$

   2. 卷积实现

      $$\begin{array}{rl}\bar{K}& =(C\bar{B},C\bar{A}\bar{B},...C{\bar{A}}^k\bar{B},...)\\ y& =x\bar{K}\end{array}$$

   通常，我们训练时使用卷积方式进行并行运算，而在infer过程中使用recurrent方式进行实现。

其他主要的SSM的结构主要包括以下一些：

1. Linear attention (Katharopoulos et al. 2020)
2. H3 (Dao, Fu, Saab, et al. 2023)
3. Hyena (Poli et al. 2023)
4. RetNet (Y. Sun et al. 2023)
5. RWKV (B. Peng et al. 2023)

2. Selective State Space Models

下面，我们来看一下文中主要的优化结构，即SSSM模型。

如前所述，虽然SSM可以通过recurrent的方式在模型中加入时间信息，但是却并没有content的信息被加入其中，导致对于以summary为代表的一些需要copy的任务当中就很难获得很好的效果，因此文中在SSM的基础上加入了内容的selection，有些类似于attention或者LSTM当中的遗忘门和输出门，负责对前文进行重点选择。

具体的selective部分的示意图和SSSM的为代码逻辑如下：

最后，文中给出了一个典型的SSSM的实现，也就是文中最终使用的Mamba的模块如下：

3. 实验考察 & 结论

有了上述Mamba模型的介绍，下面，我们来看一下文中给出的针对Mamba模型的具体实验效果。

1. 简单问题上的验证

首先，如前所述，Selective机制的加入是为了增加输出内容对于历史信息的选择性记忆，因此，文中首先在两个Syntactic问题考察了一下Selective机制是否成功实现：

1. Selective Copying
2. Induction Heads

得到结果如下：

可以看到，在两个任务上，S6较之没有selective的S4都有着显著的效果增强。

2. 实际场景效果

然后，既然我们发现selective机制确实已经按照预定的设计发挥了作用，下面我们就来看一下Mamba的整体设计是否能够真正在各类场景下发挥出好的效果，文中主要考查了三个Transformer的典型应用场景：

1. 语言模型
2. DNA模型
3. 语音模型

下面，我们来逐一看一下文中在这三个场景下的效果。

1. 语言模型

首先，文中在当前Transformer效果最显著的LM下面考察了Mamba的效果，得到结果如下表所示：

可以看到，在各个参数量级下，Mamba模型都表现出了不输于GPT的效果，并且模型总的size越小，这种趋势就越明显。

此外，传统的RNN的一大缺点就是缺乏可扩展性，当参数规模增大时模型效果无法有效地同比例增强，因此文中也进一步考察了Mamba模型的可扩展性，考察不同size下模型的效果，得到结果如下：

可以看到，随着模型size地增加，Mamba模型的效果是持续变好的。

2. DNA模型

同样的，文中先考察了一下Mamba模型和其他模型的效果对比，得到结果如下：

可以看到，同样有Mamba模型在长序列DNA模型上的表现优于HyenaDNA模型，且参数量更大的7M模型效果优于1.4M的Mamba模型。

文中同样考察了DNA模型上面Mamba模型的可扩展性，得到结果如下：

可以看到，在DNA模型上同样Mamba莫i选哪个满足可扩展性，参数量越大，模型效果越好，并没有看到显著的瓶颈现象。

3. 语音模型

文中考察的第三个使用场景是语音的场景，得到实验结果如下所示：

可以看到，Mamba模型是全面超越其他当前的语音模型的，且在Mamba系列模型内部做消融饰演的结果也同样由Mamba模型的效果是最好的。

另外，文中也在语音场景考察了Mamba模型的可扩展性，得到结果如下：

可以看到，随着序列长度的增加，Mamba的模型效果可以得到持续地提升。

3. 细节考察

除了上述在实际场景中的效果考察和对比之外，文中还针对Mamba本身的细节特质进行了考察，主要包括两部分的内容：

1. Mamba模型设计初衷的在速度和内存上是否优于Transformer
2. Mamba模型在迭代中的细节设计的消融实验

下面，我们来逐一看一下文中针对这两方面的考察。

1. 速度和内存考察

首先，如开头所说明的，Mamba的设计初衷是为了替代transformer架构，在不损失性能的情况下更加节省时间和内存，因此要证明这个目的，文中首先在相同尺寸的模型下对Transformer以及Mamba模型所需的内存和infer数据进行了一下对比，得到结果如下：

可以看到，确实在执行速度和内存使用上，Mamba都显然是更加优秀的一方。

2. 消融实验

然后，文中对Mamba模型本身的各个结构设计进行了消融实验，证明各个组成部分都是对模型效果有正向作用的，具体包括：

1. 模型结构
2. SSM Layer设计
3. Selective机制的引入部位
4. 参数初始化
5. $\Delta$的投影维度
6. SSM的维度

可以看到：

• S6的SSM结构设计是优于其他的SSM结构设计的
• 模型结构方面的话Mamba模型好于H3架构
• Selective机制的话最好在$\Delta ,A,B$上都加上
• 参数初始化的话最好还是正态分布初始化
• $\Delta$以及SSM的维度越大，模型效果越好。

4. 结论 & 思考

综上，在这篇文献当中，作者提出了一个名为Mamba的模型结构，它是在作者的前作S4以及H3模型的基础上持续优化得到的，效果上宣称可以匹敌Transformer，并且同样具有很好的可扩展性，且避开了Transformer的有限窗口长度的问题。

更牛逼的是，作者确实在LM，DNA以及语音领域展示了Mamba模型的有效性和可扩展性，真的是非常厉害。

不过可惜的是这个模型结构本身我并没有怎么完全看懂，后续可能还是得去看看作者的开源代码来好好理解一下Mamba的结构设计。

至于Mamba是否真的可以干掉Transformer，这个暂时还是让子弹再飞一会吧……