Bendi新闻
>
YOCO:打破传统Decoder-only架构,内存消耗仅为Transformer的六分之一

YOCO:打破传统Decoder-only架构,内存消耗仅为Transformer的六分之一

9天前


(本文阅读时间:8分钟)


编者按:近期,微软亚洲研究院推出了一种创新性的 Decoder-Decoder 架构 YOCO(You Only Cache Once)。通过自解码器和交叉解码器的独特架构,YOCO 仅需缓存一次键值对,从而显著降低 GPU 内存的使用。

在模型评估中,YOCO 展现出与同规模 Transformer 模型相媲美的性能,并在语言建模评估、模型大小扩展以及长上下文处理方面具有显著优势。特别是在降低 GPU 内存占用和缩短预填充延迟方面,YOCO 实现了“模型越大,内存越省”,为自然语言处理领域带来了全新的研究和应用范式。


微软亚洲研究院&清华大学最新研究,打破 GPT 系列开创的 Decoder-Only 架构——提出 Decoder-Decoder 新型架构,名为 YOCO (You Only Cache Once)。

YOCO 仅缓存一次键值对,可大幅降低 GPU 内存需求,且保留全局注意力能力。

一张图来看 YOCO 和标准 Transformer 的比较。


在处理 512K 上下文长度时,标准 Transformer 内存使用是 YOCO 的6.4倍,预填充延迟是 YOCO 的30.3倍,而 YOCO 的吞吐量提升到标准 Transformer 的9.6倍。

去年一张“大语言模型进化树”动图在学术圈疯转,模型架构还只有三大类:Decoder-Only、Encoder-Only、Encoder-Decoder。



那么这个新出的 Decoder-Decoder 架构到底长啥样?


嗯,如网友所言,要读的论文又增加了。




话不多说,一起来看。



打破Decoder-Only


YOCO 整体架构设计如下,分为自解码器(Self-Decoder)交叉解码器(Cross-Decoder)两部分。



具体来说,YOCO 由 L 个块堆叠而成,其中前 L/2 层是自解码器,其余模块是交叉解码器。


自解码器利用高效自注意(efficient self-attention)机制来获取键值(KV)缓存


接收输入序列的嵌入表示,并使用高效自注意力来生成中间向量表示;使用因果掩码(causal masking)保证解码的自回归特性;自解码器的输出用于生成全局 KV 缓存。


交叉解码器使用交叉注意力(cross-attention)来重用自解码器生成的共享 KV 缓存


在自解码器生成的 KV 缓存基础上进行堆叠,以获得最终的输出向量;同样使用因果掩码来维持自回归生成;允许交叉解码器层间高效地重用 KV 缓存,减少了对 GPU 内存的需求。


总的来说,自解码器和交叉解码器的模块设计与 Transformer 的解码器层类似,包含交错注意力和前馈网络子层。不过,研究人员还进行了预 RMSNorm、SwiGLU 和分组查询注意力等改进。


两部分之间的区别在于注意力模块。


自解码器使用高效自注意力,如滑动窗口注意力(Sliding-Window Attention)门控保留(gated retention)


而交叉解码器使用标准的多头交叉注意力,Query 向量通过注意力与自解码器产生的全局键值缓存相关联。


推理大幅度省 省 省


实验阶段,研究人员将 YOCO 模型与同体量的 Transformer 模型进行比较。

分析维度有四个:语言建模评估、与 Transformer 比较的可扩展性、长上下文评估、推理优势。

语言建模评估

研究人员训练了一个 3B 参数的 YOCO 语言模型,并根据训练 token 数量(1T 和 1.6T)进行评估。

在 LM Eval Harness 的多个下游任务上,YOCO 与 Transformer 模型 OpenLLaMA-3B-v2、StableLM-base-alpha-3B-v2、StableLM-3B-4E1T 打得有来有回。


可扩展性对比

接着,研究人员在 160M 到 13B 参数规模范围内,分别训练了 YOCO(门控保留和滑动窗口注意力版本)和 Transformer 语言模型。

对比了它们在验证集上的语言模型损失,YOCO 的表现与 Transformer 基本持平:


结果证明 YOCO 在模型大小扩展方面具有很强的可扩展性。

长上下文评估

将 3B 的 YOCO 模型扩展到上下文为 1M,在“大海捞针”等长序列的 needle retrieval 任务上,YOCO-3B-1M 的准确率接近100%。


在多针检索任务上,YOCO-3B-1M 的性能优于一些超 3B 的 Transformer 模型:


此外,YOCO 模型在长序列上的 NLL 随着上下文长度的增加而一致下降,表明 YOCO 能够有效地利用长距离依赖信息进行语言建模:


综上,可见 YOCO 在性能上完全不输 Transformer,关键来看 YOCO 在推理效率上取得的显著提升。

推理优势

研究人员评估了 YOCO 在 GPU 内存占用、prefilling 延迟、吞吐量和服务容量等方面的优势,评估上下文范围为 32K 至 1M。

如下图所示,与 Transformer 相比,YOCO 大幅度降低了 GPU 内存占用,且 YOCO 的内存消耗随上下文长度增加,增长幅度很小。

例如,在 1M 长度下,整体推理内存使用量仅为 12.4GB,而传统的 Transformer 则占用了9.38倍的 GPU 内存。


下面展示了 token 的 KV 缓存对 GPU 内存的占用情况。


YOCO 模型只缓存一层全局的键值对,因此与 Transformer 模型相比,它需要的内存约少了 L(指模型的层数)倍。


例如,YOCO 模型可以使用 1GB 的 GPU 内存来处理 128K token。而具有 GQA 的 Transformer 65B 大小模型,仅能支持 1.6K token。

也就是说,模型越大,YOCO 可以节省更多

在预填充阶段,模型并行编码输入 token。对于 512K 和 1M 长度的输入,Transformer 分别需要大约180秒和300秒。Transformer 的计算复杂度为 O(N^2),处理长上下文需要大量的浮点运算操作。

相比之下,YOCO 的预填充时间为 O(N),随序列长度线性增长。


YOCO 将 Transformer 的 512K 上下文预填充时间从180秒减少到不到6秒。

预填充阶段可以在进入交叉解码器之前提前退出。因此,即使对于短上下文,预填充延迟的加速至少是两倍。例如,对于 32K 长度,YOCO 比 Transformer 快2.87倍。


吞吐量表示模型每秒可以处理多少个 token,涵盖了预填充和生成时间。如下图所示,与 Transformer 相比,YOCO 在不同上下文长度下实现了更高的吞吐量。

以512K查询为例,Transformer 的吞吐量为4.5 token/秒,而 YOCO 达到了43.1token/秒,即实现了9.6倍的加速。

吞吐量提高的原因如前所述,YOCO 减少了预填充所需的时间。其次,由于内存消耗减少,因此可以在推理时使用更大的批量大小,这也有助于提高吞吐量。




详细细节可以查看原论文。


You Only Cache Once: Decoder-Decoder Architectures for Language Models

论文链接:

https://arxiv.org/abs/2405.05254















你也许还想看:






微信扫码关注该文公众号作者

来源:微软亚洲研究院
logo
联系我们隐私协议©2024 bendi.news
Bendi新闻
Bendi.news刊载任何文章,不代表同意其说法或描述,仅为提供更多信息,也不构成任何建议。文章信息的合法性及真实性由其作者负责,与Bendi.news及其运营公司无关。欢迎投稿,如发现稿件侵权,或作者不愿在本网发表文章,请版权拥有者通知本网处理。