论文地址： Mamba: Linear-Time Sequence Modeling with Selective State Spaces
代码： https://github.com/state-spaces/mamba
天下苦 Transformer 久矣!整整7年了，一个理论上一点也不漂亮，纯粹靠着模型堆砌和工程上的大力出奇迹，统治了整个NLP 领域和如今大模型的发展。然而，它的局限性日益凸显，计算量的庞大，数据和算力需求越来越让凡人仰望。俱往矣，数风流人物，还看今朝。Mamba 模型的破土而出，让我们看到了结束这一切的曙光，也是为什么引起 A1 界广泛兴奋的原因，尤其是它回归时空序列建模，从中寻找灵感并将其与新型注意力机制结合的思路很有启发。
原文较为晦涩，网上技术文章虽然不少，但一方面写文章的人照本宣科念经的比较多，自己压根没懂，导致读的人反而怀疑自己的智商;另一方面看到网上很多吃瓜群众跟着瞎起哄What l can say, Mamba out!你连它是啥都不知道，嘴都没亲一下，就 OUTOUT 啥?
它其实是用一个李指数映射代替了 Transformer 的非线性状态方程。估计大部分人连啥是“李指数”都没听说过。那这个模型到底是什么来头?凭什么能对transformer 的江湖地位形成挑战?原理是什么，与 RNN 等模型的联系是那些?为什么说它是用流体力学的思想指导了时序建模?牛逼之处和未来发展的看点在哪里?子日:知之为知之，不知为不知，是知也。
一、Transformer 的死穴
Transformer 结构的核心是自注意力机制层，甭管是encoder，还是 decoder，序列数据都先经过位置编码后喂给这个模块。
这里有个天然的缺陷，就是自注意力机制的计算范围仅限于窗口内，而无法直接处理窗口外的元素。像极了古语所说:两耳不闻窗外事，一心只读圣贤书。某种程度上造成视野狭窄，信息孤立，缺乏全局观一样，这种机制无法建模超出有限窗口的任何内容，看不到更长序列的世界。
本质上说，它这是通过位置编码，把序列数据空间化，然后通过计算空间相关度反向建模时序相关度。这个过程中忽视了数据内在结构的细腻关联关系，而是采取了一种一视同仁的暴力关联模式，好处是直接简单，但显然参数效率低下，兄余度高，训练起来不易。那怎么办呢?俗话说:魔鬼的归魔鬼，天使的归天使。明明是时序数据，非要用空间化实现注意力机制。这在当年是为了充分利用 GPU的并行能力，非常有效，但并不是万能的。本质上说还是有问题的。让长序列数据建模回归传统，某种程度上说，这是整个SSM 类模型思考问题的初衷和视角。而 mamba 是其中的佼佼者。因此，咱们先从时序状态空间模型开始讲起。
二、时序状态空间模型 SSM
Mamba是基于结构化状态空间序列模型 (SSMs)的，这是2021年提出的工作，算是火了两年了。需要补充读这篇文章。
实际上还是旧瓶装新酒，模型还是时序的，本身就是个 RNN 模型。我们用下面这三幅图来讲清它是怎么来的，长什么样，又怎么用。
实际上就是个RNN模型，但现在不是用神经网络实现这个模型，而是用状态方程，线性代数实现数学建模 1.连续空间的时序建模
很多实际问题都能用左图所示的连续空间模型来建模。尤其是控制理论、信号处理或者线性系统领域特别常见。我们称为 LTI，线性时不变 linear tiem-invariant 系统，用公式表示就是:
线性时不变系统（linear, time-invariant system， LTI） 就是有如下两个性质的一类系统：
线性：如果你同时向系统输入两个不同的信号，那么得到的输出和这两个信号各自向系统的输出之和相同。 用数学语言描述为，如果输入 x1 产生输出 y1 ，输入 x2 产生输出 y2，那么 输入 ax1+bx2 产生的输出为 ay1+by2 。这里 a 和 b 为常数。
时不变性：系统的作用效果不随时间的变化而变化，仅依赖于系统本身的状态。因此，如果输入变化仅是在时间上平移， 那么产生的输出变化也仅仅是相同的时间平移，其他都是一样的。
输入x乘以B+隐状态h乘以A得到隐状态导数，h乘以C得到输出 牺牲通用性，换来特定场景下的更高性能，这就是所有 SSM 模型的最底层逻辑。 (画图:表示逻辑关系)
2.时序离散化与 RNN
连续系统不方便计算机处理，中间图是对它的离散化展开，就是沿时间拉长，模型和RNN 长相上几乎一样。公式与上面类似，只是导数h'(t)改为不同时刻角标，形成递归过程， 这些完全是大学线性系统的内容。
原文在这里给出了一种从连续系统转换为离散系统的 ABC 参数对应关系，这段如果你没有控制理论基础或者线性代数不好的话，基本直接懵逼。不过它其实只是一种方法的举例，学名叫“零阶保持(Zero-Order Hold,ZOH)”。公式不用记，只要知道用了这么一个 delta 函数，经过这么一通运算能得到新的A和B就可以了。这个后面用得到。
其实，还有很多种离散化方法。离散化主要是为了方便计算机处理，同时也是 Mamba的一个技巧，后面会讲它怎么转化成类似 RNN 的门控机制。
3.并行化处理与CNN
如果只是时序建模，SSM 和 RNN 相比也就没啥意思了。它最大的特点就是通过下图卷积实现了计算上的并行化。
看图，就是把长长的链条一下子弄成了输入输出直接对应的样子，隐变量关联关系都跑到中间肚子里去了。先说结论:核心思想是用 CNN 对时序数据建模，借助不同尺度的卷积核，从不同时间尺度上捕获时序特征。 数学上一番推导猛如虎之后能得到下面的公式。
第二个式子和前面的图完全对应，就是个CNN 实现。时序递归状态方程改用卷积操作实现了，ABC参数矩阵变成了K，就是图中绿色框内的卷积核。这就牛逼了哈，拖拖拉拉的时序计算变成了能一口气的并行计算。体现在效果上，借助不同卷积核捕捉局部时间序列特征，能同时对短期和长期依赖关系学习，并行计算提高了训练和推理效率，使得 SSMs 在处理复杂的时序数据时表现出色。
离散卷积的定义：
将两个离散序列中的数，按照规则，两两相乘再相加的操作。计算卷积的过程：序列翻转，移位，相乘，取和。
你会发现，二者在结构上类似:因此，时序的状态空间模型能改写成卷积的形式啦。 到这里，不得不感慨一下。卷积可以说是个极其伟大的数学运算发明，起源于 100 多年前傅里叶变换的工作。它就像是一个“滤镜”，可以帮助我们突出信号中的某些特征，识别和提取信号中的重要信息，因此在信号处理与通信系统领域占据了核心地位。
除此之外，在上面式子中需要注意的是，实际问题中往往会对 ABC矩阵进一步简化假设成更简单的对角阵方便计算，这就是所谓的结构化 SSM，S4 模型。
到这里，再回来看看这幅图。你有没有赫然发现，历史总是重演的，三十年河东，三十年河西。大家熟悉的 CNN、RNN，眼瞅着过时了，没想到又能杀回来。我常提醒大家，想学好深度学习，一定要系统的学，这样你才能高屋建瓴，有全局观。不要一上来总盯着 transformer，但有些人不信邪，往往只想偷懒追风。其实，很多时候,“横看成岭侧成峰，远近高低各不同”。不同模型的巧妙利用和合理组合，也许会迸发新的活力，就像这个 SSM 一样，某种程度上，其实只是换了个名字的CNN 化的 RNN。
我们用生活中的例子来帮你更好理解 RNN 和 SSM 的区别。想象你在读一本书，时序嵌套的 RNN 每次只能读一行，然后把记忆传递到下一行，这种方法只适合处理短故事，故事一长，容易忘记前面的情节。而SSM 并行处理，同时打开所有页看到每行内容，这样就能快速找到和理解整本书，无需逐行传递记忆。这个系统还有两个强假设:线性+时不变。极大的限制了它的应用范围，因为实际系统大多为非线性、时变系统。Mamba本质上就是一个 SSM 模型的改进版，放开了这两个约束。
三、Mamba: 选择性 SSM
解决了 SSM 线性+时不变的缺陷
接下来咱们着重讲解什么是选择性 SSM 。Mamba 主要体现在设计了一种机制，让状态空间具备选择性，达到了 Transformers 的建模能力，同时在序列长度上实现了线性扩展，也就是克服了 Transformers 的缺陷，可处理最长达百万长度的序列，而且效率贼高，与 GPU 硬件适配，比Transformers 快5倍，准确率相当甚至更好。这就是它为啥牛逼起来的原因啦。
核心就是搞懂这幅图，其实大致看看就明白，在时间序列模型中间设计了这么一坨非常类似LSMT的门结构，实现所谓的选择性，BC都带了t变成了时变参数，A虽然没有直接含t，但由于delat函数影响A，所以A其实也是时变的了。下面的蓝色部分就是所谓的选择机制，这个delat别小瞧，它就是前面离散函数，也是一个非线性的函数，一会有大用。要真正明白其中的细节，先要从增加选择性的动机开始讲，但其实简单理解就是把整个系统该用: 一个总开关 delat+若干个旋钮Bt，Ct ==非线性时变系统
下面的解释可以认为是对这种选择的合理化。
1.要解决什么问题
从某种角度看，序列建模的核心就是研究如何将长序列的上下文信息压缩到一个较小的状态中。比如，语言模型实际上就是在一个有限的词汇集合中不断进行转换。有统计表明，3500 多个常用中文字，3000个常用英文词能覆盖90%以上的日常用语。transformer 的注意力机制虽然很有效，但效率低，因为它需要存储整个上下文，导致推理和训练时间较长。前面讲的 SSM 递归模型可以并行处理，但因为它们的状态是有限的(单纯时不变导致)，效率高但有效性受限于状态的压缩能力。
那能不能设计一种模型，平衡一下，实现这种类似的强选择能力。具体来说，本文关注两种能力:一是抓重点的能力(选择性复制任务)。从大量信息中选择和记住关键的信息，忽略不相关的部分。类似于在人群中找到你的朋友或者在一篇文章中找到关键词。比如从下面句子中找出名词。
二是上下文联想/推理能力(诱导头任务)。在处理连续的信息时，能够保持逻辑一致
这就像我们经常说一个人口才好，又能抓住重点，逻辑性又强
那具体怎么改进呢?当然还是在 SSM 基础上。前面我们说过，原来的 SSM 有个 LTI线性时不变系统的强假设(ABC是固定的)，这样就导致难以有效的选择上下文信息。以下图为例，甭管输入x是啥，B自岿然不动，这不行，得想办法让它时变。
前面的全局卷积虽然能用不同的卷积核进行时序特征捕捉，但是缺乏内容感知，也就是不知道输入的重点和逻辑。而 transformer 人家本身是没有这些限制的，别说时不变，连线性系统假设都没有。这么一分析，改进的方向就很明确了哈。放开LTI模型的时不变约束，让模型参数依赖于输入内容不就行了吗?说起来简单，具体看看是怎么实现的。
2.怎么增加选择性
说穿了就是让 B和C由固定的变成了可变的，根据输入xt和它的压缩投影学习可变参数。A为了简化，自身还是不变的。蓝色部分(包括投影及其连线)就是所谓的选择机制，目的是根据输入内容选择性地记忆和处理信息，从而提高对复杂序列数据的适应能力。你看，这个思想是不是像极了LSTM，就是增加开关。
所谓的选择，其实人话理解，就是“掺和”。一次不够，用三条线给 B_t掺和了三次，给Ct掺和了两次，给A掺和了两次，看下面的算法因为函数 \tau 是非线性激活函数因此 delta 是非线性的，所以 ABC 都是非线性时变的，整个系统两个条件都放开了加了几个开关，多接了几条水管子。
这里面的 deltat是前面离散化计算时的参数，我们刚才讲了。投影出来的三条蓝线其实就是S_B,S_C,S_delta 三个选择函数，共享一个投影模块(Project)，主要是为了实现参数共享和计算效率。右边的 GPU SPAM、HBM 我们稍等再解释，主要和硬件实现相关。
这么看 Mamba 的思想也不复杂啊，本质上就是拿 LSTM 的门控思想移植到了 SSM 模型上，放开了时不变的约束放成非线性系统。。
注意:这里的 B/L/N/D 符号乍一看让人很困惑，其实就是张量的维度。
B:批次大小(Batch size)。表示一次输入的数据量的大小。
L:序列长度(Sequence length)。表示每个序列中包含的时间步数。
N:特征维度(Feature dimension)。表示每个时间步的特征数量
D:输入特征维度(Inputfeature dimension)。
对左边的传统 SSM 模型，参数矩阵长这样，都是静止的
Mamba 的参数增加了L输入序列长度，因为是并行计算，且为时变参数，所以每个时间步的参数都不同，需要学习 L长序列时间步长的参数 其中步长大小 delta 像是个放大镜观察窗口，影响信息处理的焦点。步长较小时模型倾向于忽略具体的单词，而更多地依赖于之前的上下文信息。你可以简单的认为，就是靠着它，实现了注意力的选择。拿着放大镜忽远忽近的看。
另外，原文没有详细解释这几个选择函数，让人有些困惑。我这里都给你补上了。
经过这样的优化设计，最终希望达到一个什么样的效果呢?图2给了一个示意图
左边是 LTI的效果，输出只能对规则的输入特征进行发现，而右边上面能自己找重了，带色的尽管开始间隔大小不一，但都能找出来排好队。右下是联想能力的体现再看到黑的后就想到以前后面应该跟着蓝色的。也就是说，对于非线性时变数据备了很强的特征捕捉能力。
3.核心原理:流体力学与李指数映射
明白了上面的内容，你只能算是知其然，还没有知其所以然。梗直哥带你从更高角度来理解 Mamba 的精髓，这部分内容的理解其实是哥的发挥，已然超出原文。Transformer 描述的是粒子运动，通过自注意力机制映射动态调整每个输入的权重，类似粒子间通过牛顿力学相互作用力来动态调整自己的轨迹。训练的过程，就是在用牛顿力学拟合粒子轨迹，每个输入(粒子)独立计算与其他输入的关系。
如果把记忆的流淌比作一个水流管道系统，可以看做一个“李群”，进行各种复杂变换(比如旋转、推移等)。固定矩阵A就是主管道(全局演变路径)，类似于流体运动的全局关系，让系统状态更新有固定的全局路径和规则，因此能表现出更高的灵活性和适应性。Bt/Ct就是阀门或旋钮，delta这个离散化因子，就像是流体力学中的时间步长，决定流体运动的离散时间点。选择机制就像是根据具体情况选择和调整旋钮，控制流体在管道中的流动路径。
Transformer 的自注意力机制粒子运动强调个体的独立性和动态调整，灵活性高，能动态调整权重，捕捉复杂的上下文依赖关系，但计算复杂度高，资源消耗大。Mamba流体运动强调整体的连续性和全局关系，借用独特的矩阵 A的固定性，提供了稳定性和确定性，通过李指数映射实现高效状态更新和决策。
4.实现细节和网络结构
这部分主要研究如何充分利用 GPU 实现选择性 SSM 的并行计算，也就是前面的图和算法。 原文细节很多，大量引用了前人的工作，读起来有些费劲，尤其在没有阅读之前文章的情况下。为此，提出了三种创新的解决方案:内核融合、并行扫描和重计算。
1)所谓的内核融合:就是把离散化和循环在 GPU SRAM 内存中实现，快然后加载和存储参数 ABC矩阵都用 HBM 高带宽内存。这俩简写对 GPU 不熟的同学可能不知道，其实就是一种分层提升效率的新技术，前者快但内存带宽小，后者慢但是带宽大。这是随着 AI崛起的芯片新技术。
2)并行扫描:这是啥意思呢?本来 SSM 好不容易能用卷积并行运算了，但是放开 LTI假设后又回归 RNN 动态模型，每次 Bt不一样导致并行卷积不行了。Mamba为此做了改进，提出了并行扫描技术，scan 就是操作或处理的意思，
3)重计算:这步是为了避免存储反向传播所需的中间状态，在输入从 HBM 加载到SRAM 时在反向通道中重计算。这部分的更多细节在附录 D中。简单小结:这部分讲解的就是内存管理、并行计算和动态计算方面的执行细节。开始看不懂不影响整篇文章主体内容的理解。可以先放放。
好，到这里我们总算把 Mamba 的核心原理讲完了，那它的网络结构具体是怎么实现的呢?是不是有点“干呼万唤始出来，犹抱琵琶半遮面”。
下图是放大版，看的更清楚，修改了一下它用线性投影实现维度调整，以及部分特征提取与转换，输出时用于数据压缩。增加了卷积层捕捉局部的时间依赖性
Selective ssM(状态空间模型)相当于把时序链条又折叠回去，肚子里包含了sA(x)和 T_Δ。用来捕获更广泛的依赖性，包括长距离的时间依赖性和复杂的内容依赖关系。Silu 函数与算法中的 softplus 函数类似，都是用来增加非线性。其实这个 silu 是工程上人为进一步增加的非线性，增强模型表达力。此外，增加卷积模块是为了捕获局部时空特征。
封装后就是一个的模块，如同 transformer 的 decoder 一样可以实现进一步的堆叠，形成大模型。把这个模块和前面的框图对比学习，同时要清楚的意识到哪里是对应的哪些是工程化添加的边角料。当然，这也是为啥再单独起一个名字 Mamba 的原因毕竟和前面讲的 selective ssM 不完全一样了。
接下来看看实验内容。
四、实验
对比 transformer/RNN和Mamba 的训练与推理性能。transformer 训练快推理慢，RNN 训练慢推理快，Mamba结合了二者的优点，都快。实验部分很全面也是一大特点。我们重点看结论。
1.合成任务验证选择机制
选择性复制任务要求模型能够记住并复制序列中的特定单词。表- Mamba 架构与选择性机制结合后的表现优秀，(S6):准确率为 99.8扩展序列长度任务要求模型看到一个二元组(如“Harry Potter”)时，能够记住“Harry并在序列中再次出现时预测“Potter”。Mamba架构，也就是最上面的棕色线，比其他方法要好两倍。
2、语言模型预训练
左图是较短右图是较长的序列长度，
3.DNA 序列
由于大型语言模型的成功，人们开始探索将基础模型范式应用于基因组学。DNA被视为一种由有限词汇组成的离散序列，需要模型处理长程依赖。实验和图表展示了Mamba 架构在 DNA 建模任务中的卓越性能，特别是在处理长序列和扩展模型大小方面。
横轴为参数量，纵轴为困惑度。Mamba(橙色线)的困惑度左图随着参数数量的增加(从约 200K 到约 40M)显著下降，右图随着序列长度保持稳定。下面是物种 DNA分类任务的微调准确率和不同数据集上的扩展定律。随序列长度，Mamba性能更好。
4.音频例子
音频波形建模和生成问题也是序列建模任务。
5.训练推理效率分析
训练效率:Mamba 的扫描实现比标准实现快 40 倍，处理长序列时时间增长最慢(橙线)。对比蓝线，最右边512时，1ms/0.025ms=40倍
6.消融实验分析
主要结论包括:
五、代码实现
GitHub：https://github.com/state-spaces/mamba
Mamba 类:
六、小结与探讨
Transformer注意力机制的窗口小了效果差，大了计算复杂度平方暴涨，的死穴:两难。时序问题空间化出现了瓶颈，单纯的注意力机制有缺陷，并非万能，不是机制本身有问题，而是实现方式。这促使人们思考更换视角。 SSM 模型:从 LTI连续空间线性时不变系统讲起，类似 RNN 离散化，CNN 并行化，但都没离开线性+参数矩阵时不变两个假设。 Mamba 原理:提出了选择机制，其实就是时序门控单元。通过离散化函数 delta 这个非线性总开关控制 ABC矩阵三个小旋钮开关，相当于放开了时不变约束，实现时变。与此同时，在选择性 SSM 这个核心模块之外，还通过增加激活函数进一步提升模型表征力，利用卷积层增强空间特征捕捉力。可以认为Mamba=RNN(变形GRU)+CNN+选择性注意力机制 思想精髓:流体力学系统+李指数映射+固定矩阵A最优主管道=独特的处理方法
