Voorbij Homogene Aandacht: Geheugenefficiënte LLM's via Fourier-Benaderde KV-Cache

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

June 13, 2025

Auteurs: Xiaoran Liu, Siyang He, Qiqi Wang, Ruixiao Li, Yuerong Song, Zhigeng Liu, Linlin Li, Qun Liu, Zengfeng Huang, Qipeng Guo, Ziwei He, Xipeng Qiu

cs.AI

Samenvatting

Grote Taalmodellen hebben moeite met de geheugeneisen van de groeiende Key-Value (KV)-cache naarmate de contextlengtes toenemen. Bestaande compressiemethoden homogeniseren de dimensies van de aandachtskoppen of vertrouwen op aandacht-gestuurde tokenpruning, wat vaak ten koste gaat van nauwkeurigheid of rekenkundige overhead introduceert. Wij stellen FourierAttention voor, een trainingsvrij raamwerk dat gebruikmaakt van de heterogene rollen van transformer-kopdimensies: lagere dimensies prioriteren lokale context, terwijl hogere dimensies langeafstandsafhankelijkheden vastleggen. Door de dimensies die ongevoelig zijn voor lange context te projecteren op orthogonale Fourier-basen, benadert FourierAttention hun temporele evolutie met vaste-lengte spectrale coëfficiënten. Evaluaties op LLaMA-modellen laten zien dat FourierAttention de beste nauwkeurigheid voor lange context behaalt op LongBench en Needle-In-A-Haystack (NIAH). Daarnaast is een aangepaste Triton-kernel, FlashFourierAttention, ontworpen om het geheugen te optimaliseren via gestroomlijnde lees-schrijfbewerkingen, wat efficiënte implementatie mogelijk maakt zonder prestatieverlies.

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.