💻 computer science

DepthKV: Layer-Dependent KV Cache Pruning for Long-Context LLM Inference

DepthKV is a novel layer-dependent KV cache pruning framework that optimizes long-context LLM inference by dynamically allocating a fixed global memory budget across layers based on their individual pruning sensitivity, thereby outperforming traditional uniform pruning methods.

Original authors: Zahra Dehghanighobadi, Asja Fischer

Published 2026-04-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Zahra Dehghanighobadi, Asja Fischer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Problem: The "Overstuffed Backpack"

Imagine a Large Language Model (LLM) is a brilliant student trying to write a long story or summarize a massive book. To do this, the student needs to remember everything they've read so far.

In the computer world, this memory is called the KV Cache (Key-Value Cache). Think of the KV Cache as a backpack the student carries. Every time the student reads a new word, they put a note about that word into the backpack.

The Issue: If the student reads a 100-page book, the backpack gets huge. If they read a 1,000-page book, the backpack becomes so heavy and bulky that it breaks the student's back (runs out of computer memory).
The Current Fix: To keep the backpack from getting too heavy, the student throws away some notes. The current method is like a uniform rule: "Throw away 60% of the notes from every page of the book, no matter what."

The Discovery: Not All Pages Are Equal

The authors of this paper realized that the "uniform rule" is a mistake. They discovered that not all parts of the student's brain (or the computer's layers) are equally important.

Imagine the student's brain is a multi-story building with many floors (layers):

Floor 1 (Early layers): These handle basic things like recognizing words and grammar.
Floor 10 (Middle layers): These handle the deep meaning, the plot, and the logic.
Floor 20 (Late layers): These handle the final polish and sentence structure.

The researchers tested what happens if they throw away notes only from specific floors. They found that:

If you throw away notes from the middle floors, the student forgets the story entirely and writes nonsense.
If you throw away notes from the top or bottom floors, the student can still tell the story, maybe with a slightly different accent or style, but the core meaning remains.

The Analogy: It's like trying to save space in a moving truck. The current method is like throwing away 60% of the furniture from the living room, the kitchen, and the bedroom equally. The new discovery is that the kitchen (middle layers) holds the most critical items (the recipe books and the stove). If you throw away 60% of the kitchen, you can't cook. But if you throw away 60% of the bedroom (less critical layers), you can still sleep.

The Solution: DepthKV (The Smart Packing List)

The paper proposes a new system called DepthKV. Instead of treating every floor of the building the same, DepthKV acts like a smart packing manager.

It measures importance: It checks which floors are "sensitive" (critical for the story) and which are "robust" (can handle losing some notes).
It reallocates the budget: It keeps a fixed total size for the backpack, but it distributes the space differently:
- Critical Floors (The Kitchen): Keep almost all the notes. Don't throw anything away here.
- Less Critical Floors (The Bedroom): Aggressively throw away notes here to save space.

How They Measure "Importance"

How does the computer know which floor is the "kitchen"? The researchers used a mathematical test called InfoNCE.

The Analogy: Imagine you shake the backpack.
- If the notes on a specific floor fall out and the student immediately forgets the story, that floor is fragile (high importance).
- If the notes on another floor fall out and the student doesn't even notice, that floor is robust (low importance).
- DepthKV uses this "shake test" to decide where to keep the notes safe.

The Results: Smarter Packing

The researchers tested this on three different types of students (AI models: Gemma, LLaMA, and Qwen) and various tasks (summarizing news, answering questions, solving math).

The Old Way (Uniform Pruning): Threw away notes evenly. The student often got confused or gave short, incomplete answers.
The New Way (DepthKV): Threw away notes strategically.
- Result: The student wrote better summaries, answered questions more accurately, and solved math problems correctly, all while carrying a backpack of the exact same size.

Summary

The paper argues that one size does not fit all. By realizing that different parts of an AI model play different roles, we can be much smarter about what we delete from its memory. DepthKV is the method that keeps the most important memories safe while ruthlessly cutting the fluff, allowing AI to handle longer stories without running out of memory.

1. Problem Statement

Large Language Models (LLMs) face a critical memory bottleneck during long-context inference. While autoregressive generation relies on the Key-Value (KV) cache to avoid recomputing past tokens, the memory footprint of this cache grows linearly with sequence length ( $O(L \cdot N \cdot d)$ , where $L$ is layers, $N$ is sequence length, and $d$ is hidden dimension).

Existing post-training solutions primarily use KV cache pruning, which discards tokens with low attention scores to fit within a fixed memory budget. However, a fundamental limitation of current methods (e.g., H2O, StreamingLLM) is the assumption of uniform pruning: they apply the same pruning ratio across all transformer layers. The authors argue this is suboptimal because it implicitly assumes all layers contribute equally to model performance, ignoring the heterogeneous sensitivity of different layers to token removal.

2. Methodology: DepthKV

The proposed framework, DepthKV, challenges the uniform pruning assumption by introducing layer-dependent KV cache allocation. Instead of discarding a fixed percentage of tokens per layer, DepthKV reallocates a fixed global KV budget across layers based on their specific sensitivity to pruning.

A. Pre-study: Layer Sensitivity Analysis

Before designing the allocation strategy, the authors conducted a layer-wise ablation study:

Method: They pruned tokens in only one specific layer at a time while keeping others intact, measuring the resulting performance drop (ROUGE-1).
Finding: Performance degradation is highly non-uniform. Certain "critical" layers cause massive drops when pruned, while others show negligible impact or even redundancy.
Content Amplification: They identified "content amplification layers." Pruning these layers not only lowers accuracy but significantly suppresses output length and informativeness (measured by YapScore), indicating these layers are crucial for sustaining generation.
Statistical Validation: Permutation tests confirmed that the variation in layer importance is statistically significant across models (Gemma, LLaMA, Qwen) and datasets.

B. Representation Metrics for Importance

To automate the allocation without exhaustive ablation, the authors analyzed representation properties to find a proxy for layer importance.

Metric Selection: They evaluated six metrics (Spectral Entropy, Effective Rank, Curvature, DiME, LiDAR, and InfoNCE).
Key Finding: InfoNCE (a contrastive loss measuring representation robustness to input perturbations) showed the strongest negative correlation with performance degradation.
- Interpretation: Layers with lower InfoNCE values (more robust representations) are more critical. Pruning them causes larger performance drops.
- Usage: DepthKV uses InfoNCE scores to rank layers and determine budget allocation.

C. Allocation Strategies

DepthKV operates under a fixed global pruning ratio ( $\rho$ ) but distributes it non-uniformly. Three strategies are proposed:

Middle-Layer Protection (MLP): A heuristic approach preserving a fixed subset of middle layers (e.g., layers $\lfloor L/2 \rfloor$ and $\lfloor L/2 \rfloor + 1$ ) while pruning others uniformly.
Metric-Guided Allocation (MGA): The primary proposed method. It calculates an importance score $s^{(l)}$ based on InfoNCE. Layers with higher robustness (lower InfoNCE) receive a larger KV budget (lower pruning ratio), while less critical layers are pruned more aggressively.
Middle-Layer Metric Allocation (MLMA): A hybrid approach that protects middle layers structurally while distributing the remaining budget among other layers based on InfoNCE scores.

3. Key Contributions

Theoretical Insight: Demonstrated that transformer layers exhibit statistically significant, non-uniform sensitivity to KV cache pruning, invalidating the uniform pruning assumption.
Discovery of Content Amplification: Identified specific layers whose pruning suppresses content generation, linking representation robustness directly to output quality and length.
Framework Proposal: Introduced DepthKV, a layer-dependent allocation framework that redistributes a fixed memory budget based on layer importance (specifically using InfoNCE).
Empirical Validation: Showed that DepthKV consistently outperforms uniform pruning across diverse tasks and models without requiring additional training.

4. Experimental Results

The authors evaluated DepthKV on three model families (Gemma-7B, LLaMA-3.1-8B, Qwen2.5-7B) across seven datasets:

Tasks: Long-document summarization (arXiv, PubMed, GovReport, LegalCase), Document-grounded QA (HotpotQA, Qasper), and Mathematical Reasoning (GSM- $\infty$ ).
Setup: Fixed global KV reduction of 60% (40% retention).
Performance:
- Summarization: DepthKV (specifically MGA) significantly improved ROUGE-1 and SBERT similarity compared to uniform pruning (e.g., +3.0 ROUGE-1 on arXiv for Gemma).
- QA & Reasoning: DepthKV variants achieved higher Exact Match (EM) and F1 scores. For instance, on Qasper, MLMA-6L outperformed all baselines.
- Robustness: The benefits extended to reasoning tasks (GSM- $\infty$ ), where all DepthKV variants outperformed uniform pruning.
- LLM-as-a-Judge: Evaluations using Prometheus confirmed that DepthKV outputs were more correct, complete, and concise than uniform pruning baselines.
Ablation: The study confirmed that how the budget is distributed (layer-wise allocation) matters more than the specific token importance estimator (e.g., H2O w/ V vs. w/o V).

5. Significance and Limitations

Significance:

Efficiency: DepthKV offers a simple, post-training solution to the long-context memory bottleneck, enabling longer context windows on existing hardware without retraining.
Generalizability: The method is architecture-agnostic and effective across different model families and task types.
Paradigm Shift: It moves the field from "token-level" pruning to "layer-level" resource allocation, acknowledging the functional heterogeneity of transformer layers.

Limitations:

Non-Query-Aware: DepthKV estimates token importance during the prefill stage based on historical attention. It does not dynamically adjust to the specific decoding query, potentially missing tokens that become critical only for specific future queries.
Head-Level Granularity: The current framework allocates budgets at the layer level but aggregates attention across heads. It does not account for intra-layer variability where specific attention heads might be more critical than others.

Conclusion

DepthKV establishes that not all transformer layers are created equal in long-context inference. By reallocating memory budgets to preserve the most sensitive and robust layers (identified via InfoNCE), the method achieves superior performance-per-memory trade-offs compared to uniform pruning, offering a practical path forward for scaling LLM context windows.