균일한 주의력 너머: 푸리에 근사 KV 캐시를 통한 메모리 효율적인 LLMs

대형 언어 모델(Large Language Models)은 컨텍스트 길이가 증가함에 따라 발생하는 키-밸류(KV) 캐시의 메모리 요구량 문제에 직면해 있습니다. 기존 압축 방법들은 헤드 차원을 동질화하거나 주의(attention)를 기반으로 한 토큰 프루닝(token pruning)에 의존하여, 종종 정확도를 희생하거나 계산 부담을 초래합니다. 본 연구에서는 FourierAttention이라는 학습이 필요하지 않은 프레임워크를 제안합니다. 이 프레임워크는 트랜스포머 헤드 차원들의 이질적인 역할을 활용합니다: 하위 차원은 로컬 컨텍스트를 우선하며, 상위 차원은 장거리 의존성을 포착합니다. FourierAttention은 장컨텍스트에 덜 민감한 차원들을 직교傅리叶基底(orthogonal Fourier bases)로 투영하여, 고정된 길이의 스펙트럼 계수로 그 시간적 변화를 근사합니다. LLaMA 모델에서의 평가 결과, FourierAttention은 LongBench와 Needle-In-A-Haystack (NIAH)에서 가장 우수한 장컨텍스트 정확도를 달성했습니다. 또한, 메모리를 최적화하기 위해 스트림라인된 읽기-쓰기 연산을 통해 설계된 맞춤형 트리톤 커널인 FlashFourierAttention이 효율적인 배포를 가능하게 하면서 성능 저하 없이 작동하도록 최적화되었습니다.