NeuIPS 2024 | YOCO的高效解码器-解码器架构
该研究提出了一种新的大模型架构，名为YOCO（You Only Cache Once），其目的是解决长序列语言模型推理中的内存瓶颈。YOCO通过解码器-解码器结构的创新设计，显著减少推理时的显存占用并提升了长序列的处理效率。
现有大模型的挑战：长序列缓存的瓶颈
当前的大规模语言模型通常采用基于解码器的Transformer架构，通过缓存前序生成的键值对（KV缓存）来提升推理效率。然而，随着输入序列的增长，KV缓存占据的显存也在急剧增加。例如，一个65B参数量的语言模型，在处理512K长度的序列时，需要约86GB显存，超过了一张H100-80GB GPU的容量。此外，预填充（prefilling）长输入序列的延迟也非常高，限制了模型的实际部署。
YOCO架构：解码器-解码器的新范式
为了应对上述问题，YOCO架构应运而生。它由两个部分组成：
-
自解码器（Self-Decoder）：用于生成全局KV缓存。
-
交叉解码器（Cross-Decoder）：通过交叉注意力机制复用自解码器生成的KV缓存。
整个模型从外部看起来依然像传统的解码器模型，能够自然适应自回归生成任务。YOCO的独特之处在于，它只需要在自解码器中缓存一次KV对，然后在交叉解码器中多次复用。这种设计不仅大幅降低了显存消耗，还保留了传统模型的全局注意力能力。
YOCO的三大核心优势
1. 显存节省，支持更长序列
YOCO通过只缓存一次全局KV对，将显存需求显著降低。与传统Transformer相比，YOCO在处理长序列时的显存占用得到了大幅优化。实验数据显示，在65B参数量的模型中，YOCO将KV缓存的显存需求减少了80倍，使得在处理1M长度的输入时，YOCO仅需12.4GB显存，而Transformer则需占用超过100GB显存。
2. 预填充加速，优化用户体验
在推理阶段，YOCO利用交叉解码器复用自解码器的输出，使得模型在预填充阶段可以提前退出，节省计算资源。例如，在处理512K长度的序列时，YOCO的预填充时间仅为6秒，相较于传统Transformer的180秒，实现了约30倍的加速。
3. 提升推理吞吐量，降低部署成本
YOCO显著提升了推理的吞吐量，使每秒生成的token数量大幅增加。在512K长度下，YOCO的吞吐量为每秒43个token，而传统Transformer仅为4.5个，达到了近10倍的提升。这意味着，YOCO能够以更少的硬件资源完成更多的推理任务，从而降低了部署成本。
YOCO的实验表现
实验结果表明，YOCO在多个语言建模任务中展现了极强的竞争力。与同等规模的Transformer模型相比，YOCO在处理长序列任务时，不仅性能更优，而且显存占用和推理延迟也明显更低。
在“针尖查找（Needle-in-a-Haystack）”测试中，YOCO将上下文长度扩展至1M，依然保持了高精度的查询结果。这表明，YOCO在长序列任务中的记忆与处理能力非常出色，适合处理如长文档摘要、代码理解等需要长上下文的任务。
此外，下表展示了YOCO与其他模型在语言建模任务中的性能对比。实验数据显示，YOCO的性能与现有最优模型相当，甚至在某些任务上有所超越，进一步验证了YOCO架构的有效性。
技术实现细节：高效的自解码器模块
YOCO的自解码器模块采用了两种高效的注意力机制来减少计算复杂度：
-
滑动窗口注意力（Sliding-Window Attention）：通过限制注意力的窗口范围，将KV缓存的复杂度从线性降低到常量级，极大减少了显存的消耗。
-
门控保留机制（Gated Retention）：结合数据驱动的门控机制，提升了模型的计算效率和生成效果。
YOCO在书籍和代码数据集上的实验结果表明，其累积平均负对数似然（NLL）随着上下文长度的增加而不断降低，显示了其在长序列任务中的优越性能。
总结：未来前景
YOCO通过创新性的解码器-解码器架构，突破了长序列语言模型在推理阶段的内存和延迟瓶颈。在保持模型性能的前提下，显著提升了显存利用率和推理速度，使其成为未来大模型架构的有力候选。未来，YOCO的设计理念或许还可以拓展到更多任务场景中，如长文本生成、对话系统等需要高效长序列处理的领域。
论文链接：
https://openreview.net/forum?id=25Ioxw576r