Source-linked AI summary
ZigZagkv: Dynamic KV Cache Compression for Long-context Modeling based on Layer Uncertainty
Meizhi Zhong, Xikai Liu, Chen Zhang, Yikun Lei, Yan Gao, Yao Hu, Kehai Chen, Min Zhang
TL;DR
Long-context inference can make growing KV caches exceed available memory, while uniform per-layer compression may not preserve information efficiently. ZigZagKV analyzes layer-dependent uncertainty and dynamically allocates cache budgets, achieving strong retrieval performance with substantially less cache than FullKV. The method is evaluated on Needle-in-a-Haystack and LongBench and is reported to outperform existing compression methods.
Problem
Growing KV caches create severe memory demands, while existing methods commonly use uniform layer budgets despite layer- and model-dependent information-retention requirements.
Method
ZigZagKV is a training-free KV cache compression method that dynamically allocates each layer’s budget using layer uncertainty.
Results
ZigZagKV achieves 89.33% accuracy with an average budget of 256 on Needle-in-a-Haystack, closely matching FullKV while using a much smaller cache.
Takeaways & Limitations
The results support adaptive, layer-specific budgeting as an effective way to reduce information loss during PartialKV inference across the evaluated benchmarks and models.
Takeaways & Limitations
The study primarily analyzes LLaMA and Mistral and does not validate encoder-decoder or encoder-only architectures.
Abstract
from arXiv · showhide
Large Language models (LLMs) have become a research hotspot. To accelerate the inference of LLMs, storing computed caches in memory has become the standard technique. However, as the inference length increases, growing KV caches might lead to out-of-memory issues. Many existing methods address this issue through KV cache compression, primarily by preserving key tokens throughout all layers to reduce information loss. Most of them allocate a uniform budget size for each layer to retain. However, we observe that the minimum budget sizes needed to retain essential information vary across layers and models based on the perspectives of attention and hidden state output. Building on this observation, this paper proposes a simple yet effective KV cache compression method that leverages layer uncertainty to allocate budget size for each layer. Experimental results show that the proposed method can reduce memory usage of the KV caches to only $\sim$20\% when compared to Full KV inference while achieving nearly lossless performance.
1 Introduction
Long-context inference stores growing KV caches that can cause severe memory pressure, while existing compression methods usually assign uniform budgets across layers. ZigZagKV instead uses layer uncertainty to allocate budgets dynamically, motivated by evidence that information-retention requirements vary across layers and models.
- Over 50GB is required to maintain a 100K-token KV cache for LLaMA-2 7B, versus less than 1GB for a 2K context.
- Existing KV compression methods typically preserve the top-B tokens uniformly in every layer, but the effectiveness of this strategy remains unclear.
- Minimum budgets needed to preserve attention or hidden-state information vary across layers and models.
- ZigZagKV dynamically allocates each layer’s cache budget according to layer uncertainty, assigning more budget to diffuse-attention layers and less to concentrated-attention layers.The method first assigns initial budgets and then adjusts remaining cache based on uncertainty.
- ZigZagKV outperforms existing KV cache compression methods on Needle-in-a-Haystack and LongBench.
2 Problem Formulation
Autoregressive inference stores key-value states for prior tokens, but the resulting cache grows with sequence length and can cause memory pressure. PartialKV inference addresses this by evicting cache entries under a per-head budget before computing attention.
- 2.1 FullKV Inference: FullKV inference caches computed key and value vectors at each autoregressive step for reuse in subsequent attention computations.The cache is formed from key and value states projected from the input embeddings.
- 2.1 FullKV Inference: The KV cache is organized across h attention heads, with head index i identifying separate key and value projection matrices.The notation defines X as input embeddings and W K_i and W V_i as head-specific projections.
- 2.1 FullKV Inference: At each step, the current input is projected into query, key, and value states, which update the previous KV cache before attention weights are computed.The projections use W Q, W K, and W V for each attention head.
- 2.2 PartialKV Inference: FullKV cache size grows linearly with total sequence length, creating potential out-of-memory issues during long inference.This memory-growth problem motivates partialKV inference methods.
- 2.2 PartialKV Inference: PartialKV inference maintains a bounded cache by applying an eviction policy to key-value states under budget size B for each attention head.The retained cache is denoted by the evicted key and value states.
- 2.2 PartialKV Inference: After eviction, the retained key-value cache is used to calculate attention weights for the next-token computation.The eviction step therefore precedes attention calculation in the partialKV pipeline.
- 2.2 PartialKV Inference: Figure 2 presents LMBA across layers of Mistral and LLaMa on the 2WikiMQA dataset.The caption identifies the compared models, layer dimension, and evaluation dataset but does not state a numerical outcome.
3 Rethinking PartialKV Inference
PartialKV methods reduce memory by evicting KV-cache tokens, but uniform top-B retention may not preserve information equally across layers. The required budgets vary by layer and model for both attention retention and hidden-state output stability.
- Motivation: Uniformly retaining the top-B tokens in every layer may fail to optimize information retention because token-removal effects differ across layers.This question is examined from both attention-mechanism and hidden-state-output perspectives.
- Layer-Specific Budget Setting for Attention Retention: LMBA averages each head’s minimum budget for retaining 90% of its attention score, thereby defining the layer budget needed for attention loss of 0.1.A higher LMBA indicates that more tokens are needed in that layer.
- Layer-Specific Budget Setting: Layer-specific minimum budgets vary across models and layers for maintaining attention and hidden-state information near fullKV behavior.The analysis measures attention retention through LMBA and hidden-state output stability through LMBO.
- Layer-Specific Budget Setting for Attention Retention: LMBA varies across layers in Mistral and LLaMA, initially requiring relatively larger lower-layer budgets before decreasing in middle layers.The empirical analysis uses 200 samples from the 2WikiQA dataset.
- Layer-Specific Budget Setting for Hidden State Output Retention: LMBO also varies by layer and model: it rises with depth for Mistral, while LLaMA requires minimal cache initially and more at greater depth.LMBO is the minimum budget needed for at least 90% similarity between partialKV and fullKV hidden-state outputs.
4 ZigZagKV
ZigZagKV replaces uniform layer budgets with dynamic allocation based on layer uncertainty. It assigns more budget to uncertain layers while using a minimum bound to prevent low-uncertainty layers from receiving too little.
- 4.1 Dynamic Budget Allocation Based on Layer Uncertainty: ZigZagKV dynamically allocates KV-cache budget across layers according to uncertainty rather than using one uniform budget.Higher-uncertainty layers receive a larger budget share, while lower-uncertainty layers receive a smaller share.
- 4.1 Dynamic Budget Allocation Based on Layer Uncertainty: A fixed minimum budget bound prevents uncertainty-based allocation from starving low-uncertainty layers and risking inadequate information retention.The remaining budget is distributed dynamically according to layer uncertainty.
- 4.2 KV Cache Selection: After layer budgets are determined, ZigZagKV selects crucial tokens separately for each attention head using cumulative attention scores.The selection budget is calculated by the dynamic allocation rule.
- 4.2 KV Cache Selection: Token importance scores use the prompt sequence length and the final instruction-token segment when selecting KV-cache entries.The last segment is represented as [n −w, n].
5 Experiments
Experiments on Needle-in-a-Haystack and LongBench compare ZigZagKV with established KV-cache compression methods across two open-source LLMs. ZigZagKV improves retrieval and multi-task performance while reducing attention and hidden-state information loss, with bounded budgets improving performance and latency remaining close to PyramidKV.
- Needle-in-a-Haystack: 89.33% accuracy at mean budget 256 closely matches FullKV on Needle-in-a-Haystack, including 30K-context evaluation.ZigZagKV consistently outperforms prior methods under constrained budgets, while FullKV retains the entire cache.
- LongBench: ZigZagKV achieves higher average scores across multiple LongBench tasks than four baseline methods.The comparison covers 16 datasets for each of Mistral-7B-Instruct-v0.3 and LLaMA-3.1-8B-Instruct; Table 1 and Table 2 present the results.
- LongBench: ZigZagKV outperforms FullKV using mean KV-cache size 128 on the TriviaQA few-shot task.The result is reported as evidence that the method reduces memory overhead while retaining information from few-shot examples.
- Information retention: ZigZagKV substantially reduces attention loss and minimizes the attention-score gap from FullKV on Mistral and LLaMA.SnapKV and PyramidKV exhibit higher attention loss, associated with uniform treatment or smaller budgets in higher layers.
- Information retention: ZigZagKV has the lowest mean hidden-state output loss among evaluated methods and models.Output loss is measured as 1−similarity(y, ŷ) between PartialKV and FullKV inference outputs.
- Ablation and overhead: Bounded-budget allocation improves performance across budget sizes on both Mistral and LLaMA, while ZigZagKV latency is similar to PyramidKV.StreamLM is faster but shows a performance drop relative to the other methods.
6 Related Work
Existing KV cache compression methods use fixed or adaptive token-selection policies, but generally assign uniform budgets across layers. ZigZagKV differs by allocating budgets according to layer uncertainty rather than a heuristic layer pattern.
- KV cache compression methods include fixed policies that retain initial and local tokens, and adaptive policies that select tokens using attention weights.
- Most existing methods apply a uniform budget size across layers, despite potentially different optimal budgets.
- Some recent methods allocate budgets across layers but overlook the minimum budget needed to preserve essential information.
- Unlike PyramidKV’s arithmetic allocation from lower to higher layers, ZigZagKV uses layer uncertainty, which can allocate more cache to middle layers when uncertainty peaks there.
7 Conclusion
The paper analyzes how minimum information-preserving budgets vary across layers and uses layer uncertainty to allocate cache dynamically. Across two benchmarks and several models, ZigZagKV reduces information loss during PartialKV inference.
- Minimum budgets needed to retain attention and hidden-state information differ across layers, motivating uncertainty-aware allocation.
- ZigZagKV is a training-free method that dynamically allocates each layer’s budget based on uncertainty to reduce information loss during PartialKV inference.
- Experiments on two benchmarks and several models demonstrate the effectiveness of ZigZagKV.
Limitations
The evaluation primarily covers the LLaMA and Mistral decoder-only language-model families. Encoder-decoder and encoder-only architectures are not validated.
- The paper primarily analyzes two decoder-only language-model families, LLaMA and Mistral.
- Encoder-decoder and encoder-only architectures are outside the paper’s validation study.