Mamba模型初步解析 — Mamba : Linear-Time Sequence Modeling with Selective State Spaces

Mamba模型初步接触 — Mamba : Linear-Time Sequence Modeling with Selective State Spaces

"Mamba"是一种序列建模架构，它采用了称为选择性状态空间模型（SSMs）的结构来优化处理长序列数据的效率和性能，这在语言处理、基因组学和音频处理等领域特别有用。

Selective State Space Model with Hard ware-aware State Expansion（S6）
在这里插入图片描述

Transform的问题

Transformer 结构的核心是自注意力机制层，不管是encoder，还是decoder，序列数据都先经过位置编码后喂给这个模块。

我们在说明Transform的问题之前也对Transform的一些核心的概念来进行进一步的回顾与学习。

它创建了一个矩阵，将每个token与之前的每个token进行比较，矩阵中的权重由token对彼此的相关性决定。

在这里插入图片描述

这里就说明了他的计算过程是一个未进行压缩的形式。矩阵可以一次的进行创建从而可以支持并行化的方法。

推理的curse

在推理的时候存在的问题就是，即使我们前面已经计算了token但是在计算新的token的过程中依然需要计算整个序列的注意力。

在这里插入图片描述

Transformer的核心结构Core Attention对序列长度的内存和计算是成二次关系的，也就是说序列长度扩大为两倍，对于内存和计算量的要求会增大4倍。

在这里插入图片描述

作者将Transformers对时间序列的建模的这种特性概括为，有效但不高效，即可以有效的对长序列建模，但计算效率和内存占用上不高效。

快速推理的RNN结构

对于之前我们提到的RNN推理速度过慢的问题，我们很容易就可以联想到之前学过的RNN网络。

在生成输出时，RNN 只需要考虑前一个隐藏状态和当前输入。这样可以避免重新计算所有先前的隐藏状态，这是 Transformer 会做的事情。

在这里插入图片描述
换句话说，RNN 可以快速进行推理，因为它与序列长度成线性比例增长！理论上，它甚至可以具有无限的上下文长度。

在RNN中每个隐藏状态是所有先前隐藏状态的聚合，通常是一个压缩视图。

在这里插入图片描述

注意，在生成名称"Maarten"时，最后一个隐藏状态不再包含有关单词"Hello"的信息。RNNs 随着时间的推移往往会忘记信息，因为它们只考虑一个先前状态。

RNN这种模型的性质，就决定了RNN无法进行并行的计算只能按照顺序来进行推理。也就有了无法建模长序列关系。

总结：结合之前提到过的这些问题，其实就有了Mamba这种模型产生的想法，既可以进行并行化的计算又可以满足线性的序列关系。

时序状态空间模型SSM

状态空间模型（SSM：StateSpaceModel）20世纪60年代由卡尔曼和其他研究者提出。

在这里插入图片描述

它是一种用于描述动态系统的数学模型，广泛应用于系统控制、信号处理等多个领域。
它利用状态变量描述系统的状态，并通过状态方程和输出方程，描述系统的演化和输出。

连续的SSM

首先连续的SSM的一般形式如下：

state Equation:状态方程 描述系统的动态行为，表示系统状态随时间的变化

$x^{\prime}(t)=A x(t)+B u(t)$

在这里插入图片描述

其中：

X(t)是状态向量，表示系统在时间t的状态
X’(t)是状态向量的一阶导数，表示状态随时间的变化率
A是状态转移矩阵
B是输入矩阵，描述外部输入对系统状态的影响
u（t）是控制输入向量

Output Equation:输出方程(或者叫观测方程)描述系统的输出与状态之间的关系

$\begin{aligned} y(t) & =C x(t)+D u(t) \text { 或 } \quad y(t)=C x(t) \end{aligned}$

在这里插入图片描述

其中：

y(t)是输出向量，表示系统在时间t的输出值。
C是输出矩阵，用于控制状态对输出的影响。
D是直接传输矩阵，用于控制输入对输出值的直接影响。

同时参考一下其他的资料我们同时也给出经典的状态空间方程，SSM：

$\begin{array}{l} h^{\prime}=A h(t)+B x(t)(1 \mathrm{a}) \\ y(t)=C h(t)(1 \mathrm{~b}) \end{array}$

上式给出了-个一维函数或序列 X(t）∈ R 基于隐藏状态 h(t)∈ RN×1 向 y(t）∈ R 映射的过
程，这里的A∈RN×N，B∈RN×1，C ∈R1×N

在这里插入图片描述

注意：矩阵A、B、C和D通常被称为参数，因为它们是可学习的。回到我们的简化视角，我们现在可以将矩阵 A、B 和 C视为 SSM 的核心。

在这里插入图片描述

但是我们对于常用的文本数据来说，大多数的文本序列都是离散化的数据，因此需要引入离散化的SSM的相关概念。

离散的SSM

state Equation:状态方程 描述系统的动态行为，表示系统状态随时间的变化

$x_{k+1}=A x_{k}+B u_{k}$

其中：

Xk：表示系统在K时刻的状态向量。
Xk+1:是系统在K+1时刻的状态向量。
A是状态转移矩阵，描述系统从k时刻到k+1时刻的状态转移
B是输入矩阵，描述外部输入对系统状态的影响
uk:k时刻的输入向量。

Output Equation:输出方程(或者叫观测方程)描述系统的输出与状态之间的关系

$y_{k}=C x_{k}+D u_{k} \text { 或 } y_{k}=C x_{k}$

其中：

yk:是K时刻的输出向量。
C是输出矩阵，用于控制状态对输出的影响。
D是直接传输矩阵，用于控制输入对输出值的直接影响。

离散化的原因：

对于连续的系统，计算机没办法处理，计算机处理数据和执行算法时通常使用离散的时间步长。离散化将连续时间系统转换为离散时间系统，使其能够在计算机上进行数值模拟和计算。
通过连续SSM 的状态方程，我们可以知道任意时刻：输入、状态、状态微分(状态对时间的导数)的关系，但没办法根据上一时刻的状态推测下一时刻的状态。离散化后，可以将系统表示为状态的递推公式，可逐步递推系统状态。

离散化方法

离散化方法主要包括：

前向欧拉方法（Eulermethod)：使用简单的前向欧拉公式将连续时间系统离散化。
后向欧拉方法(Backward Euler method)：使用后向欧拉公式，将系统的状态更新方程以隐式方法表示。
双线性变换(Bilinear transform,Tustin’s method)：将连续时间系统的拉普拉斯变换映射到离散时间系统的z变换，从而保持系统的稳定性特性。
零阶保持(Zero-order hold, ZOH)：假设在每个采样周期内输入信号保持不变，从而将连续系统离散化。

离散化的推导过程

连续的SSM的状态方程为：

$x^{\prime}(t)=A x(t)+B u(t)$

两边同时在区间 [t(k-1),tk) 内求积分 :(公式1)

$x\left(t_{k}\right)-x\left(t_{(k-1)}\right)=\int_{t_{k-1}}^{t_{k}} A x(s) d s+\int_{t_{k-1}}^{t_{k}} B u(s) d s$

用uk表示每个离散时间区间的平均值：（公式2）

$u_{k}=\frac{1}{\Delta_{k}} \int_{t_{k-1}}^{t_{k}} u(s) d s$

将公式2带入到公式1中就可以得到推导后的形式。（公式3）

$x\left(t_{k}\right)-x\left(t_{(k-1)}\right)=\int_{t_{k-1}}^{t_{k}} A x(s) d s+\Delta_{k} B u_{k}$

在得到上面的形式之后就可以发现其中的Buk的部分已经得到了一个近似的解。之后就需要用到之前提到的离散化的方法来对

$\int_{t_{k-1}}^{t_{k}} A x(s) d s$

求一个近似的解。(前向欧拉法、后向欧拉法、梯形法,零阶保持法)

前向欧拉法

在这里插入图片描述

前向欧拉法：计算橙色部分的面积来进行近似的操作。

$\int_{t_{k-1}}^{t_{k}} A x(s) d s \approx \Delta A x\left(t_{k-1}\right)$

我们将其带入公式3并写为离散形式可以得到：

$x_{k}-x_{k-1} \approx \Delta_{k} A x_{k-1}+\Delta_{k} B u_{k}$

进一步进行整理可以得到：

$x_{k}=\left(I+\Delta_{k} A\right) x_{k-1}+\Delta_{k} B u_{k}$

我们根据状态空间模型：可以写出离散化之后的状态方程为：

$x_{k}=\bar{A} x_{k-1}+B u_{k} \text {, 其中: } \bar{A}=I+\Delta_{k} A, \quad B=\Delta_{k} B$

后向欧拉法

在这里插入图片描述
计算对应的橙色矩形面积，近似阴影部分的面积：

$\int_{t_{k-1}}^{\sim} A x(s) d s \approx \Delta_{k} A x\left(t_{k}\right)$

我们将其带入公式3并写为离散形式可以得到：

$x_{k}-x_{k-1} \approx \Delta_{k} A x_{k}+\Delta_{k} B u_{k}$

进一步整理可以得到：

$x_{k}=\left(I-\Delta_{k} A\right)^{-1} x_{k-1}+\left(I-\Delta_{k} A\right)^{-1} \Delta_{k} B u_{k}$

我们根据状态空间模型：可以写出离散化之后的状态方程为：

$x_{k}=\bar{A} x_{k-1}+B u_{k} \text {, 其中: } \bar{A}=\left(I-\Delta_{k} A\right)^{-1}, \quad B=\left(I-\Delta_{k} A\right)^{-1} \Delta_{k} B$

梯形法

梯形法简单的理解就是结合之前的两种方法取中间值的部分来进行求解。

$\int_{t_{k-1}}^{t_{k}} A x(s) d s \approx \Delta_{k}\left(\frac{1}{2} A x_{k-1}+\frac{1}{2} A x_{k}\right)$

带入公式写为离散的形式可以得到：

$x_{k}-x_{k-1} \approx \Delta_{k}\left(\frac{1}{2} A x_{k-1}+\frac{1}{2} A x_{k}\right)+\Delta_{k} B u_{k}$

进一步整理可以得到：

$x_{k}=\left(I-\frac{\Delta_{k}}{2} A\right)^{-1}\left(I+\frac{\Delta_{k}}{2} A\right) x_{k-1}+\left(I-\frac{\Delta_{k}}{2} A\right)^{-1} \Delta_{k} B u_{k}$

最后可以得到状态空间的表达形式为：

$x_{k}=\bar{A} x_{k-1}+\bar{B} u_{k} \text {, 其中: } \bar{A}=\left(I-\frac{\Delta_{k}}{2} A\right)^{-1}\left(I+\frac{\Delta_{k}}{2} A\right), \quad \bar{B}=\left(I-\frac{\Delta_{k}}{2} A\right)^{-1} \Delta_{k} B \quad{ }^{1} \Delta_{k} B$

零阶保持法

零阶保持（Zero-OrderHold，ZOH）用于将离散时间信号转换为连续时间信号。这种方法假设输入信号在每个采样周期内保持恒定不变，即在每个采样点之后，直到下一个采样点到来之前，信号的值不再变化。

我们保持数值的时间由一个新的可学习参数表示，称为步长△。它代表了输入的分辨率。

在这里插入图片描述

首先先写出通过该方法所得到的结论：

连续状态下的状态方程可以写为：

$x^{\prime}(t)=A x(t)+B u(t)$

离散空间状态下的方程可以写为：

$x_{k}=e^{A \Delta_{k}} x_{k-1}+A^{-1}\left(e^{A \Delta_{k}}-I\right) B u_{k}$

进一步可以写为：

$x_{k}=\bar{A} x_{k-1}+\bar{B} u_{k} , 其中:- \bar{A}=e^{A \Delta_{k}} -\bar{B}=A^{-1}\left(e^{A \Delta_{k}}-I\right) B$

其卷积形式的表示为：

$\begin{array}{l} \bar{K}=\left(C \bar{B}, C \overline{A B}, \ldots, C \bar{A}^{k} \bar{B}, \ldots\right)\\ y=x * \bar{K}(4 \mathrm{~b}) \end{array}$

循环与卷积表示形式

注意我们如何可以使用RNN的基础方法来使用这个离散化版本。这种技术给我们带来了RNN的优点和缺点，即快速推理和缓慢的训练。
在这里插入图片描述

卷积表达形式的举例：以y2输出方程来进行推导

$\begin{aligned} y_{2} & =C h_{2} \\ & =C\left(\bar{A} h_{1}+\bar{B} x_{2}\right) \\ & =C\left(\bar{A}\left(\bar{A} h_{0}+\bar{B} x_{1}\right)+\bar{B} x_{2}\right) \\ & =C\left(\bar{A}\left(\bar{A} \cdot \bar{B} x_{0}+\bar{B} x_{1}\right)+\bar{B} x_{2}\right) \\ & =C\left(\bar{A} \cdot \bar{A} \cdot \bar{B} x_{0}+\bar{A} \cdot \bar{B} x_{1}+\bar{B} x_{2}\right) \\ & =C \cdot \bar{A}^{2} \cdot \bar{B} x_{0}+C \cdot \bar{A} \cdot \bar{B} \cdot x_{1}+C \cdot \bar{B} x_{2} \end{aligned}$

$\begin{array}{l} \text { 以此内推，可得 }\\ y_{3}=\mathbf{C} \overline{\mathbf{A A A B}} x_{0}+\mathbf{C} \overline{\mathbf{A A B}} x_{1}+\mathbf{C} \overline{\mathbf{A B}} x_{2}+\mathbf{C} \overline{\mathbf{B}} x_{3} \end{array}$

我们换一种形式来看将y3改写为点积的形式可以写为下的结构，可以进一步的对之前的卷积表达形式来进行说明。

${y_{3}=\left(\begin{array}{llll} \mathbf{C} \overline{\mathrm{AAAB}} & \mathbf{C} \overline{\mathrm{AAB}} & \mathbf{C} \overline{\mathbf{A} \bar{B}} & \mathbf{C} \overline{\mathbf{B}} \end{array}\right)\left(\begin{array}{l} x_{0} \\ x_{1} \\ x_{2} \\ x_{3} \end{array}\right)}$

对于处理的一维的线性序列，我们的卷积表示图可以如下所示。
在这里插入图片描述

$\begin{array}{l} \bar{K}=\left(C \bar{B}, C \overline{A B}, \ldots, C \bar{A}^{k} \bar{B}, \ldots\right)\\ y=x * \bar{K}(4 \mathrm{~b}) \end{array}$

总结：按照规律我们就可以得到之前的卷积的形式。

在进行0填充之后我们的卷积核依次的向后进行移动。就类似于CNN的卷积操作了，这种就可以实现了CNN的卷积的并行计算，但推理的速度相比较于循环的表示形式会有所下降。

在这里插入图片描述

SSM的三种表示

对之前的推理过程和得到的结果来进行表示，其实就可以看出这三种表示，连续的、循环的和卷积的，都有不同的优缺点：

连续表示
循环表示
卷积表示

我们现在具有循环SSM的高效推断和卷积SSM的可并行训练。

在这里插入图片描述

有了这些表示，我们可以使用一个巧妙的技巧，即根据任务选择表示。在训练期间，我们使用可以并行化的卷积表示，在推断期间，我们使用高效的循环表示。

这个模型被称为线性状态空间层（LSSL)这些表示共享一个重要属性， 即线性时不变性 (LTI) 。 LTI 表示 SSM 的参数 A、B 和 C 对于所有时间步长都是固定的。这意味着矩阵 A、B 和 C对于SSM生成的每个token都是相同的。无论向 SSM 提供什么序列，A、B 和 C 的值都保持不变。我们有一个静态表示，它不具备内容感知能力。

线性状态空间层（LSSL)训练和推理使用不同的结构，个人感觉有点类似于YOLO v7中的结构化重参数卷积。

矩阵A的重要意义

矩阵A捕获了之前状态的信息，来构建一个新的状态，网上部分的文章也提到了说。本质上是由矩阵A产生出了隐状态。

那我们如何创建矩阵A？来使得模型获得更多的隐状态信息呢？

由此引出了在Mamba中的一个很重要的技术HiPPO技术。即高阶多项式投影算子。HiPPO 试图将迄今为止看到的所有输入信号压缩成系数向量

我们可以通过HiPPO压缩和重建输入信号的信息。得到重建的信号

在这里插入图片描述

它使用矩阵A来构建一个状态表示，很好地捕捉最近的token并衰减较l旧的token。其公式可以表示如下：

$\mathbf{A}_{n k}\left\{\begin{array}{l} (2 \boldsymbol{n}+1)^{1 / 2}(2 \boldsymbol{k}+1)^{1 / 2} \\ \boldsymbol{n}+1 \\ 0 \end{array}\right.$

在这里插入图片描述
使用HiPPO构建矩阵A显示出比将其初始化为随机矩阵要好得多。因此，与I旧信号（初始token）相比，它更准确地重建了更新的信号（最近的tokn)。

这个矩阵旨在将过去的历史压缩成一个状态，其中包含足够的信息以大致重构历史。

这个矩阵的解释是它产生一个记忆其历史的隐藏状态。

S4模型的引出

总结：它由三个部分组成 Structured State Spaces for Sequences（s4）

State Space Models 状态空间模型(有连续的状态空间表示)
HiPPO for handling long-range dependencies用于处理长距离依赖HiPPO
Discretization for creating recurrent and convolution representations用于创建循环和卷积表示的离散化

然后，HiPPO 被应用于我们之前看到的循环和卷积表示，以处理长距离依赖关系。结果是序列的结构化状态空间（S4），这是一类能够高效处理长序列的 SSM。

在这里插入图片描述

A Selective SSM(S6)概述

在这一部分，我们结合之前的一些概念推导和相关公式和方法的引入。来简单的概述Mamba中的一些主要的创新点（贡献）

S6主要包含的两个部分为：

A selective scan algorithm,which allows the model to filter(ir) relevant information（一种选择性的扫描算法）允许模型过滤不相关的信息
A hardware-aware algorithm that allows for efficient storage of (intermediate) results through parallel scan, kernel fusion, and Pecomtaton 一种硬件感知算法，通过并行扫描、内核融合和重新计算实现对（中间）结果的高效存储

它们共同创建了选择性SSM或S6模型，可以像自注意力一样用于创建Mamba块

A Selective SSM -> S6模型

S4（结构化的状态空间模型）存在的问题

SSM + HiPPO + Structured Matrices = S4

状态空间模型，甚至S4（结构化状态空间模型在语言建模和生成中的某些关键任务上表现不佳，即聚焦或忽略特定输入的能力。）

S4在两个任务上表现不佳，其中一个是选择性复制的任务。SSM的目标是复制输入的部分并按顺序输出它们。

也就是说距离的间隔是发生变化的，而A B C矩阵是固定的，无法执行基于内容的推理

有选择的压缩信息。

在S4的模型之后我们将HiPPO矩阵参数化为Normal Plus Low-Rank (NPLR) 形式

在S4模型中, "结构化"主要指的是将状态转移矩阵A参数化为一种特殊的结构, 即Normal Plus Low-Rank (NPLR)形式。这种结构中:Λ是一个对角矩阵(diagonal matrix)
P和Q是两个低秩矩阵(low-rank matrices), 通常秩为1或2"结构化"实际上是S4模型的核心思想, 通过巧妙地对矩阵结构进行分解和参数化, 同时在表达能力和计算效率之间取得了很好的平衡。在S4模型的NPLR(Normal Plus Low-Rank)参数化中, Normal矩阵指的是一个正规矩阵(Normal matrix)。

在这里插入图片描述