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

Estimated read time 1 min read

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

—    —

点这里关注我,记得标星哦~

一键三连「分享」、「点赞」和「在看」

科技前沿进展日日相见 ~  

推荐内容

+ There are no comments

Add yours