Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

Muhammad Adnan; Akhil Arunkumar; Gaurav Jain; Prashant J. Nair; Ilya Soloveychik; Purushotham Kamath

Source-linked AI summary

Keyformer: KV Cache Reduction through Key Tokens Selection for Efficient Generative Inference

Muhammad Adnan, Akhil Arunkumar, Gaurav Jain, Prashant J. Nair, Ilya Soloveychik, Purushotham Kamath

arXiv:2403.09054v2Mar 14, 2024cs.LGcs.AIcs.ARcs.CL

TL;DR

Long-context LLM generation is limited by KV-cache memory traffic and growing cache size. Keyformer dynamically retains important key tokens alongside recent tokens during inference, reducing the cache without fine-tuning. Across multiple models and long-context tasks, it preserves accuracy while improving latency and token-generation throughput.

Problem
Growing KV caches create memory-bandwidth and latency challenges for long-context LLM inference, while many reduction methods overlook cache expansion or require retraining.
Method
Keyformer dynamically reduces the inference-time KV cache by selecting essential key tokens with a novel score function and discarding unnecessary tokens.
Results
Keyformer preserves accuracy with a 50% KV-cache reduction, reducing inference latency by 2.1× and increasing token-generation throughput by 2.4× across multiple models and tasks.
Takeaways & Limitations
Keyformer provides an inference-time KV-cache reduction strategy that applies across GPT-J, Cerebras-GPT, and MPT without fine-tuning.
Takeaways & Limitations
Reducing attention to window or key tokens alone lowers accuracy, motivating Keyformer’s mixed strategy; score functions must also account for discarded tokens.

Abstract

from arXiv · show

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.

1 INTRODUCTION

Long-context LLM inference is constrained by KV-cache growth and memory movement, motivating Keyformer’s inference-time selection of important tokens. Keyformer combines recent tokens with selected key tokens to reduce cache demands while preserving accuracy.

Motivation: Longer sequences intensify LLM inference latency and throughput problems because autoregressive generation repeatedly accesses growing KV caches.The KV cache stores previously computed key-value pairs to avoid recomputation, but accessing it from off-chip memory is bandwidth-limited.
Motivation: More than 50× higher inference latency results when MPT-7B sequence length increases 16× from 512 to 8K.Approximately 40% of total inference time is consumed by KV-cache data movement, and the cache exceeds model size beyond 8K tokens.
Motivation: Existing system and attention techniques often overlook KV-cache expansion or require resource-intensive retraining, creating demand for accuracy-preserving inference-time reduction.The paper frames KV-cache reduction as especially constrained by the need to maintain model accuracy.
Keyformer: Keyformer dynamically discards unnecessary tokens from the KV cache during inference rather than retaining all preceding tokens.The method exploits attention sparsity by replacing a sequence of n cached tokens with a subset of k tokens, where k < n.
Keyformer: Keyformer mixes recent-window attention with preceding key tokens, unlike full attention and fixed or dilated windows.Window-based methods reduce cache size but can lose comprehensive semantic information from the past, while Keyformer combines recent and selected context.
Results: Nearly 90% of attention weight focuses on key tokens, and Keyformer preserves accuracy while reducing KV cache by 50%.Across GPT-J, Cerebras-GPT, and MPT, evaluations include long-context summarization and conversation tasks.

2 BACKGROUND AND MOTIVATION

Generative inference relies on a growing KV cache, whose memory movement and attention sparsity create both scalability challenges and opportunities for cache reduction. The section shows why retaining only accumulated-attention key tokens is insufficient and motivates combining key tokens with recent tokens while addressing score-distribution distortion.

2.2 Reducing KV Cache Size by Exploiting Sparsity: KV cache stores keys and values for prior tokens, growing from n tokens to a retained subset of k tokens where k < n.The reduction exploits inherent sparsity in the attention mechanism.
2.3.1 Leveraging Score Function to Identify Key Tokens: A small subset of tokens receives most attention during generation, motivating a score function fθ to identify k key tokens among n tokens.Accumulated attention scores are proposed as the score function for selecting consistently important tokens.
2.3.1 Leveraging Score Function to Identify Key Tokens: Retaining only key tokens or only recent window tokens reduces accuracy relative to full attention when each is reduced to n/2 tokens.Key-token attention loses recent context, whereas window attention loses crucial context; the proposed response is a mixed approach combining selected key tokens with recent tokens.
2.3.2 Problem: Uneven Score Distribution: Removing KV-cache tokens makes attention scores uneven among remaining tokens because discarded attention weight is redistributed through the softmax function.This changes the score function used to identify key tokens.
2.3.3 Motivation: Damping the Score Function: 50% KV-cache reduction leaves key tokens with approximately 90%–95% of average accumulated attention, but damping still fails to match full-attention quality.The Cerebras-GPT-6.7B experiment uses a 20% recent-token ratio; the remaining score-distribution change makes accumulated attention alone inadequate.

3 KEYFORMER: INTUITION AND DESIGN

Keyformer identifies and retains key tokens by combining recent tokens with score-selected tokens, while Gumbel-based regularization and temperature scheduling compensate for discarded context. Its accumulated score function supports consistent token selection across decoding steps while maintaining a fixed KV cache budget.

Key-token selection: Keyformer applies regularization to unnormalized logits to adjust score changes caused by discarded tokens when identifying key tokens.The method leverages decoder-layer sparsity and combines key-token selection with recent-token retention.
Gumbel regularization: Gumbel-based logit adjustment produces a more uniform normalized distribution, supporting robust key-token identification when discarded-token information is unavailable.The paper measures distribution spread with entropy and reports H (E[zGumbel]) > H (E[z]).
Keyformer score function: Keyformer’s score function integrates Gumbel noise and introduces temperature τ to account for discarded tokens in the probability distribution.This score function provides a continuous relaxation aligned with selecting a subset of tokens that preserves the semantic information of the full sequence.
Temperature scheduling: Higher τ values produce more uniform probabilities, whereas lower τ values sharpen selection toward tokens with larger unnormalized logits.Removed tokens cannot be reintroduced without recomputing their keys and values, making this control important during cache reduction.
Temperature scheduling: Keyformer increases τ from τinit to τend across decoding steps as more tokens are discarded, using τinit = 1 and τend = 2 based on empirical analysis.The schedule keeps prompt-phase scoring close to standard softmax while increasing randomization during later cache reduction.
Iterative cache reduction: The accumulated score function selects tokens according to consistent behavior across prompt processing and token generation rather than only the current token’s correlations.The cache retains a recent window of w tokens and selects k − w additional tokens from the remaining context to maintain k tokens.

4 EVALUATION

Keyformer is evaluated across diverse model families, tasks, KV-cache budgets, and inference metrics, generally preserving accuracy while reducing cache-related cost.

Evaluation setup: Keyformer was evaluated on GPT-J, Cerebras-GPT, and MPT, which use RoPE, learnable, and ALiBi positional embeddings, respectively.
Evaluation setup: The evaluation covered summarization, conversation, long-context summarization, and four few-shot lm-eval-harness tasks.
Accuracy results: Keyformer reached target ROUGE with 70% of the KV cache, while H2O missed the target with a larger cache budget.
Accuracy results: Keyformer maintained the 99% accuracy threshold with 50% KV cache on MPT-7B-storywriter long-context summarization, whereas H2O performed substantially worse.
Performance results: 2.1× lower inference latency followed a 50% KV-cache reduction, while throughput improved 2× at the same batch size and 2.4× at a larger batch size.
Performance results: KV-cache data movement decreased 2.9× and scaled-dot-product efficiency improved 1.3× for 4k-token sequences, including Gumbel-softmax overhead.
Additional evaluations: Keyformer remained close to full-attention accuracy across four few-shot tasks and both 0-shot and 5-shot settings with 50% cache reduction.
Ablations and comparisons: Per-layer scoring improved accuracy, original positional information outperformed new positional information, and Keyformer still exceeded H2O with the latter.

5 RELATED WORK

Related work reduces attention cost through architectural, sparse-attention, KV-cache, and model-training strategies, but these approaches differ in scope and requirements.

Attention speedup: Attention-speedup methods reduce transformer computation through word-vector elimination, linearization, locality-sensitive hashing, or accumulated states.
Sparse attention: Sparse-attention methods combine windowed, random, global, or task-based local patterns to reduce attention computation and memory.
KV-cache reduction: KV-cache reduction methods reuse stable vectors, retain heavy hitters, or remove heads and tokens, with some requiring task-specific predictors or fine-tuning.
Model architecture: MQA and GQA train foundation models with fewer attention heads, whereas Keyformer targets token redundancy during inference.

6 FUTURE WORK

Future work will evaluate Keyformer more fully with MQA and GQA models and integrate its sparsity into training-time attention.

Future evaluation: Detailed evaluation with MQA- and GQA-based models is deferred to future work, although Keyformer can be applied on top of them.
Training-time integration: Replacing standard softmax with Keyformer-based softmax during training is proposed to introduce sparsity and address transformer scaling challenges.

7 CONCLUSION

Keyformer reduces KV-cache size by identifying and retaining essential tokens during inference, improving latency and throughput without sacrificing accuracy.

Conclusion: Keyformer reduces KV-cache size up to 50% by discarding tokens across heads, layers, and beams according to a novel score function.
Conclusion: The inference-time method requires no fine-tuning while improving latency and token-generation throughput without sacrificing accuracy.

A.1 Qualitative Comparison of Text Generation

Keyformer produces more faithful summaries than window-based and H2O baselines at the same 50% KV-cache budget. Its output preserves more of the reference summary’s content while achieving higher ROUGE scores than H2O in the example.

Qualitative comparison: At 50% KV cache, Window Attention repeats details and omits much of the source’s transfer context.Its output focuses on Chelsea and repeated club-interest statements rather than the fuller source narrative.
Qualitative comparison: At 50% KV cache, Keyformer’s summary retains the transfer-scramble, £25million, and interested-clubs details from the source.The generated summary names Chelsea, Barcelona, and Real Madrid as interested in Costa’s signature.
Quantitative comparison: Keyformer records ROUGE-1 46.1538, ROUGE-2 28.3688, and ROUGE-L 33.5664 in the example.These scores exceed H2O’s ROUGE-1 44.6043, ROUGE-2 24.8175, and ROUGE-L 33.0935.

A.2 Extended Related Work

Related work addresses memory efficiency through KV-cache compression, quantization, and popular-embedding placement, while Keyformer studies sparsity as threshold-dependent token selection. Increasing the threshold percentage increases sparsity.

KV-cache compression and quantization: KV-cache compression and quantization target the growing cache size of autoregressive LLMs and can accelerate inference.FastGen preserves special tokens through adaptive KV-cache compression.
Popular embeddings: Popular or key embeddings have also been studied for intelligent placement across heterogeneous memory hierarchies.This line of work is described in recommendation and other large models rather than as Keyformer’s inference method.
Threshold-based sparsity: Increasing the percentage threshold increases sparsity in the evaluated attention patterns.Figure 11 reports the resulting increase in sparsity as the threshold percentage rises.

A.4 Recent Window versus Key Token Window Ratio

A sensitivity study at a fixed 70% KV-cache size varies the recent-token ratio and reallocates cache capacity between recent and key tokens. Model accuracy is best when the recent-token ratio is 20%–30%.

Sensitivity study: At 70% KV cache, varying recent-token ratio w changes the number of retained key tokens k −w.The study examines how cache capacity is divided between recent tokens and key tokens.
Accuracy trend: Model accuracy is highest when the recent-token ratio w falls between 20% and 30%.This trend is reported for all three evaluated models.

A.5 ROUGE-1 and ROUGE-L Scores

The section evaluates ROUGE-based summarization quality, attention sparsity, positional-encoding effects, temperature sensitivity, and the experimental setup. It frames reduced-cache evaluation against full-attention quality targets while examining why sparsity patterns differ across models.

ROUGE evaluation: MLPerf requires ROUGE-1, ROUGE-2, and ROUGE-L to remain between 99% and 99.9% of original summarization scores.Figure 13 compares ROUGE-1 and ROUGE-L for three models on CNN/DailyMail.
Attention sparsity: GPT-J attention sparsity is dispersed across layers and heads without a distinct pattern, complicating direct exploitation.The heat map uses 360 context tokens and produces a 65-token summary.
Positional encoding: GPT-J and MPT heat maps vary across layers and heads, mainly reflecting differences in positional encoding schemes.The MPT heat map also shows the effect of ALiBi across heads.
Attention-sink comparison: The evaluated models do not show the initial-token attention-sink trend associated with retaining the first four tokens.The discussion contrasts these observations with StreamingLLM and notes MPT’s ALiBi setup.
Temperature sensitivity: Sweeping MPT-7B’s temperature τ from 1 to 2 during token generation works better than using a fixed temperature across both phases.The ablation studies Keyformer on CNN/DailyMail summarization.
Evaluation setup: The evaluation uses validation sets from CNN/DailyMail, GovReport, and SODA across summarization and conversation tasks.GovReport contains approximately 19.5k reports with longer documents and summaries; SODA contains diverse dialogues.