Skip to content

New press: DropKV (eviction by decoupled residual-output perturbation) #284

Description

@AozhongZhang

Press

DropKV is a principled eviction framework based on Decoupled Residual-Output Perturbation.

Paper: https://openreview.net/forum?id=MqfNzH3TVH

Let $p_{ij}$ be the attention probability of query $i$ on position $j$, $\mathbf{v}_j$ the value vector at position $j$, and $\mathbf{o}i = \sum_k p{ik}\mathbf{v}_k$ the attention output of query $i$. Write $\mathbf{o}_i^{(-j)}$ for the attention output of query $i$ once pair $j$ has been evicted. Eviction renormalizes the softmax over the remaining entries, so the residual admits a closed form:

$$ \mathbf{o}_i^{(-j)} - \mathbf{o}_i ;=; \frac{p_{ij}}{1 - p_{ij}},\bigl(\mathbf{o}_i - \mathbf{v}_j\bigr) $$

DropKV accumulates the squared residual over the window $\mathcal{W}$ of the most recent window_size queries:

$$ s_j ;=; \sum_{i \in \mathcal{W}} \left(\frac{p_{ij}}{1 - p_{ij}}\right)^{!2},\bigl\lVert \mathbf{v}_j - \mathbf{o}_i \bigr\rVert_2^{2} $$

It fits the ScorerPress interface directly: one score() returning (batch, num_kv_heads, seq_len), no extra forward passes, no checkpoints, no new dependencies.

Because the score is an output deviation measured in a common unit (squared L2 in value space) rather than a per-head normalized quantity, its magnitudes are comparable across heads, so it also composes well with AdaKVPress. I would register both "dropkv": DropKVPress() and "adakv_dropkv": AdaKVPress(DropKVPress()), the latter mirroring the existing "adakv_snapkv" entry; the AdaKV variant is the stronger configuration in my evaluations.

Motivation

Minimizing the attention-output perturbation over an eviction set is NP-hard, because the evicted pairs interact through the softmax normalization. Existing methods that go beyond attention-score heuristics therefore settle for mathematically suboptimal approximations, or for formulations that are hard to kernelize. DropKV decouples the joint decision into independent per-token scores, which sidesteps the combinatorial problem and admits a constant-factor approximation guarantee under a cone condition that the paper verifies empirically on three long-context LLMs.

None of the presses currently in the library derive their criterion from the residual-output perturbation in closed form. Attention-score based scorers rank by a proxy for importance, and the scorers that do incorporate value representations approximate the joint objective heuristically. DropKV instead evaluates, per token, the exact deviation its removal induces in the attention output, which is what the eviction objective is defined over; the value vectors enter the criterion directly rather than as a separate heuristic term. The decoupling is what makes this tractable, and the approximation guarantee is what bounds the cost of decoupling.

On RULER 4096 with query-aware compression, averaged over the 13 subtasks, AdaKVPress(DropKVPress()) is ahead of every query-aware entry currently on the leaderboard at all four standard ratios on both Qwen3-8B and Llama-3.1-8B-Instruct, with the margin widening as the budget tightens (77.35 vs 61.68 for the best baseline at 0.875 on Qwen3-8B; 79.68 vs 72.39 on Llama-3.1-8B-Instruct). Full runs:
https://huggingface.co/spaces/nvidia/kvpress-leaderboard/discussions/22

I have a working implementation (6 files, make style clean, press test matrix passing) and am happy to open the PR.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions