当前位置:AIGC资讯 > AIGC > 正文

Transformer挑战者出现!斯坦福CMU联合团队,开源模型及代码,公司已创办

现在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

更新时间 2023-12-05