KV 캐시 압축을 위한 간단하고 효과적인 L_2 노름 기반 전략

A Simple and Effective L_2 Norm-Based Strategy for KV Cache Compression

June 17, 2024

저자: Alessio Devoto, Yu Zhao, Simone Scardapane, Pasquale Minervini

cs.AI

초록

대규모 언어 모델(LLM)의 배포는 특히 컨텍스트 길이가 증가함에 따라 키-값(KV) 캐시의 방대한 메모리 요구 사항으로 인해 종종 제약을 받습니다. 기존의 KV 캐시 크기를 줄이기 위한 접근 방식은 모델을 미세 조정하여 압축 전략을 학습시키거나 어텐션 점수를 활용하여 시퀀스 길이를 줄이는 방법이었습니다. 우리는 디코더 전용 트랜스포머 기반 모델에서의 어텐션 분포를 분석하고, 대부분의 레이어에서 어텐션 할당 패턴이 일관되게 유지된다는 것을 관찰했습니다. 놀랍게도, 캐시된 KV 쌍에 대한 L_2와 어텐션 점수 사이에 명확한 상관관계가 있음을 발견했는데, 여기서 키 임베딩의 낮은 L_2는 일반적으로 디코딩 중 높은 어텐션 점수로 이어집니다. 이 발견은 KV 쌍의 영향력이 쿼리되기 전에 키 임베딩 자체에 의해 결정될 가능성이 있음을 시사합니다. 이러한 관찰을 바탕으로, 우리는 키 임베딩의 L_2를 기반으로 KV 캐시를 압축합니다. 실험 결과, 이 간단한 전략은 언어 모델링 및 건초 더미 속 바늘 찾기 작업에서 KV 캐시 크기를 50% 줄이고, 패스키 검색 작업에서는 90% 줄이면서도 정확도를 잃지 않음을 보여줍니다.

English

The deployment of large language models (LLMs) is often hindered by the extensive memory requirements of the Key-Value (KV) cache, especially as context lengths increase. Existing approaches to reduce the KV cache size involve either fine-tuning the model to learn a compression strategy or leveraging attention scores to reduce the sequence length. We analyse the attention distributions in decoder-only Transformers-based models and observe that attention allocation patterns stay consistent across most layers. Surprisingly, we find a clear correlation between the L_2 and the attention scores over cached KV pairs, where a low L_2 of a key embedding usually leads to a high attention score during decoding. This finding indicates that the influence of a KV pair is potentially determined by the key embedding itself before being queried. Based on this observation, we compress the KV cache based on the L_2 of key embeddings. Our experimental results show that this simple strategy can reduce the KV cache size by 50% on language modelling and needle-in-a-haystack tasks and 90% on passkey retrieval tasks without losing accuracy.