Au-delà de l'attention homogène : des LLM économes en mémoire grâce à un cache KV approximé par transformée de Fourier

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

Résumé

Les grands modèles de langage (LLM) rencontrent des difficultés liées aux exigences de mémoire croissantes du cache clé-valeur (KV) à mesure que les longueurs de contexte augmentent. Les méthodes de compression existantes uniformisent les dimensions des têtes ou reposent sur l'élagage de tokens guidé par l'attention, sacrifiant souvent la précision ou introduisant une surcharge computationnelle. Nous proposons FourierAttention, un cadre sans entraînement qui exploite les rôles hétérogènes des dimensions des têtes de transformateur : les dimensions inférieures privilégient le contexte local, tandis que les dimensions supérieures capturent les dépendances à longue portée. En projetant les dimensions insensibles au contexte long sur des bases de Fourier orthogonales, FourierAttention approxime leur évolution temporelle avec des coefficients spectraux de longueur fixe. Les évaluations sur les modèles LLaMA montrent que FourierAttention atteint la meilleure précision en contexte long sur LongBench et Needle-In-A-Haystack (NIAH). Par ailleurs, un noyau Triton personnalisé, FlashFourierAttention, est conçu pour optimiser la mémoire via des opérations de lecture-écriture rationalisées, permettant un déploiement efficace sans compromis sur les performances.

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.