🤖 AI

SelKV: Selective KV Cache Merging with Per-Token Merge-or-Drop and Attention Compensation

SelKV is a training-free framework that compresses KV caches by selectively merging or dropping tokens based on value-vector similarity and compensating for attention imbalances via logit bias, achieving near-lossless generation quality and significant speedups while retaining only 25% of the original cache.

Original authors: Soumia Bouyahiaoui, Manel Kara laouar, Aicha Boutorh, Mohamed Hadj Ameur

Published 2026-07-21

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Soumia Bouyahiaoui, Manel Kara laouar, Aicha Boutorh, Mohamed Hadj Ameur

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

Imagine you are trying to remember a massive, 100,000-page story while writing a new chapter. Every time you think of a plot point, you have to flip back through your entire notebook to find the right details. The more pages you write, the heavier your notebook gets, until your arms are too tired to hold it. This is exactly what happens inside the "brain" of a Large Language Model (LLM) when it tries to process long texts. These models are incredibly smart, but they have a memory problem: as they read, they store a "Key-Value cache"—a digital filing system of everything they've seen so far—to help them understand what comes next. The longer the text, the bigger this filing cabinet becomes, eventually filling up all the computer's memory and slowing everything to a crawl.

To fix this, scientists have tried two main tricks. The first is "eviction," which is like throwing away old pages you think aren't important. The second is "merging," which is like gluing similar pages together to save space. But both methods have a flaw. Throwing pages away can lose crucial clues, and gluing pages together often creates a "blurry" version where the model forgets how much attention to pay to the glued group. It's like if you glued ten photos of a cat together; the model might still only give that single glued blob the same amount of attention it would have given to one photo, even though it now represents ten. This causes the model to get confused and make mistakes.

Enter SelKV, a new, clever method that acts like a smart librarian for these digital filing cabinets. Instead of blindly gluing pages together or tossing them in the trash, SelKV uses a "soft cosine gate"—think of it as a bouncer at a club who checks how much two pages look alike before deciding what to do. If two pages are very similar, they get glued together tightly. If they are totally different, the bouncer sends one away to keep the memory clean. But the real magic is the "attention compensation." Since gluing pages together usually makes the model ignore them, SelKV adds a little "volume boost" to the glued pages, ensuring the model pays the right amount of attention to them.

The researchers tested this system on three different AI models using 16 different tasks, from answering questions about multiple documents to writing code. They found that by keeping only 25% of the original memory space, SelKV performed almost as well as having the full memory, and in some tricky multi-document quizzes, it actually did better than the full version. It seems that by being selective, the AI actually focused on the most important parts of the story. Furthermore, when the text got huge (100,000 tokens), this method made the AI decode text 3.3 times faster. The paper suggests that this approach is especially good for modern AI models that use a specific type of attention (called GQA), offering a way to keep these powerful brains running fast without losing their memory.

Technical Summary: SelKV - Selective KV Cache Merging with Per-Token Merge-or-Drop and Attention Compensation

1. Problem Statement

Large Language Models (LLMs) face a critical bottleneck during autoregressive inference: the Key-Value (KV) cache. As the context length ( $n$ ) increases, the memory footprint of the KV cache grows linearly ( $O(L \cdot H \cdot d \cdot n)$ ), severely limiting batch sizes, throughput, and deployment costs for long-context applications.

While existing compression methods attempt to mitigate this, they suffer from two primary limitations:

Uniform Merge-or-Drop Decisions: Current approaches typically apply a binary strategy: either merge all evicted tokens (e.g., KVMerger, WeightedKV) or drop them entirely (e.g., H2O, SnapKV). This lack of granularity means that when an evicted token's value vector is dissimilar to its merge target, the weighted average corrupts the retained semantic representation. Conversely, indiscriminate dropping loses valuable information.
Attention Sag: When multiple tokens are merged into a single cache entry, that entry receives the same softmax attention mass as a single unmerged token. This results in "attention sag," where merged positions are systematically under-attended by a factor proportional to the number of merged tokens ( $\sim m$ ), causing the model to underutilize the preserved information.

2. Methodology: SelKV

The authors propose SelKV (Selective KV Cache Merging), a training-free, dual-component framework designed to address these limitations. SelKV operates as a pluggable module integrated into existing merging pipelines after the prefill phase.

Core Components

1. Soft Cosine Gate (Addressing Uniform Decisions)
Instead of a binary merge-or-drop decision, SelKV introduces a continuous, per-token gating mechanism based on value-vector similarity.

Mechanism: For each evicted token routed to a merge target, the system computes the cosine similarity between their value vectors.
Gating Function: A soft gate $g = \max(\cos\_sim, 0)$ $g = max (cos_s im, 0)$ modulates the merge intensity:
- $g \approx 1$ : High similarity $\rightarrow$ Full merge.
- $g = 0$ : Dissimilar (orthogonal) vectors $\rightarrow$ Drop the token.
- $0 < g < 1$ : Moderate similarity $\rightarrow$ Partial merge.
Implementation: This requires no learned parameters or threshold tuning. The merged value for a kept token $i$ is calculated as a gated, attention-weighted average of the original value and the contributions from evicted tokens, scaled by their respective gates.

2. Attention Compensation (Addressing Attention Sag)
To correct the systematic under-attention of merged positions, SelKV applies a decoding-time logit bias.

Mechanism: Unlike previous methods that use raw merge counts (which can be unstable at high compression ratios), SelKV derives an attention-ratio ( $R_{h,i}$ ) from prefill statistics. This ratio estimates the total attention mass a merged position should receive relative to its original state.
Correction: A bias term $\alpha \cdot \log(R_i)$ is added to the attention logits during decoding. This recalibrates the softmax distribution to ensure merged positions receive appropriate attention mass proportional to the information they now represent.

Pipeline Overview

The SelKV pipeline proceeds in six stages post-prefill:

Importance Scoring: Tokens are scored based on the product of attention weight and value vector magnitude (following VATP), rather than attention weight alone.
Token Selection: Top- $m_s$ tokens are selected per layer (and per head for MHA, or via union for GQA) to retain, alongside a fixed recent window.
Merge-Target Routing: Evicted tokens are routed to retained tokens within positional buckets based on prefill attention patterns.
Selective Merging: The soft cosine gate determines the merge intensity for each evicted token.
Attention Compensation: The logit bias is computed and applied.
RoPE Repositioning: Keys are re-embedded at contiguous positions (disabled for long-context evaluation to avoid positional distortion).

3. Key Contributions

Per-Token Granularity: The introduction of a soft cosine gate that enables continuous modulation of merge intensity based on semantic similarity, eliminating the "all-or-nothing" flaw of prior methods.
Stable Attention Compensation: A novel attention-ratio based logit bias that corrects attention sag more robustly than raw count-based methods, particularly under heavy compression.
Training-Free Integration: The framework requires no fine-tuning and can be integrated into existing one-shot compression pipelines.
Architecture Adaptability: The method supports both Multi-Head Attention (MHA) and Grouped-Query Attention (GQA) architectures, with specific handling for token selection unions in GQA models.

4. Experimental Results

The framework was evaluated on LongBench (16 English datasets) across three models: LongChat-7B (MHA), LLaMA-3.1-8B (GQA), and Gemma-2-9B (GQA), retaining only 25% of the KV cache.

Performance vs. Baselines: SelKV consistently outperformed representative one-shot baselines (SnapKV, LOOK-M, PyramidKV).
- On GQA models, SelKV achieved the closest performance to the full-cache baseline, with average gaps of $-0.67$ (Gemma-2) and $-0.73$ (LLaMA-3.1).
- On MHA models (LongChat), PyramidKV performed slightly better in the specific 31.5k token setting, though SelKV remained competitive.
Multi-Document QA: Notably, SelKV exceeded the full-cache baseline on several multi-document QA tasks (e.g., HotpotQA and MuSiQue on LLaMA-3.1). The authors interpret this as evidence that selective merging acts as an implicit attention filter, suppressing distracting context while preserving salient evidence.
Efficiency: At a context length of 100k tokens, the compressed cache delivered a 3.3x decoding speedup compared to the full-cache baseline.
Robustness: The method demonstrated robustness across varying compression ratios (10%–90% retention) and task categories, including summarization, few-shot learning, and code completion.

5. Significance and Claims

The paper positions SelKV as a significant step forward in making long-context LLM inference feasible without sacrificing quality. The authors claim that:

Selective merging is superior to uniform strategies: By modulating merge intensity based on representation similarity, SelKV avoids the semantic corruption caused by indiscriminate merging.
Attention compensation is critical: Without correcting for attention sag, merged tokens are under-utilized; SelKV's ratio-based correction ensures the preserved information is effectively accessed.
Implicit Filtering: The ability of SelKV to outperform full-cache baselines in complex multi-document QA suggests that the compression process itself can enhance model performance by filtering noise, a finding that challenges the assumption that compression always degrades capability.

The authors conclude that SelKV offers a practical, parameter-free solution for deploying LLMs in memory-constrained environments while maintaining near-lossless generation quality, particularly for GQA-based architectures.