微软打破Decoder-Only架构，大幅降低GPU内存需求？

微软&清华最新研究，打破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比较的可扩展性、长上下文评估、推理优势。

1.语言建模评估

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

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

2.可扩展性对比

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

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

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

3.长上下文评估

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

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

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

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

4. 推理优势

研究人员评估了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减少了预填充所需的时间。其次，由于内存消耗减少，YOCO就可以在推理时使用更大的批量，这也有助于提高吞吐量。

详细细节，感兴趣的家人们可以查看原论文。

论文链接：https://arxiv.org/abs/2405.05254

本文来自微信公众号：量子位 （ID：QbitAI），作者：西风