西风发自凹非寺
量子位公众号 QbitAI
微软&清华最新研究，打破 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 减少了预填充所需的时间。其次，由于内存消耗减少，因此可以在推理时使用更大的批量大小，这也有助于提高吞吐量。
详细细节，感兴趣的家人们可以查看原论文。
