Entropy Gated ChunkKV (EG-ChunkKV)

[ShaharBenIshay]

Press

EntropyGatedChunkKVPress extends ChunkKV, already in the
library as ChunkKVPress. It is a wrapper press - a BasePress taking any ScorerPress as its
inner token scorer - and it changes at what granularity a selected chunk
is retained.

ChunkKV partitions the context into contiguous chunks of length c, ranks them by aggregate
token importance S_i, and retains each selected chunk as an indivisible unit. That aggregate
records how much importance a chunk holds but not how that importance is distributed inside it,
so two chunks with the same score can have very different internal structure: importance spread
across many tokens (coherent), or concentrated on one or two "needle" tokens with the
remaining positions contributing little (spiky).

EG-ChunkKV adds a second, shape statistic per chunk - the normalized within-chunk Shannon
entropy of the token scores:

p_ij = s_j / (S_i + eps)
H̃_i = -sum_j p_ij * log(p_ij + eps) / log|I_i|        in [0, 1]

H̃_i near 1 means importance is spread across the chunk (coherent); near 0 means it is
concentrated (spiky). A single-token or all-zero chunk is defined to have H̃_i = 0.

Retention is then gated. Under the same retained-token budget L = max(1, floor(1 - rho) * T) that
ChunkKV uses, chunks are visited greedily in decreasing importance; with remaining budget b:

important_i = S_i >= median(S)
spiky_i     = H̃_i <  tau           # tau defaults to the per-example median of H̃

if important_i and spiky_i:   keep the chunk's top-min(r, b) tokens    <- rescue the needle
elif the whole chunk fits:    keep the chunk whole                     <- coherent, keep whole
else:                         keep its top-b tokens                    <- budget-boundary cap

Motivation

The gap. Because ChunkKV ranks on S_i alone, two chunks with S_a = S_b but very different
H̃_a != H̃_b are indistinguishable to it, and both are kept or dropped whole. For an
important-but-spiky chunk that is a poor binary choice: keeping it preserves the needle together
with its filler, dropping it discards the needle entirely. This cannot be fixed by ranking whole
chunks better - the missing decision is the granularity at which each region is retained.

The waste is measurable, and it grows as the budget tightens. On real LongBench prefills at
cr = 0.90, over the same ~700K ChunkKV-selected chunks, mean retained entropy rises from
0.947 under ChunkKV to 0.983 under EG-ChunkKV — rescue strips low-signal filler without
discarding the needle. Sweeping the compression ratio, the coherence gap widens monotonically —
mean normalized within-chunk entropy of the tokens each method actually retains:

ChunkKV's retained cache becomes less coherent and more variable as the budget tightens (std
0.076 -> 0.109), while EG-ChunkKV stays flat (std 0.054 -> 0.069).

Some results. All runs used the same model and same hyperparams (cr = fraction of KV cache
removed).

RULER-4096, mean over the 13 subtasks:

Method	cr=0.85	cr=0.90	cr=0.95
ChunkKV	45.03	31.60	21.49
EG-ChunkKV	46.12	36.45	21.99

LongBench, average over the 16 tasks:

Method	cr=0.70	cr=0.80	cr=0.90
ChunkKV	40.91	38.24	32.91
EG-ChunkKV	41.34	38.52	34.26

The gate is not tied to one scorer. Swapping the inner ScorerPress (RULER-4096, 13-subtask
mean, delta = EG-ChunkKV - ChunkKV):

inner scorer	cr=0.85	cr=0.90	cr=0.95
expected_attention	+7.51	+5.56	+4.46
snapkv	+1.09	+4.85	+0.50
tova	-1.57	+3.94	+1.76

thanks for reading and we would like to have your permission to open a PR with our work.

[SimJeg]

Hi @ShaharBenIshay, thanks for your issue. Is there any publication / pre-print related EG-ChunkKV ?

[ShaharBenishayTeads]

Hi @SimJeg,
We haven't released a preprint yet, but we have worked on one and we plan to do so in the future.

[SimJeg]

Thanks ! Feel free to open a PR. Some parts are 100% clear to me

1. You compute average score per chunk $S_i$
2. You compute entropy per chunk $H_i$
3. You keep chunks with highest scores $S_i$ (?) but for those with $H_i \lt \bar{H_i}$ and $S_i \ge \bar{S_i}$ you don't keep the whole chunk but only some tokens in it ?

[Liranitz]

Thanks @SimJeg ! You're right about all three - with one nuance on how the selection actually runs:

Sᵢ (average score per chunk): yes - the mean of the (head-summed, non-negative) token importances in chunk i, the score ChunkKV already ranks on.

H̃ᵢ (entropy per chunk): yes. Sᵢ says how much importance a chunk holds, not how it's spread. H̃ᵢ captures that: near 1 = spread evenly (coherent), near 0 = concentrated in a few tokens with the rest barely mattering (spiky). Two chunks with equal Sᵢ can differ sharply in H̃ᵢ - the extra signal we add to ChunkKV.

"keep chunks with highest Sᵢ?": Instead of picking whole chunks up front, we fix the same budget L ChunkKV uses and visit chunks greedily in decreasing Sᵢ, spending budget b as we go. An important chunk (Sᵢ ≥ median(S)) is kept whole only if coherent, if also spiky (H̃ᵢ < τ, τ = median of H̃) it's cut to its top-r tokens.

"for H̃ᵢ < H̄ and Sᵢ ≥ S̄ you keep only some tokens?": exactly. Keeping such a chunk whole burns ~c slots for one needle, dropping it loses the needle.