ChunkKV: Semantic-Preserving KV Cache Compression for Efficient Long-Context LLM Inference

• 文章标题: ChunkKV: 通过保留语义的KV缓存压缩实现高效的长上下文LLM推理
• 作者/机构: Xiang LIU♡∗, Zhenheng TANG♣∗, Peijie DONG♡, Zeyu LI♡, Yue LIU♢†, Bo LI♠♣, Xuming HU♡†, Xiaowen CHU♡†
• 机构: ♡ 香港科技大学（广州）, ♣ 香港科技大学计算机科学与工程系, ♠ 广州市香港科大霍英东研究院, ♢ 特斯联科技

A1 主要贡献

核心问题

大型语言模型（LLMs）在处理长文本（如书籍、报告）时需要大量GPU内存，其中键值（KV）缓存的消耗在推理过程中可高达总内存的70%。现有的压缩方法通过评估单个token的重要性来减少内存占用，但这种做法忽略了token之间关键的语义关系，导致上下文碎片化和性能下降。如图1所示，离散的token级重要性评估可能只关注与问题相关的词语（如“turaco”），而忽略了文档中关键的对象信息（如食物），从而丢失了必要的语义信息。

图1：令牌离散方法和分块方法对语义保留的影响说明。离散方法保留了与问题相关的词语，但常常忽略了主语。相比之下，分块方法保留了词语的主语，从而保持了更准确的语义信息。对于等式：S是得分函数，c是一个令牌块。

研究目标

本文旨在解决现有KV缓存压缩方法中孤立的token重要性度量问题，并提出一种能够保留KV缓存中语义信息的新方法。

创新点

本文提出了一种名为ChunkKV的简单而有效的KV缓存压缩方法，其核心思想是将语义块（chunks）而不是孤立的token作为基本的压缩单元。这种方法能够保留完整的语言结构和上下文完整性，确保即使在激进的压缩下也能保留基本含义。

本文的主要贡献总结如下：
1. 识别了现有问题：指出现有的离散KV缓存压缩方法会无意中修剪掉必要的语义信息。
2. 提出ChunkKV方法：这是一种简单的KV缓存压缩方法，它使用分块的方式来保留语义信息。具体而言，它将token分组为块，并作为一个整体进行保留或丢弃，从而保留了信息最丰富的语义块（如图1和表13所示）。
3. 设计层级索引重用技术：研究发现，ChunkKV保留的KV缓存索引在不同层之间表现出更高的相似性。基于此，本文设计了一种层级索引重用（layer-wise index reuse）技术，它可以在多个层之间重用压缩的token索引，从而减少了KV缓存压缩方法带来的额外计算开销，并将吞吐量提高了26.5%。
4. 全面的实验验证：在多个前沿长上下文基准测试（如LongBench, Needle-In-A-HayStack）、上下文学习任务（如GSM8K）和多步推理LLM上对ChunkKV进行了评估。实验结果表明，在相同的压缩率下，ChunkKV的精确度比现有SOTA方法高出8.7%，证明了其在效率和性能上的显著提升。

表1：KV缓存压缩方法的比较。该表显示ChunkKV在保留语义信息和高效索引重用方面优于其他方法。

A2 方法细节

3.1 KV缓存压缩的初步研究

KV缓存的内存消耗。随着LLM处理长上下文能力的增强，KV缓存对于提高推理效率至关重要。然而，在处理长输入上下文时，它会消耗大量GPU内存。解码阶段KV缓存的GPU内存成本可通过以下公式计算：

$$M_{KV} = 2 \times B \times S \times L \times N \times D \times 2$$

其中B是批量大小，S是提示和解码长度的序列长度，L是层数，N是注意力头的数量，D是每个注意力头的维度，第一个2代表KV矩阵，最后一个2代表使用float16时的精度。附录E中的表格展示了LLaMA-3-8B-Instruct【【索引30，Introducing meta llama 3: The most capable openly available llm to date，2024，https://ai.meta.com/blog/meta-llama-3/】】的配置参数。当批量大小B=1，提示序列长度S=2048时，KV缓存的GPU内存成本接近1GB。如果批量大小超过24，KV缓存的GPU内存成本将超过一 块RTX 4090 GPU的容量。为解决此问题，研究人员提出了KV缓存压缩方法，旨在保留最少量的KV缓存同时尽可能多地保留信息。关于LLM配置参数的更多细节，请参阅附录E。

3.2 基于块的KV压缩

ChunkKV的核心思想。为了解决现有KV缓存压缩方法的局限性，我们提出了ChunkKV，一种新颖的KV缓存压缩方法，它保留了信息最丰富的语义块。ChunkKV的关键思想是将KV缓存中的token分组为能够保留更多语义信息的块，例如一个包含主语、谓语和宾语的块。如图1所示，ChunkKV保留了包含更多语义信息的KV缓存块。首先，我们将一个块定义为一组包含相关语义信息的token。通过从原始KV缓存中保留信息最丰富的块，ChunkKV可以有效减少KV缓存的内存使用，同时保留必要的信息。有关ChunkKV的更多信息，请参阅附录A。

ChunkKV算法流程。算法1展示了ChunkKV的伪代码。首先，我们遵循H2O【【索引11，H2o: Heavy-hitter oracle for efficient generative inference of large language models，2023，Advances in Neural Information Processing Systems】】和SnapKV【【索引15，Snapkv: Llm knows what you are looking for before generation，2024，ArXiv preprint】】的方法，通过计算注意力分数$A \leftarrow Q_{T_q-w:T_q} K^T$来设置观察窗口，其中$Q_{T_q-w:T_q}$是观察窗口，K是键矩阵，窗口大小w通常设置为{4, 8, 16, 32}。接下来，计算块的数量C，公式为$C = \lceil \frac{T_k}{c} \rceil$，其中$T_k$是键矩阵的长度，c是块大小。然后，计算每个块的注意力分数，公式为$A_i = \sum_{j=(i-1)c+1}^{ic} A_{:,j}$，其中i = 1, 2, ..., C。我们使用top-k算法作为ChunkKV的采样策略。根据注意力分数选择top-k个块，其中$k = \lfloor \frac{L_{max}}{c} \rfloor$，而$L_{max}$是压缩后KV缓存的最大长度。最后一个块的大小设置为$min(c, L_{max} - (k - 1) \times c)$。top-k个块的索引保留了原始序列顺序。在压缩步骤中，只保留与所选索引对应的键和值矩阵，从而得到压缩后的KV缓存。最后，将原始KV缓存的观察窗口拼接到压缩后的KV缓存中，通过替换最后w个token来保留重要信息。压缩后的KV缓存随后用于后续的注意力计算。在实现中，我们增加了向量化操作、内存优化等来优化代码。更多细节请参阅附录A.5。

算法 1 ChunkKV
输入: Q ∈ R^{T_q×d}, K ∈ R^{T_k×d}, V ∈ R^{T_v×d}, 观察窗口大小 w, 块大小 c, 压缩后KV缓存最大长度 L_max
输出: 压缩后的KV缓存 K′, V′
1: 观察窗口计算:
2: A ← Q_{T_q-w:T_q} K^T ▷ 观察窗口的注意力分数
3: C ← ⌈ T_k / c ⌉ ▷ 计算块的数量
4: 块注意力分数计算:
5: for i = 1 to C do
6:   A_i ← Σ_{j=(i-1)c+1}^{ic} A_{:,j} ▷ 每个块的注意力分数总和
7: end for
8: Top-K 块选择:
9: k ← ⌊ L_max / c ⌋
10: Top_K_Indices ← 基于 A_i 的 Top-k 块的索引
11: 压缩:
12: K′, V′ ← index_select(K, V, Top_K_Indices)
13: 拼接:
14: K′ ← concat(K′_{0:L_max-w}, K_{T_k-w:T_k})
15: V′ ← concat(V′_{0:L_max-w}, V_{T_v-w:T_v})
16: 返回 K′ , V′

3.3 层级索引重用

层级索引相似性观察。此外，我们研究了由ChunkKV保留的KV缓存索引，发现与先前的方法相比，它们表现出更高的相似性。图2展示了SnapKV和ChunkKV的层级相似性热图。每个单元格代表两个层之间保留的KV缓存索引的相似性，颜色越深表示相似性越高。结果表明，由ChunkKV保留的KV缓存索引与相邻层中的索引更为相似。

如表2所示，与SnapKV相比，ChunkKV在不同模型架构中，相邻层之间的平均Jaccard相似度始终更高，这表明ChunkKV中保留的token索引彼此之间更为相似。更详细的可视化请参阅附录B.1.3。

表2：不同模型相邻层保留的KV缓存索引相似度。

图2：SnapKV（左）和ChunkKV（右）在LLaMA-3-8B-Instruct上保留的KV缓存索引的层级相似性热图。深色表示相似度更高。更多可视化内容见附录B.1.3。

层级索引重用算法。基于上述关于KV缓存索引的发现，我们提出了一种无需训练的层级索引重用方法，以进一步降低KV缓存压缩的额外成本，该方法在多个层之间重用压缩的token索引。我们在第4.3节评估了层级索引重用方法的效率和有效性。与FullKV基线相比，重用层级索引可将KV缓存压缩时间减少20%，而性能仅下降0.5%。

这种层级索引重用方法在算法2中有正式描述。ChunkKV压缩过程返回压缩后的KV缓存及其各自的token索引，表示为$I_l$。对于层级索引重用，我们定义一个层的分组，使得所有$N_{reuse}$层共享相同的ChunkKV token索引。具体来说，对于一组层$\{l, l+1, \ldots, l+N_{reuse}-1\}$，我们在第一层l上执行ChunkKV以获得token索引$I_l$，并在后续的层$l+1, l+2, \ldots, l+N_{reuse}-1$中重用$I_l$。符号$K_l[I_l]$和$V_l[I_l]$表示根据$I_l$中的索引选择键和值缓存。层级索引重用的效率分析在附录B.1.1中提供。

算法 2 ChunkKV的层级索引重用
输入: LLM中的层数 N_layers, 重用层数 N_reuse
1: 初始化: 用于存储索引的字典 I_reuse = {}
2: for l = 0 to (N_layers - 1) do
3:   if l mod N_reuse == 0 then
4:     K′_l, V′_l, I_l ← ChunkKV(K_l, V_l)
5:     I_reuse[l] ← I_l
6:   else
7:     I_l ← I_reuse[⌊ l / N_reuse ⌋ × N_reuse]
8:   end if
9:   K′ ← index_select(K_l, I_l)
10:  V′ ← index_select(V_l, I_l)
11: end for

理论解释。我们从上下文学习（ICL）【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】的角度提供了一个理论解释，来说明为什么在ChunkKV中根据连续序列维护KV缓存比根据稀疏token更好。非正式地说，连续的块级KV缓存保留了ICL中的完整示例（语义信息），从而降低了对可区分性的要求，即可区分性是示例和问题之间KL散度的下界（条件2中的方程4）。完整的分析在附录C中提供。

A4 实验环境

• 数据集：

  • 上下文学习（In-Context Learning）:

    • GSM8K【【索引27，Training verifiers to solve math word problems，2021，ArXiv preprint】】：包含超过1000个算术问题，用于评估推理能力。
    • JailbreakV【【索引28，Jailbreakv: A benchmark for assessing the robustness of multimodal large language models against jailbreak attacks，2024，First Conference on Language Modeling】】：一个用于评估语言模型安全性的越狱基准。

  • 长上下文（Long-Context）:

    • LongBench【【索引25，LongBench: A bilingual, multitask benchmark for long context understanding，2024，Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics】】：用于评估模型处理长文档和复杂信息序列能力的基准套件，包含中英文子任务。
    • Needle-In-A-HayStack (NIAH)【【索引26，Needle In A Haystack - pressure testing LLMs，2023，Github】】：评估LLM从大量文档中提取隐藏信息（长上下文检索）的能力。

• 模型架构：

  • DeepSeek-R1-Distill-Llama-8B【【索引29，Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning，2025，arXiv preprint】】
  • LLaMA-3-8B-Instruct【【索引30，Introducing meta llama 3: The most capable openly available llm to date，2024，https://ai.meta.com/blog/meta-llama-3/】 】
  • LLaMA-3.1-8B-Instruct【【索引30，Introducing meta llama 3: The most capable openly available llm to date，2024，https://ai.meta.com/blog/meta-llama-3/】 】
  • Mistral-7B-Instruct【【索引31，Mistral 7b，2023，URL https://arxiv.org/abs/2310.06825】 】
  • Qwen2-7B-Instruct【【索引32，Qwen2 technical report，2024，ArXiv preprint】】
  • LLaMA-3-70B-Instruct

• 硬件配置：

  • GPU: A40 GPU

• 软件配置：

  • 依赖库: Flash Attention 2

A4 实验结果

实验中，即使对于不同的模型架构，块大小（chunk size）也统一设置为10。所有实验均进行三次，并取平均分以确保结果的稳健性。

4.1 上下文学习（In-Context Learning）

• 实验内容：评估ChunkKV在GSM8K、多样本GSM8K（Many-Shot GSM8K）和JailbreakV基准上的性能。其中，多样本GSM8K设置了50个样本，提示长度超过4k tokens，是一项比长上下文检索更具挑战性的任务。
• 实验结果：
  • GSM8K：在LLaMA-3、Qwen2和DeepSeek等多个模型和不同压缩率下，ChunkKV的性能均优于其他KV缓存压缩方法（如StreamingLLM、H2O、SnapKV、PyramidKV）。（见表3）
  • 多样本GSM8K：ChunkKV同样表现最佳，凸显了其在通过块级KV缓存维护复杂算术推理任务关键上下文信息方面的有效性。（见表4）
  • JailbreakV：在安全性基准测试中，ChunkKV在不同压缩率下均优于其他方法，表明块级KV缓存比离散token级压缩更有效。（见表5）

表3：GSM8K性能比较。# SLM=StreamingLLM, SKV=SnapKV, PKV=PyramidKV

表4：多样本（50-shot）GSM8K性能比较。

表5：JailbreakV性能比较。

4.2 长上下文基准（Long-Context Benchmark）

• 实验内容：在LongBench和Needle-In-A-HayStack (NIAH)两个广泛使用的长上下文基准上评估ChunkKV。
• 实验结果：
  • LongBench：在LLaMA-3-8B、Mistral-7B和Qwen2-7B模型上，以10%的压缩率进行测试。结果显示，ChunkKV在英文和中文子任务上均优于其他压缩方法，其与FullKV基线的性能差距最小，甚至在Qwen2模型上实现了性能超越。（见表6）
  • NIAH：在8k和32k上下文长度下进行评估。统计结果（表7）和可视化图（图3）都表明，ChunkKV在管理不同token长度和深度百分比方面是鲁棒且有效的，其检索成功率最高。

表6：KV缓存压缩方法在LongBench基准上的表现。结果显示了与FullKV基线的性能差距（负值表示性能较差）。

表7：NIAH性能比较。

图3：在8k上下文长度下，LLaMA3-8B-Instruct在NIAH基准测试中的表现，KV缓存大小=128。

4.3 索引重用（Index Reuse）

• 实验内容：从效率和任务性能两个方面评估层级索引重用技术。
• 实验结果：
  • 效率：在A40 GPU上使用LLaMA3-8B-Instruct进行测试，重用层数设为2。结果显示，索引重用策略（ChunkKV_reuse）相比FullKV基线，最多可将延迟降低20.7%，吞吐量提升26.5%，尤其在长输入序列下效果显著。（见表8）
  • 任务性能：在LongBench和GSM8K上进行评估。结果显示，索引重用在保持模型性能的同时减少了计算需求。在LongBench上性能下降极小（<0.6%），在GSM8K上则表现出中性或略微积极的影响。这验证了ChunkKV选择的语义块在相邻Transformer层中重要性一致的假设。（见表9）

表8：不同输入输出配置下ChunkKV与FullKV的延迟和吞吐量比较。括号中的百分比表示相对于FullKV基线的提升。

表9：在LongBench和GSM8K上重用索引的性能比较。△表示索引重用相比基线的性能下降。

4.4 块大小（Chunk Size）

• 实验内容：探究块大小对ChunkKV性能的影响，测试了{3, 5, 10, 20, 30}等不同大小。
• 实验结果：当块大小在5到20之间时，性能相对稳定，其中块大小为10时效果最佳。块太小（如3）会导致上下文碎片化，而太大（如30）则会使语义粒度过粗，都可能导致性能下降。这一趋势在LongBench和NIAH基准及不同模型上都保持一致，表明最优块大小对特定任务或模型不高度敏感。因此，推荐使用块大小10作为鲁棒的默认值。（见表10）

表10：不同块大小消融实验在LongBench和NIAH上的结果。

4.5 与KV量化方法的比较

• 实验内容：将ChunkKV与KIVI量化方法【【索引45，Kivi: A tuning-free asymmetric 2bit quantization for kv cache，2024，arXiv preprint】】进行比较。两者原理不同：量化降低KV矩阵精度，而本文方法减少KV矩阵大小。
• 实验结果：在效率指标上，ChunkKV表现出显著优势。尽管两种方法都大幅减小了缓存大小，但ChunkKV在首token生成时间（TTFT）和每token处理时间（TPOT）上更优。总生成时间上，ChunkKV（164.66s）比KIVI的2-bit量化（226.52s）快27.3%。（见表11）

表11：ChunkKV和KIVI在LlaMa-3-8B-Instruct上的效率结果。

4.6 混合压缩分析：不同层深度的块级与token级压缩

• 实验内容：设计了一个混合压缩模型（LLaMA-3-8B-Instruct），底层16层使用ChunkKV（块级），顶层16层使用SnapKV（token级），并在LongBench上与纯ChunkKV和纯SnapKV进行比较。

• 实验结果：

  • 纯ChunkKV模型获得了最高的平均分，验证了即使在深度层，通过分块保留局部语义完整性也是一种鲁棒有效的策略。
  • 混合模型揭示了任务依赖的性能权衡：
    • 对于局部信息检索任务（如单/多文档问答），纯ChunkKV显著更优，因其能保留完整的语言单元。
    • 对于全局理解任务（如摘要和少样本学习），混合模型表现最好，可能是因为在形成抽象表示的深层，token级的SnapKV能更好地保留更广泛、分散的全局重要信号。（见表12）

• 结论：对于通用和鲁棒的解决方案，纯ChunkKV被证明是最有效的。

表12：混合压缩模型在LongBench上的性能。虽然混合模型在全局理解任务中显示出优势，但纯ChunkKV实现了最佳的整体性能，验证了即使在深层保留语义块的鲁棒性。

A5 结论

本文首先指出现有的KV缓存方法缺乏对数据语义信息的考量，随后提出了一种新颖的KV缓存压缩方法ChunkKV，该方法通过保留信息更丰富的块来保护语义信息。通过在多种先进的LLM（包括DeepSeek-R1、LLaMA-3、Qwen2和Mistral）和多样化的基准测试（GSM8K、LongBench、NIAH和JailbreakV）上进行的大量实验，我们证明了ChunkKV在仅使用一小部分内存的情况下，性能始终优于现有方法。我们的综合分析表明，ChunkKV基于块的方法能够维护关键的上下文信息，从而在复杂推理任务、长上下文理解和安全评估中取得更优异的性能。该方法的有效性在诸如多样本GSM8K和多文档问答等对语义连贯性要求极高的挑战性场景中尤为突出。此外，我们提出的层级索引重用技术在对性能影响极小的情况下，带来了显著的计算效率提升，实现了高达20.7%的延迟降低和26.5%的吞吐量提升。这些由详细的量化分析和消融研究支持的发现，确立了ChunkKV作为KV缓存压缩技术的一项重大进步，为在资源受限环境中部署LLM并保持高质量输出提供了有效的解决方案。

A6 附录

A ChunkKV与离散token方法的深度分析

A.1 定量分析

实验设置与指标。为严格评估ChunkKV与基于离散token方法的有效性，我们使用LLaMA-3-8B-Instruct模型进行了系统性实验。我们从LongBench数据集的每个子类别中随机选择了100个序列，并分析了两个关键指标在不同模型层的表现：KV缓存L1损失和注意力余弦相似度。对于每个序列，我们：1. 计算未压缩的完整KV缓存和注意力模式作为基准。2. 应用ChunkKV、SnapKV和H2O压缩方法，固定压缩率为10%，参数设置与表17相同。3. 测量压缩版本与未压缩版本之间的差异。

图4：在LongBench的单文档问答子类别中，ChunkKV与基于离散token的方法在L1损失和注意力余弦相似度上的逐层比较。

结果分析。如图4所示，ChunkKV在两个指标上均表现出优越性能：
- KV缓存L1损失：与SnapKV和H2O相比，ChunkKV始终实现更低的L1损失，尤其是在早期和中间层（5-25层）。这表明基于语义块的方法能更好地保留原始KV缓存信息。
- 注意力余弦相似度：ChunkKV在大多数层中展现出更高的相似度得分，在0-5层和20-30层表现尤为出色。这表明它能更好地保留token之间的注意力关系，这对于维持语义理解至关重要。

为了量化这些改进，我们计算了所有层的平均指标，如表13所示。ChunkKV同时实现了最低的L1损失和最高的注意力余弦相似度，优于两种基线方法。

结果的意义。虽然这些改进在绝对值上可能看起来不大（L1损失约2%，余弦相似度约1.5%），但其实际意义重大。这些指标反映了模型维持关键语义关系和注意力模式的能力，这对于复杂的推理任务至关重要。在不同序列上的一致改进表明，保留语义块比选择单个token能带来更好的信息保留。

在模型中间层，通常负责更高级别的语义处理，性能的提升尤为明显。这为ChunkKV在下游任务上优于基于离散token的方法提供了具体证据。

表13：LongBench中不同任务类别下KV缓存指标的详细比较。

A.2 假设场景

场景描述。为了更深入地理解ChunkKV相比于基于离散token方法的有效性，我们通过一个假设场景进行详细分析。此分析旨在阐明这些方法之间的根本差异，并解释为何ChunkKV在保留长上下文中的语义信息方面更有效。假设有一份综合性文档，详细介绍了各种动物的信息，包括它们的栖息地、饮食和行为。用户提问：“野生熊猫吃什么？” ChunkKV和基于离散token的方法都会利用这个问题来计算文档的注意力分数。然而，它们选择和保留信息的方法有显著不同。

A.2.1 基于离散token的方法

信息保留方式。一个基于离散token的方法可能会识别并保留具有高相关性分数的单个token，例如：“熊猫”、“吃”、“竹子”、“野生”、“饮食”、“食物”。尽管这些token是相关的，但它们缺乏上下文和连贯性。该方法可能会丢弃其他提供关键上下文或完善信息的必要token。

A.2.2 ChunkKV方法

信息保留方式。相比之下，ChunkKV会识别并保留有语义意义的块，例如：“在野外，熊猫主要吃竹笋和竹叶”、“它们的饮食99%是竹子，但偶尔也吃其他植物”、“野生熊猫在有机会时也可能吃小型啮齿动物或鸟类”。通过保留这些块，ChunkKV不仅保留了相关的关键词，还保留了它们的上下文关系和其他相关信息。

A.3 对比分析

ChunkKV的优势。当我们考虑这些保留的信息片段在后续处理中如何被使用时，ChunkKV的优势就变得显而易见：
1. 上下文理解：离散的token要求模型从孤立的词汇中重构意义，可能导致歧义。ChunkKV提供完整的短语或句子，可以实现即时且准确的理解。
2. 语义连贯性：ChunkKV保留了块内的语义关系，这对于理解熊猫主要食物和偶尔食物来源之间的细微差别至关重要。
3. 信息密度：与离散方法相比，单个块可以在其正确的上下文中包含多个相关token，从而在相同的压缩缓存大小内可能保留更多有用信息。
4. 减少歧义：离散方法可能会保留来自描述不同动物的句子中的“吃”这个token。ChunkKV确保“吃”这个词被特别保留在野生熊猫的上下文中。
5. 时间和逻辑流：ChunkKV能够保持原文中思想的顺序，保留任何对于理解可能至关重要的时序或逻辑进展。

A.4 对模型性能的影响

性能提升的启示。该分析揭示了对模型性能的几个关键影响：
- 提高准确性：通过保留富含上下文的信息，ChunkKV能够对查询做出更准确的回应，特别是那些需要细致理解的查询。
- 增强长上下文处理：保留语义块有助于更好地处理长距离依赖和复杂的推理任务。
- 减少计算开销：尽管两种方法都压缩了KV缓存，但ChunkKV的方法可能减少了对大量上下文重构的需求，从而可能提高推理效率。
- 通用性：基于块的方法可能在更广泛的任务和领域中更有效，因为它保留了语言的自然结构。
本深度分析展示了为什么ChunkKV在保留长上下文中的语义信息方面更有效。通过保留连贯的文本块，它为语言模型提供了更具上下文和语义完整性的信息，从而在需要深度理解和从大量文档中准确检索信息的任务中表现更佳。

A.5 实现细节

优化策略。在ChunkKV的实现中，我们专注于最大化计算效率以确保在实际应用中的可行部署。在算法3中，我们采用了几项关键的优化策略来减少推理过程中的计算开销和内存占用。

优化实现的关键改进。优化的实现具有以下几个关键改进：
- 向量化操作。我们利用批量矩阵运算来计算注意力分数，而不是显式循环，从而显著减少了计算开销。通过优化的张量操作（第7-11行）计算块分数，我们实现了比逐个token处理快得多的加速。
- 内存优化。我们的实现采用了一种二元掩码方法（第15-20行），避免了冗余的内存分配并减少了碎片化。这种单遍token选择策略整合了块保留和最近窗口保留操作，从而在推理过程中实现了更高效的内存使用。
- 边界处理。我们在整个算法中加入了鲁棒的边界检查（第10、12、18行），以处理诸如不完整块或有限压缩长度等边缘情况。这确保了在不同输入长度下的稳定性能，而无需特殊情况处理。
- 单遍压缩。我们的实现不是执行多次拼接操作，而是构建一个全面的选择掩码，该掩码包括语义上重要的块和最近的上下文token。这种方法（第21-25行）最小化了内存复制操作并减少了计算开销。

我们的PyTorch实现通过利用GPU加速所有向量和矩阵运算，进一步扩展了这些优化。将CUDA核心集成到关键操作中，如注意力分数计算和top-k选择，即使在长上下文窗口下也能实现实时性能。

1: 输入: Q ∈ R^{T_q×d}, K ∈ R^{T_k×d}, V ∈ R^{T_v×d}, 观察窗口大小 w, 块大小 c, 压缩后
   KV缓存最大长度 L_max
2: 输出: 压缩后的KV缓存 K′, V′
3: q_observe ← Q_{T_q-w:T_q} ▷ 提取观察窗口查询
4: A ← q_observe K^T ▷ R^{w×T_k} 注意力分数
5: C ← ⌈ T_k / c ⌉ ▷ 块的数量
6: 分配 chunk_scores ∈ R^C 以存储块的重要性
7: 向量化块分数计算:
8: for i = 0 to C - 1 do
9:     start ← i · c
10:    end ← min((i + 1) · c, T_k)
11:    chunk_scores[i] ← Σ_{j=start}^{end-1} Σ_{q=0}^{w-1} A_{q,j} ▷ 向量化求和
12: end for
13: k ← min(⌊ (L_max - w) / c ⌋, C) ▷ 将 k 限制在可用块数内
14: kept_chunks ← Top-K(chunk_scores, k) ▷ top-k 块的索引
15: 构建选择掩码:
16: mask ← zeros(T_k) ▷ 用于token选择的二元掩码
17: for each i in kept_chunks do
18:     start ← i · c
19:     end ← min((i + 1) · c, T_k)
20:     mask[start : end] ← 1 ▷ 将整个块标记为保留
21: end for
22: 确保保留最近的上下文:
23: mask[T_k - w : T_k] ← 1 ▷ 始终保留最新的token
24: 高效压缩:
25: indices ← nonzero(mask) ▷ 获取要保留的token的索引
26: K′ ← index_select(K, indices)
27: V′ ← index_select(V, indices)
28: 返回 K′, V′

B 附加实验

B.1 层级索引重用

B.1.1 效率分析

计算复杂性降低。层级索引重用方法显著降低了ChunkKV的计算复杂性。若不使用索引重用，ChunkKV将应用于所有$N_{layers}$层，总压缩时间为$N_{layers} \cdot T_{compress}$，其中$T_{compress}$是压缩一层的耗时。使用索引重用后，ChunkKV仅应用于$\frac{N_{layers}}{N_{reuse}}$层，总时间减少为$\frac{N_{layers}}{N_{reuse}} \cdot T_{compress} + (N_{layers} - \frac{N_{layers}}{N_{reuse}}) \cdot T_{select}$，其中$T_{select}$是选择索引的耗时，通常远小于$T_{compress}$。这带来了理论上的加速比：

$$\text{Speedup} = \frac{N_{\text{layers}} \cdot T_{\text{compress}}}{\frac{N_{\text{layers}}}{N_{\text{reuse}}} \cdot T_{\text{compress}} + (N_{\text{layers}} - \frac{N_{\text{layers}}}{N_{\text{reuse}}}) \cdot T_{\text{select}}}$$

假设$T_{select}$与$T_{compress}$相比可以忽略不计，该公式可简化为约$N_{reuse}$。在实践中，实际加速比可能因具体实现和硬件而异，但仍可节省大量时间，特别是对于层数较多的模型。

B.1.2 索引重用性能

GSM8K与LongBench性能对比。图5和表14展示了ChunkKV在GSM8K基准上不同索引重用层数下的性能。图6和表15展示了在LongBench基准上的性能。实验表明，与LongBench相比，数学问题对索引重用层数更为敏感。LLaMA3-8B-Instruct和Qwen2-7B-Instruct都表现出显著的性能下降，其中LLaMA3-8B-Instruct在两层索引重用后下降得比Qwen2-7B-Instruct更剧烈。这表明Qwen2-7B-Instruct模型可能对索引重用更具鲁棒性。

图5：不同索引重用层数下的GSM8K性能比较

表14：GSM8K上重用索引的性能比较

表15：LongBench上重用索引的性能比较

B.1.3 层级索引相似度

实验设置与结果。本节详细介绍了3.3节中描述的层级索引重用相似度实验。推理提示是从LongBench基准中随机选择的，H2O、SnapKV和ChunkKV保留的索引被保存在日志文件中。对于多头注意力，只保存第一个头的索引。PyramidKV由于在不同层保留的索引大小不同，不适用于此实验。然后我们计算了不同模型相邻层保留索引的Jaccard相似度。表16显示了不同模型相邻层保留索引的Jaccard相似度。

可视化分析。图7-9（LLaMA-3-8B-Instruct）、图10-12（Mistral-7B-Instruct）和图13-15（Qwen2-7B-Instruct）展示了H2O、SnapKV和ChunkKV在不同模型上保留的KV缓存索引的层级相似度热图。层级索引相似度热图的模式在不同模型中是一致的，这与我们在3.3节中的发现相符。

图6：在LongBench上不同索引重用层数的比较。

表16：不同模型相邻层保留的KV缓存索引相似度。

LLaMA-3-8B-Instruct with H2o
图7：H2O在LLaMA-3-8B-Instruct上保留的KV缓存索引的层级相似性热图

B.2 LongBench

英文子任务性能。表17显示了KV缓存压缩方法在LongBench英文子任务类别中的平均性能。ChunkKV在总体平均分和多文档问答（Multi-Document QA）类别中均取得了最佳性能，这支持了分块方法对于语义保留更为有效的观点。

70B模型性能。表18展示了使用LLaMA-3-70B-Instruct进行的额外实验，在LongBench数据集上比较了我们的ChunkKV方法与基线方法（KV缓存大小=256）的性能：

B.3 Needle-In-A-Haystack

性能可视化。图16和图17可视化了ChunkKV在NIAH基准测试中，使用LLaMA-3-8B-Instruct和Mistral-7B-Instruct模型，在8k和32k上下文长度下，KV缓存大小为128时的性能。随着上下文长度的增加，ChunkKV的性能持续表现更优。

LLaMA-3-8B-Instruct with SnapKV
图8：SnapKV在LLaMA-3-8B-Instruct上保留的KV缓存索引的层级相似性热图

图9：ChunkKV在LLaMA-3-8B-Instruct上保留的KV缓存索引的层级相似性热图

Mistral-7B-Instruct with H2o
图10：H2O在Mistral-7B-Instruct上保留的KV缓存索引的层级相似性热图

图11：SnapKV在Mistral-7B-Instruct上保留的KV缓存索引的层级相似性热图

Mistral-7B-Instruct with ChunkKV
图12：ChunkKV在Mistral-7B-Instruct上保留的KV缓存索引的层级相似性热图

图13：H2O在Qwen2-7B-Instruct上保留的KV缓存索引的层级相似性热图

Qwen2-7B-Instruct with SnapKV
图14：SnapKV在Qwen2-7B-Instruct上保留的KV缓存索引的层级相似性热图

图15：ChunkKV在Qwen2-7B-Instruct上保留的KV缓存索引的层级相似性热图

表17：KV缓存压缩方法在LongBench英文子任务上的综合性能比较。结果展示了不同模型和任务上的表现，突出了不同压缩技术的有效性。

表18：ChunkKV在LongBench英文子任务上使用70B模型的综合性能比较。

Token Limit

(a) ChunkKV, 准确率 73.8%
Token Limit
(b) PyramidKV, 准确率 65.1%

(c) SnapKV, 准确率 58.9%

(d) H2O, 准确率 47.9%
Token Limit

(e) StreamingLLM, 准确率 23.7%
图16：LLaMA-3-8B-Instruct在NIAH基准测试中的表现，KV缓存大小=128，上下文长度为8k

图17：NIAH基准测试中，Mistral-7B-Instruct在KV缓存大小=128，上下文长度为32k时的表现

B.4 块大小

实验结果。表19和表20展示了ChunkKV在不同压缩率和不同块大小下在LongBench和NIAH上的性能。我们进行了广泛的实验，涵盖了不同的压缩率和KV缓存大小，以证明ChunkKV的有效性，并表明块大小是鲁棒的。

表19：LLaMA-3-8B-Instruct在LongBench上不同块大小和压缩率下的性能

表20：LLaMA-3-8B-Instruct在NIAH上不同块大小和KV缓存大小下的性能

性能趋势分析。图18展示了ChunkKV在不同块大小下在LongBench和NIAH基准测试中的性能。三条彩色曲线代表三种不同块大小的LLM，彩色虚线是相应的FullKV性能。从图18可以看出，ChunkKV在LongBench上的性能不受块大小的显著影响，性能变化小于1%。三条曲线紧密对齐，表明在{10, 20}范围内的块大小表现更佳。

GSM8K性能分析。表21和图19显示，ChunkKV在不同块大小下在GSM8K上的表现呈现出与LongBench相似的曲线模式。GSM8K的CoT提示长度仅为1K tokens，因此最优块大小范围更小。

LongBench性能与块大小的关系
图18：在10%压缩率下，不同块大小在LongBench上的性能比较。

图19：不同块大小下的GSM8K性能比较

表21：不同块大小下的GSM8K性能比较

B.5 多语言

中文任务性能。表22展示了支持中文的模型Qwen2-7B-Instruct在LongBench中文子任务上的评估结果。其中，ChunkKV的性能优于其他压缩方法，甚至超过了全量KV缓存的性能。无论是在英文还是中文任务上的结果都表明，ChunkKV是维护KV缓存中关键信息的一种有前途的方法。

表22：Qwen2-7B-Instruct在LongBench中文子任务上的性能比较。

B.6 KV缓存量化

与量化方法的对比。为了全面评估ChunkKV的有效性，我们进行了实验，将ChunkKV与量化方法KIVI【【索引45，Kivi: A tuning-free asymmetric 2bit quantization for kv cache，2024，arXiv preprint】】进行了比较。虽然这两种方法都旨在优化LLM推理，但它们的基本原理完全不同：量化降低了KV矩阵的精度，而我们的逐出方法则减小了KV矩阵的大小。

实现差异与性能分析。从实现角度看，量化方法在预填充阶段需要完整的KV缓存来生成量化表示，然后在解码阶段使用这些表示。相比之下，ChunkKV在预填充之前就移除了token，使得整个推理过程都可以在压缩后的缓存上运行。这种差异导致了不同的效率特性，每种方法都有其独特的优势。

性能与效率对比。下面的表格展示了我们在LLaMA-3-8B-Instruct模型上对ChunkKV和KIVI在性能和效率指标上的详细分析。需要注意的是，由于KIVI依赖于较旧的Python版本，这里报告的效率结果可能与我们论文中表8的结果存在一些差异。表23表明，ChunkKV在30%的压缩率下（41.59）与KIVI的8位量化（41.43）在LongBench子任务上的性能相当。值得注意的是，ChunkKV在代码相关任务上表现尤为出色，而KIVI在单文档问答场景中保持了更好的性能。这些结果突显了逐出和量化方法的互补优势。ChunkKV基于语义块的方法在保持与最先进量化技术相当的性能的同时，提供了特别强的效率优势。

表23：ChunkKV与量化方法在LongBench英文子任务上的综合性能比较。

B.7 与正交及基于训练的方法的比较

实验设置与结果。为了更好地将ChunkKV置于更广泛的推理优化领域中，我们进行了额外的实验，将我们这种无需训练的逐出方法与最近的一种基于训练的压缩方法Palu【【索引46，Palu: Compressing kv-cache with low-rank projection，2024，arXiv preprint】】进行了比较。实验在LongBench基准上使用LLaMA-3-8B-Instruct模型进行。

性能分析。表24的结果突显了关键差异。Palu作为一种基于训练的方法，在多样化的LongBench基准上性能出现了显著下降，这对其训练数据而言是一个分布外（Out-Of-Distribution, OOD）的挑战。相比之下，我们这种无需训练的方法ChunkKV表现出了鲁棒的性能。

表24：在LongBench基准上（LLaMA-3-8B-Instruct），ChunkKV与Palu（基于训练）的性能比较。ChunkKV在这个多样化的基准上展示了优越的性能-效率权衡，并且与正交的量化技术相比仍具竞争力。

B.8 效率结果

效率指标。表25显示了不同KV缓存策略在不同输出长度下的效率结果，指标为首个令牌生成时间（Time to First Token, TTFT）和令牌处理时间（Token Processing Time, TPOT）。

长上下文效率。表26展示了在LLaMA-3-8B-Instruct上的延迟和吞吐量比较，包括高达16k令牌的长上下文场景。括号中的百分比表示相对于FullKV基线的提升。带有层级重用的ChunkKV（ChunkKV_reuse）始终提供最佳的效率增益，尤其是在长上下文设置中。

表25：不同KV缓存策略在不同输出长度下的效率结果

表26：LLaMA-3-8B-Instruct上的延迟和吞吐量比较，包括高达16k tokens的长上下文场景。括号中的百分比表示相对于FullKV基线的提升。带有层级重用的ChunkKV（ChunkKV_reuse）始终提供最佳的效率增益，尤其是在长上下文设置中。

C 理论理解

理论背景。在本节中，我们从上下文学习（ICL）的角度提供理论解释，以进一步理解ChunkKV为何优于token级别的KV缓存压缩。

预训练数据分布。给定一个概念集合Θ和一个概念θ ∈ Θ，我们定义预训练数据是从$p(o_1, \ldots, o_T) = \int_{\theta \in \Theta} p(o_1, \ldots, o_T |\theta)p(\theta)d\theta$中采样的【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】。每个token o从词汇表O中采样。为简化起见，我们记$o_{1:t} = o_1 \ldots o_t$。

语言建模。当前的LLMs【【索引1，Language models are few-shot learners，2020，Advances in neural information processing systems】】【【索引5，Llama 2: Open foundation and fine-tuned chat models，2023，ArXiv preprint】】【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】通常使用下一个词预测作为语言建模，即对所有t = 1, ..., T，给定前面的tokens $o_1 \ldots o_{t-1}$来预测下一个token $o_t$。形式上，语言建模可以写成分布$f(o_t|o_{1:t-1})$。它在一个从预训练分布$p(o_1, \ldots, o_{t+1})$采样的巨大语料库上进行预训练【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】。考虑到模型和数据集规模巨大，可以假设$f(o_1 \ldots o_{t+1})$已经与$p(o_1 \ldots o_{t+1})$对齐【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】。

提示分布。根据【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】，一个提示由输入token序列x后跟一个输出token y组成。那么，可以出现在整个提示$o_{1:T}$中任何位置的第i个训练样本被定义为$O_i$，它包含一个输入$x_i = O_i[1 : k-1]$（前k-1个tokens）和一个在末尾的输出$y_i = O_i[k]$，为简化起见，长度k是固定的。第i个训练样本是独立生成的，过程如下：1) 从提示开始分布$p_{prompt}$生成一个开始隐藏状态$h_{start}^i$；2) 给定$h_{start}^i$，从$p(O_i|h_{start}^i, \theta^\star)$生成样本序列$O_i = [x_i, y_i]$。测试输入$x_{test} = x_{n+1}$的采样方式类似。然后，提示由一系列训练样本($S_n$)后跟一个测试样本$x_{test}$组成：

$$[S_{n}, x_{\text{test}}] = [x_{1}, y_{1}, x_{2}, y_{2}, \dots, x_{n}, y_{n}, x_{\text{test}}] \sim p_{\text{prompt}}$$

上下文学习设置与假设。我们遵循【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】中的其他设置和假设。采用贪心解码【【索引47，Break the sequential dependency of llm inference using lookahead decoding，2024，arXiv preprint】】，从语言模型$f(o_t|o_{1:t-1})$中采样下一个token成为预测器，即$y = \arg \max_o f(o_t|o_{1:t-1})$。因此，对于$[S_n, x_{test}]$，上下文学习预测器可以写为$f_n(x_{test}) := \arg \max_y p(y|S_n, x_{test})$，它输出在给定提示分布条件下，在预训练分布上最可能的预测。其期望的0-1误差为$L_{0-1}(f_n) = E_{x_{test},y_{test}\sim p_{prompt}}[\mathbf{1}[f_n(x_{test}) \neq y_{test}]]$。我们定义第i个token的分布为$p_i^\theta(o) := p(O[i] = o|O[1:i-1], \theta)$，其中包含前面的tokens和概念θ，以及在提示分布下的类似分布$p_i^{prompt} := p_{prompt}(O[i] = o|O[1:i-1])$。根据【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】，存在一个可区分性条件，它形式化了在给定概念θ时上下文学习何时发生。可区分性条件依赖于前两个分布之间的KL散度以及由每个样本的提示分布和预训练分布之间的不匹配所产生的误差项$\epsilon_\theta$。令$p_i^\theta(o)$和$p_i^{prompt}$分别对应于概念θ和θ⋆。

条件1 (可区分性)。当且仅当对于所有$\theta \in \Omega, \theta \neq \theta^\star$，$\theta^\star$是可区分的：

$$\sum_{i=1}^{k} K L_{i}\left(\theta^{\star} \| \theta\right)>\epsilon_{\theta},$$

其中$KL_i(\theta^\star||\theta) := E_{O[1:i-1]\sim p_{prompt}} [KL(p_i^{prompt}||p_i^\theta)]$。

KV缓存压缩带来的噪声。自然地，由于KV缓存的稀疏化，在不同层中，$o_{1:t-1}$中的一些历史token失去了与下一个词预测$o_t$的注意力分数计算。我们可以将其视为添加到$o_{1:t-1}$上的噪声。因此，要从θ中区分出θ⋆需要更大的KL散度。遵循【【索引48，Can language models perform robust reasoning in chain-of-thought prompting with noisy rationales?，2024，The Thirty-eighth Annual Conference on Neural Information Processing Systems】】，我们提出了关于KV缓存稀疏性下可区分性的第二个条件。

条件2 (稀疏KV缓存下的可区分性)。在稀疏率为r的稀疏KV缓存引入噪声的情况下，由LLM近似的提示分布与预训练分布之间的不匹配被放大，导致对θ的可区分性要求发生变化，需要一个误差项$\xi_\theta(r)$。如果对于所有$\theta \in \Theta, \theta \neq \theta^*$，满足以下条件，则θ是可区分的：

$$\sum_{i=1}^{k} K L_{i}\left(\theta^{*} \| \theta\right)>\epsilon_{\theta}+\xi_{\theta}(r), \quad \text { where } \quad \xi_{\theta}(r) \propto r .$$

引理1 (噪声松弛界)。令B表示不满足条件1的θ集合。我们假设对于所有θ，$KL(p_{prompt}(y_{test}|x_{test}))||p(y_{test}|x_{test}, \theta)$是有界的，并且θ⋆最小化多分类逻辑风险，即：

$$L_{C E}(\theta)=-\mathbb{E}_{x_{test} \sim p_{prompt}}\left[p_{prompt}\left(y_{test} \mid x_{test}\right) \cdot \log p\left(y_{test} \mid x_{test}, \theta\right)\right] .$$

如果

$$\mathbb{E}_{x_{test} \sim p_{prompt}}[KL(p_{prompt}(y_{test}|x_{test})||p(y_{test}|x_{test}, \theta))] \leq (\epsilon_{\theta} + \xi_{\theta}(r)), \quad \forall \quad \theta \in \mathcal{B},$$

那么

$$\lim_{n \rightarrow \infty} L_{0-1}(f_n) \leq \inf_{f} L_{0-1}(f) + g^{-1}\left( \sup_{\theta \in \mathcal{B}}(\epsilon_{\theta}) \right),$$

其中$g(\nu) = \frac{1}{2} [(1 - \nu) \log(1 - \nu) + (1 + \nu) \log(1 + \nu)]$是多分类逻辑损失的校准函数【【索引49，How to compare different loss functions and their risks，2007，Constructive Approximation】】【【索引50，Multiclass classification calibration functions，2016，arXiv preprint】】，其中$\nu \in [0, 1]$。

泰勒展开假设。根据【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】【【索引51，The bernstein-von-mises theorem under misspecification，2012，Electronic Journal of Statistics】】，假设KL散度在概念θ上具有二阶泰勒展开。于是，我们有以下定理和证明。

定理1。令不满足条件1中方程3的θ集合为B。假设KL散度在θ⋆周围有二阶泰勒展开：

$$\forall j > 1, \quad KL_i(\theta^\star||\theta) = \frac{1}{2}(\theta - \theta^\star)^\top I_{j, \theta^\star}(\theta - \theta^\star) + O(||\theta - \theta^\star||^3)$$

其中$I_{j, \theta^\star}$是第j个token分布相对于θ⋆的费雪信息矩阵。令$\gamma_{\theta^\star} = \frac{\max_j \lambda_{max}(I_{j, \theta^\star})}{\min_j \lambda_{min}(I_{j, \theta^\star})}$，其中$\lambda_{max}, \lambda_{min}$返回最大和最小特征值。那么对于$k \geq 2$且当$n \to \infty$时，上下文学习预测器$f_n$的0-1风险有界为：

$$\lim_{n \rightarrow \infty} L_{0-1}(f_{n}) \leq \inf_{f} L_{0-1}(f) + g^{-1} \left( O \left( \frac{\gamma_{\theta^{*}} \sup_{\theta \in \mathbb{B}} (\epsilon_{\theta} + \xi_{\theta}(r))}{k-1} \right) \right)$$

证明。根据【【索引24，An end-to-end contrastive self-supervised learning framework for language understanding，2022，Transactions of the Association for Computational Linguistics】】，通过对θ进行泰勒展开，我们对B中的任何θ有：

$$\sum_{j=2}^{k} \mathrm{KL}_{i}\left(\theta^{\star} \| \theta\right) \geq \frac{1}{2} \sum_{j=2}^{k}\left(\theta-\theta^{\star}\right)^{\top} I_{j, \theta^{*}}\left(\theta-\theta^{\star}\right)+(k-1) O\left(\left\|\theta-\theta^{\star}\right\|^{3}\right)$$ $$\geq \frac{1}{2}(k-1) \lambda_{\min }\left(I_{j, \theta^*}\right)\left\|\theta-\theta^\star\right\|^2$$ $$\implies\|\theta-\theta^\star\|^2 \leq \frac{(\epsilon_\theta + \xi_\theta(r))}{\frac{1}{2}(k-1)(\min_j \lambda_{\min}(I_{j,\theta^*}))}.$$

我们可以用上述项来界定最后一个KL项（第k个token）：

$$\mathrm{KL}_{k}(\theta^{\star} \| \theta)=\frac{1}{2}(\theta-\theta^{\star})^{\top} I_{k, \theta^{*}}(\theta-\theta^{\star})+O(\|\theta-\theta^{\star}\|^{3})$$ $$\leq \frac{1}{2}(\max_{j} \lambda_{\max}(I_{j, \theta^{\star}}))\|\theta - \theta^{\star}\|^{2} + O(\|\theta - \theta^{\star}\|^{2})$$ $$\le \frac{(\epsilon_\theta + \xi_\theta(r))(\max_j \lambda_{\max}(I_{j,\theta^*}) + O(1))}{(k-1) \min_j \lambda_{\min}(I_{j,\theta^*})}.$$

重新整理上述方程，并且有$KL_k(\theta^\star||\theta) = E_{x_{test}\sim p_{prompt}} [KL(p_{prompt}(y_{test}|x_{test})\|p(y_{test}|x_{test}, \theta))]$，可以得到：

$$\mathbb{E}_{x_{\text{test}} \sim p_{\text{prompt}}}[KL(p_{\text{prompt}}(y_{\text{test}}|x_{\text{test}})\|p(y_{\text{test}}|x_{\text{test}}, \theta))] \le \frac{(\epsilon_\theta + \xi_\theta(r))(\max_j \lambda_{\max}(I_{j,\theta^*}) + O(1))}{(k-1) \min_j \lambda_{\min}(I_{j,\theta^*})}$$

将方程16与方程6结合到引理1中，即可完成证明。

KV缓存稀疏化。回到条件2中的方程4，$\xi_\theta(r)$随着稀疏率r的增大而增大。压缩率r越高（意味着丢弃的KV缓存越多），噪声$\xi_\theta(r)$就越大。这导致了引理1中方程5中$\lim_{n\to\infty} L_{0-1}(f_n)$的界限更高。接下来，我们讨论KV缓存压缩如何影响方程4。

token级重要性度量。token级的KV缓存方法通常计算不同token的重要性。然后，具有较高重要性索引的KV缓存将被保留。这些索引通常被选为注意力分数。考虑到图1中的情况，其中第i个训练样本序列($O_i = [x_i, y_i]$)中的每个token都可能被压缩，并且token是具体稀疏化的，不依赖于其他token。然而，在第i个训练样本的生成过程中，$O_i = [x_i, y_i]$是从$p(O_i|h_{start}^i, \theta^\star)$中采样的，并且第j个token的分布为$p_j^\theta(o) := p(O[j] = o|O[1:j-1], \theta)$，以及类似的分布$p_j^{prompt} := p_{prompt}(O[j] = o|O[1:j-1])$。KL散度定义为$KL_j(\theta^\star||\theta) := E_{O[1:j-1]\sim p_{prompt}} [KL(p_j^{prompt}||p_j^\theta)]$，这意味着在一个训练样本$O_i = [x_i, y_i] = O_i[1:k]$中，每个token $O_i[j]$与$O_i[1:j-1]$有很强的依赖关系，前面任何一个j-th token上的噪声都会影响后续token的可区分性，即需要更大的$\{KL_u(\theta^\star||\theta)\}_{u>j}$。另一方面，token级的稀疏性对每个样本$O_i$（如图1中的情况）的可区分性要求统一放大，这统一地放宽了方程9中$L_{0-1}(f_n)$的界限。

块级重要性度量。与token级重要性度量不同，ChunkKV将连续窗口中的token视为一个基本单元，应作为一个整体保留或丢弃。保留的窗口可以看作是保存了完整的$O_i = [x_i, y_i]$而没有噪声。因此，ChunkKV减少了保留的$O_i$的噪声$\xi_\theta(r)$，从而降低了$L_{0-1}(f_n)$的界限。

直观解释。更直观地说，ChunkKV专注于减少一些完整训练样本上的噪声，但其他整体重要性较低的样本将被丢弃。然后，模型从那些干净且更相关的训练样本$O_i$中识别$x_{test}$，并忽略那些重要性较低的$O_i$。

未来工作。需要注意的是，我们在这里没有提供关于KV缓存稀疏性如何增强可区分性要求，以及$O_i = [x_i, y_i]$上不同的$KL_j(\theta^\star||\theta)$如何影响$L_{0-1}(f_n)$界限的严格证明。我们将其作为未来工作，以分析KV缓存稀疏性如何影响上下文学习。