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
— 完 —
点这里👇关注我，记得标星哦～
微信扫码关注该文公众号作者