Ada-KV: Optimizing KV Cache Eviction by Adaptive Budget Allocation for Efficient LLM Inference

发表时间: 2024-07 · arXiv:2407.11550 (NeurIPS 2025)

Yuan Feng1,3† Junlin Lv1,3† Yukun Cao1,3 Xike Xie2,3* S. Kevin Zhou2,3
1School of Computer Science, University of Science and Technology of China (USTC) 2School of Biomedical Engineering, USTC 3Data Darkness Lab, MIRACLE Center, Suzhou Institute for Advanced Research

速读

一句话结论 本文提出了一种名为 Ada-KV 的即插即用型自适应预算分配策略，通过根据不同注意力头的分布特征动态分配 KV Cache 预算，在不增加额外开销的情况下显著降低了缓存驱逐带来的精度损失，大幅提升了长文本大模型的推理质量。

要解决什么问题 自回归大语言模型在处理长文本时，KV Cache 会急剧膨胀，导致严重的显存瓶颈，并在解码阶段引发极高的 I/O 延迟。为了缓解这一问题，现有的缓存驱逐方法通过丢弃非关键的缓存元素来压缩显存。然而，这些原生方法在分配压缩预算时，对所有的注意力头都采用了均匀分配的策略。这种做法忽略了不同注意力头在注意力分布上的独特性质：有些头非常稀疏，只关注极少数的缓存元素；而有些头则比较分散，需要关注更广泛的上下文。均匀分配会导致显存预算的低效利用——在稀疏头上浪费了宝贵的预算，而在分散头上又因为预算不足导致大量关键信息被强制驱逐，最终引发严重的精度下降。

怎么做的 核心思路是打破均匀分配的限制，将缓存预算从注意力集中的稀疏头动态转移到注意力分散的头上。作者首先从理论上定义了驱逐前后的注意力输出之间的 $L_1$ 驱逐损失，并推导出了该损失的上界 $\epsilon$：

$$ \epsilon = 2hC - 2C \sum_{i \in [1, h]} \sum_{j \in [1, n]} \mathcal{T}_i^j A_i^j $$

其中 $h$ 是注意力头的数量，$\mathcal{T}_i^j \in \{0, 1\}$ 代表是否保留第 $i$ 个头的第 $j$ 个缓存元素，$A_i^j$ 是对应的注意力权重，$C$ 是常数。现有的 Top-$k$ 驱逐方法本质上是在给定各个头预算 $B_i$ 的前提下，通过保留权重最大的元素来最小化这个上界，得到最小化上界 $\epsilon^*$：

$$ \epsilon^* = 2hC - 2C \sum_{i \in [1, h]} \sum_{j \in [1, n]}^{A_i^j \in \text{Top-}k(A_i, B_i)} A_i^j $$

为了在全局层面进一步压低这个上界，Ada-KV 设计了自适应分配机制。具体而言，它将同一层内所有注意力头的权重拼接在一起，在全局范围内选出权重最大的 $B$ 个元素（$B$ 为该层的总预算）。然后，统计这 $B$ 个元素在各个头中的分布频次，以此作为每个头的新预算 $B_i^*$。为了防止某些极端稀疏的头被分配到零预算而破坏生成稳定性，设计中还引入了一个安全垫超参 $\alpha$（默认取 0.2），将自适应预算与均匀预算进行加权融合。在工程实现上，由于各个头的缓存长度变得不一致，Ada-KV 引入了变长 FlashAttention 技术，并设计了一种扁平化的缓存存储布局，将同一层内所有头的缓存拼接成一个连续的张量，配合自定义的 CUDA kernel，确保了在非均匀缓存长度下依然能保持极高的计算和显存访问效率。此外，该方法还原生兼容了 Group Query Attention 机制，通过取组内平均注意力权重作为筛选标准，消除了冗余，在 Llama-3.1-8B 上实现了 4 倍的缓存缩减。

效果如何 实验在 Llama-3.1-8B-Instruct 和 Mistral-7B-Instruct-v0.2 两个模型上展开，使用了包含 13 个数据集的 Ruler（16K 上下文）和包含 16 个数据集的 LongBench 基准。对比基线包括代表滑动窗口路线的 StreamingLLM，以及代表 Top-$k$ 路线的强基线 SnapKV 和 Pyramid。作者将 Ada-KV 作为即插即用的模块接入后两者，构建了 Ada-SnapKV 和 Ada-Pyramid。实验不仅测试了常规的已知问题压缩场景，还重点测试了更具挑战性、更符合真实应用的未知问题压缩场景（即在不知道后续问题的情况下先对长文本进行压缩）。结果表明，在未知问题场景下，原生方法的性能会出现大幅衰减，而 Ada-KV 展现出了极大的优势。在 Ruler 基准测试中，当缓存预算压缩至 80% 时，Ada-SnapKV 将 Llama-3.1-8B 的平均分从 SnapKV 的 87.59 提升至 92.67；在极限的 20% 预算下，得分从 44.02 提升至 53.29。在难度极高的 S-NIAH-3 任务中（80% 预算），Ada-SnapKV 更是将得分从 62.4 跃升至 97.6。在 LongBench 的 2048 预算设定下，Ada-SnapKV 同样有效挽回了原生方法在摘要和问答等任务上的严重掉点。此外，效率测试表明，在固定 1024 预算下，Ada-SnapKV 的峰值显存占用和解码延迟与原生 SnapKV 几乎完全一致，证明了其在获得精度收益的同时没有付出额外的计算代价。作者也坦承，在代码补全这类对缓存压缩本身就不敏感的任务上，该方法的提升空间相对有限。

主要贡献

大型语言模型（LLMs）在处理长序列推理时，由于Key-Value（KV）缓存的不断增长，面临着严重的效率挑战（如内存占用过大和I/O延迟增加）。例如，对于一个8B参数的模型，处理2M个token的单序列可能需要高达256GB的缓存。现有的Top-$k$缓存驱逐方法（Cache Eviction Methods）旨在通过在运行时驱逐大量非关键的缓存元素来减少KV缓存大小并保持生成质量。然而，现有方法通常在所有注意力头（Attention Heads）之间均匀分配压缩预算，完全忽略了每个头部独特的注意力集中度模式（有些头部注意力高度集中在少数元素上，而有些则分布广泛）。

为了解决这一问题，本文的主要贡献如下：
1. 识别了现有方法的局限性：指出现有KV缓存驱逐方法中均匀预算分配策略的缺陷，该策略忽略了单个头部的独特注意力模式。
2. 建立理论框架：建立了一个理论损失上界，用于衡量驱逐前后注意力输出之间的差异。该框架不仅解释了先前缓存驱逐方法的优化目标，还指导了自适应预算分配的设计。
3. 提出Ada-KV策略：提出了首个头部级别的自适应预算分配策略Ada-KV。该策略动态地将预算从注意力稀疏的头部重新分配给注意力分散的头部。
4. 即插即用的经验优势：通过高效的CUDA内核实现，Ada-KV提供了即插即用的兼容性，能够无缝集成到现有的SOTA方法（如SnapKV和Pyramid）中，形成了Ada-SnapKV和Ada-Pyramid。在Ruler和LongBench基准测试的29个数据集上（涵盖问题感知和更具挑战性的问题不可知场景），证明了其相较于现有方法的显著质量提升。

图1：自适应预算分配适应了不同头部的不同注意力集中度。左图：使用Llama-3.1-8B-Instruct的分析表明，大多数头部仅用很小的缓存（例如前5%）就能保留几乎所有的注意力权重，而分散的头部需要更大的缓存比例。右图：自适应分配将预算从稀疏头部转移到分散头部，与均匀分配相比，增加了聚合的保留注意力权重（从2.26增加到2.48）并降低了驱逐损失。

背景知识与相关工作

缓存驱逐方法的分类与局限
在长序列推理中，庞大的KV缓存元素规模导致了内存受限的情况，带来了显著的内存负担和I/O延迟【13, Catalyst: Optimizing cache management for large in-memory key-value systems+2023】。现有的减少缓存大小的方法主要分为两类：
1. 滑动窗口驱逐（Sliding window eviction）：如StreamingLLM【16, Efficient streaming language models with attention sinks+2023】，简单地保留几个初始缓存元素和滑动窗口内的元素，同时驱逐其他元素。然而，这种不加区分的滑动驱逐会导致生成质量显著下降。
2. Top-$k$驱逐（Top-$k$ eviction）：如FastGen【7, Model tells you what to discard: Adaptive kv cache compression for llms+2023】、H2O【8, H2o: Heavy-hitter oracle for efficient generative inference of large language models+2024】、SnapKV【11, SnapKV: LLM knows what you are looking for before generation+2024】和Pyramid【9, PyramidInfer: Pyramid KV cache compression for high-throughput LLM inference+2024】。这些方法基于注意力权重识别并保留$k$个关键缓存元素。早期的H2O聚合所有查询状态的注意力权重，但在单向掩码下容易误删最近的KV元素；SnapKV通过使用观察窗口内的查询状态来识别关键元素，达到了SOTA性能；Pyramid进一步引入了不同层之间的预算分配。然而，现有的Top-$k$方法通常在不同头部之间均匀分配整体预算，导致预算利用率低下。

稀疏注意力与缓存驱逐的区别
稀疏注意力方法（Sparse attention methods）在概念上与KV缓存驱逐相关，但在方法上存在根本差异【19, Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention+2024】。关键区别在于，KV缓存驱逐仅保留KV缓存的子集，而稀疏注意力方法保留所有条目，但在计算期间选择性地仅使用关键子集。因此，稀疏注意力方法不会减少KV缓存的内存占用，通常需要将缓存卸载到CPU内存或磁盘存储中。这两条技术路线实际上是正交的，未来可以探索先采用KV缓存驱逐压缩缓存，然后再应用稀疏注意力进行进一步加速。

方法细节

多头自注意力层的形式化描述
LLM通过自回归生成过程运行，每一步依赖上一个token来预测下一个。设$X \in \mathbb{R}^{n \times d}$表示包含序列中所有token的嵌入矩阵，设$x \in \mathbb{R}^{1 \times d}$为当前时间步用作输入的最后一个token的嵌入。使用【25, Deja vu: Contextual sparsity for efficient llms at inference time+2023】的符号系统，假设有$h$个注意力头。对于每个头$i \in [1, h]$，变换矩阵$W_i^Q, W_i^K, W_i^V \in \mathbb{R}^{d \times d_h}$将token嵌入映射到各自的Query、Key和Value状态，而最终的输出矩阵$W_i^O \in \mathbb{R}^{d_h \times d}$将中间结果变换为输出隐藏状态。在每个时间步，头$i$的先前KV缓存初始化为：$K_i = X W_i^K$，$V_i = X W_i^V$。接着，输入token $x$的嵌入被映射到每个头各自的Query、Key和Value状态，并相应地更新先前的KV缓存：

$$q_i = x W_i^Q, k_i = x W_i^K, v_i = x W_i^V \mathrm{and} K_i = \mathrm{Cat}[K_i: k_i], V_i = \mathrm{Cat}[V_i: v_i]$$

最后，输出$y \in \mathbb{R}^{1 \times d}$使用注意力权重$A_i \in \mathbb{R}^{1 \times n}$计算如下：

$$y = \sum_{i \in [1, h]} A_i V_i W_i^O \quad \mathrm{where} \quad A_i = \mathrm{softmax}(q_i K_i^T)$$

Top-$k$方法的理论基础与驱逐损失上界
Top-$k$驱逐方法通常预设关键缓存元素在未来生成过程中的稳定性。SOTA方法【11, SnapKV+2024】【9, PyramidInfer+2024】利用观察窗口内一系列token的查询状态，来计算与过去KV缓存元素的观察注意力权重。为方便表述，假设窗口大小为1。使用一组指示变量$\{\mathcal{T}_i \in \mathbb{R}^{1 \times n}\}$表示驱逐决策，分配给所有头$\{i \in [1, h]\}$的预算为$\{B_i\}$，其中$B = \sum_{i \in [1, h]} B_i$是单层注意力层的总预算：

$$\begin{aligned} \mathrm{Eviction\ Decision:} \quad \mathcal{T}_i^j = \begin{cases} 1 & \mathrm{Retain\ } K_i^j \mathrm{\ and\ } V_i^j, \\ 0 & \mathrm{Evict\ } K_i^j \mathrm{\ and\ } V_i^j. \end{cases} \end{aligned}$$

因此，头$i$仅保留预算大小为$B_i$的缓存元素：$\sum_{j \in [1, n]} \mathcal{T}_i^j = B_i$。驱逐后多头自注意力机制的输出$\hat{y}$为：

$$\hat{y} = \sum_{i \in [1, h]} \hat{A}_i V_i W_i^O, \quad \mathrm{where} \quad \hat{A}_i = \mathrm{softmax}(-\infty \odot (\mathbf{1} - \mathcal{T}_i) + q_i K_i^T)$$

其中$\odot$表示逐元素乘法。
缓存驱逐后生成质量的下降源于注意力输出的变化。我们将驱逐损失量化为自注意力机制驱逐前后输出之间的$L_1$距离：$L_1 \mathrm{\ Eviction\ Loss} = ||y - \hat{y}||_1$。
定理3.1：$L_1$驱逐损失可以被$\epsilon$界定：

$$L_1 \mathrm{\ Eviction\ Loss} \leq \epsilon = 2hC - 2C \sum_{i \in [1, h]} \sum_{j \in [1, n]} \mathcal{T}_i^j A_i^j$$

其中$C = \mathrm{Max}\{||V_i W_i^O||_\infty\}$是一个常数，代表最大行范数。
本质上，现有的缓存驱逐方法等价于最小化定理3.1中定义的上界$\epsilon$。Top-$k$驱逐决策公式化为：如果$A_i^j \in \mathrm{Top-}k(A_i, k=B_i)$，则$\mathcal{T}_i^{*j}=1$，否则为0。
定理3.2：在给定预算分配$\{B_i\}$下，Top-$k$缓存驱逐$\{\mathcal{T}_i^*\}$最小化了$L_1$驱逐损失的上界$\epsilon$，得到$\epsilon^*$：

$$\{\mathcal{T}_i^*\} = \underset{\{\mathcal{T}_i\}}{\arg \min} \epsilon \quad \mathrm{and} \quad \epsilon^* = \underset{\{\mathcal{T}_i\}}{\min} \epsilon = 2hC - 2C \sum_{i \in [1, h]} \sum_{j \in [1, n]}^{A_i^j \in \mathrm{Top-}k(A_i, k=B_i)} A_i^j$$

通过自适应预算分配优化Top-$k$方法
为了进一步最小化定理3.2中的上界$\epsilon^*$，提出了Top-$k$驱逐方法的自适应预算分配策略（Algorithm 1）。

# Algorithm 1: Ada-KV: Adaptive Budget Allocation
Input: Total budget: B; Attention weights for head i: {Ai}
Output: Allocated budgets {Bi}
1: Concatenate attention weights across heads A = Cat({Ai})
2: Select top B weights from A: Top-k(A, k = B)
3: Count number of selected weights for each head i: {fi}
4: Set the allocated budgets as {Bi = fi}
Return allocated budgets {Bi}

该策略先识别单层内所有头部中最大的$B$个注意力权重，再根据每个头部被选中的频率动态确定预算分配$\{B_i^*\}$。这使得大部分注意力集中在少数缓存元素上的稀疏头部被分配较小的预算；节省下来的预算被重新分配给注意力分布更广泛的分散头部。
定理3.3：从算法1得出的自适应预算分配$\{B_i^*\}$实现了Top-$k$驱逐策略相关损失的最小上界$\epsilon^{**}$：$\epsilon^{**} = \min_{\{B_i\}} \epsilon^*$。
经验上，该策略有效降低了实际的驱逐损失。图2的详细可视化显示，在相同的缓存预算下，自适应预算分配在大多数样本中持续降低了实际的驱逐损失。
图2：自适应分配产生更低的实际驱逐损失（Llama3.1-8B-Instruct在Qasper数据集上）。

集成到现有的缓存驱逐方法中
将Ada-KV策略无缝集成到SnapKV【11】和Pyramid【9, 10】中，创建了Ada-SnapKV和Ada-Pyramid（Algorithm 2）。

# Algorithm 2: Ada-KV Integration
Input: total budget B, tokens in observation window X^win, cache in observation window {K_i^win, V_i^win}, cache outside observation window {K_i, V_i}
Output: retained cache \hat{K}_i, \hat{V}_i
1: for i = 1 to h do
2:    Q_i^win = X^win W_i^Q
3:    \bar{A}_i = softmax(Q_i^win K_i^T)
4:    \bar{A}_i = \bar{A}_i.maxpooling1d(kernel_size).mean(dim=0)
5: end for
6: B = B - winsize * h
7: Derive budget allocation {B_i^*} using Algorithm 1(B, {\bar{A}_i})
8: Safeguard {B_i^*} = \alpha * {B_i^*} + (1 - \alpha) * (B / h)
9: Determine the Top-k eviction decision {\mathcal{T}_i^*} based on {B_i^*}
10: Select {\hat{K}_i, \hat{V}_i} from {K_i, V_i} according to {\mathcal{T}_i^*}
11: {\hat{K}_i, \hat{V}_i} = Cat({\hat{K}_i, \hat{V}_i}, {K_i^win, V_i^win})
Return retained cache \hat{K}_i, \hat{V}_i

集成过程发生在每层的驱逐过程之前。首先使用观察窗口内的查询状态计算过去缓存元素的观察注意力权重$\bar{A}_i$，然后通过最大池化层处理这些权重以保留基本信息。接着，引入了一个保护超参数$\alpha$（默认为0.2），以防止为高度稀疏的头部分配过小的预算，从而增强了对预设关键稳定性的容错能力。最后，根据自适应分配的预算进行Top-$k$选择，保留选定的元素以及观察窗口内的元素，驱逐其余元素。

自适应分配下的计算效率实现
1. 变长注意力与变长缓存元素：自适应分配导致不同头部的缓存元素长度可变，给高效计算带来挑战。通过采用广泛应用于连续批处理的变长FlashAttention技术【26, Flashattention+2022】【27, Flashattention-2+2023】，并实现了一种扁平的缓存存储布局（将层内所有注意力头的缓存连接成单个张量结构），结合自定义的CUDA内核，实现了高效的缓存更新操作。
2. 与分组查询注意力（GQA）的兼容性：现有的SnapKV和Pyramid缺乏GQA【31, Gqa: Training generalized multi-query transformer models from multi-head checkpoints+2023】兼容性，会在头部之间冗余地复制分组的KV缓存。实现了一种简单的GQA兼容机制，使用每个组内的平均注意力权重作为选择标准，消除了冗余，使得Llama-3.1-8B等模型实现了4倍的缓存大小缩减。

实验环境

数据集：
- Ruler：包含13个长序列任务，主要是Needle-in-a-Haystack测试的变体，提供全面的难度评估。选择16K长度作为主要评估基准。
- LongBench：包含16个数据集，涵盖多个任务领域（单/多文档QA、摘要、少样本学习、合成任务、代码生成），平均长度为6711。
模型架构：Llama-3.1-8B-Instruct和Mistral-7B-instruct-v0.2。两者均采用GQA技术。
基线方法：SnapKV、Pyramid（均实现了GQA兼容性）以及代表滑动窗口方法的StreamingLLM。
参数设置：在每层的prefilling阶段之后执行缓存驱逐。保护超参数$\alpha$固定为0.2。观察窗口大小设为32，最大池化核大小设为7。
软件配置：基于FlashAttention-2和自定义CUDA内核实现。

实验结果

1. Ruler基准测试评估
- 实验内容：在20%、40%、60%和80%的原始缓存预算下，评估问题感知（question-aware，问题提前已知）和更具挑战性的问题不可知（question-agnostic，压缩后才揭示问题）场景。
- 实验结果：在问题感知场景下，Ada-SnapKV和Ada-Pyramid在小压缩预算（40%和20%）下显著降低了质量损失。在问题不可知场景下，现有的Top-$k$方法性能大幅下降，而Ada-KV方法在所有预算设置下均大幅降低了质量损失。例如，在80%和20%预算下，Ada-SnapKV将SnapKV的得分分别从87.59和44.02提升至92.67和53.29（见图3）。
- 子任务分析：在困难的Needle-in-a-Haystack任务（如S-NIAH-3和MK-NIAH-2）中，80%预算下的原始SnapKV性能显著下降，而Ada-SnapKV保持了近乎无损的性能（得分从62.4和85.2提升至97.6和99.6，见图4）。
图3：Ruler基准测试13个数据集的平均得分
图4：Ruler子任务分析（问题不可知，Llama-3.1-8B-Instruct）

2. LongBench基准测试评估
- 实验内容：在固定预算{128, 256, 512, 1024, 2048}下，评估问题感知和问题不可知场景。此外，在问题不可知场景中，还进行了基于比例的缓存预算压缩测试。
- 实验结果：在问题感知场景中，Ada-SnapKV和Ada-Pyramid持续提升生成质量，在2048固定预算下接近无损性能。在问题不可知场景中，所有方法性能显著下降（例如Llama在2048预算下，SnapKV得分从49.09降至42.86），但Ada-KV方法依然优于原始方法（见图5）。
- 基于比例的压缩分析：代码任务对缓存压缩不敏感，而摘要和QA任务退化严重。Ada-SnapKV有效缓解了这些损失。在18个领域案例和三个缓存预算中，Ada-SnapKV在15个领域实现了质量提升。在更大的Llama-3.1-70B模型上，Ada-SnapKV同样展现了强大的增益。
图5：固定预算下LongBench 16个数据集的平均得分

3. 自适应分配下的计算效率
- 实验内容：在固定预算1024下，评估Ada-SnapKV、SnapKV和完整缓存的峰值内存使用量和解码延迟。
- 实验结果：凭借扁平缓存布局和自定义CUDA内核，Ada-SnapKV实现了与原始SnapKV相当的峰值内存使用量和解码延迟，两者均显著优于完整缓存情况（见图6）。这证明Ada-KV策略在提高生成质量的同时，保持了强大的计算效率。
图6：效率（均使用FlashAttention-2）

补充细节：自适应预算分配策略的广泛优势

由于其即插即用的设计，Ada-KV的适用范围超越了本文提出的两个集成案例。截至发表时，许多后续工作（如CriticalKV【39】和DefensiveKV【40】）已经采用了Ada-KV策略并获得了增强。此外，并发方法探索了通过基于训练的分析进行预算分配（如DuoAttention将头部划分为“全注意力”和“流式”类型，HeadKV实现了更细粒度的分配）。
当与CriticalKV和DefensiveKV结合时，Ada-KV持续提高性能。值得注意的是，即插即用的CriticalKV和DefensiveKV在采用Ada-KV策略增强后，超越了基于训练的方法（如DuoAttention和HeadKV）的性能。这证明了Ada-KV保留了巨大的价值，并提供了显著的增强潜力。

结论

本文重新审视了用于高效LLM推理的缓存驱逐策略，并揭示了一个被忽视的关键因素：跨注意力头的自适应预算分配。在损失上界理论分析的指导下，提出了首个自适应预算分配策略Ada-KV，以优化KV缓存驱逐方法。通过将其无缝集成到两种现有的SOTA方法中，引入了Ada-SnapKV和Ada-Pyramid，证明了其即插即用的优势。除了常见的问题感知压缩外，还在更具挑战性、较少探索的问题不可知压缩场景中进行了评估。在Ruler和LongBench基准测试上的结果不仅暴露了当前缓存驱逐策略的局限性，还突显了自适应分配增强缓存驱逐的潜力。未来的工作可以探索将头部级别的分配机制扩展到跨层分配。

附录

其他相关工作
一些不减少缓存元素数量的方法与本工作正交。例如，Paged Attention【28】采用高效的内存管理策略来减少I/O延迟，而不改变KV缓存的大小。KV缓存量化方法【43, 44, 45】通过降低单个元素的精度来减小缓存大小。本文关注的缓存驱逐技术未来可以进一步与量化结合。此外，推测解码（Speculative decoding）【46, 47, 48】通常使用具有减少KV缓存的模型来生成草稿，未来可以集成更高级的缓存驱逐方法来进一步提高此类方法的效率。

Ruler基准测试的详细结果
图7补充了主文中由于空间限制而缺失的多跳QA子任务结果。图8至图10全面概述了Ada-SnapKV和Ada-Pyramid在两个LLM上、问题不可知和问题感知场景中所有16K Ruler子任务的性能。结果表明，无论场景、模型还是缓存预算，Ada-SnapKV和Ada-Pyramid在大多数任务中都优于原始的SnapKV和Pyramid方法，突显了自适应分配策略的广泛适用性和有效性。
图7：Ruler上问题不可知场景的M-QA子任务分析（Llama3.1-8B-Instruct）
图8：Ruler子任务分析（问题感知，Llama3.1-8B-Instruct）。
图9：Ruler子任务分析（问题不可知，Mistral-7B-Instruct-v0.2）。
图10：Ruler子任务分析（问题感知，Mistral-7B-Instruct-v0.2）。

保护参数$\alpha$的鲁棒性分析
为了避免为每个模型或预算微调保护参数$\alpha$带来的复杂性，实验选择了一个固定值。在LongBench基准测试上使用Mistral-7B对$\alpha$进行的鲁棒性分析表明，较小的$\alpha$允许更激进的预算分配，从而在有限预算下提高性能；相反，较大的$\alpha$在较高预算下表现略好。基于这种权衡，本文选择了一个平衡值$\alpha = 0.2$。

定理3.1的证明过程
考虑到softmax函数 $softmax(x)^j = \frac{exp(x^j)}{\sum_j exp(x^j)}$，驱逐过程后的注意力权重 $\hat{A}_i = \mathrm{softmax}(-\infty \odot (1 - \mathcal{T}_i) + s_i)$，其中$s_i = q_i K_i^T$。
经过推导可得：

$$\hat{A}_i^j = \frac{\mathcal{T}_i^j A_i^j}{| |A_i \odot \mathcal{T}_i| |_1}$$

因此输出为：

$$\hat{y} = \sum_{i \in [1, h]} \frac{A_i \odot \mathcal{T}_i}{| |A_i \odot \mathcal{T}_i| |_1} V_i W_i^O$$

计算输出之间的L1距离：

$$| |y - \hat{y}| |_1 = | |\sum_{i \in [1, h]} \left( \mathbf{1} - \frac{\mathcal{T}_i}{| |A_i \odot \mathcal{T}_i| |_1} \right) \odot A_i V_i W_i^O| |_1$$

$$\le C \sum_{i \in [1, h]} | | \left( \mathbf{1} - \frac{\mathcal{T}_i}{| |A_i \odot \mathcal{T}_i| |_1} \right) \odot A_i | |_1$$

其中 $C = \mathrm{Max}\{||V_i W_i^O||_\infty\}$。
考虑到 $\mathcal{T}_i^j \in \{0, 1\}$ 且 $\sum_{j \in [1, n]} A_i^j = 1$，进一步展开并化简公式，最终得到：

$$L_1 \mathrm{\ Eviction Loss} \le \epsilon = 2hC - 2C \sum_{i \in [1, h]} \sum_{j \in [1, n]} \mathcal{T}_i^j A_i^j$$

定理3.2的证明过程
给定预算分配结果$\{B_i\}$，最小化$\epsilon$的目标可以分解为$i \in [1, h]$个独立的子问题：

$$\arg \max_{\mathcal{T}_i} \sum_{j \in [1, n]} \mathcal{T}_i^j A_i^j, \quad \mathrm{s.t.} \quad \sum_{j \in [1, n]} \mathcal{T}_i^j = B_i$$

Top-$k$缓存驱逐最大化了这$h$个独立的子问题，从而最小化了上界$\epsilon$。最优解满足：

$$\min_{\mathcal{T}_i} \sum_{j \in [1, n]} \mathcal{T}_i^j A_i^j = \sum_{j \in [1, n]} \mathcal{T}_i^{*j} A_i^j = \sum_{j \in [1, n]}^{A_i^j \in \mathrm{Top-}k(A_i, B_i)} A_i^j$$

定理3.3的证明过程
算法1根据Top-$k$索引自适应分配预算$\{B_i^*\}$。由此产生的Top-$k$驱逐导致上界$\epsilon^{**}$。给定固定的总预算 $\sum_{i \in [1, h]} B_i = \sum_{i \in [1, h]} B_i^*$，可以推导出对于任何预算分配结果$\{B_i\}$，都有 $\epsilon^{**} \le \epsilon^*$。这是因为在上界中，全局Top-$k$选择的注意力权重之和必然大于等于任何局部Top-$k$（即预设$\{B_i\}$）的权重之和。这个结果也可以理解为全局最优解总是优于局部最优解。

头部集中度的详细可视化
图11补充了主文中的图1a，展示了所有层的可视化结果。可以观察到，在所有层中，不同的头部在注意力集中度上表现出显著的差异。这表明自适应分配策略在实践中具有降低驱逐损失的巨大潜力。
图11：头部集中度的可视化