Mamba：5 并行扫描

若在阅读过程中有些知识点存在盲区，可以回到如何优雅的谈论大模型重新阅读。另外斯坦福2024人工智能报告解读为通识性读物。若对于如果构建生成级别的AI架构则可以关注AI架构设计。技术宅麻烦死磕LLM背后的基础模型。当然最重要的是订阅跟随“鲁班模锤”。

Mamba自从出道就一直被拿来和Transformer对比，文章巨多，带着标题的爆炸性字样“颠覆Transfomer”，“全面包围”等等。文章深浅不一，小编认为其背后的设计思路以及演化过程更加重要。为此本专题将从山上有座庙开始，沿着RNN、SSM、S4、HIPPO一路走来。期间抛弃冗杂的文辞，以更接地气的模式协助读者打开新的逻辑思维。

前情回顾

Mamba和Transformer是不同的结构，之前文章带着读者从RNN启程，路过SSM和S4，尚未拜访S5和线性RNN，乃至本次终点站Mamba，它们都属于RNN家族。这类的模式都有相同的起手式，里面蕴含的思维也是极其简单：不在将神经网络应用于两个连续的输入向量，而是将其应用于一个输入向量和神经网络的前一个输出。 犹如一条蛇，头尾相咬。

从左图可以看出这个家族一直在稳定的发展，研究课题和突破的方向主要围绕着两个主题： 模型具备稳定训练能力和可扩展（堆叠）能力。

然而RNN有两个问题需要被解决，否则在实战中就无法很好的运用。

第一个问题是：虽然循环计算能够简化到和卷积计算相同的计算量，然而还是无法跨多个处理器进行并行化运算。即便有很多可用的处理器，任何一次处理中，只有等前面的所有步骤完成，才能处理下一任务。因为下一个任务依赖于上个任务的的输出结果。而卷积层在处理任务的时候只需要查看原始输入，只要有足够的可用处理器，就可以同时在所有输入上并行运算。由于GPU的特殊架构，实际上RNN比CNN慢很多。 即便运算简单，最早的RNN在速度层面还是低于Transformer。

第二个问题是：RNN很难以训练。从理论上讲，单个循环层可以不断地提炼任意长的信息输入，但实际上却不然，最多只学习和提炼最近的数十个输入信息。

RNN 的想法自 20 世纪 80 年代以来就已存在，但由于这两个问题，RNN 发展停滞，而卷积神经网络和Transformer在实践中更加成功。事实上，RNN 在过去十年中几乎没有被使用过。直到去年发表的一篇新论文表明线性RNN可以避免这两个问题，因此线性RNN是高效的长序列模型。事实上Mamba基本上也是属于线性RNN的体系。

S4与S5

S4其实等同于SSM+HiPPO+Structured Matrices，使用结构化状态空间序列 (S4) 层的模型在远程序列建模任务上实现了最先进的性能。下图为S4 Layer的图示，其利用了状态空间模型 (SSM)、HiPPO框架和深度学习来实现高性能。

S5是基于S4层的设计基础，S4层使用许多独立的单输入、单输出SSM，而 S5层使用一个多输入、多输出SSM。S5利用S4来设计自身的初始化和参数化。S5进行高效且大范围的并行扫描计算，然而它却能够维持与S4相同计算效率，同时能够在多个远程序列建模任务上实现最先进的性能。

并行扫描与计算

仔细看看上图的式子。到目前为止，读者应该很熟悉它的推理过程了，若还存在疑问回到之前的章节好好的复习。它是SSM的卷积表达模式。RNN的并行计算一直受到质疑，然而在S5中提出了并行扫描计算解决了RNN的并行计算问题。接下来一起看看它的原理。

首先它定义了一个计算符号：

运算的两个输入qi和qj都是拓扑结构。拓扑结构的意思就是qi这个变量是由两个元素构成，包含两个值，例如qi=（矩阵A，向量B）

上面⊙代表矩阵相乘，圆圈里面有个乘号的代表矩阵和向量相乘。+号就是矩阵相加。如下的推导证明这个符合是满足结合律。

在数学中，结合律(associative property)是二元运算的一个性质，意指在一个包含有二个以上的可结合运算子的表示式，只要运算数的位置没有改变，其运算的顺序就不会对运算出来的值有影响。重新排列表示式中的括号并不会改变其值。例如：(5+2)+1=5+(2+1)=8

先构造一个变量ck，它是个拓扑，里面包含两个元素：

初始化s0为(I, 0)，I是单位矩阵（对角线全是1，其余的为0）。然后按照操作符进行如下的迭代运算：

仔细观察以下，会发现运算输出s1，...，s4的第二项数值就是x1,x2,...,x4。

以前x1,x2,...,x4一直是串行计算的，然而通过构造新的运算符，同时利用ci的结构就可以实现并行计算，例如c1和c2，c3和c4可以同时计算，然后再将两者的计算再次计算一次，得到最终的数值。如下图r4拓扑的最后一项就是x4。

要是想计算x1和x3，可以通过下面的式子同时计算：

下图为S5的架构图，中间加入了并行扫描：

RNN的梯度

神经网络是通过从模型中的每个权重中减去损失函数的梯度来训练。什么是梯度？ 可以这么理解，当神经网络的输出结果之后，可以对比它和实际的结果。然后根据偏差的大小将权重值增加很小的量再次训练，如此反复直到神经网络达到最优。这些偏差的大小和权重梯度（成比例）。它告诉如何更改权重以使神经网络更好。这里跳过一大段的公式推理。

若梯度非常大，那么神经网络权重变化太大，神经网络已经学习的功能就会被破坏。如果梯度非常小，则权重变化不够，神经网络根本不会学到任何东西。极端的情况会导致梯度消失和梯度爆炸。这就是训练RNN变得困难的原因，因此RNN需要保持梯度正确稳定，才能够正常发挥。如何初始化和设定参数就十分关键，下篇文章将随着线性RNN一起来探索探索它。