[논문 리뷰] Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

Muhammad Adnan, Akhil Arunkumar외 4명|arXiv (Cornell University)|2024. 03. 14.

Algorithms and Data Compression인용 수 6

한 줄 요약

Keyformer는 Gumbel 기반 점수와 로짓 정규화를 사용하여 핵심 토큰과 최근 토큰을 선택해 KV 캐시를 줄이는 추론 시 메서드를 도입합니다. 최대 2.1배 레이턴시 감소와 2.4배 처리량 향상을 KV 캐시의 70%로 달성하되 정확도를 유지합니다.

ABSTRACT

Transformers have emerged as the underpinning architecture for Large Language Models (LLMs). In generative language models, the inference process involves two primary phases: prompt processing and token generation. Token generation, which constitutes the majority of the computational workload, primarily entails vector-matrix multiplications and interactions with the Key-Value (KV) Cache. This phase is constrained by memory bandwidth due to the overhead of transferring weights and KV cache values from the memory system to the computing units. This memory bottleneck becomes particularly pronounced in applications that require long-context and extensive text generation, both of which are increasingly crucial for LLMs. This paper introduces "Keyformer", an innovative inference-time approach, to mitigate the challenges associated with KV cache size and memory bandwidth utilization. Keyformer leverages the observation that approximately 90% of the attention weight in generative inference focuses on a specific subset of tokens, referred to as "key" tokens. Keyformer retains only the key tokens in the KV cache by identifying these crucial tokens using a novel score function. This approach effectively reduces both the KV cache size and memory bandwidth usage without compromising model accuracy. We evaluate Keyformer's performance across three foundational models: GPT-J, Cerebras-GPT, and MPT, which employ various positional embedding algorithms. Our assessment encompasses a variety of tasks, with a particular emphasis on summarization and conversation tasks involving extended contexts. Keyformer's reduction of KV cache reduces inference latency by 2.1x and improves token generation throughput by 2.4x, while preserving the model's accuracy.

연구 동기 및 목표

• 롱 컨텍스트를 가진 자동회귀 LLM 추론에서 KV 캐시의 대기 시간 및 대역폭 병목 현상을 설명/동기 부여합니다.
• 모델 정확도를 유지하면서 KV 캐시를 줄이는 추론 시점 메커니즘을 제안합니다.
• 최근 토큰과 함께 소수의 핵심 토큰 세트를 식별하고 유지하여 축소된 KV 캐시를 구성합니다.
• 디코딩 단계 전반에서 핵심 토큰을 식별하기 위해 Gumbel 기반 정규화를 포함한 학습에서 영감을 얻은 점수 함수를 개발합니다.
• 다른 위치 임베딩을 사용하는 모델 및 긴 컨텍스트를 포함하는 작업에서의 강건성을 입증합니다.

제안 방법

• 대략 90%의 어텐션 가중치가 '키 토큰'이라 불리는 작은 토큰 부분집합에 집중된다는 것을 관찰합니다.
• 최근 윈도우(w)와 핵심 토큰(k-w)을 혼합해 축소된 KV 캐시를 구성하도록 Keyformer를 제안합니다.
• 토큰 제거 후 분포 이동을 완화하고 핵심 토큰을 식별하기 위해 비정규화된 로짓에 Gumbel 기반 정규화를 도입합니다.
• 토큰이 버려짐에 따라 확률 분포를 조정하기 위해 디코딩 단계가 진행됨에 따라 증가하는 온도 스케줄 tau를 사용합니다.
• 프롬프트 단계와 생성 단계 전반에 걸쳐 점수 함수 f_theta를 누적해 일관되게 핵심 토큰을 식별합니다.]
• research_questions:[
• 추론 중 KV 캐시 크기를 줄이면서 생성 품질을 MLPerf 정확도 목표를 넘지 않게 저하시키지 못하는가?
• Keyformer의 핵심 토큰 선택 및 혼합 어텐션이 다양한 위치 인코딩을 가진 모델들 간에 ROUGE 기반 요약 품질과 긴 컨텍스트 성능을 유지하거나 향상시키는가?
• 다양한 KV 캐시 예산 하에서 지연 시간, 처리량 및 정확도 측면에서 Keyformer가 Window Attention 및 H2O와 어떻게 비교되는가?
• KV 캐시에서 토큰이 제거될 때 정확도 유지를 위해 Gumbel 기반 로짓 정규화가 필수적인가?

Figure 1: (a) Inference latency normalized to sequence length of 512. We measure the $\mathsf{KV}$ $\mathsf{cache}$ data movement for MPT-7B Team et al. ( 2023 ) model with varying sequence length (50% context + 50% text generation). (b) The $\mathsf{KV}$ $\mathsf{cache}$ size and model size as sequ

실험 결과

연구 질문

• RQ1
• RQ2 Do Keyformer’s key-token selection and mixed attention maintain or improve ROUGE-based summarization quality and long-context performance across models with different position encodings?
• RQ3 How does Keyformer compare to Window Attention and H2O in terms of latency, throughput, and accuracy under various KV cache budgets?
• RQ4 Is the Gumbel-based logits regularization essential for maintaining accuracy when tokens are discarded from the KV cache?

주요 결과

• Keyformer reduces KV cache by about 50% while preserving ~MLPerf-level accuracy (target ROUGE within 99–99.9% of full attention).
• Keyformer achieves 2.1x inference latency improvement and 2.4x token generation throughput gains at 50% KV cache reduction.
• Keyformer maintains baseline accuracy with only 70% of the KV cache, outperforming H2O which struggles even with larger budgets.
• Across GPT-J, Cerebras-GPT, and MPT with varied position embeddings, Keyformer consistently outperforms Window Attention and H2O in accuracy at reduced KV budgets.
• For long-context summarization (MPT-7B-storywriter on GovReport), Keyformer preserves 99% accuracy at 50% KV cache, unlike H2O.
• The approach improves speed while maintaining quality across tasks including summarization and conversation, demonstrating robustness to context length and model differences.

Figure 2: Attention block for generative inference. (a) Full attention Brown et al. ( 2020 ) with current token attending all previous tokens. (b) Window attention ( $w=4$ ): Focusing on the most recent 4 tokens. (c) Dilated window attention ( $w=4$ , dilation = 1). (d) $\mathsf{Keyformer}$ ( $w=2$