Mamba一作预告新架构！长文论述Transformer≠最终解法

Mamba一作最新大发长文！
主题只有一个，即探讨两种主流序列模型——状态空间模型（SSMs）和Transformer模型的权衡之术。


简单介绍下，Mamba就是一种典型的SSMs，它建立在更现代的适用于深度学习的结构化SSM基础上，与经典架构RNN有相似之处。
在最受关注的语言任务上，Mamba-3B超越同等规模的Transformer，与两倍大的Transformer匹敌，因此被视为Transformer架构的有力挑战者。
现在，Mamba一作将自己去年的几场演讲整合成一篇科普长文，向大众分享了如下观点：
Attention虽强，但不是万能。Transformer≠最终解法，而是阶段性最优。“让每个FLOPs都有意义”才是架构设计的最终目标。将SSM层与注意力层按一定比例混合能带来更强大的模型。
而且他还提前剧透，几天后将发布“架构领域的下一个重大进展”。


虽然还不知道具体内容，但他目前扔出来的消息已经足够大家抓耳挠腮一段时间了。
因为他提到了一个重要观点——注意力机制的缺点实际上并不是它的二次复杂度。
要知道之前大家都认为，ChatGPT等大模型之所以处理长文本算力消耗巨大，背后原因是Transformer架构中注意力机制的二次复杂度。
而现在，这样的共识或许即将被推翻~
不过好消息是，即将推出的新架构能够和Transformers兼容。


那么在迎来新架构之前，先让我们完整回顾下SSMs和Transformers的“世子之争”吧（doge）。
SSMs就像人类的大脑
一上来，作者先定义了什么是状态空间模型（SSMs）？


方程看不懂不要紧，只需要知道它可以通俗理解为循环神经网络（RNN）的现代版。
更直观的类比如下：
Transformer就像人类每写一个字之前，都把前面的所有字+输入都复习一遍，所以写的慢。
RNN每次只参考前面固定的字数，写的快，但容易忘掉更前面的内容。
而以Mamba为代表的SSMs每次参考前面所有内容的一个概括，越往后写对前面内容概括得越狠，丢掉细节保留大意。
这一工作方式有点像人类的大脑——不断接收新信息（输入），并将其压缩、总结成一个固定大小的“隐藏状态”（即模型的内部记忆），一旦模型需要处理新信息时，它只与这个总结过的“记忆”互动，而不是回顾所有旧细节。
这也意味着，SSM相比其他架构更适合处理长序列信息，而且它还具备两大优势：
第一，非常适合处理非结构化或“低分辨率”数据。实验结果表明，采用SSM架构的Mamba在语言、音频、DNA序列模态上都实现了SOTA。
第二，处理长序列信息时，其计算成本与序列长度呈线性关系（不会突然急剧上升），且无论输入序列有多长，模型在推理过程中所需的内存量都是固定的（适合资源有限的环境）。
按照作者总结，Mamba的成功得益于SSM的三个关键要素：
（1）状态大小 (State size)
传统RNN通常只有一个较小的隐藏状态，而SSMs通过允许隐藏状态成为一个更高维度的向量，能够存储比旧RNN多N倍的信息。
（2）状态表达能力 (State expressivity)
早期SSMs以固定不变的方式更新状态，适合处理音频（信息变化规律）但不适合处理语言（信息速率变化快、需要选择性记忆）数据。
而Mamba通过引入“选择性SSMs”解决了这个问题，与经典RNN“门控机制”相似，它可以根据当前输入数据来决定哪些信息该记住，哪些该遗忘。
（3）训练效率 (Training efficiency)
尽管状态更大、表达力更强会增加计算难度，但Mamba通过精心的参数化和利用经典的并行扫描算法来解决计算效率问题。
而且它和其他现代循环模型一样，都注重并行化、内存管理以及模型线性度以提高计算效率。
不过作者也提醒，SSMs缺乏对过去信息的精细回忆和精确检索能力。
Transformer模型更像一个数据库
相比之下，Transformer模型更像一个数据库——
会把收到的每一个信息（通常是经过Tokenization处理的“token”）都完整记录下来，并储存在一个叫做“KV缓存”的临时记忆区中。当模型需要处理新信息时，它会回顾并比较所有以前储存过的“token”。


其核心组件是自注意力机制，所带来的优缺点也非常明显。
一方面，Transformer模型能完美记住并精细处理序列中每一个单独的“token”。
这使得它在处理已经过预处理、每个“token”都具有明确含义的数据时表现出色。如经过分词（Tokenization）处理的文本，每个词都带有语义，Transformer就能很好地利用它们。
缺点就是计算成本高以及过于依赖高质量数据。
而针对Tokenization，作者也探讨了它是否应该存在的问题，并得出如下观点：
尽管Tokenization有实用价值，但强烈建议废除。
作者认为，虽然Tokenization能够将序列长度缩短约5倍，从而显著提高语言模型的效率，但这只是表面现象。
首先，Tokenization违背了深度学习“端到端”的自动学习精神，即模型应该从原始数据中自动学习，而不是依赖人工预处理。
更要命的是，Tokenization在多语言和多模态应用中实施起来非常困难甚至不可能，而且还可能限制模型的Scaling Law和推理能力。
比如众所周知的翻车事件，模型不会数“strawberry”这个词中有多少个R，就有Tokenization的影响。


一句话，作者坚持认为从原始数据中学习才是一种更好的模式。
而且已有实验证据表明，在未经Tokenization处理的数据上，SSMs的表现显著优于Transformer，即使Transformer被允许使用更多的计算资源。这进一步强调了Transformer在处理非语义化“token”数据时的弱点。
至此可以小结一下，SSMs和Transformer模型可谓各有千秋。


那么能不能将二者结合一下呢？
混合一下性能更佳
答案是yes！
作者发现，将两种类型的信息处理方式结合起来，可能会产生更强大的效果。
这类似于人类智能既有大脑的模糊记忆，又有外部数据库的精确检索能力。


多项独立研究表明，在这些混合模型中，SSM层与注意力层之间的最佳比例大约在3:1到10:1之间。
如此也说明，Attention并非All You Need。
作者明确表示，他的主张不仅仅关乎计算效率（尽管简化数据可以减少注意力机制的二次复杂度开销），而是一个更强烈的声明，即Transformer在建模能力上存在固有的局限性。
最后，他也提到了已经为人熟知的Scaling Law定律。
在他看来，虽然Transformer目前很流行，但它们远非计算资源转化的最优选择，而要设计新的架构，一个重要衡量标准是每个FLOPs（算力）能否物尽其用，快速转化为模型能力。


一言以蔽之，未来的方向可能是结合两者的优势，并开发能够直接处理原始数据的模型。
不知道新架构又能带来多大惊喜？