Mamba模型

问题

Transformer模型基于矩阵乘法，注意力机制中Q、K、V（N×d）会进行大量运算。

如：使用Q、K运算得到score的过程中，将Q每一行与K每一列进行运算，需要N^2次点积，每次点积包含d次乘法。即最终的复杂度为O()。考虑到序列长度（L）、batch大小（B）、多层Transformer块的堆叠（M），计算量为,参考。

RNN存在梯度消失的问题，且训练时间长。

状态空间模型

Mamba中的状态空间模型（SSM）的灵感来源于传统的状态空间模型，公式描述如下：

其中，h ( t ) h(t)h(t)是当前的状态量，A AA是状态转移矩阵；x ( t ) x(t)x(t)为输入的控制量，B BB表示控制量对状态量的影响。y ( t ) y(t)y(t)表示系统的输出，C CC表示当前状态量对输出影响，D DD表示当前控制量对输出影响。

 

这个模型是时不变(time-invariant)系统，时不变系统是输出不会直接随着时间变化的系统。任意时间延迟的输入将得到相同时间延迟的输出（ABCD其实可以带上时间参数，那就是时变系统，延时输入会导致不同的输出）。

离散化

 为了方便运算，将上述状态空间模型离散化，使用零阶保持法。微分方程的解如下：

 离散近似的结果为：

 

 

 卷积化

 Mamba中的状态空间模型计算效率最大的提升就在于其序列运算可以卷积化。从0时刻开始向后推导状态空间模型几个时刻后的输出，即可得到如下的形式：

设计适当的卷积核即可将序列运算转化为卷积运算，形式如下所示：

 选择性机制

传统状态空间模型的时序结构导致了其输出状态完全依赖有序的输入数据。对于增减过、顺序打乱过的输入，在一些不相关数据混杂在序列中出现时，状态空间模型就无法对其进行有效处理。

Mamba针对这一情况进行了改进，在对B BB，C CC矩阵进行计算时，加入了选择性机制，即在计算是引入一个额外的线性层，对输入的输入的控制量和状态量进行选择，加强模型对不同输入形式的适应能力

 选择性机制对状态空间模型的改进

 

并行累加计算流程 

mamba结构

 

与Transformer结构类似，Mamba结构也是由若干Mamba块堆叠而成。一个基本的Mamba块结构如图7所示：Mamba块由H3块以及门控MLP组合而成。H3为Hungry Hungry Hippos，是一种状态空间模型的执行方式。Mamba块简化了H3的结构，并与门控MLP结合，添加了残差项防止梯度消失。

参考：

一文通透想颠覆Transformer的Mamba：从SSM、HiPPO、S4到Mamba_mamba模型-CSDN博客

Mamba模型底层技术详解，与Transformer到底有何不同？-CSDN博客

Mamba：2 状态空间模型 - 大模型知识库|大模型训练|开箱即用的企业大模型应用平台|智能体开发|53AI

 原文代码：mamba/README.md 在 main ·状态空格/曼巴 (github.com)

论文：[2312.00752] Mamba：使用选择性状态空间的线性时间序列建模 (arxiv.org)

参考代码：mamba-minimal/README.md at master · johnma2006/mamba-minimal (github.com)