引言

本文将探讨KV-Cache如何通过在内存使用和计算时间之间进行巧妙的权衡，使像ChatGPT和DeepSeek这样的语言模型在生成文本时更快。

总结11篇最近的研究论文，归纳三大类：token选择、后处理压缩技术和架构重新设计。包括DeepSeek的多头潜在注意力（MLA），这些论文在这一基本思想的基础上，进一步提高了大型语言模型（LLM）推理的时间效率。

一、思考

为什么文本生成如此缓慢

让我们从一个简单的类比开始。想象你在写一个故事，每写一个新词，你都需要重新阅读到目前为止的整个故事以保持一致性。故事越长，重新阅读的时间就越长。这正是大型语言模型在文本生成时所面临的问题。

自注意力的基本构建块

现代语言模型的核心是一种称为自注意力的机制。对于一个由个标记（大致对应单词）组成的序列，每个标记都需要“查看”或“关注”所有其他标记以理解上下文。

这种查看一切的过程的计算成本随着序列长度的增长而增长：

• 对于个标记，每个标记都需要查看所有nn个标记

• 这意味着成本与成正比

• 用数学符号表示，我们将其写为的复杂度

真正的问题：一次生成一个标记

当语言模型生成文本时，它一次生成一个标记，这就是事情变得计算密集的地方：

• 第一个标记：查看1个标记（成本：）

• 第二个标记：查看2个标记（成本：）

• 第三个标记：查看3个标记（成本：）

• 以此类推，直到第个标记：查看个标记（成本：）

如果我们将生成长度为的序列的所有这些成本加起来，我们得到：

这种的成本意味着随着文本的增长，生成时间会极其迅速地增长。例如，生成两倍长的序列大约需要八倍的时间！显然，我们需要一个更好的方法。

解决方案：键值（KV）缓存

KV 缓存背后的关键是，我们正在做大量冗余工作。在生成每个新标记时，我们会重新计算之前已经处理过的所有先前标记。让我们看看如何解决这个问题。

什么是键值缓存？

可以将 KV 缓存想象成一个智能记事本，我们会在第一次看到每个 token 时记下有关它的重要信息。对于每个 token，我们计算并存储两件事：

• 键（k）：可以将其视为一种寻址机制——它有助于确定此标记与未来标记的相关性

• 值（v）：可以将其视为当此标记被发现相关时实际使用的信息

从数学上，我们计算这些为：

• 键：（其中是标记，是一个学习到的变换）

• 值：（其中是另一个学习到的变换）

在生成一个新标记时，我们使用它的查询（计算方式类似于键）通过将其与所有存储的键进行比较来在我们的缓存中找到相关信息。然后使用匹配的值来帮助生成标记。

KV缓存如何加速

有了KV缓存，处理过程变得更加高效：

1. 当我们遇到一个新token时，只需要计算它的key和value一次

2. 对于所有后续的token，我们可以直接从缓存中查找这些预计算的值

3. 这意味着每个新token只需要做少量新的计算，而不是重新做所有之前的计算

显然有一个权衡：

• 我们需要更多的内存来存储所有的keys和values。对于一个具有：

• 层

• 注意力头

• 序列长度

• key/value维度，总的内存开销为值（这个2是因为需要存储keys和values）。这会随着序列长度线性增长，但对于大模型来说，常数因子可能非常大。

• 但作为回报，我们将计算成本从降低到。

要理解为什么是，让我们看一下每一步的成本：

1. 第一步：处理一个token成本

2. 第二步：处理一个新token + 查找1个缓存的token成本

3. 第三步：处理一个新token + 查找2个缓存的token成本

4. 依此类推…

将这些加起来：

这相比是一个显著的改进！虽然我们仍然需要做查看所有前面的tokens的基础工作，但我们避免了每一步都进行昂贵的重新计算。

内存挑战：为什么我们需要更好的解决方案

虽然KV缓存是一个强大的优化手段，但它伴随着显著的内存开销。让我们通过一个具体的例子来看看，使用像Llama3 70B这样的现代大语言模型：

• 层

• 注意力头

• 批量大小为8个序列

• key/value维度

• 16位精度

处理一个批量（8个序列，每个序列1000个token）所需的内存为：

字节字节GB

这种巨大的内存使用带来了几个挑战：

1. 随着序列长度线性增长

2. 与批量大小成倍增长，支持并行处理

3. 限制了我们可以处理的最大上下文长度

4. 限制了在内存受限设备上的部署

这些挑战激发了研究界的一波创新，导致了各种优化KV缓存使用的技术。接下来，将探讨这些前沿的解决方案。

二、如何改善传统的KV缓存？

以下论文代表了KV缓存优化的关键创新。我们将通过三大主要方法来探索它们：token选择、后处理压缩技术和架构重设计。

2.1 Token 选择和修剪方法（Token Selection and Pruning Approaches）

1) Heavy-Hitter Oracle (H2O)

• https://arxiv.org/abs/2306.14048

H2O 引入了在KV缓存中识别和保留重要token的概念：

• 重型Token（Heavy-Hitter Tokens）：H2O 识别在生成过程中具有最高累计注意力分数的token，这些token遵循幂律分布。这些token对于模型的功能至关重要，因此在缓存中优先处理。

• 动态次模撤销（Dynamic Submodular Eviction）：该方法将缓存管理问题框架化为一个优化问题，目标函数为次模函数，用于量化token集合的重要性：

其中是token的累计注意力分数。缓存通过以下方式更新：

确保每次最多只移除一个token。这个贪心算法在计算上高效，并在次模约束下保证接近最优的性能。

• 结果：通过该方法，KV缓存大小减少了5倍，几乎没有精度损失，并且吞吐量提升了高达29倍。

2) StreamLLM

• https://arxiv.org/abs/2309.17453

• 作者观察到**注意力汇聚（Attention Sinks）**现象：解码过程中，初始token充当自然的注意力锚点。

• 如果没有这些注意力汇聚的token，传统窗口注意力方法的性能会下降。

• 基于这一观察，他们引入了滚动缓存（Rolling Cache），它保留了初始token，并处理最近的上下文，从而实现了无限长度序列的处理。

• 他们还展示了这些汇聚token可以通过训练获得，作为专用的注意力锚点，从而减少对多个初始token的依赖。

3) Value-Aware Token Pruning (VATP)

• https://arxiv.org/abs/2406.12335

VATP 扩展了 H2O 的 token 重要性概念，考虑了注意力模式和价值向量的属性：

• 重要性评分：结合了注意力分数和价值向量的信息：

其中是累计注意力分数，是价值向量的 L1 范数。

• Token修剪：根据排名token，最低分数的token被修剪，而注意力汇聚token（例如，开始或换行token）被保留，以防止性能下降。

性能与效率：

• 在16个 LongBench 任务中，VATP 在12-14个任务中超越了 H2O 和 Scissorhands 等基准。

• 在保持最小性能损失的情况下，实现了50%的有效压缩。

• 引入的计算开销几乎可以忽略不计，并且与 Scissorhands 集成时兼容 FlashAttention。

2.2 后处理压缩技术（Post-hoc Compression Techniques）

这些方法压缩或优化KV缓存，同时保持标准的Transformer架构。

4) Adaptive KV Compression (FastGen)

• https://arxiv.org/pdf/2310.01801

FastGen 通过观察运行时的注意力模式引入了自适应压缩：

• 注意力分析：在提示编码过程中，FastGen 识别注意力模式，并选择压缩策略，以最小化内存开销，同时保留注意力恢复：

自适应压缩策略：

• 压缩策略包括：

• 特殊 token()：仅保留特殊 token。

• 局部性()：逐出超过相对距离的 token。

• 频率()：保留具有高累计注意力分数的 token ()。

• 混合策略结合这些策略，首先采用，并根据每个头的需要适应性地应用：

Token 生成：

• 在解码过程中，预先选择的压缩策略有效地管理 KV 缓存：

5) 动态内存压缩（DMC）

• https://arxiv.org/pdf/2403.09636

DMC 引入了自适应的 token 合并：

• 决策机制：在时刻，预测合并决策和权重：

• 加权合并：当时，合并当前和先前的条目：

其中累积重要性权重。

训练：

• 使用 Gumbel-Sigmoid 放松来训练，支持端到端的梯度下降训练：

其中是温度参数。

• 优化组合目标：

其中是语言建模损失，第二项鼓励模型匹配目标压缩比（CR）。

• 结果：达到 8 倍的压缩率，保持性能。

6)范数基础的压缩

• https://arxiv.org/pdf/2406.11430
  

本文提出了一个令人惊讶的观察：缓存 KV 对的范数与注意力分数之间存在明确的相关性，低范数的键嵌入通常会导致解码时的高注意力分数。因此，提出了一个简单但有效的压缩目标：

• 基于范数的选择：对于一组缓存键，计算并排序键的范数：

• 排序和选择：为了压缩 KV 缓存，按范数值对所有键进行排序：

保留范数最小的前个键，其中，为压缩比。

• 压缩缓存：压缩后的键值缓存为：

• 由于其简洁性，该方法与 FlashAttention 保持兼容。

2.3 体系结构重设计

这些方法改变了 Transformer 架构，以更高效地处理 KV 缓存，通常将压缩直接集成到架构中。

7) 多查询注意力（MQA）

• https://arxiv.org/pdf/2305.13245

• 核心思想：MQA 通过共享单个键值头跨所有查询头来减少 KV 缓存大小，替代传统的多头注意力（MHA）：

其中和是共享的键和值投影。

• 优点：将 KV 缓存大小减少了（注意力头的数量），显著降低了内存带宽开销。

• 权衡：虽然 MQA 更快，但在需要多样化注意力模式的任务中，通常会遭遇质量下降。

8) 分组查询注意力（GQA）

• https://arxiv.org/abs/2305.13245

• 核心思想：GQA 在完全多头注意力和 MQA 之间进行插值，提供了推理速度和模型质量之间的可扩展权衡。它将查询头分为组，每组共享一个单独的键值头：

• GQA-1：等价于 MQA。

• **GQA-**：等价于 MHA。

• 训练：通过微调将 GQA 引入现有的预训练模型：

• 首先，将 MHA 权重通过均值池化转换为 GQA。

• 然后进行微调（“上训练”）以适应新的注意力模式。

• 该适应过程仅需原始预训练计算的 5%，使其非常高效。

• 结果模型保持质量，同时获得 GQA 的内存优势。

9) 多头潜在注意力（MLA）

• https://arxiv.org/abs/2405.04434
  

DeepSeek的多头潜在注意力（MLA）采用了一种新颖的方法来减少KV缓存开销。虽然MQA和GQA通过头共享来实现这一目标，MLA则采用低秩潜在压缩技术，在保持多头注意力的优点的同时，减少了KV缓存的大小。

• MLA通过将键（keys）和值（values）压缩成低维度的潜在向量，来减少KV缓存的大小。

• 它将键值嵌入（key-value embeddings）降投到一个压缩的潜在空间：

其中，是降投矩阵，、是键和值的上投矩阵。

• 通过压缩表示，MLA保持了每个头的灵活性，不同于MQA的完全头共享。

• 它引入了旋转位置嵌入（RoPE）来解耦位置感知的键：

这进一步减少了KV缓存的存储，仅缓存压缩的潜在向量和位置键。

10) SnapKV

• https://arxiv.org/pdf/2404.14469
  

• SnapKV引入了观察窗口（Observation Window）：使用提示结束的tokens来识别注意力模式：

其中，表示注意力权重，由压缩率决定。

• 压缩：使用池化层围绕选定位置聚类特征，以保持上下文完整性。

11) 只缓存一次（YOCO）

• https://arxiv.org/pdf/2405.05254

YOCO修改了Transformer架构以优化缓存：

• 全局缓存：使用解码器-解码器设计，只有一个共享的KV缓存。

• 复杂度减少：将内存从减少到，其中是序列长度，是层数。

• 高效注意力：自解码器采用滑动窗口注意力或门控保留机制，使内存使用保持恒定（，其中是小窗口大小）。

结论

KV-Cache技术是将Transformer模型扩展和优化到实际应用中的核心。像动态逐出、压缩和结构化近似等创新，持续推动着在长上下文或资源受限的场景中实现更高效的技术。KV-Cache仍然是一个活跃的研究领域，既提供了理论上的见解，也带来了实际的改进。

---

如何系统学习掌握AI大模型？

AI大模型作为人工智能领域的重要技术突破，正成为推动各行各业创新和转型的关键力量。抓住AI大模型的风口，掌握AI大模型的知识和技能将变得越来越重要。

学习AI大模型是一个系统的过程，需要从基础开始，逐步深入到更高级的技术。

这里给大家精心整理了一份全面的AI大模型学习资源，包括：AI大模型全套学习路线图（从入门到实战）、精品AI大模型学习书籍手册、视频教程、实战学习、面试题等，资料免费分享！

1. 成长路线图&学习规划

要学习一门新的技术，作为新手一定要先学习成长路线图，方向不对，努力白费。

这里，我们为新手和想要进一步提升的专业人士准备了一份详细的学习成长路线图和规划。可以说是最科学最系统的学习成长路线。

2. 大模型经典PDF书籍

书籍和学习文档资料是学习大模型过程中必不可少的，我们精选了一系列深入探讨大模型技术的书籍和学习文档，它们由领域内的顶尖专家撰写，内容全面、深入、详尽，为你学习大模型提供坚实的理论基础。（书籍含电子版PDF）

3. 大模型视频教程

对于很多自学或者没有基础的同学来说，书籍这些纯文字类的学习教材会觉得比较晦涩难以理解，因此，我们提供了丰富的大模型视频教程，以动态、形象的方式展示技术概念，帮助你更快、更轻松地掌握核心知识。

4. 2024行业报告

行业分析主要包括对不同行业的现状、趋势、问题、机会等进行系统地调研和评估，以了解哪些行业更适合引入大模型的技术和应用，以及在哪些方面可以发挥大模型的优势。

5. 大模型项目实战

学以致用 ，当你的理论知识积累到一定程度，就需要通过项目实战，在实际操作中检验和巩固你所学到的知识，同时为你找工作和职业发展打下坚实的基础。

6. 大模型面试题

面试不仅是技术的较量，更需要充分的准备。

在你已经掌握了大模型技术之后，就需要开始准备面试，我们将提供精心整理的大模型面试题库，涵盖当前面试中可能遇到的各种技术问题，让你在面试中游刃有余。

全套的AI大模型学习资源已经整理打包，有需要的小伙伴可以微信扫描下方CSDN官方认证二维码，免费领取【保证100%免费】