发表时间: 2025-07 · ACL 2025
原文: https://aclanthology.org/2025.acl-long.703
Minwei Zhang, Haifeng Sun, Jingyu Wang, Shaolong Li, Wanyi Ning, Qi Qi, Zirui Zhuang, Jianxin Liao
State Key Laboratory of Networking and Switching Technology, Beijing University of Posts and Telecommunications
一句话结论 提出免训练的稀疏注意力方法 ClusterAttn,通过提取提示词中固有的注意力聚类特征来压缩 KV Cache,在单张 A100 上将 128K 长上下文推理的显存占用降低最高 65% 且几乎无精度损失。
要解决什么问题 大语言模型在处理长上下文推理时,KV Cache 会占用海量显存,例如 175B 参数的 GPT-3 在处理 64 个 4096 token 的序列时需要 1208 GB 显存,是模型权重所需显存的 3.45 倍。现有的稀疏注意力机制通常在不同的注意力头和输入上应用相同的静态稀疏模式(如固定的滑动窗口),这种一刀切的做法无法捕捉模型内部固有的注意力多样性。此外,现有方法大多只关注解码阶段的 KV Cache 优化,却忽略了提示词缓存(prompt cache)的压缩。在实际的商业应用中,用户的输入通常比模型的回复长得多,提示词阶段的 KV Cache 才是真正的显存瓶颈。如果不能根据模型自然关注重要 token 的“内在注意力聚类”现象进行动态且上下文感知的压缩,就会导致长文本外推时显存不足、有效上下文长度受限,或者在强行截断时丢失关键语义信息。
怎么做的 核心思路是利用大模型在解码时注意力分布呈现出的“内在注意力聚类”现象——即高注意力分数的 token 往往会在提示词中聚集在一起,且这些聚类在整个解码过程中保持高度稳定。ClusterAttn 利用提示词末尾的观察窗口(通常是用户的 query)来探测这些聚类,从而动态提取后续解码所需的关键 token。具体流程包含三个关键部件:首先,在全量 prefill 阶段后,利用观察窗口对上下文前缀(context prefix)的注意力分数进行列求和,聚合出注意力重要性特征,定义为
$$ \mathcal { P } = \sum _ { i = - \mathcal { L } _ { w i n } } ^ { i = \mathcal { L } _ { p r o m p t } - 1 } \mathcal { A } _ { w i n } ^ { m } [ : , i , : \mathcal { L } _ { p r e } ] $$ 其中 $\mathcal { A } _ { w i n } ^ { m }$ 是经过掩码和 softmax 归一化后的观察窗口注意力分数。其次,为了将连续的 token 聚类,作者设计了一种基于密度的注意力聚类算法。该算法受 DBSCAN 启发,通过超参数 num_block 控制聚类的粒度和范围,先对 $\mathcal { P }$ 进行最大池化找到聚类中心,再通过阈值过滤掉低分位置,从而自适应地拟合出最优的注意力聚类,实现对上下文前缀的压缩。为了找到最佳的 num_block,方法会在推理前抽取 20% 的样本进行离线分析,以压缩前后注意力分布的余弦相似度为优化目标,仅需几小时即可为特定数据集确定最佳参数。最后,将压缩后的聚类 KV Cache 与完整的观察窗口 KV Cache 拼接,作为后续解码阶段的最终 KV Cache。这种设计既保留了上下文中的连续关键语义块,又保留了对生成质量至关重要的最近 token,成功绕开了固定掩码无法感知上下文内容的卡点。
效果如何 实验在 Mistral-7B-Instruct-v0.2 和 LWM-text-chat-1m 模型上进行,硬件为单张 A100-80GB 显卡。对比基线包括代表全量注意力的 FlashAttn,以及代表免训练稀疏注意力路线的 StreamingLLM(保留初始和最近 token)和 H2O(保留累计注意力得分最高的重度击中 token)。在 LongBench 的 16 个长上下文理解子任务中,ClusterAttn 仅保留 1024 个 token 就能达到与全量注意力相当甚至更好的精度,全面超越了保留 4096 个 token 的 StreamingLLM 和 H2O。在“大海捞针”检索任务中,当序列长度达到 40K 时 FlashAttn 出现显存溢出,H2O 因为生成动态掩码的额外开销在 16K 或 32K 时就会显存溢出,而 ClusterAttn 能够处理高达 128K 的上下文,且检索准确率保持在 90% 以上。量化指标方面,在 1024 压缩大小下,该方法减少了 10% 到 65% 的显存使用,降低了 12% 到 23% 的推理延迟,并将吞吐量提升了 2.6 到 4.8 倍。困惑度测试也表明该方法不会破坏模型生成的连贯性。作者承认的局限性在于:它仅针对生成式模型设计,无法提升模型原生的长文本能力;此外,该方法目前未对解码阶段的 KV Cache 进行动态更新,如果输出序列过长,可能会遇到性能瓶颈。
大型语言模型(LLMs)如 GPT 和 LLaMA 系列在处理长上下文时,KV 缓存(KV cache)的高内存需求成为了高效推理的重大挑战。例如,处理 64 个长度为 4096 tokens 的序列时,1750 亿参数的 GPT-3 模型需要约 1208 GB 的 GPU 内存来存储 KV 缓存,这是模型权重所需内存的 3.45 倍。为了应对这一挑战,KV 缓存压缩技术应运而生。然而,现有的稀疏注意力方法主要集中在解码阶段(decoding)的 KV 缓存优化,忽略了提示词缓存(prompt cache)的压缩,而这在实际应用中通常是更大的内存瓶颈。此外,现有方法通常在不同的注意力头和输入上应用相同的稀疏模式,未能捕捉到 LLMs 内部固有的注意力模式多样性——即“固有注意力聚类”(intrinsic attention clustering)。
本文的核心目标是提出一种免训练(training-free)的稀疏注意力方法,在固有注意力聚类现象的指导下,为高效的 LLM 推理提供有效的提示词缓存压缩方案。
本文的主要创新点和贡献如下:
1. 揭示固有注意力聚类现象:研究发现,在解码过程中,尽管不同注意力头的注意力分布存在差异,但大多数头会一致地关注提示词中的特定聚类(clusters)。这些具有较高注意力分数的 tokens 倾向于聚集在一起,并且这些聚类在整个解码过程中保持高度稳定。这种模式可以通过在提示词末尾设置一个观察窗口(observation window)来检测。
2. 提出 ClusterAttn 方法:这是一种简单且高效的提示词 KV 缓存压缩方案。它首先利用观察窗口从上下文前缀中聚合注意力重要性以形成聚类;接着,受 DBSCAN 启发,提出了一种基于密度的注意力聚类算法(density-based attention clustering algorithm)来自适应地拟合这些聚类,从而压缩上下文前缀;最后将压缩后的聚类与观察窗口拼接,形成用于后续解码的最终 KV 缓存。
3. 卓越的实验性能:在 Mistral-7B 和 LWM 模型的多个基准测试中,ClusterAttn 展现了极高的压缩率和效率。仅保留 1024 个 tokens 的情况下,它能够减少 10%-65% 的内存使用,降低 12%-23% 的延迟,并将吞吐量提高 2.6-4.8 倍,且几乎没有精度损失。此外,在单张 A100-80GB GPU 上,ClusterAttn 能够处理高达 128k 的上下文,显著优于现有方法。
固有聚类模式介绍。本节介绍了在解码过程中注意力分布内部出现的固有聚类模式。作者探讨了如何利用这些模式进行提示词缓存压缩。这些发现基于对各种上下文以及 LLMs 中注意力机制行为的分析。
实验分析术语与公式定义。为了结构化实验分析,作者引入了以下术语:
* Prompt (prompt):用户提供的输入,由上下文前缀(context prefix)和观察窗口(observation window)组成。
* Context Prefix (pre):提示词的一部分,为当前对话或任务提供上下文信息。
* Observation Window (win):提示词的最后一部分,通常包含任务中用户的查询(query)。该窗口对于分析上下文前缀中的注意力模式至关重要。
上述定义的长度关系表示为($\mathcal{L}$ 表示长度):
* Observation Window Mask ($\mathcal{M}_{win}$):该掩码用于遮蔽观察窗口中后续 softmax 函数的注意力权重。
* Feature Aggregation (特征聚合):在观察窗口内,对上下文前缀的注意力分数进行列求和,以促进对固有注意力聚类的观察。对于每个序列,特征聚合过程可以通过以下公式表示:
命中率(Hit Rate)定义。命中率通过测量固有注意力聚类在后续生成过程中保持重要的概率,来量化特征聚合的有效性。计算公式如下:
跨上下文的聚类普遍性。作者在长文档问答数据集(包括 QMSum【Qmsum: A new benchmark for querybased multi-domain meeting summarization + 2021 + NAACL】和 OpenReview【L-eval: Instituting standardized evaluation for long context language models + 2024 + ACL】)中长度约为 3k 的序列上进行了观察。在解码之前,使用观察窗口聚合注意力特征,并过滤掉小于 $\theta$ 的位置。结果发现,对于不同的提示词,固有注意力聚类模式在解码前的注意力头中普遍存在于上下文前缀中,具有较高注意力分数的 tokens 倾向于聚集在一起(紧跟 Fig 3)。
生成过程中的聚类一致性。作者进一步探讨了固有注意力聚类在后续生成过程中是否保持其重要性。使用 Ultrachat 数据集的样本,过滤出响应长度超过 512 且提示词长度超过 1000 的序列。计算了在生成 512 个 tokens 期间固有注意力聚类的平均命中率。结果表明,聚类在整个过程中保持一致性,高命中率证明了这一点(紧跟 Fig 4)。
注意力分数求和验证。为了进一步验证这一点,作者在生成过程中对某一层的全部注意力头中的上下文前缀注意力分数进行了求和。结果显示,尽管聚类在不同头之间有所变化,但它们在每个头内部表现出高度的一致性。由于这些聚类占据了注意力分数的很大比例,作者启发式地假设可以基于固有注意力聚类模式大幅压缩上下文前缀,并设计了一种算法来动态压缩提示词缓存(紧跟 Fig 5)。
ClusterAttn 概述。ClusterAttn 是一种简单而有效的方法,用于在完全预填充(full prefill)后利用固有注意力聚类压缩提示词缓存。与全注意力(full attention)相比,它通过仅保留最相关的信息显著减少了 KV 缓存,降低了内存使用和延迟,同时保持了准确性并提高了吞吐量(紧跟 Fig 2)。
从上下文前缀聚合注意力特征以形成聚类。基于观察到的现象,作者提出使用每个提示词的查询(query)作为观察窗口,从上下文前缀中聚合注意力特征并形成固有注意力聚类。这种方法确保了上下文感知,为后续的聚类拟合奠定了基础。对于每个样本,使用前面公式得到的对应的 $\mathcal{P}$ 作为聚类拟合的对象。
用于压缩上下文前缀的基于密度的注意力聚类算法。在特征聚合之后,需要根据聚类模式压缩上下文前缀。由于固有注意力聚类是基于注意力分数形成的,且不同于传统离散数据的聚类,这些聚类由连续的 tokens 组成,作者受 DBSCAN【A density-based algorithm for discovering clusters in large spatial databases with noise + 1996 + KDD】启发,设计了一种基于密度的注意力聚类算法来拟合固有注意力聚类(如 Algorithm 1 所示)。给定 $\mathcal{L}_{pre}$ 为上下文前缀大小,$\mathcal{L}_{cp}$ 为压缩大小,$\mathcal{L}_{win}$ 为观察窗口大小。算法中定义了 blksize 和 $r$,分别对应 DBSCAN 中使用的最小点数(minPts)和邻域大小,其中 num_block 和 $\theta$ 是管理聚类粒度的超参数。与 DBSCAN 中每个聚类所需的最小点数不同,这里使用 blksize 来控制注意力聚类的大小和范围。
算法执行流程。根据算法,首先基于 blksize 对 $\mathcal{P}$ 执行最大池化(max pooling),以获得每个聚类范围内的块级最大索引 $\mathcal{T}$。然后,过滤掉分数小于 $\theta$ 的索引,得到 $\mathcal{T}_{\theta}$。聚类过程通过 gather 操作获取聚类中心,$r$ 控制聚类邻域。$\mathcal{P}_h^{\overline{\mathcal{L}_U}}$ 表示在排除已选择的注意力聚类后剩余的注意力。值得注意的是,unique 和 topk 操作旨在消除冗余索引,同时收集聚类间隙之间的关键元素。最后,获得压缩后的上下文前缀 KV 缓存 $\{\hat{K}, \hat{V}\}$。根据压缩大小,使用不同的 $\theta$ 进行阈值过滤(如 Table 1 所示)。
Table 1: 不同压缩大小下的 $\theta$ 值
| Compression Size | 1024 | 2048 | 4096 |
|---|---|---|---|
| $\theta$ | 2e-3 | 1e-3 | 8e-4 |
选择合适的 num_block。因此,固有注意力聚类的最佳拟合主要由 num_block 决定,它决定了聚类大小和范围的粒度。为了选择最合适的 num_block,需要为每个数据集在每个压缩大小下分析 num_block。为了确保压缩前后注意力分布的一致性和稳定性,使用解码期间压缩前后对应位置的注意力分布之间的余弦相似度作为指标。然后,通过测量所有样本的所有头的平均余弦相似度来评估最佳拟合。最终的优化目标如下:
这里,$\mathcal{A}{full}^i$ 和 $\mathcal{A}^i$ 表示在生成第 $i$ 个 token 时压缩前后对应位置的注意力分布。$D$ 表示生成的 token 总数。对于每个数据集,选择 20% 的样本并分析最佳的 }_Fnum_block。选择导致最高平均相似度的 num_block 作为最佳聚类拟合。通过为每个数据集应用相应的 num_block,ClusterAttn 能够在细粒度下进行定制的聚类选择,从而实现上下文感知的 KV 缓存压缩。
将聚类与观察窗口拼接作为压缩后的提示词缓存。在使用聚类模式压缩上下文前缀之后,最后一步是将聚类与观察窗口结合起来,获得提示词最终的压缩 KV 缓存。这是因为聚类代表了上下文前缀的压缩,并不包含观察窗口内用户的查询,而这对于模型的响应至关重要。此外,先前的研究【Efficient streaming language models with attention sinks + 2024 + ICLR】、【H2O: heavy-hitter oracle for efficient generative inference of large language models + 2023 + NeurIPS】表明,最近的 tokens 对于后续生成的质量至关重要,因为它们确保了模型的稳定性和流畅性。
1. 固有注意力聚类拟合评估 (Evaluation on Intrinsic Attention Clusters Fitting)
* 实验内容:在 Qasper 数据集上,将提出的基于密度的注意力聚类算法与传统的 K-means 和 DBSCAN 进行拟合性能比较。提示词长度为 3k,ClusterAttn 设置 $\theta=2e-3$,num_block=8,压缩大小为 1024。
* 实验结果与分析:ClusterAttn 在特定头和所有头的平均拟合率上均表现极高(紧跟 Fig 6)。这证明了该算法的有效性,避免了 K-means 在长上下文中通过均值定义相关性的低效,同时比 DBSCAN 能以更细的粒度捕获聚类间隙中的元素。
2. 长上下文理解评估 (Evaluation on Long Context Understanding)
* 实验内容:在 LongBench 上测试 ClusterAttn 在不同压缩大小下的表现,并与 H2O (4096) 和 StreamingLLM (4096) 进行对比。
* 实验结果与分析:如 Table 2 所示,仅使用 1024 个 tokens,ClusterAttn 在 12 个数据集上的性能下降与全注意力相比微乎其微,部分任务甚至更优。在 LWM 和 Mistral 上,1024 tokens 实现了平均 92% 的压缩率,4096 tokens 实现了 68% 的压缩率,且精度损失极小。在 Mistral 上,ClusterAttn (1024) 在 12 个任务中有 8 个优于 H2O (4096),并全面超越 StreamingLLM。
3. 长上下文检索评估 (Evaluation on Long Context Retrieval)
* 实验内容:在 LWM 模型上进行 "Needle in a Haystack" 测试,将文档长度扩展到 128k tokens(单张 A100-80GB 在压缩大小为 1024 时的极限)。
* 实验结果与分析:如 Fig 7 所示,全注意力在 40k 时发生 OOM(内存溢出),而 ClusterAttn 能够处理 128k 的序列长度,且检索准确率仍超过 90%。这突显了其在 128倍 压缩率下处理极长输入细节的潜力及节省显存的直接效果。
4. 连贯性测试与吞吐量 (Coherence Testing & Throughput)
* 实验内容:从 Qasper、MultiNew、TREC 和 LCC 数据集中提取问答对计算困惑度(Perplexity),测试连贯性(Table 3)。测试 LWM 模型在不同 batch size 下的吞吐量(Fig 8)。
* 实验结果与分析:ClusterAttn 始终表现出较低的困惑度,表明其未对生成连贯性产生负面影响,去除冗余信息反而可能增强连贯性。吞吐量方面,随着输入长度增加,ClusterAttn 的吞吐量保持恒定(因推理时 KV 缓存大小固定),而 FlashAttn 呈指数下降。相较于 FlashAttn,吞吐量提升了 2.6 到 4.8 倍。
5. 效率评估 (Efficiency)
* 实验内容:在 batch size 为 1 的情况下,测试峰值内存和推理延迟随序列长度增加的变化,ClusterAttn 和 H2O 压缩大小设为 1024。
* 实验结果与分析:内存方面(Table 4),Mistral 上 H2O 在 32k 达到极限且消耗 2-3 倍于 ClusterAttn 的内存;LWM 上,H2O 在 16k 后 OOM,FlashAttn 在 32k 后 OOM,而 ClusterAttn 延伸至 128k。延迟方面(Fig 9),当序列达到 32k 时,ClusterAttn 至少减少了 20% 的延迟(相比 FlashAttn 和 H2O)。
本文提出了一种基于固有注意力聚类的提示词缓存压缩稀疏注意力方法——ClusterAttn。该方法简单、有效且免训练,充分利用了解码过程中的注意力聚类模式。仅使用 1024 个 tokens,它就能减少 10%-65% 的内存使用,降低 12%-23% 的延迟,并将吞吐量提高 2.6-4.8 倍,且几乎没有精度损失。此外,它能在单张 A100-80GB GPU 上处理高达 128k 的上下文,优于现有方法。
局限性:ClusterAttn 的研究范围主要限于模型的生成方面,无法扩展模型本身不具备的长上下文能力。此外,其设计未涉及解码过程的优化,解码仍依赖标准注意力或 FlashAttn,这限制了其在解码期间动态更新 KV 缓存的能力。
A. 额外相关工作。介绍了动态推理方法(如 MoE)、多查询注意力(MQA)、低秩注意力、量化技术以及替代传统注意力的长序列建模机制。
B. 与 SnapKV 的扩展比较。SnapKV【Snapkv: LLM knows what you are looking for before generation + 2024 + NeurIPS】同样是免训练的 KV 压缩方法,但它对所有提示词使用固定长度的观察窗口,并采用简单的池化方法。性能与准确度对比:在 LongBench 上的实验(Table 5)表明,ClusterAttn 整体略优于 SnapKV。为了拉开差距,作者在 LWM 模型上进行了 Needle-in-a-haystack 实验(压缩大小 1024)。检索准确率验证:结果(Fig 10)表明,随着序列长度增加,SnapKV 的检索准确率显著下降,而 ClusterAttn 凭借更细粒度的压缩算法保持了稳健的性能。
C. 在 LLaMA-3 上的扩展结果。LLaMA-3适配测试:作者将 ClusterAttn 扩展到 LLaMA3-8B-Instruct,并在 LongBench 上与全注意力和 SnapKV 进行了比较(最大序列长度设为 8192)。结果(Table 6)显示,ClusterAttn 表现符合预期,准确性和效率与其他模型一致。
D. 不同稀疏注意力对保留上下文完整性的影响。在 LLMs 中,信息检索依赖高分特征及归纳头(inductive heads)复制其余上下文。上下文完整性破坏问题:简单选择排名最高的特征(如 Top-K)可能导致只保留部分细节(如只检索到区号而捏造电话号码其余部分),破坏上下文完整性。ClusterAttn的优势:如 Table 7 所示,ClusterAttn 在所有任务上表现良好,因为每个聚类代表一个连续的片段,保留了上下文完整性,从而在后续解码中保持了与全注意力相当的注意力分数稳定性。
E. LongEval 评估。LongEval-Lines任务说明:为了进一步评估检索能力,使用了修改版的 LongEval-Lines 基准【How long can context length of open-source llms truly promise? + 2023 + NeurIPS Workshop】。该任务要求在格式相似的嘈杂上下文中识别键值对,比 "Needle in a Haystack" 更具挑战性。实验表现:在 Mistral 模型(压缩大小 1024)上的结果(Fig 11)表明,随着序列长度增加,ClusterAttn 始终与全注意力表现相当,而 H2O 和 StreamingLLM 表现较差(H2O 在 20k 后 OOM)。这证明保留固有注意力聚类能显著提高长上下文中的关键信息提取准确率。
F. 消融实验 (Ablation Study)。在 Mistral-7B-Instruct-v0.2 上分析了 $\theta$ 和 num_block 的影响。
F.1 Theta 选择对聚类拟合的影响:$\theta$ 影响聚类过程中的邻域大小。Theta的变化趋势:如图 13 所示,随着 $\theta$ 增加,余弦相似度最初逐渐上升至顶点,然后迅速下降。这是因为较大的 $\theta$ 过滤掉了较弱但关键的聚类,而过小的 $\theta$ 会包含无关 tokens。作者为每个压缩大小选择了余弦相似度开始下降的第一个 $\theta$。
F.2 BlockNum 选择对聚类拟合的影响:固定 $\theta$ 后测试 num_block(范围 1-32)。BlockNum的峰值现象:结果(图 13)显示,余弦相似度在特定 num_block 处达到峰值。较小的 blksize 可能过拟合单个 tokens,较大的 blksize 可能稀释高分 tokens,降低聚类区分度。表 8 和表 9 展示了所有实验的最佳 num_block。表 10 和表 11 显示,ClusterAttn 仅需少量样本和几小时即可完成 Profiling,证明了其实用性。
F.3 ClusterAttn 与基于 N-gram 的注意力聚类对比:N-gram算法设计:由于单词被分词为 n-gram,作者设计了一个简单的 n-gram 聚类算法进行对比:(1) 聚类属于同一单词/短语的连续 tokens;(2) 保留与已聚类单词相邻的 n-1 个聚类,直到达到压缩大小。实验对比分析:在 LWM(压缩大小 1024)上的结果(Table 12)显示,n-gram 算法在 2-gram 时表现最好,但仍不及 ClusterAttn。这说明 n-gram 更侧重于物理接近度而非语义理解,导致聚类离散、缺乏语义连续性。
G. 头级别聚类拟合中 BlockNum 的可视化。头级别拟合验证:为了证明获得的最佳 num_block 能为每个头实现合适的聚类拟合,作者展示了在相同压缩大小下不同 num_block 跨注意力头的平均余弦相似度。以 Qasper 数据集(压缩大小 1024)为例(Fig 12),当 num_block 为 8 时拟合最佳,整体相似度高于其他值,突显了该方法为每个头精细应用上下文感知压缩的能力。
H. 文本生成可视化。作者在表 13、14、15 中提供了 NarrativeQA、Qasper 和 SAMSum 数据集在 Mistral-7B-Instruct-v0.2 上的生成示例,对比了 Ground Truth、ClusterAttn (1024, 2048, 4096) 和全注意力的生成结果,直观展示了模型在压缩后的生成质量。