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

梦晨 2023-12-05 13:32:12 来源：量子位

梦晨 发自 凹非寺
量子位 | 公众号 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