You Only Cache Once：用于语言模型的 Decoder-Decoder 架构
计算与语言
2024-05-10 v2
摘要
我们为大型语言模型引入了一种 decoder-decoder 架构 YOCO，它仅缓存一次键值对。它由两个组件组成，即堆叠在 self-decoder 之上的 cross-decoder。self-decoder 高效编码全局键值（KV）缓存，供 cross-decoder 通过交叉注意力复用。尽管 YOCO 只缓存一次，整个模型的行为类似于 decoder-only Transformer。该设计大幅降低了 GPU 内存需求，同时保留了全局注意力能力。此外，计算流允许 prefilling 提前退出而不改变最终输出，从而显著加快了 prefill 阶段。实验结果表明，在扩大模型大小和训练 token 数量的各种设置下，YOCO 与 Transformer 相比取得了良好的性能。我们还将 YOCO 扩展到 1M 上下文长度，并具有近乎完美的 needle 检索准确率。性能分析结果表明，在各种上下文长度和模型大小下，YOCO 将推理内存、prefill 延迟和吞吐量提高了几个数量级。代码可在 https://aka.ms/YOCO 获取。
引用
@article{arxiv.2405.05254,
title = {You Only Cache Once: Decoder-Decoder Architectures for Language Models},
author = {Yutao Sun and Li Dong and Yi Zhu and Shaohan Huang and Wenhui Wang and Shuming Ma and Quanlu Zhang and Jianyong Wang and Furu Wei},
journal= {arXiv preprint arXiv:2405.05254},
year = {2024}
}