SALS: Sparse Attention in Latent Space for KV cache Compression

📄 arXiv: 2510.24273v1 📥 PDF

作者: Junlin Mu, Hantao Huang, Jihang Zhang, Minghui Yu, Tao Wang, Yidong Li

分类: cs.LG

发布日期: 2025-10-28

💡 一句话要点

提出SALS框架，通过潜在空间稀疏注意力实现KV缓存压缩，加速长文本LLM推理。

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

关键词: KV缓存压缩 长文本LLM 稀疏注意力 低秩分解 推理加速 潜在空间 RoPE LLM优化

📋 核心要点

  1. 现有长文本LLM推理面临KV缓存巨大和内存带宽需求高的挑战，传统的低秩压缩方法因RoPE机制导致精度下降或引入新的速度瓶颈。
  2. SALS框架通过将KV缓存投影到紧凑的潜在空间，并在该空间执行稀疏token选择，避免了完全重建KV缓存的开销，从而实现高效压缩。
  3. 实验结果表明，SALS在LLaMA2、Mistral和LLaMA3模型上均表现出色，实现了显著的KV缓存压缩和推理加速，同时保持了良好的精度。

📝 摘要（中文）

为了应对长文本大型语言模型（LLM）推理中Key-Value（KV）缓存过大和内存带宽需求高的问题，本文提出了一种名为“潜在空间稀疏注意力”（SALS）的框架，用于KV缓存压缩。SALS基于两个关键观察：一是RoPE机制增加了Key向量的方差，导致其秩增大；二是Key向量在变换到潜在空间后，在大多数层中保持其表示。SALS通过低秩投影将KV缓存投影到紧凑的潜在空间，并在此空间中使用无RoPE的Query-Key交互执行稀疏token选择。通过仅重建一小部分重要token，避免了完全重建KV缓存的开销。在LLaMA2-7b-chat、Mistral-7b以及LLaMA3.1-8B-Instruct上进行了广泛的评估，结果表明SALS在保持竞争力的同时实现了SOTA性能。在不同设置下，SALS实现了6.4倍的KV缓存压缩和5.7倍的Attention算子加速（基于4K序列）。端到端吞吐量方面，相比GPT-fast，在4K和32K序列上分别实现了1.4倍和4.5倍的提升。

🔬 方法详解

问题定义：论文旨在解决长文本LLM推理过程中KV缓存占用空间大、内存带宽需求高的问题。现有方法，特别是直接应用低秩分解的方法，在使用了RoPE（Rotary Position Embedding）的位置编码的LLM中表现不佳，因为RoPE会增加Key向量的方差，使得低秩分解效果变差，并且需要先重建KV缓存才能应用RoPE，引入了额外的计算开销。

核心思路：论文的核心思路是利用Key向量在经过线性变换进入潜在空间后，其表示在不同层之间具有较高的一致性。因此，可以在这个潜在空间中进行稀疏注意力计算，选择重要的token，并只重建这些token的KV缓存，从而避免重建整个KV缓存带来的计算负担。同时，在潜在空间中进行Query-Key交互时，避免使用RoPE，进一步降低计算复杂度。

技术框架：SALS框架主要包含以下几个阶段：1) 低秩投影：使用低秩矩阵将Key和Value向量投影到低维潜在空间。2) 稀疏Token选择：在潜在空间中，使用Query向量和Key向量进行无RoPE的交互，选择重要的token。3) KV缓存重建：只重建被选中的token的KV缓存。4) Attention计算：使用重建后的KV缓存进行标准的Attention计算。

关键创新：SALS的关键创新在于：1) 潜在空间稀疏注意力：在低维潜在空间中进行稀疏注意力计算，降低了计算复杂度。2) 无RoPE的Query-Key交互：在潜在空间中进行Query-Key交互时，避免使用RoPE，进一步降低计算复杂度。3) 选择性KV缓存重建：只重建重要的token的KV缓存，避免了重建整个KV缓存带来的计算负担。

关键设计：SALS的关键设计包括：1) 低秩投影矩阵的选择：可以使用SVD等方法进行学习，也可以直接使用随机矩阵。2) 稀疏度控制：通过调整选择token的数量来控制压缩率和计算复杂度。3) 潜在空间的维度：需要根据具体的模型和任务进行调整，以保证精度和压缩率之间的平衡。4) Token选择策略：可以使用Top-K选择、阈值选择等方法。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SALS在LLaMA2-7b-chat和Mistral-7b模型上实现了显著的KV缓存压缩和推理加速。在4K序列长度下，SALS实现了6.4倍的KV缓存压缩和5.7倍的Attention算子加速（相比FlashAttention2）。端到端吞吐量方面，相比GPT-fast，在4K和32K序列上分别实现了1.4倍和4.5倍的提升。此外，SALS在RULER-128k基准测试中也验证了其在超长文本上的可扩展性。

🎯 应用场景

SALS框架可应用于各种需要处理长文本的LLM应用场景，例如长文档摘要、机器翻译、代码生成、对话系统等。通过降低KV缓存的存储需求和加速推理速度，SALS能够提升LLM在资源受限设备上的部署能力，并降低云计算成本，具有广泛的应用前景。

📄 摘要（原文）

Large Language Models capable of handling extended contexts are in high demand, yet their inference remains challenging due to substantial Key-Value cache size and high memory bandwidth requirements. Previous research has demonstrated that KV cache exhibits low-rank characteristics within the hidden dimension, suggesting the potential for effective compression. However, due to the widely adopted Rotary Position Embedding mechanism in modern LLMs, naive low-rank compression suffers severe accuracy degradation or creates a new speed bottleneck, as the low-rank cache must first be reconstructed in order to apply RoPE. In this paper, we introduce two key insights: first, the application of RoPE to the key vectors increases their variance, which in turn results in a higher rank; second, after the key vectors are transformed into the latent space, they largely maintain their representation across most layers. Based on these insights, we propose the Sparse Attention in Latent Space framework. SALS projects the KV cache into a compact latent space via low-rank projection, and performs sparse token selection using RoPE-free query-key interactions in this space. By reconstructing only a small subset of important tokens, it avoids the overhead of full KV cache reconstruction. We comprehensively evaluate SALS on various tasks using two large-scale models: LLaMA2-7b-chat and Mistral-7b, and additionally verify its scalability on the RULER-128k benchmark with LLaMA3.1-8B-Instruct. Experimental results demonstrate that SALS achieves SOTA performance by maintaining competitive accuracy. Under different settings, SALS achieves 6.4-fold KV cache compression and 5.7-fold speed-up in the attention operator compared to FlashAttention2 on the 4K sequence. For the end-to-end throughput performance, we achieves 1.4-fold and 4.5-fold improvement compared to GPT-fast on 4k and 32K sequences, respectively.