Mamba：具有选择性状态空间的线性时间序列模型

为了提高Transformer处理长序列的效率，研究人员们提出了线性注意力、门控卷积、循环模型、以及SSMs。其中，SSMs模型虽然能够高效的处理长序列数据，但是性能却没有基于注意力的优越。Mamba作者们发现这种模型性能不足的主要原因是无法执行基于内容的推理，由此作者们对SSMs模型做出了如下创新：

• 把SSMs的参数变为输入的函数，从而允许模型选择性的传递或遗忘信息，以提升其在离散模态上的能力。实际上，之前的SSMs模型属于时间和输入不变的模型，因此可利用卷积的并行计算加快模型的学习速度。
• 由于这种改变阻碍了高效卷积的使用，因此在当前循环模式下设计了一个硬件意识的并行算法。
• 通过把之前的SSM架构与Transformer中MLP相结合形成一个块，从而构建了Mamba架构。

选择性状态空间模型

动机：选择作为压缩的一种方式

作者们序列建模的根本问题为把上下文压缩到更小的状态。从这个视角看序列模型的权衡：注意力属于有效但不高效，因为它没有确切的压缩上下文，自回归推理的时候需要存储整个上下文，因此推理速度很慢；循环模型高效，但受限于上下文压缩状态的性能。

为了理解该原则，如图1所示，作者们聚焦在两个合成任务：

• 选择性复制：该任务把复制任务改为变换需要记忆的tokens的位置。它需要上下文推理记忆相关的tokens，过滤掉不相关的tokens。
• 归纳头：需要上下文推理，知道何时产出正确的输出。

图1（左）复制任务的标准版本涉及输入和输出元素之间的恒定间距，并且可以通过线性递归和全局卷积等时不变模型轻松解决。 （右上）选择性复制任务在输入之间具有随机间距，并且需要时变模型，该模型可以根据输入的内容选择性地记住或忽略输入。 （右下）归纳头任务是联想回忆的一个例子，它需要根据上下文检索答案，这是 LLM 的一项关键能力。

这些任务揭示了线性时不变模型的失效模式。从循环的视角，不变动力学转换无法从上下文选择合适的信息，或以一个依赖输入的方式影响序列传递的隐藏状态。从卷积的视角，虽然全局卷积可以解决复制问题，但是无法解决选择性复制任务。确切的说，输入和输出之间的变化，不能被静态卷积核建模。

总的来说，序列模型的有效性和高效率性的权衡被压缩状态的性能所决定：高度有效的模型必须有一个小状态，而高效率的模型一定要拥有从上下文中包含所有必要信息的状态。由此，作者们提出构建序列模型的根本原则是选择性或上下文意识的能力，从而聚焦或过滤输入到一个序列状态。特别的，一个选择机制控制信息传播方式或按序列交互的方式。

利用选择提升SSMs

Mamba作者们通过使影响交互的参数依赖于输入的方式，把选择机制包含到模型中，例如：RNN的循环动力学或CNN中卷积核。

算法1和算法2阐述了选择机制。两种算法主要的差别是$\Delta,\mathbf{B},\mathbf{C}$为输入的函数。这种方式使SSMs模型与卷积不等价。

确切的说，$s_{B}(x)=Linear_{N}(x),s_C(x)=Linear_{N}(x),s_{\Delta}(x)=Broadcast_{D}(Linear_1(x)),\tau_{\Delta}=softplus$

其中，$s_{\Delta},\tau_{\Delta}$的选择与RNN门控机制有关。

选择性SSMs的高效实现

设计的选择机制用于克服LTI模型的局限性。同时，需要提高其计算效率。为了实现该目标，作者们提出了三种技术：核融合、并行扫描、以及再计算。这些技术背后的观测是：

• 原始的循环计算利用$O(BLDN)$FLOPS，而卷积计算利用$O(BLDlog(L))$FLOPS，且前拥有更低的常量系数。因此，对于长序列和不大的状态维度$N$，循环模型可利用更少的FLOPs。其中，$B$为batch-size，$L$为每个batch的长度，$D$为通道数量，$N$为隐藏层状态维度。
• 主要的挑战是：循环的序列本质和较高的内存利用。为了应对后者，正如卷积模式，作者们尝试不完全实现整个状态$h$。

主要的思想是利用现代GPU加速器的特性，在更高效的显存层级实现状态$h$。特别的，大多数运算受限于内存带宽。这包含扫描运算，利用核融合减少显存IOs量。

确切的说，与其扫描大小为$(B,L,D,N)$的输入$(\bar{\mathbf{A}},\bar{\mathbf{B}})$，不如直接加载参数$(\Delta,\mathbf{A},\mathbf{B},\mathbf{C})$从缓慢的HBM到快的SRAM，在SRAM中直接离散化和循环，再把最终的输出$(B,L,D)$写回到HBM。

最后，为了避免存储中间状态，利用再计算减少内存需求：在反向传播时，中间状态不是被存储而是再计算。由此，与FlashAttention相一致，融合选择扫描层拥有相同的内存需求。如图2所示，选择性状态模型概览。

图2 选择性状态空间模型概览

相关思考

本质上，Mamba是对SSMs在性能和效率之间平衡的优化，使其能够充分利用硬件提升效率，且达到性能最大化。