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

发表
Liu XiaoranLiu Xiaoran 提交
作者: Liu XiaoranXiaoran Liu, Siyang He, Qiqi Wang, Ruixiao Li, Yuerong Song, Zhigeng Liu, Linlin Li, Qun Liu, Zengfeng Huang, Qipeng Guo, Ziwei He, Xipeng Qiu

摘要

随着上下文长度的增加,大型语言模型在不断增长的键值(KV)缓存中面临内存需求挑战。现有压缩方法往往同质化注意力头维度或依赖注意力引导的标记剪枝,这通常会牺牲准确性或引入计算开销。我们提出了FourierAttention,一个免训练的框架,它利用了Transformer注意力头维度的异构作用:较低维度优先处理局部上下文,而较高维度则捕获长程依赖。通过将对长上下文不敏感的维度投影到正交傅里叶基上,FourierAttention使用固定长度的频谱系数近似它们的时序演化。在LLaMA模型上的评估表明,FourierAttention在LongBench和Needle-In-A-Haystack (NIAH) 上实现了最佳长上下文准确性。此外,我们设计了一个定制的Triton内核FlashFourierAttention,旨在通过简化的读写操作优化内存,从而实现高效部署而不影响性能。
查看 arXiv 页面查看 PDF

评论

Liu XiaoranLiu Xiaoran
论文作者
论文提交者

大语言模型随着上下文长度的增加,在处理不断增长的键值(KV)缓存所带来的内存需求方面面临挑战。现有的压缩方法通常使头部维度同质化或依赖于注意力引导的令牌修剪,这往往会牺牲准确性或引入额外的计算开销。我们提出 FourierAttention,这是一个无需训练的框架,它利用了变压器(Transformer)头部维度的异构作用:较低维度优先处理局部上下文,而较高维度则捕获长程依赖性。通过将对长上下文不敏感的维度投影到正交傅里叶基上,FourierAttention 利用固定长度的频谱系数来近似它们的时序演变。在 LLaMA 模型上的评估表明,FourierAttention 在 LongBench 和 Needle-In-A-Haystack (NIAH) 上实现了最佳的长上下文准确性。此外,我们设计了一个定制的 Triton 内核 FlashFourierAttention,旨在通过简化读写操作来优化内存,从而实现高效部署而不影响性能。

Wenhao LiWenhao Li

请问这篇文章的代码什么时候可以释放出来?

Liu XiaoranLiu Xiaoran
论文作者
论文提交者

您好,中稿之后会更新并发布的