Beyond Homogeneous Attention: Memory-Efficient LLMs via Fourier-Approximated KV Cache

论文主要内容总结

研究背景与问题

大语言模型（LLMs）在处理长上下文时，键值（KV）缓存的内存需求呈线性增长，成为资源受限环境下部署的瓶颈。现有KV缓存压缩方法（如标记淘汰、量化、低秩投影等）多将头部维度同质化处理，或依赖注意力引导策略，导致准确性下降或计算开销增加。

核心方法

• FourierAttention框架：利用Transformer头部维度的异质性，发现低维度专注于局部上下文，高维度捕获长距离依赖。通过将长上下文不敏感的维度投影到正交傅里叶基上，用固定长度的频谱系数近似其时间演化，避免存储完整序列的KV缓存。
• HiPPO-FourierT压缩：在预填充阶段保留初始和局部标记，压缩中间标记的KV状态；解码阶段通过逆傅里叶变换重建缓存，结合定制的Triton内核FlashFourierAttention优化读写操作，减少内存开销。
• 细粒度压缩方案：根据KV缓存各层维度的标准差动态调整压缩比例，对变化更平缓的V缓存和低层缓存压缩更多维度，上层缓存和K缓存保留更多维度，形成非对称倒金字塔压缩模式。

实验结果

在LLaMA3.1-8B和LLaMA3.2-3B模型