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

这篇论文介绍了一种名为YOCO（You Only Cache Once）的新型解码器-解码器架构，专为大型语言模型设计，以提高推理效率和性能。以下是其核心内容的总结：

YOCO架构关键特点：

• 双重解码器结构：YOCO由自解码器和交叉解码器两部分组成，自解码器生成全局键值（KV）缓存，交叉解码器通过交叉注意力机制重用这些缓存。
• 单次缓存：与标准Transformer相比，YOCO只缓存一次KV对，显著减少了GPU内存需求。
• 推理效率：YOCO的计算流程允许预填充阶段提前退出，而不影响最终输出，从而显著加快预填充阶段的速度。

实验结果：

• YOCO在不同模型规模和训练令牌数量的设置下，与Transformer相比展现出了有利的性能。
• 通过扩展到1M的上下文长度，YOCO实现了接近完美的“针检索”准确性。
• 在不同上下文长度和模型尺寸下，YOCO在推理内存、预填充延迟和吞吐量方面均有数量级的提升。

设计选择：

• 自解码器：使用高效的自注意力机制，如滑动窗口注意力或门控保留机制。
• 交叉解码器：利用自解码器生成的全局KV缓存，通过交叉注意力机制进行解码。

推理优势：

• GPU内存节省：YOCO的KV缓存内存复杂度为O((N + L)D)，其中N是序列长度，L是常数，D是隐藏维度，远低于Transformer的O(LND)。
• 预填充时间减少：YOCO的预填充时间复杂度为O(LND)，相较于Transformer的O(LN^2D)有显著降低。

扩展性和多模态应用：

• YOCO架构具有扩展性，可以处理更长的上下文长度，适用于多模态大型语言模型和实时应用，如机器人技术。

未来工作：

• 结合YOCO与BitNet和Groq技术，以实现更高效的内存使用和降低大型语言模型的部署成本。
• 为KV缓存模块开发优化机制，如缓存压缩和索引构建。

关键概念：

训练令牌（Training Tokens）

• 定义：训练令牌是指在训练语言模型时使用的数据量度量，通常以令牌（token）数量来表示。一个令牌可以是一个词、一个字符或一个子词，依据所使用的分词器（tokenizer）而定。
• YOCO的应用：在YOCO模型中，通过增加训练令牌的数量来扩展模型的训练数据规模。实验结果显示，YOCO能够处理多达万亿级别的训练令牌，并且在这种扩展下，模型的语言建模性能与Transformer语言模型相当。

针检索（Needle Retrieval）

• 定义：针检索是一种测试模型长上下文能力的方法，即在大量数据中检索特定信息的能力。在论文中，这通常指在长文本中检索特定的、具有区分性的短语或单词（即“针”），类似于在干草堆中找针的比喻。
• YOCO的性能：YOCO模型在扩展到1M令牌的上下文长度时，进行了针检索测试，并展现出接近完美的准确性。这表明YOCO能够有效地处理和记忆长文本中的特定信息。

自解码器（Self-Decoder）

• 定义：自解码器是YOCO架构中的一个组件，负责生成全局键值（KV）缓存。它使用高效的自注意力机制来编码输入序列，并生成中间向量表示，这些表示随后被用于交叉解码器。
• 工作机制：自解码器利用诸如滑动窗口注意力或门控保留（gated retention）等技术，以减少计算和内存需求。这些技术允许自解码器在保持性能的同时，只使用常数大小的缓存。
• 设计选择：论文中提到了两种自解码器的设计选择：
  • 门控保留（Gated Retention）：通过数据依赖的门控机制增强保留网络，实现训练并行性、良好性能和低推理成本。
  • 滑动窗口注意力（Sliding-Window Attention）：限制注意力范围在一个固定窗口大小内，减少了推理时KV缓存的内存复杂度。

交叉解码器（Cross-Decoder）是YOCO架构中的关键组件之一，它负责利用自解码器生成的全局键值（KV）缓存来进行进一步的处理和生成输出。以下是关于交叉解码器的详细解释：

交叉解码器的功能：

• 重用KV缓存：交叉解码器通过交叉注意力机制重用由自解码器生成的全局KV缓存。这种设计允许模型在不同的解码层之间共享和传递信息。
• 生成输出向量：交叉解码器堆叠在自解码器之上，利用共享的KV缓存来计算最终的输出向量，这些输出向量随后用于预测下一个令牌。

交叉解码器的工作流程：

1. 生成KV缓存：自解码器首先处理输入序列，生成中间向量表示，并由此产生全局KV缓存。
2. 交叉注意力：交叉解码器的每层都会使用这些KV缓存，通过计算查询（Q）、键（K）和值（V）之间的注意力分数来更新表示。
3. 输出向量计算：通过标准的多头注意力机制，交叉解码器将自解码器的输出与KV缓存相结合，生成每个位置的输出向量。

交叉解码器的设计优势：

• 减少内存需求：由于KV缓存在交叉解码器的所有层中重用，因此大大减少了模型在推理时所需的GPU内存。
• 提高效率：交叉解码器的设计允许模型在预填充阶段提前退出，从而加快了处理速度，特别是在长序列输入时。

交叉解码器与自解码器的关系：

• 自解码器和交叉解码器共同构成了YOCO的解码器-解码器架构。自解码器负责生成全局KV缓存，而交叉解码器则利用这些缓存来生成最终的输出。
• 这种分离允许模型在保持解码器的自回归特性的同时，通过交叉注意力机制捕获更广泛的上下文信息。

交叉解码器在YOCO中的重要性：

• 交叉解码器是YOCO实现高效推理的关键，它使得模型在处理大型语言模型任务时，能够以较低的内存占用和更快的速度运行。

全局键值（KV）缓存是YOCO架构中的一个核心概念，它在模型的自解码器和交叉解码器之间起到了桥梁的作用。以下是关于全局KV缓存的详细解释：

什么是全局KV缓存？

• 键值对：在Transformer模型中，“键”（Key，K）和"值"（Value，V）是注意力机制中的概念，用于计算模型对输入序列中每个元素的关注程度。
• 全局性：在YOCO中，全局KV缓存指的是由自解码器生成的KV对，这些缓存在整个交叉解码器中被共享和重用，而不是在每个解码层重新计算。

全局KV缓存的作用：

• 信息共享：全局KV缓存允许模型在不同的解码层之间共享信息，这有助于捕捉长距离依赖关系。
• 效率提升：通过重用这些缓存，YOCO减少了计算量和内存需求，因为不需要在每一层都重新计算KV对。

全局KV缓存的生成：

• 自解码器：自解码器使用高效的自注意力机制处理输入序列，并生成中间向量表示。这些中间表示用于生成全局KV缓存。
• 计算过程：自解码器的输出通过特定的权重（例如，(W_K) 和 (W_V)）来生成全局的K和V，这些权重是可学习的参数。

全局KV缓存的使用：

• 交叉解码器：交叉解码器使用这些全局KV缓存，通过交叉注意力机制来更新其对输入序列的理解，并生成最终的输出向量。
• 减少内存占用：由于全局KV缓存在交叉解码器的所有层中重复使用，YOCO显著减少了所需的内存，与传统的Transformer模型相比，内存占用可以减少大约L倍，其中L是层数。

YOCO中的全局KV缓存优势：

• 内存效率：YOCO的内存复杂度由于单次缓存KV对而降低，这使得模型能够处理更长的序列，同时保持较低的内存占用。
• 推理速度：减少内存占用也有助于提高模型的推理速度，尤其是在处理长序列时。

结论：

全局KV缓存是YOCO架构中实现高效语言模型的关键技术之一。它通过减少重复计算和内存占用，使得YOCO能够在保持Transformer模型性能的同时，显著提高推理效率和扩展性。