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

我们介绍了一种名为YOCO的解码器-解码器架构，用于大型语言模型，它仅在第一次缓存键值对。它由两个组件组成，即一个叉式解码器和一个自解码器。自解码器有效地编码全局键值（KV）缓存，通过交叉注意力被叉式解码器重复使用。整个模型的行为类似于仅有解码器的Transformer，尽管YOCO仅缓存一次。该设计大大降低了GPU内存需求，同时保留了全局注意力能力。此外，计算流程使得可以在不改变最终输出的情况下预填充以提前退出，从而显著加快预填充阶段的速度。实验结果表明，YOCO在各种模型大小和训练标记数量的扩展设置中，与Transformer相比，实现了良好的性能。我们还将YOCO扩展到1M上下文长度，具有接近完美的针检索准确性。性能分析结果显示，YOCO在各种上下文长度和模型大小下，通过数量级提高了推理内存、预填充延迟和吞吐量。代码可在https://aka.ms/YOCO获得。

论文旨在提出一种新的大型语言模型架构，名为YOCO，该架构只需缓存一次键值对。研究人员试图解决当前大型语言模型需要大量GPU内存的问题。

YOCO架构由两个组件组成，即交叉解码器和自解码器。自解码器可以高效地编码全局键值（KV）缓存，而交叉解码器可以通过交叉注意力重复使用这些缓存。整个模型的行为类似于仅解码器的Transformer，但YOCO只需缓存一次。这种设计显著降低了GPU内存需求，同时保留了全局注意力能力。此外，计算流程使得可以在不改变最终输出的情况下提前退出预填充，从而显著加速预填充阶段。

论文使用了多种数据集进行实验，证明了YOCO在不同规模的模型大小和训练令牌数量的情况下都能取得良好的性能。此外，论文还扩展了YOCO到1M上下文长度，并实现了接近完美的针筒检索准确性。研究人员还开源了代码，可以在https://aka.ms/YOCO上找到。

最近在这个领域中，还有一些相关的研究，例如：《Scaling Laws for Neural Language Models》、《Reformer: The Efficient Transformer》、《Longformer: The Long-Document Transformer》等。