潜在空间稀疏注意力用于KV缓存压缩的SALS框架
机器学习
2025-10-29 v1
摘要
能够处理扩展上下文的大型语言模型需求日益增长，但其推理因KV缓存大小大及内存带宽要求高而具有挑战性。先前研究表明KV缓存在隐藏维度上具有低秩特征，暗示可实现有效压缩。然而，由于现代大型语言模型中广泛采用的旋转位置嵌入机制，简单低秩压缩会导致严重的精度下降或产生新的速度瓶颈，因为低秩缓存必须先重建才能应用RoPE。本文提出两个关键见解：首先，RoPE应用于关键向量会增加其方差，从而导致更高秩；其次，在关键向量转换到潜在空间后，它们在大多数层面上保持其表示。基于这些见解，我们提出了潜在空间稀疏注意力框架（SALS）。SALS通过低秩投影将KV缓存投影到紧凑的潜在空间，并在该空间中执行稀疏token选择，使用RoPE-free查询-键交互。通过仅重建少量重要token，避免了完整KV缓存重建的开销。我们在多个任务上全面评估了SALS，使用两个大规模模型：LLaMA2-7b-chat和Mistral-7b，并在RULER-128k基准测试中验证了其可扩展性。实验结果表明，SALS在保持竞争精度方面实现了SOTA性能。在不同设置下，SALS在注意力算子中实现了6.4倍KV缓存压缩和5.7倍加速，相较于FlashAttention2。对于端到端吞吐量性能，在4K和32K序列上分别实现了1.4倍和4.5倍的改进。
引用
@article{arxiv.2510.24273,
title = {SALS: Sparse Attention in Latent Space for KV cache Compression},
author = {Junlin Mu and Hantao Huang and Jihang Zhang and Minghui Yu and Tao Wang and Yidong Li},
journal= {arXiv preprint arXiv:2510.24273},
year = {2025}
}