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

今天终于有点时间来看看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后续的发展只能让子弹再飞一会了。

5.结语

更多细节可以看原文，这里主要还是说的一些大框架的东西和一些神经网络发展历史相关的东西。看不懂也可以学学这篇论文的行文逻辑。