TreeKV: Smooth Key-Value Cache Compression with Tree Structures
This addresses the challenge of scaling LLMs in resource-limited settings by improving cache compression, though it is incremental as it builds on existing compression methods.
The paper tackles the problem of efficiently compressing key-value (KV) cache in transformer-based Large Language Models for long sequences by proposing TreeKV, a training-free method using tree structures, which achieves a 16x cache reduction while maintaining high performance on benchmarks like PG19 and OpenWebText2 and outperforms baselines on Longbench with only 6% of the budget.
Efficient key-value (KV) cache compression is critical for scaling transformer-based Large Language Models (LLMs) in long sequences and resource-limited settings. Existing methods evict tokens based on their positions or importance scores, but position-based strategies can miss crucial information outside predefined regions, while those relying on global importance scores resulting in strong regional biases, limiting the KV cache's overall context retention and potentially impairing the performance of LLMs on complex tasks. Our wavelet analysis reveals that as tokens approach the end of sequence, their contributions to generation gradually increase and tends to diverge more from neighboring tokens, indicating a smooth transition with increasing complexity and variability from distant to nearby context. Motivated by this observation, we propose TreeKV, an intuitive, training-free method that employs a tree structure for smooth cache compression. TreeKV maintains a fixed cache size, allowing LLMs to deliver high-quality output even in long text scenarios. Unlike most compression methods, TreeKV is applicable to both the generation and prefilling stages. TreeKV consistently surpasses all baseline models in language modeling tasks on PG19 and OpenWebText2, allowing LLMs trained with short context window to generalize to longer window with a 16x cache reduction. On the Longbench benchmark, TreeKV achieves the best performance with only 6\% of the budget at optimal efficiency.