【花师小哲】鉴定网络热门（？）AI（26）——Mamba：被拒稿的黑马，优美的网络架构

作者：花师小哲-中二发布时间：2024-03-01

今天终于有点时间来看看Mamba的论文了。我在之前一些专栏中也提到过这个模型，不过一直没时间来看。可能非AI行业的朋友都不知道这个模型，相比于Sora之类的，Mamba确实没有那么出圈，但不能说这个模型不重要。

我看不少业界大佬都表扬过这个模型，一些大佬还指出Mamba的论文写的非常漂亮，包括现在已经有很多方法在使用Mamba做事情了。但是很可惜的是，这篇文章被ICLR给拒稿了。

作为专业性比较强的一篇论文，这里不深入讲解Mamba的细节设计，一方面是这个系列的目的还是给大众科普的，另一方面是我也没吃透，虽然不深入讲解，但涉及到的名词还是非常多的，很多也是我之前专栏没提到过的，还是有一定难度的。本文的重点还是说一说Mamba设计理念上优秀的点。

1.从不可能三角说起

对于神经网络模型，长期以来存在着一个“不可能三角”，这个三角是在RetNet中提出的，如图所示：

简单来说，就是并行训练（这往往意味着高效训练）、低成本推理和良好的扩展性能三者不可得兼，往往只能够照顾到两个。

包括图中的著名的Transformer模型就舍弃了低成本推理，当然这是合适的，毕竟如果要做权衡，那良好的扩展性能肯定是优先考虑，然后考虑成本的话肯定还是训练成本高，所以推理被往后放放真的是情有可原。现在Transformer依然是主流模型，包括之前讲的MoE也不是要取代Transformer，而是和Transformer融合，或者说MoE和Transformer只是正交架构而已。

但是人总是想要完美的，所以要把握三者的想法一直都有，包括RetNet、RWKV等方法都试图解决掉“不可能三角”，本文的Mamba也是一样，Mamba修改自状态空间模型SSM

2.状态空间模型SSM

SSM虽然不是主流，但也是在Transformer的“威压”下一直在被研究的模型。

状态空间模型实际上并不是AI领域独有的概念，在其他领域也很常见。举个物理方面的例子，就大家中学就经常见到的物理模型，例如一个用力拉一辆小车在平面上移动吧。

我们将其建模成具有系统输入、系统输出和状态量的一个系统。其中，系统输入就是拉力，系统输出是位移量，状态可以是位移、速度、加速度等等。我们怎么解题呢？一般就是根据输入算出状态，并且根据状态计算输出，简单来说是一个两步走的过程。并且在整个计算过程中状态是会变化的，我们计算完后可以更新这个状态。

神经网络的SSM模型也是类似的方式。可以看如下的图：

当然传统的SSM虽然相比Transformer有一定的优势，例如它可以以CNN的方法进行训练，以RNN的方法进行推理，以此来保证并行训练和低成本推理（为什么会这样之后稍微讲一下，不过不理解也没关系，知道这个结论就好）。理所当然，良好的扩展性能就难以保证了。特别是SSM并不擅长做离散和信息密集数据的建模，也就是不适合作为像ChatGPT这样的语言模型来使用。

3.压缩、选择、上下文与RAG

我们简要讲讲RNN和Transformer等基本网络

RNN，循环神经网络，也算是自然语言处理（NLP）领域的元老级架构了，它的特点是每次输入一个序列元素，经过计算并输出，也就是无论输入有多长，每次的计算量可以看做是恒定的，或者说RNN对不论多长的序列都压缩到恒定的大小

Transformer则不然，其核心的attention机制就是“面面俱到”，所以Transformer是完全不压缩，所以大多数Transformer其实都是豆腐块，方方正正的。

例如我们有一句话“旺财是一只狗”，我们先分词，假设拆分成“旺财”“是”“一只”“狗”，分词的每个结果可以叫做token，那么RNN每次只接受一个token，即先接受“旺财”，然后接受“是”，每次计算时间是差不多的。而Transformer会接受全部输入，然后计算attention分数，即计算“旺财”和“是”之间的分数、“旺财”和“一只”之间的分数...，每两个token都要算，也就是说，Transformer的推理时间是和输入长度的二次方成正比的，这样推理起来就慢得多。

当然，我们注意到，实际上“旺财”和“狗”之间的“相关性”才是主要的，其他分数反而没有那么重要，也就是说Transformer中的attention分数中实际上是存在很多冗余的。

CNN，卷积神经网络。虽然大家熟悉的CNN可能主要是用在计算机视觉（CV）领域的，但在自然语言处理领域依然能用，就是输入一段序列就可以了。传统SSM因为具有线性时不变性（LTI，简单理解就是状态中有些量是不变化的），所以训练时可以用CNN的方式做并行训练。线性时不变性恰恰是影响SSM性能的重要因素，但是如果不具有LTI，往往就不能做并行训练了。

OK，上面没看懂不要紧，这实际上是在说压缩的越厉害往往推理越快，但是效果往往会变差，这也很好理解，你不能既要求马儿跑，又要求马儿不吃草。然后另一个结论是强性能和并行训练似乎是矛盾的。这似乎是一个死局，但Mamba还是想办法尝试解决这些问题。

Mamba设计了选择性状态空间模型。

我猜测这也是借鉴了现在LLM（Large language model）研究的很多结果才能想到这种方法，简单来说，并不是所有输入都是有用的，那我们加一个门电路只选择最有用的信息不就好了吗？这种想法基本是借鉴了现在的ICL（上下文学习）、RAG（检索增强）的思路，这些方法强调在众多上下文中找到真正有用的东西。

对于第二个矛盾，论文作者选择了硬件解法，在抛弃了LTI的情况下对架构块进行了修改，并且设计了硬件并行训练方法。

4.实验结果与未来

单纯架构设计的漂亮没有用，还是需要训练一个模型出来和Transformer之类的模型进行比较。

实验表明Mamba的效果是非常不错的，有高质量、快速训练和推理、长上下文三个优点，很多地方是优于很多同规模或更大规模Transformer架构的LLM的，并且还可以做DNA、音频等序列建模。

之前我也说过，架构设计的优秀不一定能推广开，之前的RetNet不也是一定程度上解决了不可能三角嘛。仍然是生态，因为Transformer已经把生态搭建好了，新人想要用Transformer很简单，有大量的教程、套件等可以用，出问题也有能找到人进行请教，新诞生出来的架构需要很长时间来构建自己的生态，如果不是效果真的很好加上转架构成本够低，很多人实际上是不愿意转架构的。你像RetNet现在就没什么消息了，RWKV也是基本上靠大佬站台才能撑起来一篇好论文（RWKV的原本设计者是有些民科性格的程序员，就是整天吹自己的模型多么多么牛的，虽然RWKV也确实有意思）

好在买账Mamba的人还是不少的，很短的时间内就有各种机遇Mamba的研究。不过可惜的是，Mamba本家的论文被ICLR给拒稿了，关于Mamba后续的发展只能让子弹再飞一会了。