[논문 리뷰] Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs

Suyu Ge, Yunan Zhang외 4명|arXiv (Cornell University)|2023. 10. 03.

Topic Modeling인용 수 5

한 줄 요약

FastGen은 경량적 어텐션 프로파일링을 활용하여 LLM의 동적 적응형 KV 캐시 압축을 도입함으로써, 장기적 맥락 제거, 특수 토큰만 유지, 낮은 빈도 토큰 제거 등의 맞춤형 압축 정책을 동적으로 적용함으로써 최대 36%의 메모리 절감을 달성하며, 미세조정이나 재학습 없이도 생성 품질 손실이 거의 없음.

ABSTRACT

In this study, we introduce adaptive KV cache compression, a plug-and-play method that reduces the memory footprint of generative inference for Large Language Models (LLMs). Different from the conventional KV cache that retains key and value vectors for all context tokens, we conduct targeted profiling to discern the intrinsic structure of attention modules. Based on the recognized structure, we then construct the KV cache in an adaptive manner: evicting long-range contexts on attention heads emphasizing local contexts, discarding non-special tokens on attention heads centered on special tokens, and only employing the standard KV cache for attention heads that broadly attend to all tokens. Moreover, with the lightweight attention profiling used to guide the construction of the adaptive KV cache, FastGen can be deployed without resource-intensive fine-tuning or re-training. In our experiments across various asks, FastGen demonstrates substantial reduction on GPU memory consumption with negligible generation quality loss. We will release our code and the compatible CUDA kernel for reproducibility.

연구 동기 및 목표

• 모델의 미세조정이나 재학습 없이도 자동적 생성 LLM 추론에서 KV 캐시의 메모리 사용량을 줄이는 것.
• 목적에 맞는 구조적 패턴을 식별하고, 이를 통해 타겟팅된, 맥락 인식 기반의 KV 캐시 압축을 가능하게 하는 어텐션 모듈의 특성 분석.
• 각 어텐션 헤드의 기능적 역할에 따라 압축 전략을 적응적으로 조정할 수 있는 즉시 사용 가능한 방법 개발.
• 지능적이고 구조 인식 기반의 정제를 통해 높은 메모리 절감을 달성하면서도 생성 품질을 유지하는 것.
• 대규모 언어 모델의 효율적 온디바이스 배포를 가능하게 하여 GPU 메모리 압박을 최소화하는 것.

제안 방법

• 이 방법은 어텐션 맵을 분석하고 각 헤드를 다섯 가지 구조적 유형 중 하나로 분류하기 위해 경량적 어텐션 프로파일링 단계를 활용함.
• 프로파일링 결과에 기반해 FastGen은 다섯 가지의 압축 정책 중 하나를 적용함: 국소 맥락 제거, 특수 토큰 유지, 빈도 기반 정제, 열 기반 희소성, 또는 표준 캐싱.
• 시스템은 빈도, 특수 토큰,标점, 국소 어텐션 패턴을 포함한 정책 세트에 따라 각 헤드당 관련 있는 토큰만 선택하여 적응형 KV 캐시를 구성함.
• 프로파일링 오버헤드는 극히 작음—전체 KV 캐시 메모리의 0.78% 이내—실시간 추론에 실현 가능함.
• 압축 정책를 적용하는 데에는 우선순위가 높은 정책(예: 특수 토큰, 빈번한 토큰)을 먼저 적용하는 탐욕적 정책 순서 전략을 사용함.
• 이 방법은 기존의 LLM 추론 파이프라인과 완전히 호환되며, 모델 재학습이나 아키텍처 변경이 필요 없음.

Figure 1: Different attention heads usually have different structures. Left: Four common attention structures (more details are elaborated in Section 3 and Section 4 ). Right: Attention map compositions of three attention heads that are in the same layer.

실험 결과

연구 질문

• RQ1 LLM의 어텐션 모듈을 효율적이고 타겟팅된 KV 캐시 압축을 가능하게 하는 고유한 구조적 패턴으로 분류할 수 있는가?
• RQ2 생성 품질 저하 없이 각 어텐션 헤드별로 적응형으로 KV 캐시를 압축할 수 있는가?
• RQ3 메모리 절감을 극대화하면서도 모델 성능을 유지하기 위해 최적의 압축 정책 순서와 조합은 무엇인가?
• RQ4 즉시 사용 가능한, 미세조정이 없는 방법으로 KV 캐시에서 상당한 메모리 절감을 달성하고 품질 손실을 거의 느끼지 못할 수 있는가?
• RQ5 빈도 및 국소 어텐션 임계값과 같은 하이퍼파라미터 설정에 따라 압축 성능은 얼마나 민감한가?

주요 결과

• FastGen은 Llama 1-65B에서 KV 캐시 메모리 사용량을 최대 36.04%까지 절감하면서 AlpacaEval에서 49.75%의 승리율 유지를 달성하여 전체 캐시 기준 라인과 거의 동일한 성능 유지를 함.
• 특수 토큰 정책(${\bm{C}}_{\text{special}}$) 제거 시 성능 저하가 가장 심각함—승리율이 2.11% 감소하여 이 정책의 핵심적 역할을 확인함.
• 빈도 기반 정책(${\bm{C}}_{\text{frequent}}$)이 가장 큰 기여를 함—기타 정책과 조합했을 때 캐시 크기를 36.04%까지 줄임.
• 정책 적용 순서가 성능에 큰 영향을 미침—기본 탐욕적 순서(special → punct → frequent → local)가 가장 높은 승리율과 압축 비율 달성.
• 하이퍼파라미터 변화에 대해 강건함—모든 테스트 설정에서 45% 이상의 승리율 유지—압축 비율은 설정에 따라 크게 다름.
• 프로파일링 오버헤드는 극히 미미함—전체 KV 캐시 크기의 0.78% 이내—실시간 배포에 실현 가능함.

Figure 2: Performance of Adaptive KV Cache (FastGen) and Fixed KV Cache (Frequency, Local, and Frequency+Local; Zhang et al. , 2023 and Liu et al. , 2023a ) on AlpacaEval.