DeepSeek的多头潜在注意力（MLA）和及其11种KV-Cache技巧演进大总结 原创

引言

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

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

一、思考

为什么文本生成如此缓慢

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

自注意力的基本构建块

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

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

• 对于n个标记，每个标记都需要查看所有nn个标记
• 这意味着成本与n x n = n²成正比
• 用数学符号表示，我们将其写为O(n²)的复杂度

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

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

• 第一个标记：查看1个标记（成本：O(1²)）
• 第二个标记：查看2个标记（成本：O(2²)）
• 第三个标记：查看3个标记（成本：O(3²)）
• 以此类推，直到第n个标记：查看n个标记（成本：O(n²)）

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

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

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

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

什么是键值缓存？

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

• 键（k）：可以将其视为一种寻址机制——它有助于确定此标记与未来标记的相关性
• 值（v）：可以将其视为当此标记被发现相关时实际使用的信息

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

• 键：k = xW_k（其中是x标记，W_k是一个学习到的变换）
• 值：v = xW_v（其中W_v是另一个学习到的变换）

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

KV缓存如何加速

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

1. 当我们遇到一个新token时，只需要计算它的key和value一次
2. 对于所有后续的token，我们可以直接从缓存中查找这些预计算的值
3. 这意味着每个新token只需要做少量新的计算，而不是重新做所有之前的计算

显然有一个权衡：

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

     L层

     H注意力头

     序列长度n

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

• 但作为回报，我们将计算成本从O(n³)降低到O(n²)。

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

1. 第一步：处理一个token ->成本O(1)
2. 第二步：处理一个新token + 查找1个缓存的token -> 成本O(2)
3. 第三步：处理一个新token + 查找2个缓存的token -> 成本O(3)
4. 依此类推...

将这些加起来：

   O(1 + 2 + 3 + ... + n) = O(n²)

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

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

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

• L = 80层
• H = 64注意力头
• B = 8批量大小为8个序列
• d_k= 128key/value维度
• 16位精度

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

L x H x B x n x d_k x 2 x 2字节=80 x 64 x 8 x 1000 x 128 x 2 x 2字节=20.97GB

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

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）：该方法将缓存管理问题框架化为一个优化问题，目标函数为次模函数F(S)，用于量化token集合的S重要性：

确保每次最多只移除一个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 重要性概念，考虑了注意力模式和价值向量的属性：

性能与效率：

• 在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 通过观察运行时的注意力模式引入了自适应压缩：

自适应压缩策略：

5) 动态内存压缩（DMC）

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

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

6)L₂范数基础的压缩

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

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

2.3 体系结构重设计

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

7) 多查询注意力（MQA）

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

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

K = XW_K ， V = XW_V

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

• 优点：将 KV 缓存大小减少了H（注意力头的数量），显著降低了内存带宽开销。
• 权衡：虽然 MQA 更快，但在需要多样化注意力模式的任务中，通常会遭遇质量下降。

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

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

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

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

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

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

• MLA通过将键（keys）和值（values）压缩成低维度的潜在向量，来减少KV缓存的大小。
• 它将键值嵌入（key-value embeddings）降投到一个压缩的潜在空间：

10) SnapKV

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

11) 只缓存一次（YOCO）

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

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

• 全局缓存：使用解码器-解码器设计，只有一个共享的KV缓存。
• 复杂度减少：将内存从O(N x L)减少到O(N + L)，其中N是序列长度，L是层数。
• 高效注意力：自解码器采用滑动窗口注意力或门控保留机制，使内存使用保持恒定（O(C)，其中C是小窗口大小）。

结论

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

公众号大模型自然语言处理  作者：余俊晖

原文链接：​​https://mp.weixin.qq.com/s/7j9sVIlJQDPnji9bas09ig​​​