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

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

本论文旨在解决大型语言模型中的内存问题，通过引入一种新的解码器-解码器架构YOCO，实现只缓存一次键值对，从而显著减少GPU内存需求。

YOCO架构由两个组件组成，即自解码器和交叉解码器。自解码器高效地编码全局键值（KV）缓存，交叉解码器通过交叉注意力重复使用全局键值缓存。YOCO只缓存一次，但整个模型的行为类似于仅解码器的Transformer，保留了全局注意力能力。

论文通过实验表明，YOCO在各种模型大小和训练令牌数量的情况下，与Transformer相比具有优越的性能。此外，论文还将YOCO扩展到1M上下文长度，实现了近乎完美的针筒检索准确性。YOCO还改善了推理内存、预填充延迟和吞吐量等方面的性能，具有很高的应用价值。代码已经开源。

近期在这个领域中，还有一些相关的研究，比如《Attention is All You Need》、《BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding》等。