用傅里叶基近似降低大模型长上下文记忆开销
Beyond Homogeneous Attention: Memory-Efficient LLMs via Fourier-Approximated KV Cache
- 将不同注意力头按功能分层，低维聚焦局部，高维捕捉远距离
- 在长文本任务上准确率优于现有方法，如LongBench和NIAH
- 无需训练，适配推理部署，特别适合资源受限场景
大语言模型因上下文长度增加导致键值（KV）缓存内存需求激增。现有压缩方法通常均一化头维度或依赖注意力引导的令牌剪枝，常以牺牲精度或引入计算开销为代价。本文提出FourierAttention，一种无需训练的框架，利用Transformer头维度的异质性：低维头侧重局部上下文，高维头捕捉长程依赖。通过将对长上下文不敏感的维度投影到正交傅里叶基上，用固定长度的频谱系数近似其时间演化。在LLaMA模型上的评估显示，FourierAttention在LongBench和Needle-In-A-Haystack（NIAH）任务中达到最优长上下文准确率。此外，设计了定制Triton内核FlashFourierAttention，通过优化读写操作减少内存占用，实现高效部署且无性能损失。
原文摘要 · Abstract (English)
Large Language Models struggle with memory demands from the growing Key-Value (KV) cache as context lengths increase. Existing compression methods homogenize head dimensions or rely on attention-guided token pruning, often sacrificing accuracy or introducing computational overhead. We propose FourierAttention, a training-free framework that exploits the heterogeneous roles of transformer head dimensions: lower dimensions prioritize local context, while upper ones capture long-range dependencies. By projecting the long-context-insensitive dimensions onto orthogonal Fourier bases, FourierAttention approximates their temporal evolution with fixed-length spectral coefficients. Evaluations on LLaMA models show that FourierAttention achieves the best long-context accuracy on LongBench and Needle-In-A-Haystack (NIAH). Besides, a custom Triton kernel, FlashFourierAttention, is designed to optimize memory via streamlined read-write operations, enabling efficient deployment without performance compromise.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。