Más allá de la atención homogénea: LLMs eficientes en memoria mediante caché KV aproximada por Fourier

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

June 13, 2025

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

Resumen

Los Modelos de Lenguaje de Gran Escala enfrentan dificultades con las demandas de memoria del creciente caché Clave-Valor (KV) a medida que aumentan las longitudes de contexto. Los métodos de compresión existentes homogenizan las dimensiones de las cabezas o dependen de la poda de tokens guiada por atención, lo que a menudo sacrifica la precisión o introduce sobrecarga computacional. Proponemos FourierAttention, un marco libre de entrenamiento que aprovecha los roles heterogéneos de las dimensiones de las cabezas del transformador: las dimensiones inferiores priorizan el contexto local, mientras que las superiores capturan dependencias de largo alcance. Al proyectar las dimensiones insensibles al contexto largo sobre bases de Fourier ortogonales, FourierAttention aproxima su evolución temporal con coeficientes espectrales de longitud fija. Las evaluaciones en modelos LLaMA muestran que FourierAttention logra la mejor precisión de contexto largo en LongBench y Needle-In-A-Haystack (NIAH). Además, se ha diseñado un núcleo Triton personalizado, FlashFourierAttention, para optimizar la memoria mediante operaciones de lectura-escritura simplificadas, permitiendo un despliegue eficiente sin comprometer el rendimiento.

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.