【论文笔记】Mamba: Linear-time sequence modeling with selective state spaces

【引用格式】：Gu A, Dao T. Mamba: Linear-time sequence modeling with selective state spaces[J]. arXiv preprint arXiv:2312.00752, 2023.

【网址】：https://arxiv.org/pdf/2312.00752

【开源代码】：https://github.com/state-spaces/mamba

【参考】：一文通透想颠覆Transformer的Mamba：从SSM、HiPPO、S4到Mamba(被誉为Mamba最佳解读)_mamba模型-CSDN博客（强推）

【PPT】：https://download.csdn.net/download/Word_And_Me_/90409454

一、瓶颈问题

二、本文贡献

三、解决方案

1、State Space Models（SSM）

2、Selective State Space Models

3、Hardware-aware Algorithm For Selective SSMs

4、简化SSM及整体流程

四、实验结果

一、瓶颈问题

Transformer的注意力机制在处理长序列时，计算复杂度为 $O(n^{2})$ ，对内存和计算资源的需求非常大，限制了在长序列任务中的应用。

二、本文贡献

将动态选择机制与状态空间模型相结合，提出了选择性SSM的概念【动态调整状态矩阵和输入矩阵，能够更灵活地选择性保留和遗忘序列中的重要信息，来提升序列模型的表现力】
去掉了传统Transformer中的注意力机制，在处理长序列时，将计算复杂度降低为O(n)；【将先前的 SSM 架构设计（Dao、Fu、Saab 等人，2023 年）与 Transformers 的 MLP 模块合并为一个模块】
设计了一种基于状态更新的高效并行算法，将状态更新过程划分为多个小块进行并行处理

三、解决方案

1、State Space Models（SSM）

状态空间模型（SSM）是一种经典的动态系统建模方法，广泛用于信号处理、时间序列预测等任务。SSM的核心思想是使用一个隐藏状态来描述系统在不同时间步的内部状态，输入信号和状态的相互作用决定了模型的输出。S4 模型【即SSM】由四个参数（Δ、𝑨、𝑩、𝑪）定义，分两个阶段进行序列到序列的转换：

第一阶段（1a 1b）：通过公式 $\overline{A}=f_{A}(\Delta ,A)$ 和 $\overline{B}=f_{B}(\Delta ,A,B)$ ，将连续参数∆，A，B转化为离散参数 $(\overline{A}, \overline{B})$ ，其中 $(f_{A},f_{B})$ 为离散化规则，可用多种规则来实现这一转换。例如下述方程中定义的零阶保持（ZOH）：

$\overline{A} = exp(\Delta A)$

$\overline{B}=\Delta A^{-1}(exp(\Delta A)-I)\cdot \Delta B$

第二阶段（2a 2b和3a 3b）：在参数离散化之后，模型可以使用两种方式计算，即线性递归（2a 2b）或全局卷积（3a 3b）。对于SSM模型，通常是使用全局卷积进行训练【卷积训练可以进行并行训练，提高训练的效率】，使用线性递归进行推理【推理时给定文本是时间延迟的，卷积（全局）需要等到所有内容输入完毕再推理，而线性递归，可以边输入边推理】

上述两个方程可做拆解如下：

a) 假设有一些输入信号 $x(t)$ ，该信号首先乘以矩阵 $B$

b) 上面第一步的结果，加上：上一个状态与矩阵A相乘(矩阵A描述了所有内部状态如何连接)的结果，用来更新状态state

c) 然后，使用矩阵C来将状态转换为输出

2、Selective State Space Models

文中认为，序列建模的一个基础问题是把上下文压缩成更小的状态，从这个角度看：

Transformer注意力机制虽然有效果但是效率不高，其需要显式地存储整个上下文（KV缓存），会导致训练和推理消耗算力大【例如人类在每写一个字之前都要把前面所有的字和输入都复习一遍】
RNN的推理和训练效率高，但性能容易受到对上下文压缩程度的限制
SSM推理中，ABC矩阵不会随着输入的不同而不同，即无法针对不同的输入进行针对性的推理

针对SSM中的问题，文中设计了一个选择机制，通过“参数化SSM的输入” 让模型对信息有选择性处理，以便关注或忽略特定的输入。

文中通过 $s_{B}(x)=Linear_{N}(x)$ ， $s_{C}(x)=Linear_{N}(x)$ ， $s_{\Delta }(x)=Broadcast_{D}(Linear_{1}(x))$ 和 $\tau_{\Delta }=softplus$ 来对SSM中的三个参数进行处理，是其可以根据输入可以动态变化。【Broadcast是广播操作，softplus是一种激活函数】

矩阵B、C的大小从原来的(D,N) 变成了(B,L,N)【B为batch size，C为sequence length】

3、Hardware-aware Algorithm For Selective SSMs

由于改进之后的SSM中的B、C矩阵是动态的了，故不能与S4一样进行卷积训练（卷积需要固定的内核），因此其训练中也失去了并行训练的优势。为此，文中提出了并行扫描算法。

扫描操作：

每个状态比如 $H_{1}$ 都是前一个状态比如 $H_{0}$ 乘以 $\overline{A}$ ，加上当前输入 $X_{1}$ 乘以 $\overline{B}$ 的总和，这就叫扫描操作(scan operation)，可以使用 for 循环轻松计算，然这种状态之下想并行化是不可能的(因为只有在获取到前一个状态的情况下才能计算当前的每个状态)

文中假设系统执行的顺序与关联属性无关，故分段计算序列并迭代地组合他们，即动态矩阵B和C以及并行扫描算法一起创建：选择性扫描算法。【可以将时间复杂度由O(n)变为O(n/t)，其中t代表用于执行任务的处理器或计算单元的数量】

此外，为了让传统的SSM在现代GPU上也能高效计算，Mamba中也使用了Flash Attention技术

简而言之，利用内存的不同层级结构处理SSM的状态，减少高带宽但慢速的HBM内存反复读写这个瓶颈
具体而言，就是限制需要从 DRAM 到 SRAM 的次数(通过内核融合kernel fusion来实现)，避免一有个结果便从SRAM写入到DRAM，而是待SRAM中有一批结果再集中写入DRAM中，从而降低来回读写的次数

4、简化SSM及整体流程

将大多数SSM架构比如H3【Hungry hungry hippos: Towards language modeling with state space models】的基础块，与现代神经网络比如Transformer中普遍存在的Gated MLP相结合，组成新的Mamba块，然后重复这个块(且与归一化和残差连接结合)，便构成了Mamba架构。