TreeKV: Smooth Key-Value Cache Compression with Tree Structures

TL;DR

TreeKV introduces a tree-structured, training-free cache compression method that enhances language model performance on long sequences by maintaining high quality with significantly reduced cache size.

Contribution

It presents a novel, smooth cache compression technique using tree structures, enabling effective long-sequence processing without additional training.

Findings

Outperforms baseline models on PG19 and OpenWebText2.

Enables models trained on short contexts to generalize to longer contexts with 16x cache reduction.

Achieves top performance on Longbench with only 6% of the original cache budget.

Abstract

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.