Oltre l'attenzione omogenea: LLM efficienti in memoria tramite cache KV approssimata con trasformata di Fourier

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

June 13, 2025

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

Abstract

I grandi modelli linguistici (Large Language Models) affrontano difficoltà legate alle esigenze di memoria derivanti dalla crescente cache Chiave-Valore (KV) all'aumentare delle lunghezze del contesto. I metodi di compressione esistenti omogeneizzano le dimensioni delle testine (head dimensions) o si basano sulla potatura dei token guidata dall'attenzione, spesso sacrificando l'accuratezza o introducendo un sovraccarico computazionale. Proponiamo FourierAttention, un framework senza necessità di addestramento che sfrutta i ruoli eterogenei delle dimensioni delle testine nei trasformatori: le dimensioni inferiori privilegiano il contesto locale, mentre quelle superiori catturano le dipendenze a lungo raggio. Proiettando le dimensioni insensibili al contesto lungo su basi di Fourier ortogonali, FourierAttention approssima la loro evoluzione temporale con coefficienti spettrali a lunghezza fissa. Le valutazioni sui modelli LLaMA dimostrano che FourierAttention raggiunge la migliore accuratezza su contesti lunghi in LongBench e Needle-In-A-Haystack (NIAH). Inoltre, è stato progettato un kernel Triton personalizzato, FlashFourierAttention, per ottimizzare la memoria attraverso operazioni di lettura-scrittura semplificate, consentendo un'implementazione efficiente senza compromettere le prestazioni.

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.