【论文速读】解码器-解码器架构 YOCO，《You Only Cache Once: Decoder-Decoder Architectures for Language Models》

bylander

已于 2024-07-04 15:51:46 修改

阅读量1.4k

点赞数 19

分类专栏： AI论文阅读文章标签：架构语言模型人工智能自然语言处理

于 2024-07-04 10:52:05 首次发布

本文链接：https://blog.csdn.net/bylander/article/details/140087916

版权

AI论文阅读专栏收录该内容

54 篇文章 2 订阅

订阅专栏

文章链接：https://arxiv.org/html/2405.05254

文章的标题是“You Only Cache Once: Decoder-Decoder Architectures for Language Models”，作者介绍了一种名为YOCO的新型解码器-解码器架构，专为大型语言模型设计。

对于大语言模型，目前比较普遍的采用解码器架构，另外还有编码器架构、编码器-解码器架构。这篇文章提出了一种解码器-解码器架构，非常新颖！

摘要：

文章为大型语言模型引入了一个解码器-解码器架构 YOCO，它只缓存一次键值对。它由两个组件组成，即堆叠在自解码器上的交叉解码器。自解码器有效地对全局键值（KV）缓存进行编码，交叉解码器通过交叉注意力重用这些缓存。整个模型的行为类似于仅解码器的 Transformer，尽管 YOCO 只缓存一次。该设计大大降低了 GPU 内存需求，同时保留了全局注意力能力。此外，计算流程可以在不改变最终输出的情况下提前预填充，从而显著加快预填充阶段。实验结果表明，YOCO在模型规模和训练令牌数量等各种设置下均比Transformer具有较好的性能。文章还将 YOCO 扩展到 1M 上下文长度，具有近乎完美的针头检索精度。分析结果表明，YOCO 在上下文长度和模型大小上将推理内存、预填充延迟和吞吐量提高了几个数量级。

解决的主要问题：
大型语言模型在进行自回归生成时，随着服务Token数量的增加，键值缓存占用大量GPU内存，导致推理过程受到内存限制。
长序列输入的预填充延迟非常高，这使得在实际部署中难以使用长上下文语言模型。

主要方法：
YOCO架构：提出了一种新型的解码器-解码器架构，该架构只在自解码器中缓存键值对一次，然后由交叉解码器重用这些缓存。
自解码器：使用高效的自注意力机制来生成全局键值缓存，利用了如滑动窗口注意力等技术来降低内存占用。
交叉解码器：利用交叉注意力机制来重用自解码器生成的共享键值缓存，从而避免了对每个令牌重复编码历史信息的需要。
计算流程优化：允许在不改变最终输出的情况下，在自解码器之前提前退出预填充阶段，显著加快了预填充阶段的速度。
实验验证：通过大量实验展示了YOCO在不同设置下的语言建模性能，并证明了其在推理效率方面的多个优势。

在这里插入图片描述
文章中的图1提出了YOCO（You Only Cache Once）架构的概念，并展示了其相对于传统Transformer架构的优势。下面是对图1内容的解释：

架构图示：图中展示了YOCO的解码器-解码器架构，包括自解码器（Self-Decoder）和交叉解码器（Cross-Decoder）两个主要部分。自解码器生成全局键值缓存，然后交叉解码器使用这些缓存进行交叉注意力机制的操作。

性能对比：图中通过柱状图展示了YOCO与Transformer在几个关键性能指标上的对比：

GPU内存消耗（GPU Memory）： YOCO显著降低了GPU内存的使用量，这是通过减少键值缓存的内存占用实现的。
吞吐量（Throughput）： YOCO提高了模型的吞吐量，即每秒钟可以处理的单词数（wps，words per second）。
预填充延迟（Prefilling Latency）： YOCO减少了预填充阶段所需的时间，这是模型在生成输出之前对输入进行编码的时间。
性能提升倍数：图中还标注了YOCO在不同性能指标上相对于Transformer的性能提升倍数。例如，预填充延迟减少了6.4倍，吞吐量提高了30.3倍。

模型架构

模型架构叫做YOCO（You Only Cache Once)，专为自回归建模而设计，例如大型语言模型（LLM）。如图2所示，其解码器-解码器（decoder-decoder）架构由两部分组成，即自解码器和交叉解码器。具体来说，YOCO堆叠了𝐿块，其中第一个L/2层是自解码器，而其余模块是交叉解码器。
在这里插入图片描述
文章中的图2提供了YOCO架构的概览，展示了其解码器-解码器（decoder-decoder）结构的工作原理。以下是对图2内容的详细解释：

1、自解码器（Self-Decoder）：自解码器负责生成全局键值（KV）缓存。它使用高效的自注意力机制来处理输入序列，并生成中间向量表示，这些表示随后用于创建键值缓存。
2、交叉解码器（Cross-Decoder）：交叉解码器堆叠在自解码器之上，使用自解码器生成的全局KV缓存。它通过交叉注意力机制来重用这些共享的KV缓存，从而避免了对每个令牌重复编码历史信息的需要。
3、计算流程：图中展示了YOCO的计算流程，包括预填充（Prefilling）和生成（Generation）两个阶段：
预填充：在这个阶段，模型并行地对输入Token进行编码，生成KV缓存。
生成：在这个阶段，模型逐个解码输出Token。由于计算流程的设计，预填充阶段可以在不改变最终输出的情况下提前退出，这显著加快了预填充速度。
4、注意力掩码（Causal Masking）：自解码器和交叉解码器都使用因果掩码，确保在生成Token时，只能关注到之前的位置，保持了自回归生成的特性。
5、架构行为：尽管YOCO采用了解码器-解码器结构，但从外部视角看，它的行为类似于仅解码器的Transformer模型，自回归地生成Token。

图2通过可视化的方式，清晰地展示了YOCO架构的组件、计算流程以及其如何通过仅缓存一次全局键值对来提高效率。这种设计使得YOCO能够在保持Transformer模型自回归生成能力的同时，显著降低内存需求和提高处理速度。

Self-Decoder

自解码器包含L/2层，过程如下：
在这里插入图片描述
ESA是表示高效自注意力，LN采用RMSNorm。注意力采用因果掩蔽。
高效的自我注意力模块的关键属性是𝒪⁢(1)推理复杂度，即恒定数量的KV缓存。例如，滑动窗口注意的缓存大小取决于窗口大小，而不是输入长度。
SwiGLU的说明：
在这里插入图片描述

Cross-Decoder

对于交叉解码器，自解码器的输出生成全局 KV 缓存：
在这里插入图片描述
Wk和Wv是可学习的权重矩阵。
交叉解码器包含L/2层，过程如下：

Attention⁡(⋅)是标准的多头注意力，WQ是可学习的权重矩阵。因果掩蔽也用于交叉注意力。
因为交叉注意力与组查询注意力（group query attention）兼容，可以进一步节省KV缓存的内存消耗。
获得最后的𝑋𝐿后, 再用softmax分类器执行下一个标记预测。

推理优势

1、GPU内存节省（Saving GPU Memory）：
YOCO通过仅缓存一次全局键值对（KV caches），显著减少了GPU内存的使用量。与传统的Transformer模型相比，YOCO在处理长序列时尤其节省内存，因为它避免了对每个Token重复存储键值对。
2、更高的吞吐量（Higher Throughput）：
由于YOCO的内存占用较低，它可以使用更大的批量大小（batch size）进行推理，这有助于提高模型的吞吐量。吞吐量是指模型每秒可以处理的令牌数量，这对于实时应用和大规模部署至关重要。
3、减少预填充时间（Reducing Prefilling Time）：
YOCO的计算流程允许在不改变最终输出的情况下，在自解码器之前提前退出预填充阶段。由于交叉解码器重用了自解码器的输出，因此可以在预填充阶段在进入交叉解码器之前提前退出。这种设计显著加快了预填充速度，减少了推理阶段的延迟。
YOCO在保持与传统Transformer模型相当的性能的同时，显著降低了推理成本。这包括减少内存占用、加快预填充速度和提高吞吐量。

在文中提到的“预填充”（Prefilling）是指在生成模型开始生成输出之前，对输入序列进行处理的阶段。这个阶段是必要的，因为它允许模型利用输入数据构建所需的内部状态或缓存，这些状态或缓存对于后续生成任务至关重要。
在这里插入图片描述
在YOCO（You Only Cache Once）架构中，预填充阶段涉及到以下几个关键步骤：
1、并行编码输入token：在预填充阶段，模型会并行处理输入序列中的所有token，生成中间表示，这些表示随后会被用于生成阶段。
2、构建KV缓存：自解码器（Self-Decoder）利用高效的自注意力机制来生成全局键值（KV）缓存。这些缓存包含了输入序列的重要信息，并且在生成阶段被重复使用。
3、提前退出：YOCO架构的一个关键特性是，在预填充阶段可以提前退出，而不会改变最终的输出。就是在预填充阶段在进入交叉解码器之前提前退出。

在传统的Transformer模型中，每个token的生成都需要访问完整的键值对缓存，这会导致内存使用量大和延迟高。然而，YOCO通过仅在预填充阶段创建一次这些缓存，并在生成阶段重用它们，显著减少了内存需求和预填充时间。这样，即使对于长序列，模型也能快速进入生成阶段，提高了整体的推理效率。

试验结果

文章第四部分的实验，验证了YOCO（You Only Cache Once）架构在多个方面的表现，并与传统的Transformer模型进行了比较。以下是对文章第四部分中一些重要实验结果的分析：

语言模型评估（Language Modeling Evaluation）：
- 作者训练了一个3B大小的YOCO语言模型，并将其与其他基于Transformer的语言模型进行了比较。实验结果表明，YOCO在多种下游任务上的零样本（zero-shot）性能与之前训练有素的Transformer模型相当，显示出良好的语言建模能力。
模型扩展性（Scalability Compared with Transformers）：
- 实验展示了YOCO在不同模型大小（从160M到13B参数）上的扩展性。通过训练不同大小的语言模型，作者发现YOCO在验证损失上与传统Transformer架构相当，甚至在某些情况下更优。这表明YOCO在不同规模的模型上都能保持有效的性能。
长上下文评估（Long-Context Evaluation）：
- 作者将YOCO模型的上下文长度扩展到1M个令牌，并评估了其在长上下文建模任务上的能力。YOCO在“Needle in a Haystack”测试中表现出近乎完美的准确性，这证明了YOCO在处理长文档和检索特定信息方面的强大能力。
多针检索（Multi-Needle Retrieval）：
- 在多针检索评估中，YOCO与其他长上下文语言模型进行了比较。YOCO即使在模型大小减半的情况下，也展现出与LWM-1M-text等其他模型相当的性能，甚至在某些情况下更优。
长序列上的困惑度（Perplexity over Long Sequences）：
- 作者评估了YOCO在长序列上的语言建模性能，结果显示随着上下文长度的增加，YOCO的困惑度（Perplexity）持续下降，这表明YOCO能够更有效地利用长距离依赖信息。
推理优势（Inference Advantages）：
- YOCO在推理效率方面表现出显著的优势，特别是在GPU内存占用、预填充延迟、吞吐量等方面。例如，与传统Transformer相比，YOCO在512K上下文长度的预填充时间上实现了从180秒减少到不到6秒的巨大改进。
部署成本降低（Deployment Cost Reduction）：
- 实验结果表明，YOCO通过减少内存占用和提高处理速度，显著降低了部署成本，尤其是在需要处理长序列推理的场景中。

这些实验结果证明了YOCO架构在处理大型语言模型时的有效性和效率，特别是在需要长上下文信息和高推理速度的应用场景中。通过这些实验，作者展示了YOCO作为未来大型语言模型架构的潜力，特别是在需要原生长序列支持的情况下。

结论

在这项工作中，作者提出了一种用于大型语言模型的解码器-解码器架构（YOCO），与Transformers相比，YOCO的推理效率和竞争性能明显更好。实验结果表明，在增加训练token数量、扩大模型尺寸以及将上下文长度扩展到100万个token的各种设置下，YOCO都能为大型语言模型带来有利的结果。分析结果还表明，YOCO在推理效率方面提高了数个数量级，尤其是在长序列建模方面。

这项工作可以从以下几个角度进一步发展：
1、YOCO + BitNet + Groq：Groq通过将所有内容放在SRAM中实现了非常高的吞吐量。然而，内存容量瓶颈限制了模型尺寸和输入token数量。现在，数百个芯片连接在一起仅托管一个模型。作为解决方案，YOCO减少了键值缓存内存，而BitNet减少了模型权重内存。使用上述组合，预期可以将大型语言模型的部署成本降低数个数量级。
2、YOCO用于多模态大型语言模型：YOCO的布局通常适用于使用多个自解码器。交叉注意力层自然适用于多模态融合。自解码器的因果依赖性也非常适合流式视频。异步多模态大型语言模型可以避免不同数据流相互阻塞，这对于实时应用（如机器人技术）至关重要。
3、优化的KV缓存模块机制：图2明确突出了KV缓存，这为开发原生内存机制开辟了新的机会。首先，可以集成一个缓存压缩机制以获得更紧凑的内存。其次，可以构建一个索引以实现高效的键值检索。由于YOCO重用缓存，它使我们能够只维护一个索引，而不是为每一层创建一个索引。第三，解耦建模支持预缓存上下文，这对于原生RAG和LLM原生搜索引擎可能非常有用。