Transformer挑战者出现！斯坦福CMU联合团队，开源模型及代码_风闻

　　梦晨 发自 凹非寺

　　量子位 | 公众号 QbitAI

　　现在ChatGPT等大模型一大痛点：

　　处理长文本算力消耗巨大，背后原因是Transformer架构中注意力机制的二次复杂度。

　　FlashAttention作者Tri Dao参与提出的新架构，成为有力挑战者，引起大量关注：

　　Mamba（曼巴，一种蛇），在语言任务上击败/匹配Transformer性能，具有线性复杂度和5倍推理吞吐量。

　　具体来说，Mamba在语言、音频、DNA序列模态上都实现SOTA。

　　在最受关注的语言任务上，Mamba-3B超越同等规模的Transformer，与两倍大的Transformer匹敌。

　　并且相关代码、预训练模型checkpoint都已开源。

　　两位作者的解读都获得大量转发。

　　有网友发现，连在线预测平台上的“Transformer在2027年还是SOTA吗？”都在这一天出现明显下降。

　　有选择处理信息+硬件感知算法。

　　Mamba是一种状态空间模型（SSM，State Space Model）。

　　建立在更现代的适用于深度学习的结构化SSM（S4, Structured SSM）基础上，与经典架构RNN有相似之处。

　　在先前研究的Mamba主要有三点创新：

　　对输入信息有选择性处理

　　硬件感知的算法

　　更简单的架构

　　选择性状态空间模型

　　作者认为，序列建模的一个基础问题是把上下文压缩成更小的状态。

　　从这个角度来看，注意力机制虽然高性能但低效率，需要显式地存储整个上下文（也就是KV缓存），直接导致训练和推理消耗算力大。

　　类RNN的循环神经网络具有有限的状态，高效，但性能受到对上下文压缩程度的限制。

　　Mamba的解决办法，是让模型对信息有选择性处理，可以关注或忽略传入的内容，即使状态大小固定也能压缩上下文。

　　一个直观的类比：

　　Transformer就像人类每写一个字之前，都把前面的所有字+输入都复习一遍，所以写的慢。

　　RNN每次只参考前面固定的字数，写的快，但容易忘掉更前面的内容。

　　Mamba每次参考前面所有内容的一个概括，越往后写对前面内容概括得越狠，丢掉细节保留大意。

　　在其前身结构化状态空间模型（S4）中，四个参数A、B、C、∆都是固定的，不随输入变化。

　　在Mamaba中，作者让这些参数B、C、∆成为输入的函数，让模型能够根据输入内容自适应地调整其行为。

　　硬件感知的状态扩展

　　为了让古老的SSM在现代GPU上也能高效计算，Mamba中使用了FlashAttention同款技术。

　　核心思想是利用内存的不同层级结构处理SSM的状态，减少高带宽但慢速的HBM内存反复读写这个瓶颈，具体来说：

　　在更高速的SRAM内存中执行离散化和递归操作，再将输出写回HBM。

　　通过并行扫描算法实现并行化。

　　当输入从HBM加载到SRAM时，中间状态不被保存，而是在反向传播中重新计算。

　　简化的SSM架构

　　将大多数SSM架构的基础块，与现代神经网络中普遍存在的门控MLP相结合，组成新的Mamba块。

　　重复这个块，与归一化和残差连接结合，构成Mamba架构。

　　实验结果

　　Mamba在Chinchilla缩放定律下预训练时，语言任务优于同类开源模型。

　　对比对象中的**Transformer++**为标准GPT-3架构加上谷歌PaLM和Meta Llama中的改进方案，也就是已知最强Transformer配方。

　　下游任务上，每个规模尺寸的Mamba都是同类最佳，并且通常与两倍规模的基线性能匹配。

　　特别是当序列长度增加到512k时，相比使用FlashAttention-2的Transformer快几个数量级，而且不会内存不足。

　　Transformer的下一步？

　　最终，Mamba是第一个真正实现匹配Transformer性能的线性时间序列模型，无论是在预训练困惑度还是下游任务评估方面。

　　并且在音频和DNA序列建模上也优于之前的SOTA模型，表现出一定的通用性。

　　作者在结论中提出，Mamba是通用序列模型骨干的有力候选者。

　　Stability AI创始人当即表示关注。

　　英伟达科学家Jim Fan也对Transformer的挑战者出现感到兴奋。

　　论文两位作者Albert Gu和Tri Dao，博士都毕业于斯坦福大学，导师为Christopher Ré。

　　Albert Gu现在是CMU助理教授，多年来一直推动SSM架构发展。

　　他曾在DeepMind 工作，目前是Cartesia AI的联合创始人及首席科学家。

　　Tri Dao，以FlashAttention、FlashDecoding系列工作闻名，现在是普林斯顿助理教授，和Together AI首席科学家，也在Cartesia AI担任顾问。

　　Cartesia AI公司介绍中提到致力于基于新架构构建下一代基础模型，现在看来主要就是指创新的SSM架构。

　　联创及CEO Karan Goel同为斯坦福博士毕业，也是Mamba的前身S4论文作者之一。

　　对于Mamba的下一步，在论文中有提到“探索新架构是否能适用于Transformer已建立起的丰富大模型生态”。

　　其中包括微调、自适应、提示学习、上下文学习、指令微调、RLHF、量化……也就是要把基础模型发展成GPT-3.5、Llama同类的助手模型了。

　　但作者也提到，目前的实验规模较小，要全面评估SSM是否能与Transformer和其他架构如RWKV、微软RetNet竞争，至少还需要验证7B规模。

　　在扩展SSM的过程中，还会遇到新的工程挑战和对模型的调整，论文中没有涉及。

　　最后，Albert Gu还分享了为什么把新架构起名为一种毒蛇的名字：

　　速度快、对序列建模问题很致命、前身S4是SSSS（嘶嘶嘶嘶）。

　　论文：

　　https://arxiv.org/abs/2312.00752

　　参考链接：

　　[1]https://twitter.com/_albertgu/status/1731727672286294400

　　[2]https://twitter.com/_albertgu/status/1731727672286294400