동종 어텐션을 넘어서: 푸리에 근사 기반 KV 캐시를 통한 메모리 효율적 대형 언어 모델

초록

대규모 언어 모델은 컨텍스트 길이가 증가함에 따라 키-값(KV) 캐시의 메모리 요구량이 증가하는 문제에 직면해 있습니다. 기존의 압축 방법들은 헤드 차원을 균일하게 처리하거나 주의 기반 토큰 제거에 의존하는 경우가 많아, 정확도를 희생하거나 계산 오버헤드를 초래하는 경우가 많습니다. 우리는 퓨리에 어텐션(FourierAttention)을 제안합니다. 이는 트랜스포머 헤드 차원의 이질적인 역할을 활용하는 훈련이 필요 없는 프레임워크로, 하위 차원은 지역 컨텍스트를 우선시하고 상위 차원은 장거리 의존성을 포착합니다. 장거리 컨텍스트에 민감하지 않은 차원을 직교 퓨리에 기저에 투영함으로써, 퓨리에 어텐션은 이들의 시간적 변화를 고정 길이 스펙트럼 계수로 근사합니다. LLaMA 모델에 대한 평가 결과, 퓨리에 어텐션은 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.

동종 어텐션을 넘어서: 푸리에 근사 기반 KV 캐시를 통한 메모리 효율적 대형 언어 모델

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

초록

Support