Beyond 84×: A Vector Quantization Quality Framework
for Transformer KV Cache Compression

1. Background and Motivation

Our Part I report demonstrated that TSC achieves 63× lossy compression on GPT-2 KV caches with zero measurable quality loss — top-1 token prediction rate of 1.0000, KL < 10⁻⁴, and perplexity delta within ±0.1. Stacking an H2O-style semantic eviction layer at 75% token retention raised this to 84× at KL=0.08 — still within our "defensible" threshold of KL < 0.1.

TSC stores KV caches as: one keyframe every 64 tokens (fp32) + 4-bit delta residuals for the tail. The keyframes represent ~95% of compressed storage. The obvious next question: can we compress those keyframes too?

Vector Quantization (VQ) maps each small block of values to the nearest centroid in a learned codebook. With a block size of 8 values and 256 centroids, each block is stored as 1 byte — a 16× reduction vs fp16 (2 bytes/value). Applied to TSC keyframes, the combined ratio is:

We achieved 1034×. The problem is that it came with KL=1.09 — a completely unusable output distribution for live inference. This paper documents our systematic investigation of why, and how far we can close the gap.

2. Root Cause Analysis: Why 1034× Is Lossy

Before building solutions, we needed to understand exactly why naive VQ fails. We profiled the per-head KV magnitude distribution for both GPT-2 (124M) and TinyLlama-1.1B across 512-token WikiText-2 sequences.

2.1 KV Magnitude Distribution

Figure 1. Box plot of per-head max absolute KV value across all layers and heads. Both models show median ≈ 4–6 with tails extending past 20. A uniform 256-entry codebook in 8D covering this range has step size ≈ max_abs/16 ≈ 0.37 — catastrophic for generation quality.

The fundamental issue: with a shared 256-centroid codebook covering the full range [−max_abs, +max_abs], the quantization step size is roughly max_abs / 128 per dimension. At the median magnitude of 5.88, that's 0.046 per value per dimension. But in the 8-dimensional block, errors accumulate: the centroid assignment error in 8D can reach 0.37 per value at p90 magnitudes. This is enough to flip attention scores and corrupt the output distribution.

2.2 The "Attention Sink" Problem

A single shared codebook must cover all heads across all layers. Attention sinks — heads with max|v| > 15 — dominate codebook placement, leaving low-magnitude heads (max|v| ≈ 0.5–1.0) with effectively only 2–3 distinct codes. The mixed-distribution codebook reconstructs neither the high-magnitude nor low-magnitude heads well.

3. The VQ Quality Framework

We built a three-stage encoding pipeline in kv_vq_quality.py (NormalizedKVQuantizer, PQNormalizedQuantizer) that progressively improves quality at the cost of compression ratio.

3.1 Stage 1 — Per-Head Normalization

Before quantization, each attention head's values are divided by their maximum absolute value:

The scale factor max|v_h| is stored as a single fp32 scalar per head per keyframe. With GPT-2's 12 layers × 12 heads = 144 heads, scale storage is 576 bytes — negligible vs the keyframe size.

After normalization, every head's values live in [−1, 1] regardless of original magnitude. The VQ codebook now uses its 256 codes to cover this uniform range, giving equal precision to every head.

3.2 Stage 2 — Selective fp16 Fallback

After encoding each 8-value block with VQ and measuring its max absolute reconstruction error, any block exceeding an error_threshold is stored as raw fp16 instead. A bitmask (1 bit per block) marks which blocks use which path.

Storage cost per block: VQ path = 1 byte/8 values = 0.125 bytes/value. fp16 path = 16 bytes/8 values = 2 bytes/value. As threshold decreases, more blocks fall back to fp16 and the effective compression drops but quality improves.

3.3 Stage 3 — Product Quantization (PQ)

Rather than a single 256-centroid codebook in 8D, Product Quantization splits each 8-value block into 2 sub-vectors of 4 values, each quantized with its own 256-code codebook:

Storage: 2 bytes per 8-value block = 8× vs fp16 (vs 16× for flat VQ). The effective codebook size is 256² = 65,536 centroids, dramatically reducing quantization error for the same storage budget. Combined ratio with TSC:

PQ is implemented in PQCodebook (kv_codebook.py) with independent GPU k-means per sub-space, avoiding Windows OpenMP deadlocks via a pure-torch chunked distance computation.

4. Experimental Results

4.1 GPT-2 VQ Threshold Sweep

We swept error_threshold from 1.0 (pure VQ = 1034×) to 0.0 (pure fp16 = 1×), measuring compression ratio and KL divergence on WikiText-2 (n=5 sequences, seq_len=512). A fresh codebook is trained per sweep point.

Figure 2. Threshold sweep on GPT-2. The elbow (KL < 0.1) appears at threshold=0.07 → 123× compression. Below the elbow, nearly all blocks fall back to fp16 (99.4% fp16 at threshold=0.07).

4.2 TinyLlama-1.1B — Product Quantization Sweep

PQ provides a better tradeoff: 8× per-keyframe gain instead of 16×, but with 65,536 effective centroids vs 256, the quality is substantially better for the same storage budget. Results on TinyLlama-1.1B (n=10 sequences, seq_len=256, WikiText-2):

4.3 Stability Analysis: 429× Under the Microscope

The 429× point (VQ bs=2, pure, no fallback) shows mean KL=0.113 across 10 sequences — borderline defensible. But the variance tells a different story:

Figure 3. Left: per-sequence KL values at the 429× operating point. Two sequences show KL ≈ 0.60 — well outside defensible range. Right: achievable compression ratios vs KL threshold, with the KL<0.1 defensible boundary marked.

5. The Complete Compression Tier Ladder

6. What Would Make 1034× Defensible

The 1034× operating point requires a model architecture where all KV heads have uniformly small magnitudes — specifically, max|v| < ~0.5 across all layers and heads. With that constraint, 4-bit quantization has a step size of ~0.03, which keeps per-block VQ error below our threshold.

The most promising architectural path is models that explicitly normalize the KV projections before caching — for example, architectures that apply RMSNorm to key and value outputs similar to how some models normalize query vectors. Gemma-style and some Falcon variants may exhibit this property and are the next candidates to test.

7. Implementation Details

7.1 GPU K-Means (Torch Backend)

scikit-learn's KMeans deadlocked on Windows due to OpenMP conflicts with PyTorch. We replaced it with a pure-torch GPU k-means:

Runtime: <2 seconds for 20,000 × 8-dimensional vectors on CUDA. This enabled the full threshold sweep (11 operating points × 5 sequences) to complete in under 5 minutes.

7.2 DynamicCache Compatibility (TinyLlama)

TinyLlama uses the newer transformers DynamicCache API rather than raw tuple[tuple[Tensor]]. Passing a decoded pkv tuple back to the model caused:

Fix: wrap the decoded tuple through DynamicCache.update() before passing to the model for the quality measurement forward pass. The codebook training and encoding operates on raw numpy arrays extracted from the cache, avoiding the Cache API entirely.

7.3 Storage Model

The combined compression ratio is modeled conservatively:

The 0.05 factor accounts for delta/tail storage (4-bit, not VQ-compressed). The 0.95 factor represents the fraction of TSC-compressed storage that lives in keyframes. The TSC ratio of 63× is a proven empirical measurement from the Part I long_seq_bench.

8. Comparison to Prior Work

9. Honest Limitations

9.1 The 84× Eviction Top-1 Caveat

At 84× (75% token retention), top-1 match rate drops to 0.60. This means 40% of tokens predicted by the compressed cache differ from the exact prediction. However, KL=0.08 indicates the output distribution is very close — the second-best prediction in the compressed case is typically the first-best in the exact case. This is acceptable for most inference scenarios but not for strict next-token guarantees.

10. Files and Reproducibility

Reproduce the quality sweep (GPT-2)

Reproduce the multi-model defensible sweep