文献阅读：Mamba: Linear-Time Sequence Modeling with Selective State Spaces
- 文献阅读：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 x_t xt输出 y t y_t yt以及态函数 h t h_t ht，且训练过程可以并行，然后infer过程可以做到迭代串行的模型结构都可以称之为SSM。
下面就是一个SSM的典型case，它参数主要包括4个部分 ( Δ , A , B , C ) (\Delta, A, B, C) (Δ,A,B,C)，而整体的操作则是包括两个部分：
-
Discretization
A ˉ = e x p ( Δ A ) B ˉ = ( Δ A ) − 1 ( e x p ( Δ A ) − I ) ( Δ B ) \begin{aligned} \bar{A} &= exp(\Delta A) \\ \bar{B} &= (\Delta A)^{-1} (exp(\Delta A) - I) (\Delta B) \end{aligned} AˉBˉ=exp(ΔA)=(ΔA)−1(exp(ΔA)−I)(ΔB)
-
Computation
Computation部分则可以有以下两种等价地表述：
-
线性recurrence实现
h t = A ˉ h t − 1 + B ˉ x t y t = C h t \begin{aligned} h_t &= \bar{A}h_{t-1} + \bar{B} x_{t} \\ y_t &= C h_t \end{aligned} htyt=Aˉht−1+Bˉxt=Cht
-
卷积实现
K ˉ = ( C B ˉ , C A ˉ B ˉ , . . . C A ˉ k B ˉ , . . . ) y = x K ˉ \begin{aligned} \bar{K} &= (C\bar{B}, C\bar{A}\bar{B}, ... C\bar{A}^k \bar{B}, ...) \\ y &= x \bar{K} \end{aligned} Kˉy=(CBˉ,CAˉBˉ,...CAˉkBˉ,...)=xKˉ
通常，我们训练时使用卷积方式进行并行运算，而在infer过程中使用recurrent方式进行实现。
-
其他主要的SSM的结构主要包括以下一些：
- Linear attention (Katharopoulos et al. 2020)
- H3 (Dao, Fu, Saab, et al. 2023)
- Hyena (Poli et al. 2023)
- RetNet (Y. Sun et al. 2023)
- 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机制是否成功实现：
- Selective Copying
- Induction Heads
得到结果如下：
可以看到，在两个任务上，S6较之没有selective的S4都有着显著的效果增强。
2. 实际场景效果
然后，既然我们发现selective机制确实已经按照预定的设计发挥了作用，下面我们就来看一下Mamba的整体设计是否能够真正在各类场景下发挥出好的效果，文中主要考查了三个Transformer的典型应用场景：
- 语言模型
- DNA模型
- 语音模型
下面，我们来逐一看一下文中在这三个场景下的效果。
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本身的细节特质进行了考察，主要包括两部分的内容：
- Mamba模型设计初衷的在速度和内存上是否优于Transformer
- Mamba模型在迭代中的细节设计的消融实验
下面，我们来逐一看一下文中针对这两方面的考察。
1. 速度和内存考察
首先，如开头所说明的，Mamba的设计初衷是为了替代transformer架构，在不损失性能的情况下更加节省时间和内存，因此要证明这个目的，文中首先在相同尺寸的模型下对Transformer以及Mamba模型所需的内存和infer数据进行了一下对比，得到结果如下：
可以看到，确实在执行速度和内存使用上，Mamba都显然是更加优秀的一方。
2. 消融实验
然后，文中对Mamba模型本身的各个结构设计进行了消融实验，证明各个组成部分都是对模型效果有正向作用的，具体包括：
- 模型结构
- SSM Layer设计
- Selective机制的引入部位
- 参数初始化
- Δ \Delta Δ的投影维度
- SSM的维度
可以看到：
- S6的SSM结构设计是优于其他的SSM结构设计的
- 模型结构方面的话Mamba模型好于H3架构
- Selective机制的话最好在 Δ , A , B \Delta, A, B Δ,A,B上都加上
- 参数初始化的话最好还是正态分布初始化
- Δ \Delta Δ以及SSM的维度越大，模型效果越好。
4. 结论 & 思考
综上，在这篇文献当中，作者提出了一个名为Mamba的模型结构，它是在作者的前作S4以及H3模型的基础上持续优化得到的，效果上宣称可以匹敌Transformer，并且同样具有很好的可扩展性，且避开了Transformer的有限窗口长度的问题。
更牛逼的是，作者确实在LM，DNA以及语音领域展示了Mamba模型的有效性和可扩展性，真的是非常厉害。
不过可惜的是这个模型结构本身我并没有怎么完全看懂，后续可能还是得去看看作者的开源代码来好好理解一下Mamba的结构设计。
至于Mamba是否真的可以干掉Transformer，这个暂时还是让子弹再飞一会吧……