超越同质注意力：通过傅里叶近似KV缓存实现内存高效的LLM模型

大型语言模型在上下文长度增加时面临着来自不断增长的键值（KV）缓存的内存需求挑战。现有的压缩方法通常通过统一头部维度或依赖注意力引导的标记剪枝来应对这一问题，但这些方法往往以牺牲准确性为代价或引入额外的计算开销。我们提出了一种名为傅里叶注意力（FourierAttention）的无训练框架，该框架利用了变压器头部维度的异质角色：较低维度优先关注局部上下文，而较高维度则捕捉长距离依赖关系。通过将对长上下文不敏感的维度投影到正交傅里叶基上，傅里叶注意力使用固定长度的频谱系数来近似其时间演化。在LLaMA模型上的评估结果显示，傅里叶注意力在LongBench和针尖寻觅（Needle-In-A-Haystack, NIAH）任务中实现了最佳的长上下文准确性。此外，还设计了一个定制的Triton内核——快速傅里叶注意力（FlashFourierAttention），通过优化读写操作来减少内存占用，从而实现高效部署而不影响性能。