12. Information-Theoretic Token Eviction and KV-Cache Compression for Memory-Bounded Long-Context LLM Inference

Executive Summary Working Title: Value-Norm and Attention-Dispersion Bounded KV-Cache Eviction for Long-Context LLMs

In long-context Large Language Model inference ($16\text{K}\text{–}128\text{K}$ tokens), the Key-Value (KV) cache quickly surpasses the model weights themselves as the primary GPU VRAM and memory-bandwidth bottleneck. Training-free token eviction methods such as Heavy-Hitter Oracle ($\text{H}_2\text{O}$, Zhang et al., NeurIPS 2023) and SnapKV (Li et al., 2024) retain only a small budget $B \ll T$ of cached tokens by keeping "Heavy Hitters" with the largest cumulative attention scores $\sum_{t} \alpha_{t,j}$. However, when I analyzed why $\text{H}_2\text{O}$ fails on multi-needle retrieval and code summarization at tight budgets ($10\text{–}20\%$ KV capacity), I found three mathematical defects: (1) Exposure Bias: token $j=10$ has been attended to $(T - 10)$ times whereas token $j = T - 20$ has only had $20$ opportunities to accumulate attention; (2) Value-Norm Blindness: a token with attention $\alpha_{t,j} = 0.05$ and value norm $\|\mathbf{v}_j\|_2 = 10.0$ contributes $5\times$ more to the output $\sum_j \alpha_{t,j}\mathbf{v}_j$ than a sink token with $\alpha_{t,j} = 0.20$ and $\|\mathbf{v}_j\|_2 = 0.5$; and (3) Redundancy Splitting: two semantically identical tokens split their softmax probability in half and risk both being evicted. This project develops Exposure-Corrected Value-Aware KV Eviction & Merging with provable single-head output perturbation bounds.

Tags: #KV-Cache Compression #Token Eviction #Long-Context LLMs #H2O / SnapKV #Information Theory #VRAM Optimization

Research Motivation & Feasibility

Valuable ML-systems direction for long-context LLM serving on memory-bounded GPUs. A wave of 2024–2026 papers (H2O, SnapKV, PyramidKV, DuoAttention) makes a sharp, defensible theoretical contribution essential.

Core Identified Research Gap

Leading training-free KV-cache eviction methods (H2O, Zhang et al., NeurIPS 2023; SnapKV, Li et al., 2024) evict tokens based solely on cumulative attention weights A_ij summed over past queries. This suffers from three mathematical flaws: (1) early tokens accumulate higher sums purely by being present for more steps (position bias), (2) it ignores the L2 norm and directional uniqueness of the Value vector ||v_j||_2, and (3) redundant duplicate tokens split attention mass and both get evicted!

Basic Information

Field	Details
Core Focus	High-Impact ML Systems & Long-Context LLM Topic
Final Research Title	Information-Theoretic Token Eviction and KV-Cache Compression for Memory-Bounded Long-Context LLM Inference
Working Title	Bounding Attention Perturbation via Value-Norm Exposure Scoring and Redundant KV Merging
Research Area	Machine Learning Systems, Efficient LLM Inference & Linear Algebra
Subfield	KV-Cache Compression & Long-Context Serving
Target Academic Level	MSc Thesis / PhD Track (100% Training-Free Evaluation!)
Current Status	Mathematical Perturbation Bound Formulated

Idea Summary

2.1 One-Sentence Core Thesis

Bound the $L_2$ perturbation error of KV-cache eviction by replacing raw cumulative attention sums with exposure-normalized attention $\times$ value norm, while merging evicted tokens into their nearest retained key neighbor with log-count softmax compensation so total attention mass is conserved.

2.2 Research Motivation & Strategic Fit

Huge Practical Advantage: Because KV-cache eviction is a training-free inference algorithm, you do not need to train or fine-tune any LLM! You can implement the custom KV-cache manager in PyTorch and benchmark Llama-3.1-8B or Qwen2.5-7B on LongBench in a single afternoon on one 24GB GPU.
Why #12: This is a fiercely competitive topic (with dozens of papers in ICML/NeurIPS/ICLR 2024–2026). To stand out, the paper must lead with a clean mathematical perturbation theorem and show decisive gains on Multi-Needle RULER and LongBench at extreme compression ($10\text{–}15\%$ KV budget).

Problem Exploration & Identified Research Gap

3.1 Why Cumulative Attention ($\text{H}_2\text{O}$) Evicts the Wrong Tokens

Consider the exact attention output at decoding step $t$:

$$ \mathbf{o}_t = \sum_{j=1}^{t} \alpha_{t,j}\, \mathbf{v}_j, \qquad \alpha_{t,j} = \frac{\exp(\mathbf{q}_t^\top \mathbf{k}_j / \sqrt{d})}{\sum_{u=1}^{t} \exp(\mathbf{q}_t^\top \mathbf{k}_u / \sqrt{d})} $$

If we evict index $j$ from the cache, the single-step contribution lost is proportional to $\|\alpha_{t,j} \mathbf{v}_j\|_2 = \alpha_{t,j} \|\mathbf{v}_j\|_2$, not $\alpha_{t,j}$ alone! Empirically (see Attention Sink literature, Xiao et al., ICLR 2024), the initial sink tokens $j \in \{0, 1, 2, 3\}$ receive massive attention $\alpha_{t,j}$ precisely because the model learns to assign them near-zero value norms $\|\mathbf{v}_j\|_2 \approx 0$ as a no-op bias absorber. Furthermore, if we simply drop evicted indices $E$, the denominator of the softmax shrinks, artificially inflating the weights of all remaining tokens and shifting the output distribution!

3.2 Explicit Research Gap

Most KV eviction policies ($\text{H}_2\text{O}$, SnapKV, PyramidKV) ignore $\|\mathbf{v}_j\|_2$ during importance scoring, suffer from step-count exposure bias during long autoregressive decoding, and discard evicted KV pairs completely rather than conserving softmax partition mass and value subspaces via weighted token merging (CAM / D2O style).

Proposed Research Direction

4.1 Mathematical Formulation: Perturbation-Bounded Eviction + Partition-Conserving Merge

1
Exposure-Corrected Value-Impact Score: Over an observation window of recent queries, each cached token $j$ is scored by its expected output norm contribution corrected for step exposure $(t - j + 1)^\gamma$:

$$ I(j) = \left( \frac{1}{(t - j + 1)^\gamma} \sum_{\tau = \max(j, t-W)}^{t} \alpha_{\tau, j} \right) \cdot \|\mathbf{v}_j\|_2 $$

1
Partition-Conserving Key-Value Merge: When token $e \in E$ is selected for eviction and has high cosine similarity $\text{sim}(\mathbf{k}_e, \mathbf{k}_{r^\star}) > \tau_{\text{merge}}$ with a retained token $r^\star \in S$, simply dropping $e$ biases the softmax denominator. Instead, we merge $e$ into $r^\star$ and add a multiplicity log-bias $m_{r^\star} \leftarrow m_{r^\star} + 1$ so future attention preserves the exact partition function:

$$ \tilde{\alpha}_{t, r} = \frac{\exp\big(\mathbf{q}_t^\top \tilde{\mathbf{k}}_r / \sqrt{d} + \log m_r\big)}{\sum_{u \in S} \exp\big(\mathbf{q}_t^\top \tilde{\mathbf{k}}_u / \sqrt{d} + \log m_u\big)} $$

Interactive Diagram SpecificationMermaid Flow / Schedule

flowchart LR
    KV["Full KV Cache (K, V) at Limit B"] --> SCORE["Compute Exposure-Corrected<br/>Value-Impact I(j) = Ā_j · ||v_j||_2"]
    SCORE --> SPLIT["Partition into Retained S (Top B) & Evicted E"]
    SPLIT --> MERGE["Cosine-Matched Evicted Tokens e ∈ E<br/>Merge into Nearest Retained Key r* ∈ S"]
    MERGE --> BIAS["Update Multiplicity Log-Bias +log(m_r)<br/>Conserves Softmax Denominator & Output Norm"]

Preliminary Technical Plan

Models: Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct (using Grouped-Query Attention GQA).
Benchmarks:

1
LongBench (HotpotQA, 2WikiMultihopQA, GovReport, LCC, RepoBench-P).
2
Multi-Needle RULER ($4\text{K}, 8\text{K}, 16\text{K}, 32\text{K}$ context lengths) at KV cache budgets $B \in \{10\%, 20\%, 35\%\}$.

Baselines: Full KV Cache, StreamingLLM (Sink + Sliding Window), $\text{H}_2\text{O}$ (NeurIPS 2023), SnapKV (2024), and PyramidKV (2024).
Metrics: Task Accuracy / F1 / Rouge-L vs. KV Cache Memory (MB), Decode Throughput (tokens/sec), and Layer-wise Output L2 Reconstruction Error $\|\mathbf{o}_{\text{full}} - \mathbf{o}_{\text{comp}}\|_2$.

Expected Contributions

1
Theoretical Upper Bound on Attention Eviction Error: Proving that value-norm weighting + multiplicity log-bias minimizes worst-case single-layer output drift under bounded key clustering radius.
2
Training-Free Plug-and-Play KV Manager: Achieves near-lossless LongBench accuracy at $20\%$ KV budget while recovering $15\text{–}25\%$ lost accuracy on Multi-Needle retrieval compared to $\text{H}_2\text{O}$.

Feasibility, Scope, and Honest Risks

Feasibility: $100\%$ training-free! All experiments are forward-pass inference runs on pretrained open-weight models.
Main Risk & Mitigation: In FlashAttention-2, explicit attention weights $\alpha_{t,j}$ are not materialized over the full sequence during prefill. Mitigation: Follow the SnapKV pattern of computing exact attention scores only over a compact local observation window $W=32$ at the end of the prompt (which costs $<1\%$ overhead and remains compatible with FlashAttention).

Initial Literature Notes

[1]
Z. Zhang et al., "$\text{H}_2\text{O}$: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models," NeurIPS, 2023.
[2]
Y. Li et al., "SnapKV: LLM Knows What You are Looking for Before Generation," NeurIPS, 2024.
[3]
G. Xiao et al., "Efficient Streaming Language Models with Attention Sinks (StreamingLLM)," ICLR, 2024.
[4]
Z. Cai et al., "PyramidKV: Dynamic KV Cache Compression based on Pyramidal Information Funneling," arXiv:2406.02069, 2024.

Idea Validation Checklist

Identified 3 concrete mathematical flaws in $\text{H}_2\text{O}$ (exposure bias, value-norm blindness, softmax denominator drift)Verified
Compatible with FlashAttention via $W=32$ observation windowVerified
Zero training cost—pure inference evaluation on 1 GPUVerified
Plot $\|\mathbf{v}_j\|_2$ vs. $\alpha_{t,j}$ across layers of Llama-3.1-8B to verify the inverse sink-norm phenomenonPending

Current Assessment

Criterion	Score	My Assessment Notes
Problem clarity	9 / 10	KV-cache VRAM bottleneck is crystal clear.
Novelty potential	8 / 10	Crowded 2025–2026 field, so our perturbation bound + log-multiplicity merge is key.
Technical feasibility	9 / 10	Training-free; only modifies `Cache` class in HuggingFace `transformers`.
Resource feasibility	9 / 10	16K–32K inference in FP16/BF16 fits on 1x 24GB GPU.
Dataset availability	10 / 10	LongBench and RULER are standard open suites.
Publication potential	8 / 10	Good target for MLSys, EMNLP, ACL, ICML, or NeurIPS.
Personal interest	8 / 10	Strong linear algebra and systems flavor.

Next Actions

1
Subclass transformers.cache_utils.DynamicCache to log $\|\mathbf{v}_j\|_2$ and attention scores on 50 LongBench examples.
2
Compare output reconstruction MSE $\|\mathbf{o}_{\text{full}} - \mathbf{o}_{\text{evict}}\|_2^2$ between $\text{H}_2\text{O}$, Value-Weighted $\text{H}_2\text{O}$, and Multiplicity-Merged Eviction.

Open Questions and Working Notes

Under Grouped-Query Attention (GQA, where 4 or 8 query heads share 1 key-value head), should importance scores $I(j)$ be aggregated across the group using $\max_{h \in G}$ (preserving head-specific needles) rather than mean $\frac{1}{|G|}\sum_{h \in G}$?