Mamba

1、背景

         从早期的循环神经网络（RNN）和三维卷积神经网络（3D CNN），到目前广受瞩目的 Transformer 模型，每一次技术的飞跃都极大地拓宽了我们对视频数据的理解和应用。

         特别是 Transformer 模型，以其卓越的性能在视频理解的多个领域 —— 包括但不限于目标检测、图像分割、以及多模态问答等 —— 取得了显著成就。然而，面对视频数据固有的超长序列特性，Transformer 模型也暴露出了其固有的局限性：由于其计算复杂度呈平方增长，使得对超长视频序列的直接建模变得异常困难。

        在这样的背景下，状态空间模型架构 —— 以 Mamba 为代表 —— 应运而生，以其线性计算复杂度的优势，展现出处理长序列数据的强大潜力，为 Transformer 模型的替代提供了可能。

2、对比

（1）RNN（循环网络）

      RNN在序列中的每个时间步需要两个输入，即时间步 t 的输入和前一个时间步 t-1 的隐藏状态，以生成t时的隐藏状态，最终预测输出。因此RNN只需要考虑当前的输入和上一时刻的隐藏状态。 RNN和SSM是一个本质。

    RNN主要存在两个问题：

      第一个问题在于，虽然每个隐藏状态都是所有先前隐藏状态的聚合，然随着时间的推移，RNN 往往会忘记某一部分信息。限制了学习长距离依赖的能力。（在实践中，一般只包含前面若干步而非之前所有步的隐藏状态）

     第二个问题在于，RNN没法并行训练，相当于推理快，训练慢。

     为什么RNN不能并行训练: 在反向传播过程中，RNN需要计算每个时间步的梯度，并且这些梯度需要沿着时间步依次传播回去。如果同时更新所有时间步的参数，会导致梯度混乱和不稳定。同时，RNN的时间递归特性导致隐藏状态之间存在直接的依赖关系。因此不能并行训练。

（2）Transformer（自注意力）

      <1> 位置编码

         因为Transformer可以并行运算，它并没有按照RNN一样进行排列，然后一步一步计算。因此它不具备天然的位置关系，所以引入了位置编码，来对每个词进行位置处理。

    <2>并行运算以及二次复杂度

     

Transformer主要存在问题：

       现有的ChatGPT等大模型处理长文本算力消耗巨大，背后原因就是Transformer架构中注意力机制的二次复杂度。训练快，推理慢。

3、状态空间

（1）状态空间模型（State Space Model, SSM）

     <1>  SSM 是用于描述这些状态表示并根据某些输入预测其下一个状态可能是什么的模型。

      

​

​

      A表示系统状态如何从一个时间步转移到下一个时间步。它描述了系统内部状态的动态变化。

      B表示了外部输入如何影响系统的状态。它描述了输入信号如何转换为状态的变化。

      C 矩阵定义了系统状态如何映射到观测值。它描述了状态向量如何转换为可观测的输出。

      D 矩阵定义了输入如何直接影响观测值，而不通过状态变量。在某些系统中，输入可以直接影响输出。

    <2> SSM的离散化：

                                                           图  SSM离散化

                                                           图  零阶保持法

                                                       图 连续化到离散化的图像表示 

     <3>  SSM的并行化

       SSM使用CNN的卷积操作来实现并行化，绕开了状态计算。由下图可以得到SSM的输出方程：

       对比离散卷积的定义：

          

         

（2）Mamba

      Mamba是一种状态空间模型(SSM)，建立在更现代的适用于深度学习的结构化SSM (简称S6)基础上，与经典架构RNN有相似之处。

     在SSM中，矩阵A，B，C的值是不随着输入x变化的，即SSM是静态的。这就使得SSM无法进行内容感知推理。相比之下，这些任务对transformer来说相对容易，因为它们根据输入序列动态改变注意力。他们可以选择性地“看”或“关注”序列的不同部分。

    

                                                              图 Mamba架构 

      

与先前的研究相比，Mamba主要有三点创新：

      （1）对输入信息有选择性处理(Selection Mechanism)：  通过有选择地将数据压缩到状态中来实现这一点。（当有一个输入句子时，通常会有一些信息，比如标点，没有太多意义。这些无意义的信息就可以被忽略掉。）

    SSM的循环表示创建了一个非常高效(more efficient)的小状态，因为它压缩了整个历史状态。然而，与没有压缩历史状态(通过attention map)的Transformer模型相比，它的性能要差的多（less powerful）。

    Mamba通过合并输入的序列长度和批次大小，使矩阵B和C，甚至步长∆依赖于输入。

      （2）扫描操作（The Scan Operation）：该算法采用“并行扫描算法”而非“卷积”来进行模型的循环计算(使得不用CNN也能并行训练)。

   上述的选择性保留信息也带来了一些问题：由于这些矩阵(B,C,∆)现在是动态的，它们不能使用卷积表示进行计算。我们只能使用递归表示，而失去了卷积提供的并行化。

                                                              图 扫描操作

 首先可以看到， Ht = Ht-1 * A  + xt-1 * B 。 相当于SSM中的状态方程。在这里称为扫描操作，可以用for循环轻松计算。相比之下，并行化似乎是不可能的。因为只有在我们拥有前一个状态的情况下，每个状态才能计算出来。然而，Mamba通过并行扫描算法使这成为可能。

                                                               图 并行化 

        这里运用了Blelloch工作高效（Work-Efficient）并行扫描。因为A B C是动态的， 且通过X就能得到，所以可以两两并行计算，逐步得到最终解。

     （3）硬件感知算法 （Hardware-aware Algorithm）：GPU的一个缺点是它们在小型但高效的SRAM和大型但略低效率的DRAM之间的传输(IO)速度有限。Mamba通过核融合限制从DRAM切换到SRAM的次数。

     将大多数SSM架构比如H3的基础块，与现代神经网络比如transformer中普遍存在的门控MLP相结合，组成新的Mamba块，重复这个块，与归一化和残差连接结合，便构成了Mamba架构。

 ​

​