Jenseits homogener Aufmerksamkeit: Speichereffiziente LLMs durch Fourier-approximierten KV-Cache

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

June 13, 2025

Autoren: 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

Zusammenfassung

Große Sprachmodelle kämpfen mit den steigenden Speicheranforderungen durch den wachsenden Key-Value (KV)-Cache, wenn die Kontextlängen zunehmen. Bestehende Kompressionsmethoden homogenisieren die Kopf-Dimensionen oder verlassen sich auf aufmerksamkeitsgesteuertes Token-Pruning, was oft zu Genauigkeitseinbußen oder zusätzlichem Rechenaufwand führt. Wir schlagen FourierAttention vor, ein trainingsfreies Framework, das die heterogenen Rollen der Transformer-Kopf-Dimensionen ausnutzt: niedrigere Dimensionen priorisieren den lokalen Kontext, während höhere Dimensionen langreichweitige Abhängigkeiten erfassen. Durch die Projektion der langkontextunempfindlichen Dimensionen auf orthogonale Fourier-Basen approximiert FourierAttention deren zeitliche Entwicklung mit festen spektralen Koeffizienten. Evaluierungen an LLaMA-Modellen zeigen, dass FourierAttention die beste Langkontext-Genauigkeit auf LongBench und Needle-In-A-Haystack (NIAH) erreicht. Zusätzlich wurde ein benutzerdefinierter Triton-Kernel, FlashFourierAttention, entwickelt, um den Speicher durch optimierte Lese-Schreib-Operationen zu optimieren, was eine effiziente Bereitstellung ohne Leistungseinbußen ermöglicht.

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.