Mamba: Linear-Time Sequence Modeling with Selective State Spaces

论文地址: https://arxiv.org/abs/2312.00752
代码: https://github.com/state-spaces/mamba
主要学习Mamba的结构，因此不包含实验及后面的部分

摘要

1. 基础模型几乎普遍基于Transformer架构及其核心注意力模块
2. 许多新框架(如线性注意力、门控卷积和循环模型以及结构化状态空间模型(SSM))被开发出来以解决Transformer在长序列上的计算效率低的问题。但是他们在重要模式如作为语言的一个关键弱点是它们无法执行执行基于内容的推理。
3. 让SSMs参数作为输入函数可以解决其离散模态的弱点，允许模型根据当前字符选择性地沿着序列长度维度传播或忘记信息。
4. 在循环模式下设计了一种硬件感知的并行算法，可以将简化后的SSMs集成到简单的端到端神经网络架构中，不需要注意力，甚至不需要MLP
5. Mamba推理速度比Transformer高5x的吞吐量并且序列长度可以线性缩放，它的性能在高达百万长度序列的实际数据上得到提高。
6. 作为通用序列模型backbone, Mamba在语言、音频和基因组学等多种模态上实现了最先进的性能。

引言

1. 结构化状态空间序列模型(SSMs)有效地成为一类有前途的序列建模架构。它可以解释为由循环神经网络(RNN)和卷积神经网络(CNN)的组合，其灵感来自经典状态空间模型。它可以非常有效地计算为递归或卷积，并且序列长度具有线性或近线性缩放。此外，还拥有在某些数据模式下对远程依赖关系进行建模的原则性机制。然而，他们在对离散和信息密集的数据(例如文本)进行建模方面效果交叉。
2. 提出了一种新的选择性状态空间模型，在多个先前的工作的多个维度上进行了改进，以实现Transformer的建模能力，同时线性缩放序列长度。
   • 选择机制: 确定了先前模型的一个关键限制——以依赖输入的方式有效选择数据的能力(即关注或忽略特定输入)。基于选择性复制和感应头等重要合成任务的直觉，根据输入参数化基于输入的SSMs参数来设计一种简单的选择机制。这使得模型能够过滤掉不相关的信息并无限期地记住相关信息。
   • 硬件感知算法: 所有先前的SSM模型都必须具有时间和输入不变性，才能提高计算效率。硬件感知算法通过扫描而不是卷积来计算模型，但不会具体扩展状态，以避免GPU内存层次结构的不同级别之间的IO访问。
   • 结构: 将先前的SSMs架构的设计与Transformer的MLP块组合成一个块来简化先前的深度序列模型架构，从而形成包含选择性状态空间的简单并且同质的架构设计。
3. 选择性SSM以及扩展的Mamba架构师完全循环模型，其关键属性使其适合在序列上运行的通用基础模型的backbone
   1. 高质量: 选择性作为语言和基因组学等密集模式带来了强大的性能
   2. 快速训练和推理: 训练期间，计算和内存与序列长度呈线性关系，推理期间自回归展开模型只需要每一步恒定的时间，因为它不需要先前元素的缓存。
   3. 长上下文: 质量和效率共同提高了序列长度高达1M的实际数据的性能。
4. 验证了Mamba作为序列FM backbone的潜力
   1. 合成: 在合成任务中，Mamba不仅可以轻松应对，还可以无限长地推理解决方案(>1M的token)
   2. 音频和基因组学: 在音频波形和DNA序列建模方面，Mamba在预训练质量和下游指标方面均优于最先进的模型。在这两种设置中，它的性能随着更长的上下文(达到百万级长度的序列)而提高。
   3. 语言建模: 无论是在训练前困惑度还是下游评估方面, Mamba 是第一个真正实现 Transformer 质量性能的线性时间序列模型。通过高达1B参数的缩放法则，Mamba超过了大范围baseline的性能。

相关工作(SSMs)

为了更好的了解SSMs，该部分在原文的基础上会做一些补充，补充部分为博主不了解部分(博主是CV方向的)

结构化状态空间序列模型(S4)是最近一类用于深度学习的序列模型，与RNN、CNN和经典状态空间模型广泛相关。他们接收特定连续系统的启发，该连续系统通过隐式潜在状态$h(t)\in {\mathbb{R}}^N$映射一维函数或序列$x(t)\in \mathbb{R}\to y(t)\in \mathbb{R}$

具体来说，S4模型由四个参数$(\Delta ,A,B,C)$定义，它们定义了两个阶段的序列到序列的转换
$$\begin{array}{rlrlrlr}& h^{\prime }(t)=\mathbf{A}h(t)+\mathbf{B}x(t)& (1a)& h_t=\bar{\mathbf{A}}{h}_{t-1}+\bar{\mathbf{B}}{x}_t& (2a)& \bar{\mathbf{K}}=(\mathbf{C}\bar{\mathbf{B}},\mathbf{C}\overline{\mathbf{AB}},\dots ,\mathbf{C}{\bar{\mathbf{A}}}^{\mathbf{k}}\bar{\mathbf{B}},\dots )& (3a)\\ & y(t)=\mathbf{C}h(t)& (1b)& y_t=\mathbf{C}{h}_t& (2b)& y=x\ast \bar{\mathbf{K}}& (3b)\end{array}$$

离散化

第一阶段通过固定公式$\bar{\mathbf{A}}=f_A(\Delta ,\mathbf{A})$和$B=f_b(\Delta ,\mathbf{A},\mathbf{B})$将"连续参数"$(Delta,\mathbf{A},\mathbf{B})$转换为"离散参数"$(\bar{\mathbf{A}},\bar{\mathbf{B}})$, 其中$(f_A,f_B)$成为离散化规则，可以使用各种规则，，例如零阶保持(ZOH):
$$\begin{array}{rl}\bar{\mathbf{A}}& =exp(\Delta \mathbf{A})\\ \bar{\mathbf{B}}& =(\Delta \mathbf{A}{)}^{-1}(exp(\Delta \mathbf{A}-\mathbf{I}))\cdot \Delta \mathbf{B}\end{array}$$
离散化与连续时间系统有很深的联系，可以赋予它们额外的属性。然而从机械角度来看，离散化可以简单地视为SSMs前向传递中计算图的第一步。SSMs的替代风格可以绕过离散化步骤并直接参数化$(\mathbf{A},\mathbf{B})$, 这可能更容易推理。

计算

参数从$(\Delta ,\mathbf{A},\mathbf{B},\mathbf{C})\to (\mathbf{A},\mathbf{B},\mathbf{C})$转换后，可以通过两种方式计算模型: 线性递推(公式1)或全局卷积(公式2)。
通常，模型使用卷积训练模式(公式3)来实现高效地并行化(提前看到整个输入序列)，并切换到循环模式(公式2)用于高效地自回归推理(一次看到一个时间步的输入)

线性时间不变性(LTI)

公式1到公式2的一个重要属性是模型的动态随时间保持恒定。换句话说，$(\Delta ,\mathbf{A},\mathbf{B},\mathbf{C})$以及$(\mathbf{A},\mathbf{B})$对于所有时间步都是固定的。该属性称为线性时间不变性(LTI) ，它与递归和卷积密切相关。作者认为LTI SSMs等同于任何线性递归(公式2a)或卷积(公式3b)，并使用LTI作为此类模型的总称。

到目前为止，由于基本效率限制，所有结构化SSMs都是LTI(例如computed as convolutions)。然而，这项工作的核心见解是, LTI模型在对某些类型的数据进行建模时存在根本局限性，而本文的技术贡献涉及消除LTI约束，同时克服效率瓶颈。

结构和尺寸

最后，作者注意到结构化SSMs之所以如此命名，是因为有效地计算它们需要在矩阵$\mathbf{A}$上施加结构。最流行的结构形式是对角线，本文也使用了它。

在这种情况下，$\mathbf{A}\in {\mathbb{R}}^{N\times M},\mathbf{B}\in {\mathbb{R}}^{N\times 1},\mathbf{C}\in {\mathbb{R}}^{1\times N}$矩阵都可以用$N$个数表示。为了对具有D个通道的batch size为B, 长度为L的输入序列$x$进行操作, SSMs独立应用于每个通道。值得注意的是，在这种情况下，总隐藏状态的每个输入的维度为$DN$, 并且在序列长度上计算需要$O(BLDN)$的时间和内存。

一般状态空间模型

可以注意到状态空间模型是一个非常广泛的含义，它简单地表示任何具有潜在状态的循环过程的概念。它被用来指代不同学科中的许多不同概念，包括马尔科夫决策过程(MDP)、动态因果模型(DCM)、卡尔曼滤波器、隐马尔可夫模型(HMM)和线性动力系统(LDS)(机器学习)以及一般地循环模型(有时是卷积)(深度学习)。

在整篇论文中，使用术语"SSMs"来专门指代结构化SSMs类或S4模型，并且这些术语可以互换使用。为了方便起见，还可能包含此类模型的衍生模型，例如那些侧重于线性递归或全局卷积观点的模型，并且在必要时澄清细微差别。

SSMs架构

SSMs是独立的序列转换，可以合并到端到端神经网络架构中(有时也将SSMs架构称为SSNN, 其与SSMs层的关系就像CNN与线性卷积层的关系一样)。作者讨论了一些最著名的SSMs架构，其中许多架构也将作为本文的主要baseline

• Linear attention: 是自注意力的近似，涉及递归，可以被视为退化的线性SSMs
• H3: 推广递归来使用S4, 它可以被看做是一种由两个门控连接SSM的架构。H3还在主SSM层前插入一个标准局部卷积，将其构造呈一个shift-SSM
• Hyena: 使用与H3相同的架构，但用MLP参数化全局卷积代替了S4层。
• RetNet: 在架构中添加一个额外的们，并使用简单的SSM，从而允许使用多头注意力(MHA)的变体代替卷积来提供可并行化的路径计算
• RWKV: 是一种近期的RNN，基于另一种现行注意力近似(attention-free Transformer)设计用于语言建模。其主要"WKV"机制设计LTI循环，可以视为两个SSM的比率。

S4(补充)

参考博客: https://blog.csdn.net/m0_63745825/article/details/138136975

S4在SSMs的基础上有三个升级

• 离散化SSM
• 循环/卷积表示
• 基于HiPPO处理长序列

离散数据的连续化: 基于零阶保持技术做连续化并采样

1. 首先，每次收到离散信号时，都会保留该值直到收到新的离散信号，如此操作导致的结果就是创建了SSMs可以使用的连续信号
2. 保持该值的时间由一个新的可学习的参数表示，成为步长$\Delta$
3. 有了连续的输入信号后，就可以生成连续的输出，并且根据输入的时间步长对值进行采样
   

这样采样值就是离散输出，并且可以针对$\mathbf{A}、\mathbf{B}$按如下方式做零阶保持:
$$\begin{array}{rl}\bar{\mathbf{A}}& =exp(\Delta \mathbf{A})\\ \bar{\mathbf{B}}& =(\Delta \mathbf{A}{)}^{-1}(exp(\Delta \mathbf{A})-\mathbf{I})\cdot \Delta \mathbf{B}\end{array}$$
其中, $\mathbf{A},\mathbf{B}$为离散化矩阵。最终从连续SSMs转变为离散SSM, 使得不再是函数导函数$x(t)\to y(t)$, 而是序列到序列$x_k\to y_k$。所以使用的是$k$表示离散的时间步长而不是$t$。

循环结构表示: 方便快速推理

每个时间步，都会涉及到隐藏状态的更新。比如$h_k$取决于$\bar{\mathbf{B}}{x}_k$和$\bar{\mathbf{A}}{h}_{k-1}$的共同作用，然后通过$\mathbf{C}{h}_k$预测输出$y_k$

模型展开后如下图所示($h_k$始终是$\bar{\mathbf{B}}{x}_k$和$\bar{\mathbf{A}}{h}_{k-1}$的共同作用下进行更新)。

卷积结构表示: 方便并行训练

在经典的图像识别任务中，使用过滤器(卷积核)来导出聚合特征，而SSMs也可以表示呈卷积的形式

而用来表示这个过滤器的SSMs公式为

长距离依赖问题的解决之道——HiPPO

HiPPO: Recurrent Memory with Optimal Polynomial Projections

方法(选择性状态空间模型)

动机: 选择性作为压缩手段

作者认为序列建模的一个基本问题是: 将上下文压缩为更小的状态。事实上，作者从这个角度来看待流行序列模型的权衡。例如，注意力既有效又低效，因为它明确地根本不压缩上下文。这可以从以下事实看出: 自回归推理需要显式存储整个上下文(即KV缓存)，这直接导致Transformers的线性时间推理和二次时间训练缓慢。另一方面，循环模型的高效的，因为它们具有有限状态，这意味着恒定时间推理和线性时间训练。然而，它们的有效性受到该状态压缩上下文程度的限制。

为了理解这一原理，作者重点关注两个合成任务的运行示例

• 选择性复制 任务通过改变要记忆的标记的位置来修改流行的复制任务。它需要内容感知推理才能记住相关标记(彩色)并过渡不相关标记(白色)。
• Induction Heads 任务是一种众所周知的机制，据推测可以解释LLM(大语言模型)的大部分情景学习能力。它需要上下文感知推理才能知道何时再适当的情境中产生正确的输出(黑色)

这些任务揭示了LTI模型的故障模式。从循环试图来看，它们的恒定动态(公式2中的$\bar{\mathbf{A}},\bar{\mathbf{B}}$转换)不能让它们从上下文中选择正确的信息，也不能以输入相关的方式影响沿序列传递的隐藏状态。从卷积的角度来看，全局卷积可以解决普通复制任务，因为它只需要时间感知，但由于缺乏内容感知，他们很难完成选择性复制任务(上面的图)。更具体地说，输入到输出之间的间距变化，无法通过静态卷积核建模。

总之，序列建模的效率与有效性权衡的特征在于它们压缩状态的程度: 高效模型必须具有较小的状态，而有效模型必须具有包含上下文所有必要信息的状态。反过来，文章提出构建序列模型的基本原则是选择性: 或上下文感知能力，以将输入集中或过滤到顺序状态中。

通过选择性改进SSM

将选择机制纳入模型的一种方法是让影响序列交互的参数(例如RNN的循环动态或CNN的卷积核)依赖于输入。

上面两个算法说明了文章使用的主要选择机制。主要区别在于简单地输入的 多个参数$\Delta ,\mathbf{B},\mathbf{C}$成为了函数，以及整个张量形状的改变。

作者特别选择$S_B(x)=Linea{r}_N(x),S_c(x)=Linea{r}_N(x),S_{\Delta }(x)=Boradcas{t}_D(Linea{r}_1(x))$和${\tau }_{\Delta }=softplus$, 其中$Linea{r}_d$是维度d的参数化投影。$s_{\Delta }$和${\tau }_{\Delta }$的选择在后面解释。

选择性SSM的高效实施

卷积核Transformer等硬件友好型架构得到了广泛应用。本节的目标是让选择性SSM在现代硬件(GPU)上也能高效运行。选择机制可以很自然的结合特殊的选择情况。然而，使用SSM的核心限制是其计算效率。

选择性扫描概述: 硬件感知状态扩展

选择机制的目的是克服LTI模型的局限性；与此同时，需要重新审视SSM的计算问题。文章通过三种经典技术来解决这个问题: 内核融合、并行扫描和重新计算:

• 朴素循环计算使用$O(BLDN)$ FLOPs, 而卷积计算使用$O(BLD\log (L))$ FLOPs，前者的常数因子较低。因此，对于长序列和不太大的维度$N$, 循环模式实际上可以使用更少的FLOPs。
• 两个挑战是循环的顺序性和大量的内存使用。为了解决后者，就想卷积模式一样，可以尝试不真正的实现完整状态$h$。

主要思想是利用GPU的特性，仅在内存层次结构的更高效级别中实现状态h。基于扫描操作，使用内核融合来减少内存IO量，从而进行加速。

最后，应用经典的重新计算技术来减少内存需要。

简化SSM架构

选择性SSM与结构化SSM一样是独立的序列转换，可以灵活地合并到神经网络中。Mamba通过将由线性启发的块组成的注意力和MLP合并为一个组件来简化H3架构。如下图所示。这是收到门控注意力单元(GAU)的启发，它对注意力做了类似的事情。

该架构涉及通过可控扩展因子$E$来扩展模型维度$D$。对于每个块，大多数参数$(3E{D}^2)$都在线性投影中($2E{D}^2$用于输入投影, $E{D}^2$用于输出投影)，而内部SSMs贡献较少。相比之下, SSMs参数量($\Delta ,\mathbf{B},\mathbf{C}$的投影和菊展$\mathbf{A}$)要少得多。重复这个块并与标准化和残差连接交织形成Mamba架构。

选择机制的属性

选择机制是一个更宽泛的概念，可以以不同的方式应用，例如应用于更传统的RNN或CNN, 应用于不同的参数或使用不同的Transformer$s(x)$。

与门控机制的链接

最重要的联系: RNN的经典门控机制是SSMs选择机制的一个实例。作者注意到RNN门控和连续时间系统的离散化之间的联系中得到了很好的建立。SSMs中的$\Delta$可以被视为发挥了RNN门控机制的通用作用。作者认为SSM的离散化是启发式门控机制的原则基础。

$$\begin{gathered} 定理1:WhenN=1,A=-1,B=1,s_{\Delta }=Linear(x)并且{\tau }_{\Delta }=softplus,那么选择性SSM递归(算法2)采用形式: \\ \begin{array}{rl}{g}_t& =\sigma (Linear(x_t))\\ h_t& =(1-g_t)h_{t-1}+g_t{x}_t\end{array} \end{gathered}$$
$s_{\Delta },{\tau }_{\Delta }$的具体选择来自这个链接。需要注意的是，如果给定的输入$x_t$应该被完全忽略(合成任务中需要)，那么所有通道$D$都应该忽略它，因此将输入投影投影到1维，然后使用$\Delta$进行重复/广播。

选择机制解读

• 可变间距 选择性允许过滤掉感兴趣的输入之间可能出现的不相关的噪声标记。模型可以机械地过滤掉任何特定输入$x_t$, 例如在门控RNN情况下点那个$g_t\to 0$时。
• 过滤上下文 根据之前的经验，尽管更多的上下文能带来更好的性能，但许多序列并没有随着更长的上下文而得到改善。一个解释是许多序列模型在必要时无法有效地忽略不相关的上下文。另一方面，选择性模型可以随时简单地重置其状态以删除无关的历史记录，因此原则上它们的性能随着上下文长度的增加而单调提高。
• 边界重置: 选择性SSM还可以再边界处重置其状态(如${\Delta }_t\to \infty$或定理1$g_t\to 1$)。
• $\Delta$的解释: 一般来说，$\Delta$控制者关注或者忽略当前输入$x_t$之间的平衡。它概括了RNN门: 一个大的$\Delta$重置状态$h$并关注当前输入$x$, 而小的$\Delta$则保留状态并忽略当前输入。公式1和2可以解释为按时间不长$\Delta$离散化的连续系统，在这种情况下，较大的$\Delta \to \infty$表示系统更长时间关注当前输入(从而忘记当前状态)，而较小的$\Delta \to 0$表示忽略当前的输入。
• $\mathbf{A}$的解释: 虽然参数$\mathbf{A}$也可以是选择性的，但它最终通过$\bar{\mathbf{A}}=exp(\Delta \mathbf{A})$与$\Delta$交互来影响模型(离散化公式4)。因此$\Delta$中的选择性足以保证$(\bar{\mathbf{A}},\bar{\mathbf{B}})$中的选择性，并且是改进的主要来源。
• $\mathbf{B}$和$\mathbf{C}$的解释: 选择性最终的属性是过滤掉不相关的信息，以便序列模型的上下文可以被压缩为有效地状态。在SSMs中，将$\mathbf{B}$和$\mathbf{C}$修改为选择性的可以更细粒度地控制是否让输入$x_t$进入状态$h_t$, 或让状态进入输出$y_t$。这些可以解释为允许模型分别根据内容(输入)和上下文(隐藏状态)来调节循环动态。

其他模型详细信息

• Real vs Complex: 大多数先前的SSMs在状态$h$下使用复数，这对于在感知模式中有效地执行许多任务的强大性能是必要的。然而，完全实值的SSMs在某些设置下似乎工作得很好，甚至可能更好。Mamba使用实际值作为默认值，这对于除一项任务外的所有任务都适用。假设复数与实数的权衡与数据模态中的连续离散谱有关，其中复数对于连续模态(如音频、视频)有帮助，但对离散模态(如文本、DNA)没有帮助。
• 初始化: 之前大多数的SSMs还建议特殊的初始化。Mamba对复杂情况的默认初始化是S4D-Lin, 对真实情况的默认初始化是S4D-Real。
• $\Delta$的参数化: Mamba将$\Delta$的选择性调整定义$s_{\Delta }(x)=Broadcas{t}_D(Linea{r}_1(x))$, 这是由$\Delta$的机制驱动的。作者观察到它可以从1维扩广到更大的维度$R$。作者将其设置为$D$的一小部分，与块中的主要线性投影相比，它使用的参数数量可以忽略不计。广播操作可以被视为另一个投影，初始化为1和0的特定模式。如果这个投影是可训练的，这会导致替代的$s_{\Delta }(x)=Linea{r}_D(Linea{r}_R(x))$被视为低接头应。在实验中, 根据之前的SSM工作将参数$\Delta$(可以视为偏差项)初始化为${\tau }_{\Delta }^{-1}(Uniform([0.001,0.1]))$