ThinK: 쿼리 기반 가지치기를 통한 경량화된 키 캐시

ThinK: Thinner Key Cache by Query-Driven Pruning

July 30, 2024

저자: Yuhui Xu, Zhanming Jie, Hanze Dong, Lei Wang, Xudong Lu, Aojun Zhou, Amrita Saha, Caiming Xiong, Doyen Sahoo

cs.AI

초록

대형 언어 모델(LLM)은 모델 크기와 시퀀스 길이의 증가를 활용하여 다양한 응용 분야에서 전례 없는 성능을 달성하며 자연어 처리 분야를 혁신했습니다. 그러나 이와 관련된 계산 및 메모리 비용의 증가는 특히 트랜스포머 어텐션 메커니즘의 2차 복잡성으로 인해 긴 시퀀스를 관리하는 데 있어 상당한 과제를 제기합니다. 본 논문은 장문맥 시나리오에 초점을 맞추어, 추론 과정에서 KV 캐시 메모리 소비의 비효율성을 해결합니다. 기존의 접근 방식들이 시퀀스 길이에 기반하여 메모리를 최적화하는 것과 달리, 우리는 KV 캐시의 채널 차원이 불균형한 크기 분포와 어텐션 가중치의 저랭크 구조로 특징지어지는 상당한 중복성을 보인다는 점을 발견했습니다. 이러한 관찰을 바탕으로, 우리는 ThinK라는 새로운 쿼리 의존적 KV 캐시 프루닝 방법을 제안합니다. 이 방법은 가장 중요하지 않은 채널을 선택적으로 제거하면서 어텐션 가중치 손실을 최소화하도록 설계되었습니다. 우리의 접근 방식은 모델 정확도를 유지하거나 향상시킬 뿐만 아니라, 기존의 KV 캐시 제거 방법에 비해 메모리 비용을 20% 이상 절감합니다. LLaMA3 및 Mistral 모델을 다양한 장문 시퀀스 데이터셋에 대해 광범위하게 평가한 결과, ThinK의 효율성을 확인하며 성능 저하 없이 LLM 배포의 새로운 기준을 제시했습니다. 또한, 우리는 이 방법을 값 캐시 프루닝으로 확장할 가능성을 제시함으로써 ThinK의 다재다능성과 메모리 및 계산 오버헤드를 줄이는 데 있어 광범위한 적용 가능성을 입증했습니다.

English

Large Language Models (LLMs) have revolutionized the field of natural language processing, achieving unprecedented performance across a variety of applications by leveraging increased model sizes and sequence lengths. However, the associated rise in computational and memory costs poses significant challenges, particularly in managing long sequences due to the quadratic complexity of the transformer attention mechanism. This paper focuses on the long-context scenario, addressing the inefficiencies in KV cache memory consumption during inference. Unlike existing approaches that optimize the memory based on the sequence lengths, we uncover that the channel dimension of the KV cache exhibits significant redundancy, characterized by unbalanced magnitude distribution and low-rank structure in attention weights. Based on these observations, we propose ThinK, a novel query-dependent KV cache pruning method designed to minimize attention weight loss while selectively pruning the least significant channels. Our approach not only maintains or enhances model accuracy but also achieves a reduction in memory costs by over 20% compared with vanilla KV cache eviction methods. Extensive evaluations on the LLaMA3 and Mistral models across various long-sequence datasets confirm the efficacy of ThinK, setting a new precedent for efficient LLM deployment without compromising performance. We also outline the potential of extending our method to value cache pruning, demonstrating ThinK's versatility and broad applicability in reducing both memory and computational overheads.