$\mathrm { { \cal D } _ { 2 } O }$ : Dynamic Discriminative Operations for Efficient Long-Context Inference of Large Language Models

发表时间: 2025-04 · arXiv:2406.13035 (ICLR 2025)

原文: https://arxiv.org/abs/2406.13035

Zhongwei Wan, Xinjian Wu, Yu Zhang, Yi Xin, Chaofan Tao, Zhihong Zhu, Xin Wang, Siqi Luo, Jing Xiong, Longyue Wang, Mi Zhang / The Ohio State University, University of Chinese Academy of Sciences, Tongji University, Nanjing University, The University of Hong Kong, Peking University, Alibaba International Digital Commerce

速读

一句话结论 本文提出了一种名为 $\mathrm { D _ { 2 } O }$ 的免微调双层动态键值缓存压缩方法，通过在层级动态分配缓存并在词元级融合被剔除的上下文，在将大语言模型推理吞吐量提升 3 倍以上的同时，依然保持了高质量的长文本生成能力。

要解决什么问题 大语言模型在处理长上下文时，键值缓存（KV Cache）的显存占用会随序列长度线性爆炸，例如 30B 参数模型在 128 批次大小下处理 1024 个词元就需要高达 180 GB 的显存。为了缓解显存卡点，现有的做法通常是基于注意力分数直接剔除不重要的 KV 对，但这会带来严重的上下文信息丢失。这种信息丢失源于两个层面的机制缺陷。首先是层级的无差别对待：浅层网络主要通过全局注意力捕捉句法结构，注意力权重分布密集；而深层网络主要关注局部语义，注意力分布呈现阶梯状的稀疏模式。现有方法在所有层使用相同的剔除比例，会破坏浅层密集的关键信息。其次是词元级的不可逆剔除：由于很难准确预测当前被判定为不重要的词元在未来的文本生成中是否会被用到，直接丢弃它们会导致不可逆的信息损失，进而引发模型幻觉、上下文不一致以及长文本逻辑断裂等问题。

怎么做的 $\mathrm { D _ { 2 } O }$ 的核心思路是将 KV Cache 的压缩过程分为层级动态分配和词元级动态融合两步，以此绕开传统方法一刀切导致的上下文丢失卡点。首先在层级，方法利用注意力权重的方差来衡量每一层的注意力密度，方差越小代表注意力越密集。基于此，设计了一种逆方差 Softmax 的动态分配策略，为注意力密集的浅层分配更大的缓存预算，为稀疏的深层分配较小的预算。第 $l$ 层的缓存大小 $S_l$ 定义为：

$$S _ { l } = \alpha _ { l } \cdot S , \ \mathrm { w h e r e } \ \alpha _ { l } = { \frac { \exp ( - F _ { v } ^ { l } ) } { \sum _ { l = 1 } ^ { L } \exp ( - F _ { v } ^ { l } ) } } \cdot L \cdot \rho$$

其中 $F _ { v } ^ { l }$ 是该层累积注意力分数的方差，$\rho$ 是整体压缩率。其次在词元级，为了弥补被剔除词元带来的信息损失，$\mathrm { D _ { 2 } O }$ 引入了基于指数移动平均（EMA）阈值的动态补偿机制。在每一层确定了要保留的注意力沉淀词元、最近窗口词元以及累积注意力分数最高的 $N$ 个词元后，剩下的词元本该被直接丢弃。但 $\mathrm { D _ { 2 } O }$ 会计算这些待丢弃词元的 Key 与保留词元的 Key 之间的余弦相似度，并维护一个 EMA 阈值 $\tau_t$ 来判断当前待丢弃词元是否还有价值：

$$\tau _ { t } = \beta \mathbf { M a x } ( \mathbf { U } _ { t } [ : ] ) + ( 1 - \beta ) \tau _ { t - 1 }$$
如果待丢弃词元的最大相似度高于该阈值，说明它与某个保留词元高度相关，此时不直接丢弃，而是采用类似图注意力网络的加权机制，将该待丢弃词元的 Key 和 Value 融合进与其最相似的保留词元中；如果低于阈值，则永久丢弃。这种设计既维持了严格的缓存大小限制，又以极小的计算代价把潜在有用的历史信息压缩到了保留的词元中。

效果如何 实验在 NVIDIA A100 80GB 硬件上进行，评估了 Llama-2、Llama-3、Falcon、Mistral 和 Code-Llama 等 7B 到 34B 参数规模的模型。对比基线包括代表保留最近窗口和初始词元路线的 StreamingLLM、代表基于累积注意力分数剔除路线的 $\mathrm { H _ { 2 } O }$、代表基于平均注意力分数剔除路线的 RoCo、代表层级剔除路线的 PyramidKV 以及代表 Value 词元融合路线的 CaM。在 LongBench 长文本基准测试中，当 KV Cache 预算被压缩至原长度的 20% 时，$\mathrm { D _ { 2 } O }$ 在各项任务上均击败了所有基线方法，且相比全缓存模型的精度下降极小。在 Needle In A Haystack 长上下文事实检索任务中，当上下文长度为 50K 且缓存预算限制为 8192 时，$\mathrm { D _ { 2 } O }$ 达到了 94.48% 的准确率，远超 StreamingLLM 的 62.84% 和 $\mathrm { H _ { 2 } O }$ 的 82.32%。在吞吐量方面，由于大幅节省了显存，$\mathrm { D _ { 2 } O }$ 允许使用更大的批次大小，在 Llama-3-8B 模型上处理 2048 个输入词元并生成 8192 个词元的设置下，实现了比原始模型高 3.04 倍的吞吐量，且推理延迟与纯剔除路线的 $\mathrm { H _ { 2 } O }$ 几乎持平。作者也承认该方法的局限性在于目前仅在原始模型架构上验证，尚未与量化、剪枝或高效注意力架构等其他压缩技术进行叠加测试以探索极限压缩能力。

主要贡献

大型语言模型（LLMs）在处理长上下文任务（如对话系统、文档摘要、问答和代码补全）时，其生成推理过程严重受制于键值（KV）缓存不断增长的内存需求。传统的KV缓存驱逐策略通常基于注意力分数丢弃被认为不太关键的KV对，但这种不可逆的丢弃往往会降低生成质量，导致上下文丢失或幻觉问题。

为了解决这一问题，本文提出了一种名为动态判别操作（Dynamic Discriminative Operations, $\mathrm { D _ { 2 } O }$）的无须微调的KV缓存压缩方法。该方法能够在保留关键上下文的同时，在两个层级上动态且具判别性地优化KV缓存大小：
1. 层级（Layer level）：$\mathrm { D _ { 2 } O }$ 利用浅层和深层之间注意力权重密度的变化，通过一种新颖的动态分配策略，动态决定哪些层应避免过度驱逐，从而最小化信息损失。
2. Token级（Token level）：$\mathrm { D _ { 2 } O }$ 引入了一种补偿机制，该机制维护一个相似度阈值，用于重新判别当前被丢弃Token的重要性，从而决定是否应将它们召回并与当前保留的相似Token进行合并。

在多种基准测试和LLM架构上的实验表明，$\mathrm { D _ { 2 } O }$ 不仅实现了显著的内存节省，将推理吞吐量提高了3倍以上，而且还保持了高质量的长文本生成能力。

图1：LLaMA-2-7B在GSM8K数据集上浅层（第0、1层）和深层（第16、30、31层）的注意力图密度比较。我们使用了每层所有注意头的平均值。
图1：LLaMA-2-7B在GSM8K数据集上浅层（第0、1层）和深层（第16、30、31层）的注意力图密度比较。我们使用了每层所有注意头的平均值。
图2：在Token级别上，现有方法（a）与 $\mathrm { D _ { 2 } O }$（b）的对比。
图2：在Token级别上，现有方法（a）与 $\mathrm { D _ { 2 } O }$（b）的对比。

背景知识与设计原则

Prompt编码阶段。在Prompt编码阶段，模型利用提示序列为LLM中的每个Transformer层生成KV缓存。给定输入提示张量 $\mathbf { X } \in \mathbb { R } ^ { L _ { \mathrm { p r o m p t } } \times D }$（其中 $L _ { \mathrm { p r o m p t } }$ 为提示长度，$D$ 为隐藏层维度），键（Key）和值（Value）张量的计算方式为 $\mathbf { K } = \mathbf { X } \mathbf { W } _ { K } , \mathbf { V } = \mathbf { X } \mathbf { W } _ { V }$，其中 $\mathbf { W } _ { K } , \mathbf { W } _ { V } \in \mathbb { R } ^ { D \times D }$ 分别代表键和值层的权重。计算完成后，$\mathbf { K }$ 和 $\mathbf { V }$ 被存储在KV缓存中以加速后续的Token生成过程。

Token生成阶段。在Token生成阶段，模型利用并更新KV缓存来逐个自回归地生成Token。在每个时间步 $i$，模型仅计算新Token $\mathbf { x } _ { i }$ 的键和值，而 $\mathbf { x } _ { < i }$ 的键和值则直接从缓存中检索。随后缓存被更新，新生成Token的输出计算如下：$\mathbf { K } = [ \mathbf { K } , \mathbf { x } _ { i } \mathbf { W } _ { K } ] , \mathbf { V } = [ \mathbf { V } , \mathbf { x } _ { i } \mathbf { W } _ { V } ]$，且 $\mathbf { x } _ { i , o u t } = \mathrm { S o f t m a x } \left( \mathbf { q } _ { i } \mathbf { K } ^ { \top } / \sqrt { D } \right) \mathbf { V } , \mathbf { q } _ { i } = \mathbf { x } _ { i } \mathbf { W } _ { Q }$，其中 $\mathbf { W } _ { Q } \in \mathbb { R } ^ { D \times D }$ 是查询层的权重矩阵。随着每个新Token的生成，KV缓存的线性增长会显著增加内存消耗和延迟，这凸显了压缩KV缓存的必要性。

图3：$\mathrm { D _ { 2 } O }$ 概览。在层级的层缓存大小分配中，$\mathrm { D _ { 2 } O }$ 通过在每层引入动态缓存来解决高低层之间注意力密度不一致的问题。缓存大小由注意力的方差指标和动态分配策略决定。在Token级别，$\mathrm { D _ { 2 } O }$ 通过结合Token驱逐方案和动态Token合并技术（包含补偿机制）来解决长上下文信息丢失的问题。
图3：$\mathrm { D _ { 2 } O }$ 概览。在层级的层缓存大小分配中，$\mathrm { D _ { 2 } O }$ 通过在每层引入动态缓存来解决高低层之间注意力密度不一致的问题。缓存大小由注意力的方差指标和动态分配策略决定。在Token级别，$\mathrm { D _ { 2 } O }$ 通过结合Token驱逐方案和动态Token合并技术（包含补偿机制）来解决长上下文信息丢失的问题。

方法细节

动态层级判别操作。现有方法如 $\mathrm { H _ { 2 } O }$ 【索引编号：Zhang 等人，H2O: Heavy-hitter oracle for efficient generative inference of large language models+2024+NeurIPS】和 StreamingLLM 【索引编号：Xiao 等人，Efficient streaming language models with attention sinks+2023+arXiv】在所有层采用统一的缓存大小，这可能会损害模型性能。为了解决这个问题，本文提出使用特定指标 $F _ { v } ^ { l }$ 来评估每个层 $l$ 的注意力密度：$F _ { v } ^ { l } = \mathrm { V a r } \left( \sum _ { i = 0 } ^ { L _ { \mathrm { p r o m p t } } } \mathbf { A } _ { p } ^ { l } [ i , : ] \right) , \mathbf { A } _ { p } ^ { l } = \mathrm { S o f t m a x } \left( \mathbf { Q } _ { p } ^ { l } {\mathbf { K } _ { p } ^ { l }} ^ { \top } / \sqrt { D } \right)$，其中 $\mathbf { A } _ { p } ^ { l }$ 表示每层Prompt编码的注意力分数。模型对 $\mathbf { A } _ { p } ^ { l }$ 中每列的元素求和以建立累积注意力序列的初始状态。随后，每层的注意力密度由该序列的方差量化，因为更密集的注意力权重对应更小的方差（如图4所示）。作者在GSM8K数据集上的所有模型中观察到一个一致的现象：浅层（如0、1、2）和中间层（如13、14）的注意力分数方差较低，表明注意力权重密集；而在深层中，方差增加，注意力权重呈现出更稀疏的模式。

图4：不同模型在不同层上的注意力分数方差。
图4：不同模型在不同层上的注意力分数方差。

动态分配策略。利用上述一致性现象，本文提出了一种基于逆方差Softmax的新型动态分配策略，用于调整每层的KV缓存大小。具体而言，具有较高方差 $F _ { v } ^ { l }$ 的层被分配较小的缓存大小，而具有较低方差的浅层则获得较大的缓存分配。对于给定的压缩率 $\rho$，每层的缓存大小 $S _ { l }$ 计算如下：$S _ { l } = \alpha _ { l } \cdot S$，其中 $\alpha _ { l } = \frac { \exp ( - F _ { v } ^ { l } ) } { \sum _ { l = 1 } ^ { L } \exp ( - F _ { v } ^ { l } ) } \cdot L \cdot \rho$。这里 $\rho$ 表示缓存大小的压缩比，$S$ 是原始缓存大小（在Prompt编码阶段等于 $L _ { \mathrm { p r o m p t } }$），$L$ 表示模型层数。该方法采用类似Softmax的函数来动态分配缓存比例 $\alpha _ { l }$。

动态Token级判别操作。在执行了层级判别操作之后，为了补偿长上下文信息的丢失，本文在Token级判别操作中引入了两个关键策略：Token驱逐和动态Token合并。虽然 $\mathrm { D _ { 2 } O }$ 兼容任何Token驱逐技术，但本文提出了一种基于累积注意力的驱逐策略来动态修剪生成任务中的KV缓存。对于Token合并，本文引入了一种新策略，该策略利用基于指数移动平均（EMA）的相似度阈值，通过加权合并动态决定是否将丢弃的Token合并回保留的KV缓存中。

Token驱逐。Token驱逐的核心概念是利用累积注意力分数动态更新KV缓存。该过程系统地排除最不重要的KV对，以维持每层压缩后的缓存大小 $S _ { l }$。与传统基于累积的方法不同，本文的策略通过保留输入前 $T$ 个Token中的关键注意力汇聚（attention sink）Token，并将它们与大小为 $M$ 的近期窗口整合来提高性能。注意力分数的计算公式如下：如果Token $i <= L _ { \mathrm { p r o m p t } }$，则 $\mathrm { A t t n S c o r e } = \sum _ { i = 0 } ^ { L _ { \mathrm { p r o m p t } } } \mathbf { A } _ { p } [ i , : ]$；否则在Token生成阶段，$\mathrm { A t t n S c o r e } = \mathrm { S o f t m a x } \left( \mathbf { q } _ { i } \mathbf { K } ^ { \top } / \sqrt { D } \right) + \sum _ { i = 1 } ^ { L _ { \mathrm { p r o m p t } } } \mathbf { A } _ { p } [ i , : ]$。在获得当前的累积注意力分数后，保留大小为 $M$ 的最新窗口并包含 $T$ 个注意力汇聚Token。接着，从剩余的KV缓存中选择得分最高的 $N$ 个Token来完成驱逐过程。过程定义为：$\mathbf { K } _ { c } = [ \mathbf { K } [ : T , : ] , \mathbf { K } [ I , : ] , \mathbf { K } [ - M : , : ] ]$，$\mathbf { V } _ { c } = [ \mathbf { V } [ : T , : ] , \mathbf { V } [ I , : ] , \mathbf { V } [ - M : , : ] ]$，其中 $I = \mathrm { T o p } _ { N } \left( \mathrm { A t t n S c o r e } [ T : - M ] , N \right)$。此处 $\mathrm { T o p } _ { N } \left( \cdot , N \right)$ 选出在AttnScore中最重要的 $N$ 个Token的索引 $I$，$( \mathbf { K } _ { c } , \mathbf { V } _ { c } )$ 表示驱逐后保留的KV缓存，当前缓存大小为 $S = T + N + M$。

动态Token合并。直接丢弃被驱逐的Token（即 $\mathbf { K } _ { e } = \mathbf { K } - \mathbf { K } _ { c }$）可能会破坏长上下文的完整性。为了减轻信息损失，本文提出了一种动态Token合并方法，以极低的计算成本检索仍然包含潜在价值的Token，并将这些选定的Token与相似的保留Token进行整合。考虑到KV对的对齐特性，模型仅在Key Token上计算相似度矩阵，并在Value Token上共享相似度度量和加权合并权重。该方法分为三个关键步骤：
* 最近邻匹配：利用多对一最近邻匹配算法【索引编号：Dang 等人，Nearest neighbor matching for deep clustering+2021+CVPR】计算 $\mathbf { K } _ { e }$ 和 $\mathbf { K } _ { c }$ 之间的相似度矩阵 $\mathbf { U }$。然后从 $\mathbf { K } _ { c }$ 中识别出最相似的Token作为合并候选。设 $I ^ { e }$ 和 $I ^ { c }$ 为索引，$L ^ { e }$ 和 $L ^ { c }$ 分别代表 $\mathbf { K } _ { e }$ 和 $\mathbf { K } _ { c }$ 中的Token长度。$\mathbf { U }$ 中的每个元素 $u _ { i , j }$ 表示用于匹配的Token之间的交互，其中 $i \in I ^ { e }$ 且 $j \in I ^ { c }$。接着为每个被驱逐的Token $\mathbf { k } _ { i }$ 确定 $\mathbf { K } _ { c }$ 中最接近的Token $\mathbf { k } _ { * } ^ { \mathrm { n e a r e s t } }$。公式为：$\mathbf { k } _ { * } ^ { \mathrm { n e a r e s t } } = \mathop { \mathrm { A r g m a x } } _ { j \in I ^ { c } } \left( u _ { i , j } \right)$，其中 $u _ { i , j } = \frac { \mathbf { k } _ { i } ^ { \top } \mathbf { k } _ { j } } { \| \mathbf { k } _ { i } \| \ \| \mathbf { k } _ { j } \| }$。这里采用余弦相似度。由于相似度矩阵 $\mathbf { U }$ 在提示编码期间直接从输入提示中派生，且在Token生成期间 $\mathbf { U } \in \mathbb { R } ^ { L ^ { c } }$，因此它不引入额外参数并确保了计算效率。
* EMA阈值：在计算相似度并识别出候选Token $\mathbf { K } _ { * } ^ { \mathrm { n e a r e s t } }$ 后，直接对Token对应用平均加权融合可能会导致特征分散。受时间序列任务中使用的指数移动平均（EMA）【索引编号：Hunter，The exponentially weighted moving average+1986+Journal of Quality Technology】启发，本文提出了用于Token级操作的EMA阈值。该方法强调当前被驱逐Token与保留Token之间近期相似度的重要性，同时平滑先前被驱逐Token与保留Token之间的历史相似度信息。EMA阈值公式为：如果 $t = 0$（提示编码阶段），$\tau _ { t } = \frac { 1 } { L ^ { e } } \sum _ { i = 0 } ^ { L ^ { e } } \mathbf { M a x } ( \mathbf { U } _ { t } [ i , : ] )$；否则（Token生成阶段），$\tau _ { t } = \beta \mathbf { M a x } ( \mathbf { U } _ { t } [ : ] ) + ( 1 - \beta ) \tau _ { t - 1 }$。初始阈值 $\tau$ 设置为被驱逐Token与保留集之间最高相似度值的平均值。平滑常数 $\beta$ 调节当前相似度矩阵与前一相似度阈值之间的平衡。如果给定被驱逐Token的最大相似度低于 $\tau _ { t }$，则将其永久丢弃；否则，应用加权合并策略。
* 加权合并：对于一个保留的Token，具有较高相似度的被驱逐Token应被分配更大的权重。受图注意力网络（GAT）【索引编号：Velickovic 等人，Graph attention networks+2017+arXiv】启发，本文采用基于相似度的加权策略进行加权合并。加权合并公式定义为：$\mathbf { k } _ { c j } = \mathbf { w } _ { c j } \mathbf { k } _ { c j } + \sum _ { \mathbf { k } _ { e i } \in \mathbf { K } _ { e } } \mathbf { w } _ { e i } \mathbf { k } _ { e i }$，$\mathbf { v } _ { c j } = \mathbf { w } _ { c j } \mathbf { v } _ { c j } + \sum _ { \mathbf { v } _ { e i } \in \mathbf { V } _ { e } } \mathbf { w } _ { e i } \mathbf { v } _ { e i }$。其中权重的计算方式为：$\mathbf { w } _ { c j } = \frac { e } { \sum _ { \mathbf { k } _ { e i } \in \mathbf { K } _ { e } } \exp ( \mathbf { u } _ { i j } ) \mathbf { m } _ { i j } + e }$，$\mathbf { w } _ { e i } = \frac { \sum _ { \mathbf { k } _ { e i } \in \mathbf { K } _ { e } } \exp ( \mathbf { u } _ { i j } ) \mathbf { m } _ { i j } } { \sum _ { \mathbf { k } _ { e i } \in \mathbf { K } _ { e } } \exp ( \mathbf { u } _ { i j } ) \mathbf { m } _ { i j } + e }$。这里 $\mathbf { m } _ { i , j }$ 表示 $\mathbf { U }$ 的掩码矩阵。如果 $\mathbf { x } _ { j } \in \mathbf { K } _ { c }$ 是最相似的Token，则 $\mathbf { m } _ { i , j } = 1$；否则为0。每个保留的Token $\mathbf { k } _ { c j }$ 保持最高的融合权重（自身相似度为1），因此保留的Token本身不会被修改，而被驱逐的Token被整合到它们最相似的对应物中。

实验环境

  • 数据集：评估分为标准上下文和长上下文两类。标准上下文任务包括 LM-Eval 中的常识和数学推理（CoQA、TruthfulQA、GSM8K）。长上下文任务使用 LongBench（包括单文档问答、多文档问答、摘要、合成任务和代码补全）、用于长上下文事实检索的 Needle-In-A-Haystack 测试，以及用于长序列语言建模评估的 PG-19。
  • 模型架构：评估使用了三个不同LLM家族的四种模型：Llama-2 (7B, 13B)、Llama-3 (8B)、Falcon (7B) 和 Mistral (7B)。其中 Llama-2 和 Mistral 使用多头注意力，Falcon 使用多查询注意力，Llama-3 使用分组查询注意力。
  • 硬件配置：所有实验均在 NVIDIA A100 80GB GPU 上进行。
  • 软件配置：使用 Hugging Face Transformers 代码库实现 $\mathrm { D _ { 2 } O }$。超参数设置中，EMA阈值平滑常数 $\beta$ 默认设为 0.7，重要Token大小 $N$ 与近期窗口大小 $M$ 的比例典型设置为 $N:M = 3:1$，整体KV缓存压缩比 $\rho$ 设置为 0.2、0.4 或 0.8。

实验结果

不同KV缓存压缩比下的性能对比
在 GSM8K、CoQA 和 TruthfulQA 数据集上，将配备完整KV缓存的模型与使用 $\mathrm { D _ { 2 } O }$ 压缩技术的模型（涵盖 Llama-1-7B、Llama-2-7B、Llama-2-13B 和 Llama-3-8B）进行了对比。结果（图5）显示，$\mathrm { D _ { 2 } O }$ 在所有配置下均持续优于其他KV压缩方法（如 StreamingLLM、$\mathrm { H _ { 2 } O }$）。特别是在预算减少的情况下，$\mathrm { D _ { 2 } O }$ 显著提升了性能，防止了推理能力的下降。有趣的是，在 TruthfulQA 数据集上，$\mathrm { D _ { 2 } O }$ 在大多数预算比例下甚至超越了完整KV缓存模型，证明其统一驱逐和动态合并策略能够修剪输入文本中的无关Token，保留基本上下文从而提高推理准确性。

图5：$\mathrm { D _ { 2 } O }$ 和其他方法在 Llama 骨干网络上针对 CoQA、GSM8K 和 TruthfulQA 推理数据集的性能。
图5：$\mathrm { D _ { 2 } O }$ 和其他方法在 Llama 骨干网络上针对 CoQA、GSM8K 和 TruthfulQA 推理数据集的性能。

长上下文任务的准确性比较
* LongBench 结果：在 Falcon-7B、Mistral-7B、Llama-2-7B/13B 和 Llama-3-8B 上，默认KV缓存预算比设为 $\rho = 0.2$。表1表明，与完整模型相比，$\mathrm { D _ { 2 } O }$ 在对准确性影响极小的情况下有效管理了KV缓存压缩，并成功捕获了长文本中的关键信息。$\mathrm { D _ { 2 } O }$ 显著优于 StreamingLLM、$\mathrm { H _ { 2 } O }$、RoCo 和 CaM 等其他基于驱逐的方法。此外，在 Code-Llama-34B 模型上的代码任务评估（表2）表明，$\mathrm { D _ { 2 } O }$ 同样优于其他基线，并紧跟完整KV缓存模型的性能，验证了其在更大模型规模上的可扩展性。
* 长上下文事实检索任务：使用 Llama-2-7B-80k 模型在 Needle In A Haystack 任务上进行验证。在 4096 和 8192 的KV缓存预算下，测试最大长度为 50k 和 100k 的上下文。表3显示，$\mathrm { D _ { 2 } O }$ 不仅优于其他基于驱逐的方法，而且与完整模型相比性能准确度下降最小。
* 长序列建模困惑度：在 PG-19 数据集上，将KV缓存容量设定为2048。图6显示了累积平均负对数似然（NLL）随上下文长度的变化。$\mathrm { D _ { 2 } O }$ 实现了比其他基线更低的困惑度，证明其即使在有限的KV缓存下也能有效利用长距离依赖关系。

图6：长序列建模PPL。
图6：长序列建模PPL。

补充细节

动态分配策略分析
为了评估层级动态分配策略的影响，实验比较了缓存分配因子 $\alpha _ { l }$ 的各种设计（表4）。结果表明，“逆方差Softmax（本文方法）”优于“方差倒数”、“指数衰减分配”和“均匀分配”，证明了基于方差有效分配缓存的优越性。

吞吐量分析
在 Llama-3-8B 架构上（无CPU卸载）测试了吞吐量（表5）。结果表明，$\mathrm { D _ { 2 } O }$ 减少了内存使用，允许更大的批处理大小。随着文本长度增加，$\mathrm { D _ { 2 } O }$ 相比完整模型的吞吐量优势也随之增长，在 $2048+8192$ 设置下达到了 3.04 倍的吞吐量提升，同时保持了与 $\mathrm { H _ { 2 } O }$ 相当的效率，但由于减少了信息损失，推理准确度更高。

消融分析
* 各组件的作用：表6显示，移除 $\mathrm { D _ { 2 } O }$ 的层级操作或EMA阈值都会导致性能下降，证明每个组件都能有效缓解信息损失。
* Token相似度指标：表7表明，基于注意力键（Key）的指标性能显著高于值（Value）或独立的键值指标。
* 合并策略：表8表明，加权合并策略的性能优于平均合并。
* 重要Token与近期Token大小的平衡：表9表明，$N:M = 3:1$（强调历史上下文中的重要Token）对性能的提升最大。

附录细节

超参数确定
在 LongBench (TREC) 和 LM-Eval (COQA, GSM8K, TruthfulQA) 上进行了超参数搜索。EMA 阈值参数 $\beta$ 在 $0.5 \sim 0.9$ 之间搜索，发现 $\beta = 0.7$ 产生最佳结果。$N:M$ 比例在 $3:1$ 到 $1:3$ 之间搜索，发现 $3:1$ 性能最优。

多轮对话生成样本
使用 MT-bench 数据集（2048 Token 预算）测试了多轮对话能力。如图7所示，在对话后半段，$\mathrm { H _ { 2 } O }$ 和 StreamingLLM 开始产生不相关的内容，失去了对话连贯性；而完整模型遇到内存溢出（OOM）错误。$\mathrm { D _ { 2 } O }$ 则有效地保留了上下文并生成了准确的响应。

图7：$\mathrm { D _ { 2 } O }$、$\mathrm { H _ { 2 } O }$、StreamingLLMs 和 Full Cache 生成文本的可视化比较。
图7：$\mathrm { D _ { 2 } O }$、$\mathrm { H _ { 2 } O }$、StreamingLLMs 和 Full Cache 生成文本的可视化比较。

计算成本分析
表12和表13分析了推理时间开销。Prompt编码过程（包含Token驱逐和动态合并）仅占总推理时间的一小部分（在长生成中降至 $0.224\%$）。与 $\mathrm { H _ { 2 } O }$ 相比，$\mathrm { D _ { 2 } O }$ 的推理时间没有显著差异，证明了相似度计算和合并策略是轻量级的。

长上下文事实检索任务可视化
图8展示了 Needle-in-a-Haystack 测试的可视化比较。在 100k 长度下，基于驱逐的 $\mathrm { H _ { 2 } O }$ 丢失了大量上下文信息，而 $\mathrm { D _ { 2 } O }$ 凭借动态合并策略有效保留了被驱逐Token的信息。

图8：多种方法在长上下文事实检索任务上的可视化比较。图(a)(b)为完整模型，(c)(d)为 $\mathrm { H _ { 2 } O }$ 方法，(e)(f)为本文的 $\mathrm { D _ { 2 } O }$ 方法。
图8(b)
图8(c)
图8(d)
图8(e)
图8(f)

注意力权重可视化
图9可视化了 Llama-1-7B 和 Llama-3-8B 等模型在不同层上的注意力权重。一致的模式表明，模型的较低层比较高层呈现出更高的注意力密度，这有力地证实了层级判别操作（为不同密度的层采用不同驱逐比例）的合理性。

图9：不同模型和数据集上的注意力权重可视化。
图9：不同模型和数据集上的注意力权重可视化。

动态分配策略的理论分析
附录利用信息论证明了动态KV缓存分配策略优于统一压缩率策略。通过构建拉格朗日函数求解在总缓存大小约束下最大化总信息保留量 $Q$ 的最优分配。理论证明表明，当所有注意力头的边际信息增益平衡时，总信息保留量最大。动态分配策略通过与注意力方差成反比地分配缓存，确保了方差较小（熵较高）的注意力头获得更多缓存，从而实现了比统一压缩策略更优的总信息保留。

在100K+超长上下文及其他基线上的比较
在 InfiniteBench 和 RULER（均包含超过100K Token的实例）上使用 Llama-3.1-8B-Instruct (128K) 模型进行了额外评估。结果表明 $\mathrm { D _ { 2 } O }$ 一致优于 $\mathrm { H _ { 2 } O }$、PyramidKV 和 CaM。此外，在与 SnapKV、SirLLM、LoCoco 等最新基线的比较中，$\mathrm { D _ { 2 } O }$ 同样在大多数数据集上展现出最佳性能。

结论

本文提出了动态判别操作（$\mathrm { D _ { 2 } O }$），该方法无需微调即可通过动态合并Token来有效应对LLMs中的KV缓存管理挑战，以保持基本的上下文信息。通过利用跨层注意力特征密度的变化，$\mathrm { D _ { 2 } O }$ 最大限度地减少了驱逐期间的信息损失，并显著降低了计算和内存需求。实验证实，$\mathrm { D _ { 2 } O }$ 不仅在长文本场景中保持了生成质量，而且在KV缓存压缩与性能之间实现了最佳平衡。未来的研究方向包括探索将 $\mathrm { D _ { 2 } O }$ 与量化、蒸馏以及高效注意力架构等其他压缩方法相结合。