TreeKV: Smooth Key-Value Cache Compression with Tree Structures

TL;DR

TreeKV introduces a training-free, tree-structured cache compression for transformer KV caches guided by a wavelet-based frequency analysis of information flow. By using a moving eviction scope and averaged attention-based importance scores, TreeKV achieves a smooth left-to-right retention that reduces memory and computation while preserving long-context capability. It applies to both decoding and prefilling stages, delivering state-of-the-art perplexity on long-text benchmarks (PG19 and OpenWebText2) and strong long-context performance on Longbench, with up to a $16\times$ cache reduction and effective extensions to very long sequences. Overall, TreeKV provides a practical, efficient solution for scalable long-context generation and understanding in LLMs.

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.

Figures (5)

• Figure 1: Distribution map of tokens selected by H2O, TOVA and TreeKV, using a cache size of 128 on a 512-length sequence from PG19. We quantify the values of selected token positions as 1 (dark color) and 0 (light color) otherwise, then average across the 32 heads to reduce noise. Note that sinks and recent tokens are not particularly kept.
• Figure 2: The average magnitude of gathered information's high-frequency components given by wavelet decomposition. We illustrate the top 5 frequency components at the 512th generation step, excluding the first 32 tokens and the last 32 tokens.
• Figure 3: Applying TreeKV cache compression method on a 17-length sequence with cache size 4. Each subplot demonstrates the eviction process and state of the cache at the end of step 4, 5, 6 and 17, respectively, while steps 1-4, showing the cache filling process, are omitted. Variable idx and selection scope are explained in Section \ref{['sec: decoding stage']}.
• Figure 4: We create a 10M length sequence by concatenating the first 50 books from the PG19 test set and applying 4 efficient decoding methods: TOVA, H2O, StreamingLLM, and TreeKV on it. We then plot the negative log-likelihoods against sequence length.
• Figure 5: We plot the log mean perplexity of the first book from PG19, which has a length of 65k. In the TreeKV_Select_Left_Token experiment, we consistently select the left token within the eviction scope instead of the one with the highest attention weight.