Além da Atenção Homogênea: LLMs Eficientes em Memória via Cache KV Aproximado 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

Resumo

Modelos de Linguagem de Grande Escala enfrentam desafios com as demandas de memória do crescente cache Chave-Valor (KV) à medida que os comprimentos de contexto aumentam. Os métodos de compressão existentes homogeneizam as dimensões das cabeças ou dependem da poda de tokens guiada por atenção, muitas vezes sacrificando precisão ou introduzindo sobrecarga computacional. Propomos o FourierAttention, uma estrutura livre de treinamento que explora os papéis heterogêneos das dimensões das cabeças do transformador: dimensões inferiores priorizam o contexto local, enquanto as superiores capturam dependências de longo alcance. Ao projetar as dimensões insensíveis ao contexto longo em bases de Fourier ortogonais, o FourierAttention aproxima sua evolução temporal com coeficientes espectrais de comprimento fixo. Avaliações em modelos LLaMA mostram que o FourierAttention alcança a melhor precisão de contexto longo no LongBench e no Needle-In-A-Haystack (NIAH). Além disso, um kernel Triton personalizado, FlashFourierAttention, foi projetado para otimizar a memória por meio de operações simplificadas de leitura e escrita, permitindo implantação eficiente sem comprometer o desempenho.

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.