提升LLM推理效率新方法：POD，让长文本处理更高效！

🕙发布时间：2025-02-24

更多LLM架构文章：LLM架构专栏
近日热文：
1. 全网最全的神经网络数学原理（代码和公式）直观解释
 2. 大模型进化史：从Transformer到DeepSeek-R1的AI变革之路
 3. 2W8000字深度剖析25种RAG变体：全网最全~没有之一
 4. 3W6000字了解大模型LLM：部署、优化与框架
 知乎【柏企】
公众号【柏企科技说】【柏企阅文】

在处理更长上下文的LLM推理时，现有的方法，像选择性令牌保留和基于窗口的注意力机制，虽然能提高效率，但却存在丢弃未来文本生成所需重要令牌的风险。在本文中[1]，作者基于近距离令牌比远距离令牌更重要这一观察，提出了POD（Proximal Tokens over Distant Tokens ，近端优先于远端令牌）方法。它通过在相似层之间为远距离令牌共享注意力来分配更少的资源，不是直接丢弃不太重要的令牌，而是减少它们的内存和计算负载，进而在不丢失令牌的情况下提高LLM的效率。

它解决了两大挑战：

调查上下文中重要标记的分布：发现上下文中最近的标记比远距离的标记更重要。
通过跨层共享注意力分数来优化远程令牌的资源：

主要观察

核心动机是为不太重要的令牌分配更少的资源，而不是完全丢弃它们。这就带来两个挑战：

一个令牌需要关注的重要令牌分布在哪里？
如何为不太重要的令牌优化内存和计算？

本文试图通过两个关键观察来解决这两个挑战：

观察1：邻近标记（初始标记 + 最近的标记）比远距离标记更重要。
下面的图展示了不同窗口大小下相同预测的比例：窗口与密集注意力。

上图表明，即使只关注256个近端标记，对于相同的输入序列，在80%的情况下，模型预测的下一个标记与关注所有标记时的预测相同，这为这一观察提供了支持。

观察2：连续层之间的注意力得分相似。
下面的图展示了Llama3 - 8B - 32K中第14个头在不同层的注意力得分相似性。

在下图中，红框内各层之间的注意力得分表现出很强的相似性。

POD（远距离标记上的邻近标记）

基于上述观察，作者提出了POD来优化解码阶段的推理效率。它仅为远距离标记共享层间注意力分数，同时基于上述两个观察结果，保持近端标记不变。

该方法包含三个主要阶段：

离线层间注意力共享探索：确定哪些层可以共享注意力分数。
轻量级训练适应：基于已识别的层间注意力共享模式，使用有限的数据对密集模型进行后训练。
高效推理：在层之间共享远程令牌的注意力分数，这使我们能够在KV缓存中保留单个层的关键状态。

方法论

如上文所述，该方法包含三个主要步骤，如下图所示：

i) 离线层间注意力共享探索
- 分析给定长上下文LLM中层之间注意力分数的相似性，并将连续的相似层分组到块中。为了评估各层之间注意力分数的相似性，我们在LLM中输入了几个标记，并收集了一系列注意力分数。随后，计算层之间的注意力相似度，并将连续的相似层分组到块中，以此作为提高LLM推理效率的准备工作。
- a) 注意力分数计算：对于输入到模型M中的N个样本 $s_i = (x_1, x_2, \ldots, x_n)_{i = 1}^N$，收集每个样本的最后 $q$（$1 \leq q \leq n$）个词元关注其相应前序词元的注意力分数。其中 $L$、$H \in N^+$ 分别表示模型中的层数和注意力头数，$S_{\ell, h}^i \in R^{q×n}$ 表示在第 $h$ 个注意力头的第 $\ell$ 层收集的注意力分数。
- b) 注意力相似度评估：对于任意两个不同的层 $\ell_a$ 和 $\ell_b$（$1 \leq \ell_a, \ell_b \leq L$ 且 $\ell_a \neq \ell_b$），它们之间第 $h$ 个头的注意力相似性定义为所有 $N$ 个样本中最后 $q$ 个标记的平均Jensen - Shannon（JS）散度。公式为：
  $$sim_h(\ell_a,\ell_b)=\frac{1}{Nq}\sum_{i = 1}^N\sum_{j = 1}^qJS(S_{\ell_a,h}^i,j,S_{\ell_b,h}^i,j)$$
  其中 $S_{\ell, h}^i,j$ 表示 $S_{\ell, h}^i$ 的第 $j$ 行，$0 \leq sim_h(\cdot, \cdot) \leq 1$。
- c) 图层分组：在计算了层之间的头部注意力相似性后，我们将连续的相似层分组为头部块，作为准备工作。分组策略基于同一区块中的任何两层都应该足够相似的思想。当 $sim(\ell_a, \ell_b) \geq \delta$ 时，$\ell_a$ 和 $\ell_b$ 被视为相似，其中 $0 \leq \delta \leq 1$ 是超参数。我们采用自下而上的贪婪算法，将连续的相似层迭代合并为区块，详见下面的算法

ii) 轻量级训练适应
- 在每个模块中应用注意力共享，并对LLM进行后训练
- a) 每个区块内的注意力共享：首先将前面的标记分为两组：近端标记和远距离标记。考虑到 “注意力沉” 现象，将几个初始标记分类为近端标记。然后，令牌 $x_i$ 会关注两组令牌，但远程令牌的注意力分数将使用相应区块分组的最低层的注意力分数。对于任何注意力头，设 $Q_{\ell},K_{\ell},V_{\ell} \in R^{n×d}$ 分别表示第 $\ell$ 层的查询、键和值状态。层 $\ell$ 属于块 $B_{\ell} = \{\bar{\ell} | \ell_a \leq \bar{\ell} \leq \ell_b\}$，该块由连续的层组成。最近标记的注意力输出和标记 $x_i$ 的远距离标记的计算方法如下：
- b) 将注意力输出聚合到近端和远距离标记：无参数的门控机制可以通过（原文应该有具体方式，需补充完整）整合对近端和远端标记的注意力。
iii) 高效推理
- 使用后训练的LLM进行高效推理
- a) KV缓存内存占用优化：如前面章节的公式所示，对于远距离的tokens，查询和键状态在同一块内的层之间共享。在推理过程中，缓存查询状态是不必要的，因为它们不会被重用。这里提到的方法会减少KV缓存中key state的内存消耗，如上图所示，解码时共享注意力分数的层只保留一次远距离的token；例如，只有第1层保留 $x_2$ 和 $x_3$ 的关键状态，而第2层和第3层则不保留。
- b) 远距离代币的计算优化：经验证据表明，在许多情况下，无需关注遥远的代币即可有效地完成对下一个代币的预测。如上图所示，对于块中不是最低的层，我们可以抢先评估 $g_{\ell,i}$ 的值。如果 $g_{\ell,i} \geq \tau$（$0 \leq \tau \leq 1$ 是一个超参数），则可以省略远距离标记的注意力计算，从而减少远距离标记的计算。

实验

i) 性能评估
- a) 大海捞针：下图展示了不同方法的搜索结果。结果显示，当指针超出其预定义窗口时，StreamingLLM和H2O会失败。相比之下，我们的方法避免了标记丢失，其性能与密集模型相似，几乎可以定位所有针。
- b) 长上下文基准测试：下表显示了不同方法对两个著名的长背景基准的评估结果。POD优于基于代币驱逐的方法，表明我们不丢失代币的方法确实有效。凭借少量的训练后数据，POD击败了经典的基于层共享的方法CLA，这表明我们的模型在适应现有LLM方面具有优势。基于PoD和标记选择的方法都可以实现与标准密集模型相当的性能。此外，POD与基于令牌选择的方法正交，将它们结合起来可以进一步减小KV缓存的大小，同时保持模型性能。
ii) 效率评估
- a) 内存占用：下表显示了内存消耗结果。结果表明，POD在不同输入文本长度下的最大批处理大小增加了30%以上，这与我们理论上35%的KV缓存节省率非常一致，这表明POD有效地减少了内存使用。
- b) 远距离令牌的计算：下图显示了LEval上的计算节省与性能损失之比与 $\tau$ 的值之间的关系。观察到，随着 $\tau$ 的减小，更容易忽略对遥远标记的计算，从而节省更多的计算成本，但会损失一些性能。然而，当 $\tau < 0.7$ 时，性能下降速度减慢，而计算节省变得更加明显。当 $\tau = 0.7$ 时，计算成本降低了25%，而性能仅下降了5%。
iii) 分析
- a) 扩展到更长的上下文和其他LLMs：下表显示了InfiniteBench中5个子任务在不同上下文大小下的评估结果。与基于令牌驱逐的方法相比，POD方法导致的性能下降较小。然而，一个显著的区别是，基于标记选择的方法似乎难以在较长的上下文场景中保持模型性能。
- b) 模型性能与POD中两个关键超参数之间的关系：近端标记的数量和KV缓存保存率：下图显示，随着近端标记数量的增加，POD的性能稳步提高。下图显示，性能会随着节省率的提高而降低。为了平衡性能和效率，将KV缓存压缩到35%，在相同的训练数据下保持与LLaMA3 - 8B - 32K相当的性能。
结论
基于近端标记比远距离标记更重要的观察，本文提出了POD方法，它通过在相似层之间为远距离标记分担注意力来分配更少的资源。评估表明，POD可以在不牺牲模型性能的情况下节省35%的KV缓存。