Equal-Volume KV Vector Quantization
Implementation & benchmark of arXiv:2608.04074 — Spend Bits Where Queries Look: KV Cache Vector Quantization with Attention-Preserving Transforms
Source paper: Spend Bits Where Queries Look: KV Cache Vector Quantization with Attention-Preserving Transforms arXiv:2608.04074 ⓘ · analyzed Aug 31, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Replace consecutive or randomly assigned transformed KV coefficients with groups whose variance-volume is approximately equal. Train one equal-size vector-quantizer codebook per group, so a fixed-width cache does not waste its low-rate budget by forcing high-variance and low-variance coordinates into badly mismatched groups. This is a drop-in quantization-layout change that can be applied to keys, values, or both.
Formulas
Mathematical statement
The paper gives the fixed-width high-resolution distortion bound $C_g 2^{-2b}\sum_{\ell}v_{\ell}(\pi)^{1/g}\geq D^{*}(b)$, with equality if and only if all group volumes coincide. Here $b$ is the average number of bits per scalar coefficient, $g$ is the number of coefficient groups, $\pi$ is a partition of the $d$ transformed coordinates into groups, $v_{\ell}(\pi)$ is the volume of group $\ell$, $C_g$ is a quantizer-dependent constant, and $D^{*}(b)$ is the variable-rate high-resolution optimum. Estimate the variance of transformed coordinate $i$ as $\sigma_i^2$ from calibration data and define $v_{\ell}(\pi)=\prod_{i\in G_{\ell}}\sigma_i^2$ for group $G_{\ell}$. Equation (9) normalizes these volumes by the equalizing target $v_{\mathrm{equal}}=\left(\prod_{i=1}^{d}\sigma_i^2\right)^{g/d}$. The inequality predicts that reducing the spread of $v_{\ell}$ lowers distortion, with equality when all groups have equal volume. For transformed vectors $z=Ax$, quantize each group $z_{G_{\ell}}$ using an equal-cardinality codebook and reconstruct with $\widehat{x}=A^{-1}\widehat{z}$.
Implementation notes
(1) Integration point: modify the KV-cache quantizer immediately after key and value projection tensors are produced and before they are written to GPU memory. During calibration, collect several thousand representative key/value vectors from long-context prompts for every selected layer and attention head. Apply the existing transform $A$, initially identity, Hadamard, or PCA/KLT, and estimate each transformed-coordinate variance $\sigma_i^2=\operatorname{Var}(z_i)$ with a small floor $\epsilon$. (2) Pseudocode: Z = A @ X; logvar = log(var(Z)+eps); target_log_volume = (sum(logvar)/d) * (d/g); groups = g empty lists; for each coordinate i sorted by descending logvar: assign i to the currently smallest-log-volume group that has fewer than d/g coordinates; train an M-entry k-means codebook C_l on Z[:, G_l] for every group; cache each group codeword index. Each group must have exactly $d/g$ coordinates, and at $b$ bits per scalar its codebook has $M=2^{b(d/g)}$ entries; if this is too large, use residual or product quantization while preserving the same grouping. At inference, transform each new vector, gather its group subvectors, find the nearest codeword, store only the fixed-width indices, dequantize, concatenate groups, apply $A^{-1}$, and use reconstructed K/V in ordinary attention. (3) Computed from the mathematics: $v_l$, $v_{equal}$, the volume-spread proxy $\sum_l v_l^{1/g}$, and the greedy partition. Estimated empirically: codebooks, attention-logit error, attention-output MSE, retrieval accuracy, and actual memory bandwidth. (4) First experiment: use a 1B-3B decoder-only model on LongBench and synthetic needle-in-a-haystack prompts at 2 bits per scalar, comparing equal-volume, consecutive eigenvalue-order, random grouping, scalar quantization, and unquantized KV. Measure perplexity, retrieval accuracy, attention-output MSE, decode tokens per second, and peak KV memory at 4K, 16K, and 64K contexts. Success means lower attention-output MSE and higher retrieval accuracy than the same-bit baselines without reducing decode throughput; additionally test whether the predicted volume proxy correlates with measured distortion.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Not reproduced ✗
Stage 2 · Practical benchmark result: Inconclusive
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 8/10
Built equal-volume grouping, volume-proxy evaluation, balanced-partition enumeration, and a held-out grouped vector-quantization toy experiment. The exact 8D/4-group enumeration found the greedy equal-volume partition at the global minimum (4.00313, among 630 partitions); at zero anisotropy the layout ratio was exactly 1.0, while the consecutive/equal proxy ratio grew from 1.09 at spread 0.5 to 7.94 at spread 3.0. In the empirical sweep, equal grouping had lower MSE than consecutive grouping at every tested spread (0.872 vs 0.908 at 0.5; 1.696 vs 3.105 at 3.0), and generally beat random grouping, so the mechanism manifested; proxy-to-MSE correlation was mixed rather than universally strong.
- Agent confidence
- 8/10
- Baseline
- Consecutive grouping held-out MSE, g=4: 0.907777, 1.044027, 1.946931, 3.105135 for log-variance spreads 0.5, 1, 2, 3; random grouping: 0.883286, 0.999074, 1.609225, 1.922891.
- Idea
- Equal-volume grouping held-out MSE, g=4: 0.871968, 0.925959, 1.531733, 1.695754 for spreads 0.5, 1, 2, 3. Exact proxy minimum matched the enumerated global minimum; consecutive/equal proxy ratios were 1.093, 1.398, 3.064, and 7.939.
Limitations: This is a synthetic diagonal-Gaussian experiment, not a decoder model or real K/V cache. It does not measure perplexity, retrieval, attention-output error, throughput, GPU memory, transform/inverse-transform effects, calibration drift, or long-context behavior. The vector codebook uses a fast 16-prototype PCA-sorted approximation rather than full k-means with the mathematically prescribed 2^(b*d/g) entries, and the proxy-MSE correlation was inconsistent across individual spreads.
How to run: OMP_NUM_THREADS=1 /home/maxwelhelp/main/bin/python3 equal_volume_experiment.py && OMP_NUM_THREADS=1 /home/maxwelhelp/main/bin/python3 verify_predictions.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2608.04074)
Ahead of the baseline (−4.2%) but not statistically significant (p=0.4698, wins 4 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.4476
- Idea mean
- 0.4286
- Effect (Δ)
- -0.019 (−4.2%; negative = idea better)
- Wins
- 4 / 8 paired seeds
- p-value
- 0.4698 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±11.8%
- Mechanism
- Not confirmed ✗
- Practical effect
- Inconclusive
- Baseline tuning
- swept over 3 configs
Only the registered sequence forecasting track and small transformer_tiny were tested. The implementation uses a bounded product-code scalar approximation rather than full vector codebooks, and does not test a pretrained decoder, long-context KV memory, retrieval accuracy, throughput, bandwidth, or multiple layers/heads.
How to run: python3 sequence_bench.py
Verdict computed by deterministic test code from paired-seed statistics — not by the language model.
Stage 2 — Benchmark latest Failed ✗
Agent confidence: 9/10
Equal-volume KV quantization was implemented as a local modification of the matched transformer_tiny sequence system and evaluated on the registered sequence track. It reduced mean test MSE from 0.44760169 to 0.42862934, but the paired permutation p-value was 0.46985, so the required significant-win criterion was not met. The mechanism signature was also negative, with predicted volume-CV versus observed KV quantization MSE correlation 0.02722.
- Baseline
- Mean test MSE 0.44760169088840485 (lr=0.0015), std 0.0483900971330472.
- Idea
- Mean test MSE 0.4286293387413025 (lr=0.0015), std 0.06739731912758126; paired delta -0.018972352147102356; p=0.46985; 4/8 wins.
Benchmark result
Ahead of the baseline (−4.2%) but not statistically significant (p=0.4698, wins 4 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.4476
- Idea mean
- 0.4286
- Effect (Δ)
- -0.019 (−4.2%; negative = idea better)
- Wins
- 4 / 8 paired seeds
- p-value
- 0.4698 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±11.8%
- Mechanism
- Not confirmed ✗
- Practical effect
- Inconclusive
- Baseline tuning
- swept over 3 configs
Only the registered sequence forecasting track and small transformer_tiny were tested. The implementation uses a bounded product-code scalar approximation rather than full vector codebooks, and does not test a pretrained decoder, long-context KV memory, retrieval accuracy, throughput, bandwidth, or multiple layers/heads.
How to run: python3 sequence_bench.py
Artifacts
- bench_report.json 2.3 KB View Raw
- bench_report_rerun.json 2.3 KB View Raw
- equal_volume_experiment.py 5.2 KB View Raw
- report.md 2.0 KB View
- report_bench_2026-08-31T181031.md 4.0 KB View
- results.json 9.6 KB View Raw
- sequence_bench.py 6.7 KB View Raw
- verification.json 1.6 KB View Raw
- verify_predictions.py 2.2 KB View Raw