[논문 리뷰] Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time

Zichang Liu, Aditya Desai외 6명|arXiv (Cornell University)|2023. 05. 26.

Topic Modeling인용 수 10

한 줄 요약

이 논문은 LLM 추론을 위한 KV 캐시 압축 방법인 Scissorhands를 제안하며, 중요성 가설의 지속성을 활용해 핵심 토큰만 유지하고 미세 조정 없이 KV 캐시 메모리 사용량을 최대 5배까지 줄입니다. 또한 4비트 양자화와의 호환성을 보여주고 모델 품질을 보존합니다.

ABSTRACT

Large language models(LLMs) have sparked a new wave of exciting AI applications. Hosting these models at scale requires significant memory resources. One crucial memory bottleneck for the deployment stems from the context window. It is commonly recognized that model weights are memory hungry; however, the size of key-value embedding stored during the generation process (KV cache) can easily surpass the model size. The enormous size of the KV cache puts constraints on the inference batch size, which is crucial for high throughput inference workload. Inspired by an interesting observation of the attention scores, we hypothesize the persistence of importance: only pivotal tokens, which had a substantial influence at one step, will significantly influence future generations. Based on our empirical verification and theoretical analysis around this hypothesis, we propose Scissorhands, a system that maintains the memory usage of the KV cache at a fixed budget without finetuning the model. In essence, Scissorhands manages the KV cache by storing the pivotal tokens with a higher probability. We validate that Scissorhands reduces the inference memory usage of the KV cache by up to 5X without compromising model quality. We further demonstrate that Scissorhands can be combined with 4-bit quantization, traditionally used to compress model weights, to achieve up to 20X compression.

연구 동기 및 목표

• LLM 추론에서 KV 캐시 메모리 병목 현상을 동기화하고, KV 캐시 메모리가 모델 크기를 초과하고 배치 처리량을 제한할 수 있음을 지적합니다.
• 중요성 지속성 가설(Persistence of Importance Hypothesis)을 도입하여 시간 단계 간에 여전히 영향력이 있는 토큰을 식별합니다.
• 재훈련 없이 정확성을 보존하는 예산 친화적 KV 캐시 압축 방법을 개발합니다.
• 압축된 어텐션이 원본 어텐션 출력에 근사할 수 있다는 이론적 근거를 제공합니다.
• 추가적으로 압축을 통해 양자화와의 호환성을 확인하여 압축을 더욱 강화합니다.

제안 방법

• 반복적인 어텐션 패턴을 경험적으로 관찰하고 Persistence of Importance Hypothesis를 정의합니다.
• Algorithm 1과 Algorithm 2(Scissorhands)를 개발하여 높은 어텐션 토큰을 유지하고 비영향적 토큰을 버리는 고정 예산 하의 KV 캐시를 유지합니다.
• 히스토리 창(history window)을 사용해 토큰의 중요성을 추정하고 Reservoir sampling과 LRU에서 영감을 받은 예산 기반 KV 캐시 교체 전략을 활용합니다.
• 원본 KV-캐시 출력과 예산화된 KV-캐시 출력 간의 차이를 기대하는 근사 오차의 상한을 보이는 이론적 한계(정리 4.1)를 제공합니다.
• 레이어와 헤드를 모델의 메모리 예산에 분배하고, 후반부 레이어를 우선시하여 지속성 감소를 보완합니다.

실험 결과

연구 질문

• RQ1 Persistence of Importance Hypothesis가 서로 다른 LLM 및 레이어에서도 성립하는가?
• RQ2 재훈련 없이도 고정 예산 하에서 KV 캐시 메모리를 줄이고 생성 품질을 보존할 수 있는가?
• RQ3 예산화된 KV-캐시가 어텐션 출력과 다운스트림 태스크에 미치는 이론적·실험적 영향은 무엇인가?
• RQ4 Scissorhands가 4비트 양자화와 같은 포스트 트레이닝 양자화와 호환되는가?

주요 결과

• Scissorhands는 OPT 모델에서 언어 모델링 및 다운스트림 태스크에서 모델 품질 저하 없이 KV 캐시 메모리를 최대 5×까지 감소시킵니다.
• 지속성 비율은 일반적으로 레이어를 넘어 95%를 초과하는 경향이 있어 초기의 핵심 토큰이 후속 중요도를 크게 좌우합니다.
• 어텐션-스코어 기반 토큰 dropping은 원본 모델과 비슷한 출력을 유지하며, OPT-66B에서 최대 5× 압축 시에도 정확도 손실은 미미합니다.
• 양자화(4-bit)는 Scissorhands와 여전히 호환되어 평가된 태스크에서 성능을 유지합니다.
• 이론적 분석(Theorem 3.1 및 4.1)은 어텐션 가중치, MLP 스킵 연결, 예산 토큰 보존을 근사 동작 및 오차 한계와 연결합니다.