SparK: Query-Aware Unstructured Sparsity with Recoverable KV Cache Channel Pruning

📄 arXiv: 2508.15212v3 📥 PDF

作者: Huanxuan Liao, Yixing Xu, Shizhu He, Guanchen Li, Xuanwu Yin, Dong Li, Emad Barsoum, Jun Zhao, Kang Liu

分类: cs.CL, cs.AI, cs.LG

发布日期: 2025-08-21 (更新: 2025-11-12)

备注: accepted to AAAI 2026

🔗 代码/项目: GITHUB


💡 一句话要点

SparK:查询感知的非结构化稀疏化KV缓存通道剪枝,提升长文本LLM推理效率。

🎯 匹配领域: 支柱九:具身大模型 (Embodied Foundation Models)

关键词: 长文本推理 KV缓存 非结构化剪枝 通道剪枝 查询感知 大型语言模型 模型压缩

📋 核心要点

  1. 现有长文本LLM推理受限于KV缓存,其内存占用随序列长度线性增长,计算复杂度呈平方增长,亟需高效压缩方案。
  2. SparK提出一种查询感知的非结构化稀疏化方法,在通道级别剪枝KV缓存,并在计算时动态恢复,降低冗余。
  3. 实验表明,SparK在相同内存预算下可处理更长序列,或在相同序列长度下减少30%以上KV缓存存储,且性能下降小于5%。

📝 摘要(中文)

大型语言模型(LLM)中的长上下文推理越来越多地受到KV缓存瓶颈的限制:内存使用量随序列长度线性增长,而注意力计算则呈二次方增长。现有方法通过token驱逐或合并等策略沿时间轴压缩KV缓存来解决此问题,以减少内存和计算开销。然而,这些方法通常忽略了特征维度(即通道轴)上细粒度的重要性变化,从而限制了它们有效平衡效率和模型准确性的能力。实际上,我们观察到通道显著性在查询和位置之间变化很大:某些特征通道对于给定查询携带的信息几乎为零,而另一些通道的相关性则急剧上升。为了解决这个疏忽,我们提出SPARK,这是一种免训练的即插即用方法,它通过在通道级别剪枝KV来应用非结构化稀疏性,同时在注意力得分计算期间动态恢复被剪枝的条目。值得注意的是,我们的方法与现有的KV压缩和量化技术正交,使其可以与它们集成以实现进一步的加速。通过减少通道级冗余,SPARK能够在相同的内存预算内处理更长的序列。对于相同长度的序列,SPARK不仅保持或提高了模型准确性,而且与基于驱逐的方法相比,KV缓存存储减少了30%以上。此外,即使采用80%的激进剪枝率,与基线驱逐方法相比,SPARK也能保持性能,且性能下降小于5%,证明了其鲁棒性和有效性。

🔬 方法详解

问题定义:现有长文本LLM推理中,KV缓存成为性能瓶颈。现有方法主要集中在时间维度上的压缩,例如token驱逐或合并,忽略了通道维度上细粒度的重要性差异。不同查询和位置对不同通道的依赖程度不同,存在大量冗余信息,现有方法无法有效利用这些冗余。

核心思路:SparK的核心思路是利用通道维度的稀疏性,对KV缓存进行非结构化剪枝。通过分析不同查询对不同通道的依赖程度,动态地剪枝不重要的通道,从而减少KV缓存的存储空间和计算量。同时,在计算注意力得分时,动态恢复被剪枝的条目,以保证模型的准确性。

技术框架:SparK是一个即插即用的模块,可以方便地集成到现有的LLM架构中。其主要流程包括:1) 离线分析:对训练数据进行分析,确定每个通道的重要性;2) KV缓存剪枝:根据通道重要性,对KV缓存进行非结构化剪枝;3) 注意力计算:在计算注意力得分时,动态恢复被剪枝的条目;4) 模型推理:使用剪枝后的KV缓存进行模型推理。

关键创新:SparK的关键创新在于:1) 提出了一种查询感知的通道级非结构化剪枝方法,能够更精细地利用KV缓存的稀疏性;2) 提出了一种动态恢复机制,能够在保证模型准确性的前提下,最大限度地减少KV缓存的存储空间和计算量;3) SparK与现有的KV压缩和量化技术正交,可以与它们集成以实现进一步的加速。

关键设计:SparK的关键设计包括:1) 通道重要性评估:使用某种指标(例如梯度或激活值)来评估每个通道的重要性;2) 剪枝策略:根据通道重要性,选择合适的剪枝比例;3) 动态恢复机制:在计算注意力得分时,根据查询和位置,动态恢复被剪枝的条目。具体的参数设置和损失函数未知,论文中可能未详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,SparK在保持或提高模型准确性的前提下,能够显著减少KV缓存的存储空间。对于相同长度的序列,SparK与基于驱逐的方法相比,KV缓存存储减少了30%以上。即使采用80%的激进剪枝率,与基线驱逐方法相比,SparK也能保持性能,且性能下降小于5%,展示了其鲁棒性和有效性。

🎯 应用场景

SparK可应用于各种需要处理长文本的LLM应用场景,例如长文档摘要、机器翻译、对话系统等。通过减少KV缓存的存储空间和计算量,SparK可以提高LLM的推理速度和效率,使其能够处理更长的文本序列,并降低部署成本。该技术还有助于在资源受限的设备上部署大型语言模型。

📄 摘要(原文)

Long-context inference in large language models (LLMs) is increasingly constrained by the KV cache bottleneck: memory usage grows linearly with sequence length, while attention computation scales quadratically. Existing approaches address this issue by compressing the KV cache along the temporal axis through strategies such as token eviction or merging to reduce memory and computational overhead. However, these methods often neglect fine-grained importance variations across feature dimensions (i.e., the channel axis), thereby limiting their ability to effectively balance efficiency and model accuracy. In reality, we observe that channel saliency varies dramatically across both queries and positions: certain feature channels carry near-zero information for a given query, while others spike in relevance. To address this oversight, we propose SPARK, a training-free plug-and-play method that applies unstructured sparsity by pruning KV at the channel level, while dynamically restoring the pruned entries during attention score computation. Notably, our approach is orthogonal to existing KV compression and quantization techniques, making it compatible for integration with them to achieve further acceleration. By reducing channel-level redundancy, SPARK enables processing of longer sequences within the same memory budget. For sequences of equal length, SPARK not only preserves or improves model accuracy but also reduces KV cache storage by over 30% compared to eviction-based methods. Furthermore, even with an aggressive pruning ratio of 80%, SPARK maintains performance with less degradation than 5% compared to the baseline eviction method, demonstrating its robustness and effectiveness. Our code will be available at https://github.com/Xnhyacinth/SparK.