AdaKV-proxy: default config is non-adaptive (target==lo_bit), clamp-before-normalize saturates allocation, importance proxy sign inverted vs paper

[rajveer43]

Summary

Reviewing AdaKV-proxy against Ada-KV (arXiv:2407.11550, NeurIPS 2025) turns up four defects in the budget allocator, two of which make the shipped default configuration fully non-adaptive — i.e. method="adakv" with stock settings is provably identical to plain KIVI, while the docs advertise per-head adaptation.

To be clear about scope up front: this repo documents AdaKV-proxy as a deliberate proxy — bit budget instead of eviction budget, attention-free importance signal. That framing is not what I'm challenging. Findings 1, 2 and 4 are bugs within the proxy's own stated design: the allocator does not do what adakv.py's docstring and adakv.md say it does. Finding 3 is a fidelity gap where the proxy's signal is anti-correlated with the paper's criterion, which is worth documenting even if we keep the proxy.

All findings below were verified by execution against master (7c9db57). All 14 existing tests pass throughout — the gaps are in coverage, not in a broken tree.

---

Finding 1 — The default config is non-adaptive: target_avg_bits == lo_bit collapses to uniform

Severity: high. This is the shipped default.

adakv_target_avg_bits defaults to 2.0 and adakv_lo_bit defaults to 2 (cache/base.py:79-82). When the target equals the floor, the per-head budget H × target is exactly satisfiable only by giving every head lo_bit. Any head raised above lo must be paid for by another head going below it — impossible, since lo is the floor.

Verified:

imp=[100,1,1,1,1,1,1,1], allowed={2,3,4}
  target=2.0  -> [2,2,2,2,2,2,2,2]  avg=2.000   FLAT (== KIVI)
  target=2.1  -> [3,2,2,2,2,2,2,2]  avg=2.125
  target=2.5  -> [4,3,3,2,2,2,2,2]  avg=2.500
  target=3.0  -> [4,3,3,3,3,3,3,2]  avg=3.000
  target=4.0  -> [4,4,4,4,4,4,4,4]  avg=4.000   FLAT (== KIVI)

Adaptivity exists only on the open interval (lo_bit, hi_bit), and is strongest mid-range. The two endpoints — one of which is the default — are degenerate.

This directly contradicts adakv.md:12 ("configurable target (default 2.0 …)") and the tuning table at adakv.md:115-119, whose first row claims target=2.0 yields "mostly 2-bit, a few 3/4-bit". It yields all-2-bit, always, for every possible importance vector.

Where to change: allocate_head_bits in quantizers/adakv.py:72-175, plus the default in cache/base.py:79.

Suggested fix: raise the default target to a value strictly inside (lo, hi) — 2.5 with {2,3,4} is the natural pick and preserves a meaningful 6.4× key compression. Additionally, allocate_head_bits should detect the degenerate case (target <= lo or target >= hi) and emit a one-time warning that allocation is uniform, so the mode is never silently entered.

---

Finding 2 — Clamp-before-normalize saturates the allocation and discards importance ordering

Severity: high.

At adakv.py:124:

real = [min(max(w * budget_total, float(lo)), float(hi)) for w in norm]

norm sums to 1, so w * budget_total has mean target. But the distribution of w is set by raw importance ratios, which for realistic key-norm variances span orders of magnitude. The clamp then flattens nearly everything onto the two endpoints, destroying the interior ordering that the subsequent snap and greedy correction are supposed to act on.

Verified — importance ratio changes by 4 orders of magnitude, allocation does not move at all:

target=3.0, allowed={2,3,4}, H=4
  imp=[10,   1, 1, 1]     -> [4,3,3,2]
  imp=[100,  1, 1, 1]     -> [4,3,3,2]
  imp=[1000, 0.1,0.1,0.1] -> [4,3,3,2]

And the intermediate state showing the saturation, for H=8, target=2.0:

raw w*H*target : [14.953, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15, 0.15]
after clamp    : [ 4.0,   2.0,  2.0,  2.0,  2.0,  2.0,  2.0,  2.0 ]

Head 0's real budget of 14.95 carries no more information than one of 4.01 once clamped — the allocator cannot distinguish "somewhat important" from "overwhelmingly important". Note also that the clamped vector sums to 18 against a target_total of 16; the greedy pass then has to claw back 2 bits, and which heads it takes them from is driven by real[h] values that have already been flattened.

Where to change: adakv.py:115-127.

Suggested fix: normalize the importance vector through a bounded, order-preserving map before scaling — e.g. rank-based or softmax-with-temperature over log1p(importance) — so that the mean lands on target without the extremes saturating. A clean formulation: map importance to a spread around the target rather than a scale of it, i.e. real[h] = target + spread * z[h] where z is the standardized (or rank-normalized) importance, then clamp. That keeps the mean at target by construction and makes the clamp a genuine edge case rather than the common path.

---

Finding 3 — The importance proxy is anti-correlated with the paper's criterion

Severity: medium (fidelity/documentation, not a crash).

The paper's entire thesis (§3.3, Fig. 1a-1b) is that budget should shift from attention-sparse heads to attention-dispersed heads. A head whose attention concentrates on a few tokens needs little budget; a head spreading attention broadly needs more. Fig. 1b makes this concrete: heads 2 and 3 (dispersed) gain budget, head 1 (sparse, 0.89 on one token) loses it.

compute_head_norm_variance (adakv.py:50-69) uses Var_t(||k_t||₂) and gives more bits to higher variance. But high key-norm variance is the signature of a few outlier-norm tokens dominating the logits — that is precisely a sparse head. So the proxy systematically funds the heads the paper defunds.

Verified with two synthetic heads at S=512, D=64 (max possible entropy ln 512 = 6.238):

head                                  norm-var   attn-entropy   paper wants
SPARSE attn (few dominant keys)         0.3675      5.9024       LESS budget
DISPERSED attn (equal norms)            0.0000      6.2383       MORE budget

repo proxy gives MORE bits to: SPARSE
paper                        : DISPERSED

The proxy's sign is inverted relative to the criterion it claims to approximate. adakv.py:17-20 asserts the opposite without evidence:

"High variance ⇒ the head's key magnitudes spread widely across tokens — a proxy for high attention entropy"

Wide spread in magnitudes is not wide spread in attention — softmax over q·k concentrates on large-norm aligned keys, so magnitude spread drives entropy down.

Where to change: the docstring claim at adakv.py:17-20, the doc claim at adakv.md:53 and adakv.md:149, and ideally the signal itself in compute_head_norm_variance.

Suggested fix. Two defensible options, and I'd recommend doing both:

1. Minimum, honest: keep the signal but correct the documentation — state that it targets quantization sensitivity (high norm-variance ⇒ min/max group quant has wider dynamic range ⇒ higher quantization error ⇒ more bits help), which is a genuinely sound rationale for a bit-allocation method, and note plainly that this is a different criterion from the paper's attention-dispersion, not an approximation of it.
2. Better: add an entropy-flavored proxy that is actually computable from keys alone. SnapKV-adapted already computes an observation-window attention distribution from keys-as-proxy-queries (quantizers/snapkv.py:60-89). Reusing that to compute per-head attention entropy would give a signal with the paper's sign, at a cost the repo already pays elsewhere.

Option 1 is a docs-only change and should land regardless. Option 2 is the real fix.

---

Finding 4 — Ties broken by head index, so allocation depends on head ordering

Severity: low.

The greedy correction at adakv.py:152-162 scans heads in index order and keeps the first strict minimum (cost < best_cost). Heads with identical importance therefore receive different bit-widths, decided purely by position:

imp=[10,1,1,1] -> [4,3,3,2]     heads 1,2,3 identical, get 3/3/2
imp=[1,1,1,10] -> [3,3,2,4]     same multiset, different head loses

Deterministic (so test_determinism passes), but arbitrary — a permutation of the head axis changes which head is starved. test_equal_importance_uniform_allocation misses this because it only tests the fully uniform vector at target == lo_bit, where Finding 1 forces the flat answer anyway.

Where to change: adakv.py:152-162.

Suggested fix: break ties on a stable secondary key — e.g. prefer the head with lower current bits[h] when stepping up and higher when stepping down, falling back to importance. This makes the result invariant to head permutation.

---

Test coverage gaps

All 14 tests in tests/cache/test_adakv_cache.py pass on master despite Findings 1 and 2. Why:

• test_high_importance_heads_get_more_bits (L120) constructs its cache with adakv_target_avg_bits=3.0 — mid-range, where adaptation does work. Nothing exercises adaptation at the default 2.0.
• test_average_bits_matches_target (L136) loops (2.0, 2.5, 3.0) but only asserts the average is within ±0.5. A fully flat allocation trivially satisfies that.
• test_equal_importance_uniform_allocation (L150) asserts [2,2,2,2] at target=2.0 — it is pinning the degenerate behavior as correct, which is why Finding 1 never surfaced.

Tests to add:

1. Adaptation is non-trivial at the default config (fails today).
2. Allocation is monotone in importance — raising one head's importance never lowers its bits (guards Finding 2).
3. Allocation is permutation-equivariant in the head axis (guards Finding 4).
4. Degenerate target <= lo / target >= hi is explicitly flagged, not silently uniform.
5. Importance-proxy directionality is pinned to whatever criterion we settle on in Finding 3, with the rationale in the test docstring.

---

Proposed changes, in order

#	File	Change	Priority
1	cache/base.py:79	Default adakv_target_avg_bits 2.0 → 2.5	High
2	quantizers/adakv.py:72-175	Warn on degenerate target <= lo / >= hi	High
3	quantizers/adakv.py:115-127	Replace clamp-then-scale with spread-around-target	High
4	quantizers/adakv.py:17-20, adakv.md:53,149	Correct the entropy claim; state the real criterion	High (docs-only, ship first)
5	quantizers/adakv.py:152-162	Permutation-invariant tie-breaking	Medium
6	quantizers/adakv.py	Optional entropy proxy reusing SnapKV's obs-window scoring	Medium
7	tests/cache/test_adakv_cache.py	5 tests above; retire the degenerate-pinning assertion	High
8	adakv.md:115-119	Fix the tuning table — current first row is wrong	High

Findings 1, 2 and 4 are contained in allocate_head_bits and are straightforwardly fixable. Finding 3 is a design question; the docs correction should land regardless of whether we adopt the entropy proxy.

---

Reviewed against Ada-KV (arXiv:2407.11550v5, NeurIPS 2025) §3.3 Algorithm 1, §3.5, Fig. 1a-1b. All behavioral claims verified by execution on master @ 7c9db57.