За пределами однородного внимания: энергоэффективные большие языковые модели через кэш ключей и значений, аппроксимированный с помощью преобразования Фурье

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

June 13, 2025

Авторы: 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

Аннотация

Крупные языковые модели сталкиваются с проблемами, связанными с растущими требованиями к памяти из-за увеличения размера кэша ключей-значений (KV) при увеличении длины контекста. Существующие методы сжатия унифицируют размерности голов или полагаются на обрезку токенов, управляемую вниманием, что часто приводит к потере точности или увеличению вычислительных затрат. Мы предлагаем FourierAttention, фреймворк, не требующий обучения, который использует разнородные роли размерностей голов трансформера: нижние размерности фокусируются на локальном контексте, а верхние — на долгосрочных зависимостях. Проецируя размерности, нечувствительные к длинному контексту, на ортогональные базисы Фурье, FourierAttention аппроксимирует их временную эволюцию с помощью спектральных коэффициентов фиксированной длины. Оценки на моделях LLaMA показывают, что FourierAttention достигает наилучшей точности на длинных контекстах в тестах LongBench и Needle-In-A-Haystack (NIAH). Кроме того, разработано специальное ядро Triton, FlashFourierAttention, которое оптимизирует использование памяти за счет упрощённых операций чтения-записи, обеспечивая эффективное развёртывание без ущерба для производительности.

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.