DynamicKV: Task-Aware Adaptive KV Cache Compression for Long Context LLMs

2024-12-19

arXiv.org

AI Summary

Existing KV cache compression strategies for long-context large language models are rigid and neglect task-specific and layer-wise heterogeneity. Method: This paper proposes a dynamic hierarchical adaptive compression mechanism featuring (1) task-aware dynamic budget allocation, enabling per-layer, online adjustment of retained token counts; and (2) lightweight cache reconfiguration guided by inter- and intra-layer activation pattern analysis, under global and layer-specific budget constraints. The method requires no fine-tuning and performs periodic optimization during inference. Contribution/Results: With only 1.7% of the original KV cache retained, our approach achieves 85% of full-cache LongBench performance; under extreme compression (0.9% cache), it surpasses state-of-the-art methods by 11% accuracy on the Needle-in-a-Haystack task. The method significantly improves the trade-off between inference efficiency and accuracy for long-context processing.

Technology Category

Data Mining & Knowledge Management: Data CompressionMachine Learning: Learning on the Edge & Model CompressionNatural Language Processing: (Large) Language Models

Application Category

Economics, Online Markets and Human Computation: Cost models of using LLMs in production systemsSemantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMsSearch and Retrieval-Augmented AI: Large language models for search

Abstract

Efficient KV cache management in LLMs is crucial for long-context tasks like RAG and summarization. Existing KV cache compression methods enforce a fixed pattern, neglecting task-specific characteristics and reducing the retention of essential information. However, we observe distinct activation patterns across layers in various tasks, highlighting the need for adaptive strategies tailored to each task's unique demands. Based on this insight, we propose DynamicKV, a method that dynamically optimizes token retention by adjusting the number of tokens retained at each layer to adapt to the specific task. DynamicKV establishes global and per-layer maximum KV cache budgets, temporarily retaining the maximum budget for the current layer, and periodically updating the KV cache sizes of all preceding layers during inference. Our method retains only 1.7% of the KV cache size while achieving ~85% of the Full KV cache performance on LongBench. Notably, even under extreme compression (0.9%), DynamicKV surpasses state-of-the-art (SOTA) methods by 11% in the Needle-in-a-Haystack test using Mistral-7B-Instruct-v0.2. The code will be released.

Problem

• Adaptive KV cache compression
• Task-specific token retention optimization
• Efficient long-context LLM performance

Innovation

• DynamicKV adapts token retention per layer
• Global and layer-specific KV cache budgets
• Achieves high performance with minimal cache size