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

（本文阅读时间：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倍。