Accurate KV Cache Eviction via Anchor Direction Projection for Efficient LLM Inference

文章标题：通过锚定方向投影实现精确的KV缓存淘汰以进行高效的LLM推理
作者/机构：Zijie Geng, Jie Wang, Ziqi Liu, Feng Ju (中国科学技术大学); Yiming Li, Xing Li, Mingxuan Yuan, Jianye Hao (华为诺亚方舟实验室); Defu Lian, Enhong Chen, Yongdong Zhang, Feng Wu (中国科学技术大学); Jianye Hao (天津大学)

A1 主要贡献

大型语言模型（LLMs）通过自注意力机制彻底改变了信息处理方式，但该机制的二次计算复杂度成为瓶颈。键值（KV）缓存机制通过存储和重用中间结果，将计算复杂度从二次降至线性，但随着处理序列长度的增加，KV缓存的体积也成比例增长，导致GPU内存消耗和I/O延迟显著增加。

核心问题：现有的KV缓存淘汰策略大多依赖于简单的启发式方法（如注意力权重）来衡量token的重要性，忽略了token值状态在向量空间中的空间关系，这常常导致次优的token选择和性能下降。因此，需要开发一种能明确考虑向量空间几何关系的评分函数来改进KV缓存管理。

研究目标：提出一种新颖的方法，名为AnDPro（锚定方向投影），引入一种基于投影的评分函数，以更精确地衡量token的重要性，从而指导更准确的token选择。

创新点：
1. 问题建模：将KV缓存淘汰问题形式化为一个组合优化问题，并通过适度松弛将其转化为一个稀疏优化问题。
2. 锚定方向投影：受理论和直观分析的启发，在值向量空间中引入一个“锚定方向”（即淘汰前输出的方向），并利用token值状态在该方向上的投影来衡量其重要性。这种方法能更好地保留最重要的语义信息。
3. 性能优势：实验证明，AnDPro仅用3.44%的KV缓存预算就能保持96.07%的全缓存精度，与之前的SOTA方法相比，在不牺牲模型质量的情况下，将KV缓存预算大小减少了46.0%。

下图展示了基于投影的KV缓存淘汰机制的直观示例。

A2 方法细节

本节介绍我们提出的方法AnDPro（锚定方向投影）。第3.1节首先将KV缓存淘汰问题形式化，为深入分析提供理论基础。第3.2节对该优化问题进行理论探索，从而引出基于投影的评分函数的一般形式。第3.3节通过直观分析进一步推导出评分函数的实用形式。最后，在第3.4节中，我们提供了AnDPro的实现细节。所有证明都可以在附录A中找到。

3.1 问题形式化

Q1：如何将KV缓存淘汰问题进行形式化，以便能够深入分析设计合适的KV缓存淘汰评分函数？
我们的目标是从输入提示中选择一个token子集予以保留，丢弃其余部分，同时尽可能精确地逼近原始输出。本节我们考虑单个注意力头。给定一个输入序列 $X = [x_1, \dots, x_n]$，其中n代表上下文长度，即提示中的token数量。令 $x_i \in \mathbb{R}^d$ 表示第i个token的嵌入向量，其中d表示嵌入维度。令 $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$ 分别为查询、键和值的投影参数矩阵。对于第i个token，其查询 $q_i$、键 $k_i$ 和值状态 $v_i$ 由以下公式给出：

$$ \boldsymbol{q}_{i}=\boldsymbol{W}^{Q} \boldsymbol{x}_{i}, \quad \boldsymbol{k}_{i}=\boldsymbol{W}^{K} \boldsymbol{x}_{i}, \quad \boldsymbol{v}_{i}=\boldsymbol{W}^{V} \boldsymbol{x}_{i} . $$

淘汰前输出的定义。令 $q$ 为从观察窗口给出的一个查询，具体细节将在3.4节中进一步阐述。该注意力头的淘汰前输出由以下公式给出：

$$\boldsymbol{y} = \sum_{i=1}^{n} a_i \boldsymbol{v}_i, \quad \text{where } a_i = \mathop{\mathrm{Softmax}}_{1 \leq i \leq n}(\boldsymbol{q}^\top \boldsymbol{k}_i).$$

淘汰后输出的定义。令 $S \subset [n]$ 为被选中保留的token集合。那么，仅保留与token $i \in S$ 对应的KV对时得到的淘汰后输出由以下公式给出：

$$\hat{\boldsymbol{y}}=\sum_{i \in S} \hat{a}_{i} \boldsymbol{v}_{i}, \quad \text { where } \hat{a}_{i}=\underset{i \in S}{\operatorname{Softmax}}\left(\boldsymbol{q}^{\top} \boldsymbol{k}_{i}\right) .$$

优化目标的构建。我们的目标是使淘汰后的输出尽可能精确地逼近淘汰前的输出。很自然地，我们使用平方误差作为量化淘汰损失的度量：

$$\mathcal{L}(y, \hat{y}) \triangleq \frac{1}{2}\|y-\hat{y}\|_{2}^{2} .$$

约束与问题重述。此外，我们对保留的token数量施加稀疏性约束 $|S| \le k$。下面的命题提供了一个更便于进一步分析的淘汰后输出表达式。

命题1。令 $S \subset [n]$ 为保留的token集合，并定义指示函数

$$\begin{aligned} z_i \triangleq \begin{cases}1, & i \in S, \\ 0, & i \notin S\end{cases} \end{aligned}$$

来表示第i个token是否被保留。那么，淘汰后的输出可以表示为：

$$\hat{\boldsymbol{y}}=\frac{\sum_{i=1}^{n} z_{i} a_{i} \boldsymbol{v}_{i}}{\sum_{i=1}^{n} z_{i} a_{i}}$$

KV缓存淘汰问题的初始形式化。基于命题1，我们可以将KV缓存淘汰问题形式化如下：

$$\min_{\boldsymbol{z} \in\{0,1\}^{n}} \frac{1}{2}\left\|\boldsymbol{y}-\frac{\sum_{i=1}^{n} z_{i} a_{i} \boldsymbol{v}_{i}}{\sum_{i=1}^{n} z_{i} a_{i}}\right\|_{2}^{2}, \quad \text { s.t. } \quad \sum_{i=1}^{n} z_{i} \leq k.$$

问题的正则化形式。在稀疏优化领域，处理稀疏性约束的一个普遍方法是引入 $l_1$-范数正则化【50，Lasso regression，2018，Journal of British Surgery】。当处理二元变量时，$l_1$-范数等价于 $l_0$ 范数，即计算非零元素的数量。问题(P0-C)的正则化形式表示为：

$$\min_{\boldsymbol{z} \in\{0,1\}^{n}} \quad \frac{1}{2}\left\|\boldsymbol{y}-\frac{\sum_{i=1}^{n} z_{i} a_{i} \boldsymbol{v}_{i}}{\sum_{i=1}^{n} z_{i} a_{i}}\right\|_{2}^{2}+\lambda \sum_{i=1}^{n} z_{i} .$$

问题求解的挑战与松弛。虽然问题(P0-R)形式简洁，但它仍然是NP-hard问题，使得精确求解在计算上不可行。考虑到LLM推理的背景，需要在毫秒内找到解决方案，我们寻求快速找到高质量的解并逼近问题的最优解。为此，我们探索对(P0-R)进行适度松弛，以推导出一个能够促进快速筛选以实现高效KV缓存淘汰的评分函数。需要注意的是，直接将(P0-R)中的二元变量 $z_i$ 松弛为连续值会导致问题退化并产生平凡解。具体来说，当 $z_i$ 被松弛后，它们可以被任意缩放而不影响目标函数的第一项。如果 $z_i$ 的值被缩放得足够小，目标函数的第二项将趋近于零，从而产生平凡解。

变量替换与问题重构。我们引入新变量 $s$ 和 $\beta_i$，定义为

$$s \triangleq \sum_{i=1}^{n} z_{i} a_{i} \quad \text{and} \quad \beta_{i} \triangleq \frac{z_{i}}{s}.$$

然后，我们得到：

$$\frac{\sum_{i=1}^{n} z_{i} a_{i} \boldsymbol{v}_{i}}{\sum_{i=1}^{n} z_{i} a_{i}}=\frac{\sum_{i=1}^{n} \beta_{i} a_{i} \boldsymbol{v}_{i}}{\sum_{i=1}^{n} \beta_{i} a_{i}} \quad \text { and } \quad \sum_{i=1}^{n} \beta_{i} a_{i}=1 .$$

问题的再形式化。因此，问题(P0-R)可以被重新表述为：

$$\begin{aligned} \begin{aligned} \min _{\boldsymbol{z} \in\{0,1\}^{n}} & \frac{1}{2}\left\|\boldsymbol{y}-\sum_{i=1}^{n} \beta_{i} a_{i} \boldsymbol{v}_{i}\right\|_{2}^{2}+\lambda s \sum_{i=1}^{n} \beta_{i}, \\ \text { s.t. } & \sum_{i=1}^{n} \beta_{i} a_{i}=1, \quad s \beta_{i}=z_{i}, \quad i=1, \ldots, n. \end{aligned} \end{aligned}$$

凸松弛与最终优化问题。在这个重新表述中，我们可以选择一个任意小的 $s$，这将导致极小的 $z_i$。结果，正则化项将趋近于零，更小的 $s$ 值总是会产生更好的目标函数值。此外，如果我们对 $s$ 施加一个下界，最优解将总是使 $s$ 取这个下界。因此，我们将 $s$ 固定为一个满足 $0 < s \le 1$ 的常数，并对二元变量 $z_i$ 进行凸松弛。这引出了以下优化问题：

$$\min_{\boldsymbol{\beta} \in [0, \frac{1}{s}]^n} \quad \frac{1}{2} \left\| \boldsymbol{y} - \sum_{i=1}^n \beta_i a_i \boldsymbol{v}_i \right\|_2^2 + \lambda s \sum_{i=1}^n \beta_i, \quad \text{s.t.} \quad \sum_{i=1}^n \beta_i a_i = 1.$$

A1: 通过适度松弛，我们推导出问题(P1)，在此基础上，我们将在后续章节中进一步研究评分函数的设计。

备注1。我们想再次强调，我们的目标不是获得原始问题的精确解，该问题足够复杂以至于精确解是难以处理的。相反，我们旨在利用这一分析为评分函数的设计提供有见地的动机。

3.2 基于投影的评分函数：一般形式

问题与目标。虽然问题(P1)提供了一个简洁的公式，但在LLM推理过程中直接解决它仍然是不可行的。受稀疏优化中“筛选规则”【51，Lasso screening rules via dual polytope projection，2013，Advances in neural information processing systems】的启发，我们现在研究设计一种闭式评分函数，该函数可以识别问题(P1)中的活动变量，即那些在最优解中具有非零系数的变量。通过这种方式，我们可以保留活动变量而无需显式解决整个优化问题。因此，本节旨在回答以下问题：

Q2：是否可能找到一个闭式评分函数来识别问题(P1)的活动变量？如果可以，这种评分函数的理想数学形式是什么？
引入新变量并重构问题。令 $\beta^*$ 表示问题(P1)的一个最优解。在本节中，我们将展示在某些条件下，可以确定对于某些 $i \in [n]$ 有 $\beta^*_i = 0$，这意味着相应的token可以从缓存中移除。为了进一步研究该问题，我们引入一组新变量 $u \triangleq y - \sum_{i=1}^n \beta_i a_i v_i$。然后我们可以将问题(P1)重新表述为：

$$ \min _{\boldsymbol{\beta} \in\left[0, \frac{1}{s}\right]^{n}} \frac{1}{2}\|\boldsymbol{u}\|_{2}^{2}+\lambda s \sum_{i=1}^{n} \beta_{i}, $$

推导对偶问题。现在我们可以推导出(P-Primal)的对偶问题，如下定理所示。

定理1。问题(P-Primal)的对偶问题由下式给出：

$$\sup_{\boldsymbol{\eta}, \zeta} -\frac{1}{2} \|\boldsymbol{\eta}\|^2 + \boldsymbol{\eta}^\top \boldsymbol{y} - \zeta + \frac{1}{s} \sum_{i=1}^n \rho_i(\boldsymbol{\eta}, \zeta),$$

其中 $\eta^* \in \mathbb{R}^n$ 和 $\zeta \in \mathbb{R}$ 是对偶变量，且

$$\rho_i(\boldsymbol{\eta}, \zeta) \triangleq \min \left(0, \lambda s-a_i\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_i-\zeta\right)\right) .$$

强对偶性与KKT条件。问题(P-Primal)是凸的，其约束是仿射的。根据Slater条件【52，Convex optimization，2004，Cambridge UP】，只要(P-Primal)是可行的，强对偶性就成立。我们关心(P-Dual)的原因在于利用最优性条件，这可以为最优解 $\beta^*$ 提供有价值的见解。具体来说，通过应用Karush-Kuhn-Tucker (KKT)最优性条件，我们可以推导出以下定理：

定理2。令 $\beta^*$ 是问题(P-Primal)的一个最优解，$(\eta^*, \zeta^*)$ 是问题(P-Dual)的一个最优解。那么我们有：

$$\begin{aligned} \begin{aligned} a_{i}\left(\boldsymbol{\eta}^{* \top} \boldsymbol{v}_{i}-\zeta^{*}\right)<\lambda s & \Rightarrow \quad \beta_{i}^{*}=0, \\ a_{i}\left(\boldsymbol{\eta}^{* \top} \boldsymbol{v}_{i}-\zeta^{*}\right)>\lambda s & \Rightarrow \quad \beta_{i}^{*}=\frac{1}{s} . \end{aligned} \end{aligned}$$

推论与评分函数形式。以下推论总结了本节内容。

推论1。存在 $\theta \in \mathbb{R}^n$ 和 $b \in \mathbb{R}$ 使得：

$$\begin{aligned} \begin{aligned} a_{i}\left(\boldsymbol{\theta}^{\top} \boldsymbol{v}_{i}+b\right)<1 & \Rightarrow \beta_{i}^{*}=0, \\ a_{i}\left(\boldsymbol{\theta}^{\top} \boldsymbol{v}_{i}+b\right)>1 & \Rightarrow \beta_{i}^{*}=\frac{1}{s} . \end{aligned} \end{aligned}$$

评分函数的揭示。推论1揭示了存在一种形式为：

$$s_{i} \triangleq a_{i}\left(\boldsymbol{\theta}^{\top} \boldsymbol{v}_{i}+b\right).$$

的评分函数。

评分函数的直观解释。如图2(a)所示，这个评分函数可以精确地识别(P1)中的活动变量和非活动变量。得分低的token在(P1)中将是非活动的。因此，这些token可以安全地从缓存中淘汰。

A2：我们可以定义一个形如公式(13)的评分函数。

备注2。在定理2中，我们已经表明 $\eta^*$ 和 $\zeta^*$ 对应于我们感兴趣的参数 $\theta$ 和 $b$。然而，在不显式求解问题(P-Dual)的情况下，直接推导 $\eta^*$ 和 $\zeta^*$ 是不可行的。因此，我们在下一节转向对 $\theta$ 和 $b$ 进行直观分析，以推导出一个更实用的形式。

3.3 锚定方向投影：实用形式

Q3：在实践中，我们如何确定通常较好的参数 $\theta$ 和 $b$？
3.2节的理论分析激励我们定义一个形如公式13的评分函数。尽管理论上很完善，但参数 $\theta$ 和 $b$ 的确定仍然具有挑战性，因为在求解问题(P-Dual)之前，对偶变量 $\eta^*$ 和 $\zeta^*$ 是无法获得的。因此，本节旨在回答上述问题。

特殊情况分析：注意力权重。我们首先分析一个特殊情况，当 $\theta = 0$ 时。在这种情况下，评分函数简化为 $s_i = a_i b$，其中 $b$ 是一个常数。设置 $b=1$ 给了我们最常用的分数，即注意力权重。从这个角度看，现有仅依赖注意力权重的方法可以被视为我们框架的一个特例。

一般情况分析：锚定方向。对于更一般的情况，当 $\theta \neq 0$ 时，不失一般性，我们可以假设 $\|\theta\|_2 = 1$。在此假设下，方程 $H: \theta^\top x + b = 0$ 表示 $\mathbb{R}^n$ 中的一个超平面，而 $\theta^\top v_i + b$ 表示点 $v_i$ 到 $H$ 的带符号距离（一个可以取负值的距离度量）。因子 $a_i$ 对这些带符号距离进行缩放。这个计算过程如图2(b)所示。

锚定方向的选择。当应用此评分函数进行top-k token选择时，那些与超平面H的相对较大带符号距离（由 $a_i$ 缩放）的token被保留，而其他token则被丢弃。这个过程实质上保留了输出 $\hat{y}$ 在 $\theta$ 方向上的投影。我们将 $\theta$ 的方向称为“锚定方向”，因为它作为引导保留token选择的参考。一个自然且直观的方法是将 $y$ 的方向设置为锚定方向。这一选择有助于加强沿淘汰前输出方向的语义信息，确保最重要的语义特征得以保留。

最终评分函数形式。在这种情况下，提出的评分函数定义为 $s_i = a_i(y^\top v_i + b)$，其中 $b$ 作为一个偏置项，用于平衡投影和注意力权重的相对重要性。当 $b \to \infty$ 时，评分函数逐渐收敛到传统的基于注意力的分数。结果表明，精心选择的 $b$ 会带来轻微的性能提升，而设置 $b=0$ 也能获得稳定的性能。因此，我们在主要实验中设置 $b=0$，即 $s_i = a_i y^\top v_i$。我们在4.5节中对 $\theta$ 和 $b$ 的选择进行了实验。

A3：我们可以简单地设置 $\theta = y$ 和 $b = 0$。

3.4 实现细节

观察窗口与查询生成。本节我们补充AnDPro的一些实现细节。遵循SnapKV【16，Snapkv: Llm knows what you are looking for before generation，2024，arXiv preprint】的方法，我们保留一个观察窗口W来累积来自最后几个prompt token的注意力权重，从而生成查询 {$q_t$}$_{t \in W}$。具体来说，我们保留这个窗口内的token以及第一个token。在观察窗口的每个时间步，第i个token的注意力和相应的淘汰前输出计算如下：

$$a_{i}^{t}=\underset{1 \leq i \leq t}{\operatorname{Softmax}}(\boldsymbol{q}^{t \top} \boldsymbol{k}_{i}) \quad \text { and } \quad \boldsymbol{y}^{t}=\sum_{i=1}^{n} a_{i}^{t} \boldsymbol{v}_{i} .$$

评分函数计算。每个token的评分函数随后计算为：

$$s_i \triangleq \sum_{t \in \mathcal{W}} s_i^t = \sum_{t \in \mathcal{W}} a_i^t \boldsymbol{y}^{t \top} \boldsymbol{v}_i.$$

跨头预算分配与分块。我们基于分数在每层的所有注意力头之间进行top-k选择。这允许在头之间进行灵活的预算分配，显著提升了生成质量。此外，我们引入了一种将token合并为块的有用技术。注意到SnapKV和Ada-KV【17，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference，2024，arXiv preprint】都采用了最大/平均池化操作，确保了相邻token具有相似的分数。这导致了token的连续选择，显著提高了性能。

分块技术。受此启发，我们将token合并为块，每个块包含预定义数量的token，以保持语义连贯性。块大小被视为一个超参数。合并后的块的值状态是原始token值的加权和，其注意力分数是各个注意力权重的总和。AnDPro的伪代码在附录B.1中提供。

A4 实验环境

• 数据集:

  • LongBench: 一个长序列基准测试，包含覆盖多任务领域的16个数据集，用于评估方法的综合性能。
  • Needle-in-a-Haystack: 用于评估不同预算策略在长上下文检索任务中的能力。

• 模型架构:

  • Mistral-7B-Instruct-v0.2
  • Llama-3.1-8B-Instruct

• 硬件配置:

  • GPU: 单张 A800-80G GPU

• 软件与超参数配置:

  • 实现: 基于Ada-KV框架。
  • 预算大小 (B): 在 {128, 256, 512, 1024} 范围内进行评估。
  • 观察窗口大小: 32
  • 块大小 (Chunk size): 4
  • 代码: https://github.com/MIRALab-USTC/LLM-AnDPro

A4 实验结果

LongBench 测试评估

• 实验内容: 在LongBench基准上，使用Mistral和Llama模型，比较AnDPro与H2O、StreamingLLM、SnapKV、Ada-KV、CriticalKV等基线方法在不同预算（128, 256, 512, 1024）下的性能。

• 实验结果: 如表1和图3所示，AnDPro在所有模型和预算配置下都展现了最先进（SOTA）的性能。

  • Mistral模型: 在预算为256时，AnDPro的平均得分为40.84，达到全缓存精度的96.07%，而仅使用了3.44%的KV缓存预算。与之前的SOTA方法CriticalKV相比，AnDPro在达到相似性能的情况下，预算规模减少了46.0%。
  • Llama模型: 结果与Mistral一致，AnDPro同样表现出色（详见附录C.1）。

• 分析结论: AnDPro能以更少的KV缓存预算达到甚至超越现有SOTA方法的性能，证明了其高效性和优越性。在MF-en, SAMSum, Lcc等多个数据集上，AnDPro的表现甚至优于全缓存。
  

Needle-in-a-Haystack 测试评估

• 实验内容: 使用“大海捞针”测试评估AnDPro在长上下文检索任务中的能力。实验设置预算大小为k=128，与Ada-KV和CriticalKV进行比较。
• 实验结果: 如图4所示，AnDPro显著优于基线方法，取得了97.37的高分，非常接近全缓存的性能。
• 分析结论: AnDPro在需要精确检索信息的长上下文任务中表现出色，证明其保留关键信息的能力很强。
  

内存与延迟评估

• 实验内容: 在8K到256K的不同上下文长度下，测量各种KV淘汰策略的峰值内存使用和解码延迟。实验使用Mistral-7b-instruct-v0.2模型，固定预算k=128，生成1K个token并计算每个token的解码延迟。
• 实验结果: 如图5所示，所有淘汰方法的峰值内存使用和解码延迟大致相同，且相比全缓存模型都有显著优化。附录C.4的进一步分析表明，AnDPro引入的额外计算（Update KV Phase）所带来的时间开销可以忽略不计。
• 分析结论: AnDPro在实现更高准确率的同时，保持了与现有SOTA淘汰方法相当的内存和时间效率。
  

分析

• 消融研究: 附录C.5的消融研究证实了核心贡献——基于投影的评分函数——的有效性，它始终优于基于注意力的评分。同时，分块、保留首个token和跨头预算分配等设计都对最终性能有积极贡献。
• θ和b的选择: 附录C.6的实验表明，设置θ=y和b=0能实现优越且稳健的性能。
• 损失和分数可视化: 附录C.7和C.8的结果表明，AnDPro能有效降低淘汰损失，并且其token分数分布具有稀疏性。
• 值向量分析: 附录C.9的PCA分析和幅度直方图直观地解释了为什么基于投影的分数有效，因为它保留了沿原始输出方向的关键token，并利用了值向量本身丰富的语义信息。
• 案例研究: 附录C.10的案例表明AnDPro能比基于注意力的方法更准确地保留重要token。
• 更长上下文: 附录C.11的实验在高达384k token的极长上下文中进行，结果显示AnDPro即使在更长的上下文中也保持最佳性能。

A5 结论

本文将KV缓存淘汰建模为一个组合优化问题，并将其松弛为一个稀疏优化问题。受理论和直观分析的启发，我们提出了AnDPro（锚定方向投影）方法，该方法利用基于投影的评分函数来实现更精确的KV缓存淘汰。在LongBench和Needle-in-a-Haystack基准上的大量实验证明了我们方法的有效性。

A6 附录

A 证明

命题1。令 $S \subset [n]$ 为保留的token集合，并定义指示函数

$$\begin{aligned} z_i \triangleq \begin{cases}1, & i \in S, \\ 0, & i \notin S\end{cases} \end{aligned}$$

来表示第i个token是否被保留。那么，淘汰后的输出可以表示为：

$$\hat{\boldsymbol{y}}=\frac{\sum_{i=1}^{n} z_{i} a_{i} \boldsymbol{v}_{i}}{\sum_{i=1}^{n} z_{i} a_{i}}$$

证明。我们有

$$\begin{aligned} \begin{aligned} \hat{\boldsymbol{y}} & = \sum_{i \in S} \operatorname{Softmax}_{i \in S}(\boldsymbol{q}_n^{\top} \boldsymbol{k}_i) \boldsymbol{v}_i = \sum_{i \in S} \frac{\exp(\boldsymbol{q}_n^{\top} \boldsymbol{k}_i)}{\sum_{j \in S} \exp(\boldsymbol{q}_n^{\top} \boldsymbol{k}_j)} \boldsymbol{v}_i \\ & = \frac{\sum_{i=1}^n z_i \frac{\exp(\boldsymbol{q}_n^{\top} \boldsymbol{k}_i)}{\sum_{k=1}^n \exp(\boldsymbol{q}_n^{\top} \boldsymbol{k}_k)} \boldsymbol{v}_i}{\sum_{j=1}^n z_j \frac{\exp(\boldsymbol{q}_n^{\top} \boldsymbol{k}_j)}{\sum_{k=1}^n \exp(\boldsymbol{q}_n^{\top} \boldsymbol{k}_k)}} = \frac{\sum_{i=1}^n z_i a_i \boldsymbol{v}_i}{\sum_{i=1}^n z_i a_i}. \end{aligned} \end{aligned}$$

定理1。问题(P-Primal)的对偶问题由下式给出：

$$\sup_{\boldsymbol{\eta},\zeta}-\frac{1}{2}\|\boldsymbol{\eta}\|^2+\boldsymbol{\eta}^\top\boldsymbol{y}-\zeta+\frac{1}{s}\sum_{i=1}^n\rho_i(\boldsymbol{\eta},\zeta),$$

其中 $\eta^* \in \mathbb{R}^n$ 和 $\zeta \in \mathbb{R}$ 是对偶变量，且

$$\rho_{i}(\boldsymbol{\eta}, \zeta) \triangleq \min \left(0, \lambda s-a_{i}\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_{i}-\zeta\right)\right) .$$

证明。回想一下，问题(P-Primal)的形式为：

$$\min_{\boldsymbol{\beta} \in\left[0, \frac{1}{s}\right]^{n}} \quad \frac{1}{2}\|\boldsymbol{u}\|_{2}^{2}+\lambda s \sum_{i=1}^{n} \beta_{i}$$

拉格朗日函数的构建。通过引入新变量 $\eta \in \mathbb{R}^n$ 和 $\zeta \in \mathbb{R}$，我们得到其拉格朗日函数：

$$L(\beta, \boldsymbol{u}, \boldsymbol{\eta}, \zeta)=\frac{1}{2}\|\boldsymbol{u}\|_{2}^{2}+\lambda s \sum_{i=1}^{n} \beta_{i}+\boldsymbol{\eta}^{\top}\left(\boldsymbol{y}-\sum_{i=1}^{n} \beta_{i} a_{i} \boldsymbol{v}_{i}-\boldsymbol{u}\right)+\zeta\left(\sum_{i=1}^{n} \beta_{i} a_{i}-1\right) .$$

对偶函数的推导。原始变量是 $\beta \in [0, 1/s]^n$ 和 $u \in \mathbb{R}^n$。对偶函数 $g(\eta, \zeta)$ 于是为：

$$\begin{aligned} \begin{aligned} g(\boldsymbol{\eta}, \zeta) & =\inf _{\boldsymbol{\beta}, \boldsymbol{u}} L(\boldsymbol{\beta}, \boldsymbol{u}, \boldsymbol{\eta}, \zeta) \\ & =\boldsymbol{\eta}^{\top} \boldsymbol{y}-\zeta+\inf _{\boldsymbol{u}}\left(\frac{1}{2}\|\boldsymbol{u}\|^2-\boldsymbol{\eta}^{\top} \boldsymbol{u}\right)+\inf _{\boldsymbol{\beta} \in[0, \frac{1}{s}]^n}\left(\lambda s \sum_{i=1}^n \beta_i-\sum_{i=1}^n \beta_i a_i \boldsymbol{\eta}^{\top} \boldsymbol{v}_i+\zeta \sum_{i=1}^n \beta_i a_i\right) \\ & =\boldsymbol{\eta}^{\top} \boldsymbol{y}-\zeta+\inf _{\boldsymbol{u}}\left(\frac{1}{2}\|\boldsymbol{u}\|^2-\boldsymbol{\eta}^{\top} \boldsymbol{u}\right)+\sum_{i=1}^n \inf _{\beta_i \in[0, \frac{1}{s}]} \beta_i\left(\lambda s-a_i\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_i-\zeta\right)\right) . \end{aligned} \end{aligned}$$

求解下确界。很容易推导出

$$\inf_{\boldsymbol{u}} \left( \frac{1}{2} \| \boldsymbol{u} \|^2 - \boldsymbol{\eta}^\top \boldsymbol{u} \right) = - \frac{1}{2} \| \boldsymbol{\eta} \|^2,$$

其最优解为 $u^* = \eta$，以及

$$\inf_{\beta_{i} \in[0, \frac{1}{s}]} \beta_{i}\left(\lambda s-a_{i}\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_{i}-\zeta\right)\right)=\frac{1}{s} \min \left(0, \lambda s-a_{i}\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_{i}-\zeta\right)\right).$$

其最优解为

$$\begin{aligned} \beta_{i}^{*}=\begin{cases} 0, & \lambda s-a_{i}\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_{i}-\zeta\right)>0, \\ \in\left[0, \frac{1}{s}\right], & \lambda s-a_{i}\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_{i}-\zeta\right)=0, \\ \frac{1}{s}, & \lambda s-a_{i}\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_{i}-\zeta\right)<0. \end{cases} \end{aligned}$$

对偶函数的最终形式。因此，我们有：

$$g(\boldsymbol{\eta}, \zeta)=-\frac{1}{2}\|\boldsymbol{\eta}\|^{2}+\boldsymbol{\eta}^{\top} \boldsymbol{y}-\zeta+\frac{1}{s} \sum_{i=1}^{n} \min \left(0, \lambda-a_{i}\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_{i}-\zeta\right)\right) .$$

对偶问题的最终形式。将以上结合起来，我们得到对偶问题：

$$\sup_{\boldsymbol{\eta}, \zeta} \quad g(\boldsymbol{\eta}, \zeta)=-\frac{1}{2}\|\boldsymbol{\eta}\|^{2}+\boldsymbol{\eta}^{\top} \boldsymbol{y}-\zeta+\frac{1}{s} \sum_{i=1}^{n} \min \left(0, \lambda s-a_{i}\left(\boldsymbol{\eta}^{\top} \boldsymbol{v}_{i}-\zeta\right)\right) .$$

定理2。令 $\beta^*$ 是问题(P-Primal)的一个最优解，$(\eta^*, \zeta^*)$ 是问题(P-Dual)的一个最优解。那么我们有：

$$\begin{aligned} \begin{aligned} a_i\left(\boldsymbol{\eta}^{* \top} \boldsymbol{v}_i-\zeta^*\right)<\lambda s \quad \Rightarrow \quad \beta_i^*=0, \\ a_i\left(\boldsymbol{\eta}^{* \top} \boldsymbol{v}_i-\zeta^*\right)>\lambda s \quad \Rightarrow \quad \beta_i^*=\frac{1}{s} . \end{aligned} \end{aligned}$$

证明。根据强对偶性，很容易证明 $(\eta^*, \zeta^*)$ 是一个几何乘子。因此，从最优性条件，我们有

$$(\boldsymbol{\beta}^*, \boldsymbol{u}^*) \in \inf_{\boldsymbol{\beta}, \boldsymbol{u}} L(\boldsymbol{\beta}, \boldsymbol{u}, \boldsymbol{\eta}^*, \zeta^*).$$

然后我们根据公式21完成证明。

推论1。存在 $\theta \in \mathbb{R}^n$ 和 $b \in \mathbb{R}$ 使得：

$$\begin{aligned} \begin{aligned} a_{i}\left(\boldsymbol{\theta}^{\top} \boldsymbol{v}_{i}+b\right)<1 & \Rightarrow \quad \beta_{i}^{*}=0, \\ a_{i}\left(\boldsymbol{\theta}^{\top} \boldsymbol{v}_{i}+b\right)>1 & \Rightarrow \quad \beta_{i}^{*}=\frac{1}{s} . \end{aligned} \end{aligned}$$

证明。我们可以取

$$\boldsymbol{\theta} \triangleq \frac{\boldsymbol{\eta}^*}{\lambda s}, \quad b \triangleq-\frac{\zeta^*}{\lambda s},$$

其中 $\eta^*$ 和 $\zeta^*$ 是问题(P-Dual)的最优解。这完成了证明。

B 实现细节

B.1 算法

AnDPro算法伪代码。AnDPro的一个简化伪代码在算法1中提供。为了简洁，我们省略了一些不重要的细节（例如保留第一个token）。

算法1 AnDPro算法在单层中的实现

输入：所有头的集合 H，观察窗口的集合 W，观察窗口外的token数量 L，
      查询 {qh,t}h∈H,t∈W，键 {khi}h∈H,i∈[L]，值 {vhi}h∈H,i∈[L]，
      块大小 C，每个头的预算大小 B
输出：保留的KV缓存 {Sh}h∈H

1: ah,ti |h∈H,t∈W,i∈[L] ← Softmaxi∈[L](qh,t vhi)
2: vˆh,ti |h∈H,t∈W,i∈[L/C] ← Σj∈Chunki ah,tj vhj
3: aˆh,ti |h∈H,t∈W,i∈[L/C] ← Σj∈Chunki ah,tj
4: Wˆ ← 块的观察窗口
5: yh,t |h∈H,t∈Wˆ ← Σi∈[L/C] vˆh,ti
6: sh,ti |h∈H,t∈Wˆ,i∈[L/C] ← ⟨yh,t, vˆh,ti⟩
7: shi |h∈H,i∈[L/C] ← Σt∈Wˆ sh,ti
8: s ← Concath∈H,i∈[L/C](shi)
9: {Sh}h∈H ← TopK(s, k = |H| × B)
10: Sh |h∈H ← Sh ∪ W
11: return {Sh}h∈H

B.2 数据集

LongBench数据集。Longbench是一个全面的基准测试，由16个数据集组成，由于其提示种类繁多，涵盖不同领域、类型和长度，因此可作为一个强大的评估工具。它专为长序列任务设计，涵盖了单文档问答、多文档问答、摘要、少样本学习、合成任务和代码生成等多个领域。表2提供了LongBench中16个数据集的详细信息。这些数据集的平均输入长度从1,235到18,409个token不等，平均token长度为7,425。对于我们的评估，我们遵循了每个数据集推荐的评估程序，分数范围从0到100。

B.3 代码和超参数

代码框架与超参数设置。我们的代码框架改编自Ada-KV【17，Ada-kv: Optimizing kv cache eviction by adaptive budget allocation for efficient llm inference，2024，arXiv preprint】，特别利用了其简洁且用户友好的自定义类，以促进更高效、便捷的预算筛选和分配。主要更新包括一个跨头筛选策略。我们采用核大小为7的最大池化，并将观察窗口的大小设置为32。代码发布于 https://github.com/MIRALab-USTC/LLM-AnDPro。

C 附加结果

C.1 Llama在LongBench上的主要结果

Llama模型性能。表3显示了基于Llama模型在16个数据集上不同方法的得分。总体而言，结果与Mistral模型的结果一致，AnDPro在缓存淘汰后也带来了质量的提升。

C.2 LongBench中文数据集上的结果

中文数据集性能。表4显示了基于Llama模型在LongBench的5个中文数据集上不同方法的得分。尽管Llama模型不支持中文，AnDPro在缓存淘汰后仍然带来了质量的提升。

C.3 Needle-in-a-Haystack测试结果

长上下文检索性能。图6展示了所有方法在“大海捞针”（NIAH）测试中的得分。AnDPro显著优于所有基线，取得了97.37的高分，非常接近全缓存的性能。

C.4 运行时分析

各阶段耗时对比。我们在图7中可视化了推理过程中不同阶段的时间消耗的详细比较。结果表明，AnDPro引入的额外计算——特别是Update KV阶段——所带来的时间开销可以忽略不计。

定量运行时数据。为了提供更详细的定量分解，具体的运行时测量数据如下（单位：秒/毫秒）：

与基线的比较。尽管AnDPro的更新延迟相比SnapKV和AdaKV略高，但额外的开销相对于总的预填充时间来说是微不足道的。

解码延迟测试。我们进一步使用Mistral-7B-Instruct-v0.2模型在“大海捞针”（NIAH）基准上进行了实验。我们测量了在各种输入长度（从8K到256K个token）下，预算大小为128时，输出长度为1K的解码延迟（单位：秒）。结果见表6。

C.5 消融研究

C.5.1 基于投影的评分函数消融研究

核心贡献验证。我们进行了全面的消融研究，以分离出基于投影的评分函数（本文的核心贡献）的有效性，并将其与我们方法中的其他组件区分开来。具体来说，我们进行了5组实验，比较了基于投影的评分函数（ProjScore）和基于注意力的评分函数（AttnScore）在不同组件组合下的表现。具体来说，我们考虑了A-E五个实验组：
- A: CrossHead + Chunk + FirstToken，ProjScore vs. AttenScore
- B: Chunk + FirstToken with ProjScore vs. Pooling + FirstToken with AttenScore
- C: CrossHead + FirstToken，ProjScore vs. AttenScore
- D: CrossHead + Chunk with ProjScore vs. CrossHead + Pooling with AttenScore
- E: 单独的ProjScore vs. AttenScore比较

实验结果。我们在Mistral模型的单文档问答任务的三个数据集上进行了这些消融研究。结果如表7所示，并在图8中进行了可视化。结果表明，基于投影的评分函数在所有不同组件配置下都持续优于基于注意力的评分函数。

模型泛化性验证。我们进一步在另外两个基础模型Qwen和Llama上验证了这一结论。表8和表9显示AnDPro在其他模型上也有良好表现，证实了其有效性和鲁棒性。

C.5.2 其他采用技术的消融研究

各组件贡献分析。我们进一步进行消融研究，以调查我们算法中三个不同组件的贡献：（1）将token合并为块，（2）保留第一个token，以及（3）跨头预算分配。结果如图9所示。它直观地比较了每个组件对基于投影方法所贡献的性能提升，表明所有组件对于实现最终的高性能结果都是不可或缺的。

发现总结。我们的发现如下。(1) 分块：为了实现更好的语义连续性，我们将一定数量的相邻token合并为一个块。我们测试了不同的块大小C=1, 2, 4, 8, 16。结果如图10所示。我们发现在预算较小时，较小的块性能更好，反之亦然。主实验中设为4。(2) 保留首个token：遵循SnapKV和Ada-KV，我们默认保留第一个token，这带来了非常轻微的性能提升。(3) 跨头预算分配：我们遵循Ada-KV，根据分数在每层的所有头之间分配预算。结果显示此操作对整体性能增益至关重要。这些技术的整合带来了AnDPro的高性能。值得注意的是，当预算大小k=2048时，AnDPro的性能超过了全缓存。

C.6 θ和b的不同选择

参数选择实验。回想一下，我们的基于投影的评分函数定义为 $s_i = a_i(y^\top v_i + b)$。本节我们研究了不同选择的 $\theta$ 和 $b$ 的影响。我们使用Mistral-7B-Instruct-v2.0在单文档问答任务的数据集上进行了这些实验。

锚定方向θ的选择。为了研究锚定方向 $\theta$ 的影响，我们将其参数化为 $\theta = \sum_{i=1}^n \tilde{a}_i v_i$，其中 $\tilde{a}_i = \text{Softmax}(q^T k_i/\alpha)$，$\alpha$ 是一个温度系数。这个泛化涵盖了几个情况：$\alpha=1$ 精确对应我们提出的锚定方向 $\theta=y$；$\alpha \to \infty$ 对应所有值向量的平均值；$\alpha \to 0$ 对应具有最高注意力权重的值向量。结果如表10所示，表明将锚定方向设置为 $\theta=y$ 可以实现稳健的性能。

偏置项b的选择。然后我们研究了 $b$ 的影响，结果如图11所示。我们发现虽然仔细选择 $b$ 可能会带来更好的结果，但简单地设置 $b=0$ 就足以获得稳定的性能。因此，我们在主要实验中设置 $b=0$，无需进行耗时的超参数调整。此外，当 $b$ 趋近于无穷大（例如，$b=1000$）时，评分函数收敛到基于注意力的函数，性能接近Ada-KV的性能。

C.7 淘汰损失和余弦相似度

性能解释。图12展示了两种方法在不同预算分配下产生的淘汰损失和余弦相似度，其中淘汰损失量化为 $||y-\hat{y}||/||y||_2$。结果显示，AnDPro始终产生较低的淘汰损失和较大的余弦相似度，这在一定程度上解释了AnDPro的高性能。

C.8 Token可视化

分数分布。图13展示了我们定义的评分函数在不同层级的分布情况（使用Mistral-7B-instruct-v0.2在Qasper数据集的一个样本上），表明该分布呈现出与注意力权重相似的稀疏性特征。

C.9 值向量可视化与幅度直方图

直观解释。我们进行了两个实验来直观地解释为什么基于投影的方法优于基于注意力的方法。首先，我们在图14中提供了值向量的PCA可视化。它揭示了投影方法保留了沿原始输出方向的几何上关键的token，从而最小化了淘汰后的语义漂移。相比之下，基于注意力的选择破坏了潜在的空间分布，导致输出发散。其次，图15中的补充分析表明，值向量的幅度存在显著差异。这表明值向量包含丰富的语义信息，而这些信息在基于注意力的方法中被忽略了。这进一步强调了基于投影策略的必要性。

C.10 案例研究

定性分析。我们对一些简单的测试用例进行了案例研究。我们基于AnDPro和注意力权重计算了预填充阶段每个token的总分。一些结果呈现在图16中。这些token级别的分数被缩放到0到1的范围内，并按降序排列。具体来说，分数越高表示该token被认为越重要。

观察结论。通过观察以红色标记的关键token（由人工识别），AnDPro相比基于注意力的方法更好地强调了关键token，从而直观地展示了我们方法的有效性。

图16：案例研究。我们可视化了不同方法（AnDPro和基于注意力的方法）给出的重要性排序。红色token是人类识别的关键token。

C.11 长上下文

极长序列实验。我们使用“大海捞针”基准测试，在非常长的序列上进行了额外的实验，将测试序列长度从原来的32K token扩展到高达384K token。具体来说，我们使用了Llama-3.1-8B-Instruct模型（预训练上下文长度128K），预算大小为128。结果如图17所示。

结果分析。结果表明，AnDPro即使在长序列上也持续取得优越性能。值得注意的是，当序列长度超过256K token时，全缓存因内存限制而失败，而AnDPro仍然有效。尽管所有方法在超过512K token时都失败了，AnDPro在可行的长度范围内保持了最佳性能。

C.12 大尺寸LLMs

模型规模泛化性。为了评估AnDPro在更大模型上的泛化性和适用性，我们使用Qwen2.5-32B-Instruct模型在LongBench单文档问答基准上进行了额外的实验。结果如表11所示。

C.13 长解码任务

任务泛化性。为了进一步证明我们方法的更广泛适用性，我们进行了长解码任务的额外实验，特别是在LLM推理设置中。为了更好地与基线进行比较，我们额外实现了AnDPro和SnapKV的解码版本，而AdaKV没有解码版本。具体来说，我们在AIME24数据集上使用DeepSeek-R1-Distill-Qwen-14B评估了我们的方法。

性能与吞吐量。表12报告了不同KV缓存预算设置（2K和4K）下的推理性能（准确率，%）。表13报告了在10K-token生成过程中的吞吐量（tokens/s）和最大支持的批次大小（OOM阈值）。