Joint Tucker-rank and residual-bit KV allocator
Implementation & benchmark of arXiv:2607.12550 — A JoLT for the KV cache: Near-lossless KV cache compression via joint Lagrangian allocation of Tucker ranks and a rotated residual for llms
Source paper: A JoLT for the KV cache: Near-lossless KV cache compression via joint Lagrangian allocation of Tucker ranks and a rotated residual for llms arXiv:2607.12550 ⓘ · analyzed Aug 30, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Compress each layer's KV tensor with a partial Tucker approximation over token and feature axes, then encode the truncation residual with a rotated uniform quantizer. Select token rank, feature rank, and residual bit-width jointly under a global byte budget, allowing values with flat spectra to receive residual bits while keys may receive more low-rank capacity.
Formulas
Mathematical statement
For a cache tensor X in R^{L x H x T x D}, L is the number of layers, H is the number of heads or KV heads, T is the cached-token count, and D is the head-feature width. Group layers and retain the layer and head axes explicitly while approximating each group by a partial Tucker reconstruction ilde{X}(r_T,r_d), where r_T and r_d are the token- and feature-axis ranks. The truncation residual is R=X- ilde{X}(r_T,r_d). Let Q be an orthogonal rotation satisfying Q^TQ=I, and let Q_b denote entrywise uniform quantization at b bits; the decoded cache is \hat{X}= ilde{X}+Q^TQ_b(QR). For layer group g and cache type c in {K,V}, choose z=(r_T,r_d,b) with byte cost C_{g,c}(z) and normalized distortion D_{g,c}(z). The joint Lagrangian is L(z,lambda)=sum_{g,c}D_{g,c}(z_{g,c})+lambda(sum_{g,c}C_{g,c}(z_{g,c})-B), where B is the total byte budget and lambda>=0 is the byte price. For fixed lambda, each group independently chooses argmin_z[D_{g,c}(z)+lambda C_{g,c}(z)]; binary search on lambda finds a budget-feasible allocation.
Implementation notes
Integrate this at KV-cache admission time, before storing K and V tensors for each layer. Divide layers into groups, for example four consecutive layers, and maintain separate candidate tables for K and V. From a calibration run, collect each group tensor X with shape [heads, tokens, features]. Enumerate r_T in {4, 8, 16, 32, T}, r_d in {4, 8, 16, 32, D}, and b in {0, 2, 4, 8}. Compute a partial Tucker approximation using sequential SVD or randomized SVD on the token and feature unfoldings, reconstruct X_tilde, and set R=X-X_tilde. Apply a fixed orthogonal transform, quantize the residual, and measure normalized Frobenius distortion. Calculate actual storage including Tucker factors, core tensors, residual codes, scales, transform metadata, and alignment padding; do not rely only on the asymptotic cost expression. Store candidate pairs (D,C). For a trial lambda, choose the candidate minimizing D+lambda*C independently for every group and cache type. Binary-search lambda until total storage is within one alignment block of budget B, then optionally improve the result by swapping candidates with the best distortion reduction per added byte. At inference, reconstruct the low-rank part and residual immediately before attention, preferably in a fused kernel. First test Mistral-7B at 8k and 32k context against FP16 KV, global-rank Tucker, and uniform 4-bit KV at equal 2x and 3x memory reduction. Record perplexity, needle retrieval accuracy, peak KV bytes, decode tokens/sec, and decompression bandwidth. Success is near-baseline quality with lower memory and better distortion than fixed-rank or fixed-bit baselines.
Verification
Stage 1 · Toy mechanism gate: Passed ✓
Stage 2 · Mechanism transferred to benchmark: Confirmed ✓
Stage 2 · Practical benchmark result: No measurable effect
Methodology: 8 paired seeds, permutation test (p<0.05) against a learning-rate-tuned baseline. How verification works
Stage 1 — Mechanism check agent confidence 7/10
Built and ran a NumPy MVP for partial Tucker reconstruction, rotated residual quantization, explicit storage costs, and joint budget allocation. Orthogonal rotation preserved residual norm exactly, Tucker distortion decreased monotonically with rank, and quantization MSE decreased 16.25–25.65x per two added bits versus the high-rate 16x prediction. At a 4,000-byte budget, joint allocation reached distortion 0.030199 versus 0.032503 for fixed 4-bit allocation, a modest 7.1% improvement; the mechanism shows a real toy-scale signal.
- Agent confidence
- 7/10
- Baseline
- Fixed 4-bit allocation, 3,984 total bytes, normalized distortion 0.032503; all four groups used (rT,rd,b)=(2,2,4).
- Idea
- Joint allocation, 3,916 total bytes, normalized distortion 0.030199; structured K groups received low-rank/no-residual choices while flatter V groups received residual bits. Orthogonal relative norm error was 0.0; adjacent quantization MSE ratios were 25.65, 17.69, and 16.25 versus predicted 16.0.
Limitations: This is a synthetic four-group NumPy experiment, not a transformer KV-cache evaluation. It does not test perplexity, needle retrieval, decode speed, GPU memory, fused reconstruction, real layer grouping, alignment padding, per-channel scales, or Mistral contexts. The dense random orthogonal transform is used only for mathematical verification, not efficient inference.
How to run: python3 joint_tucker_kv.py
Stage 2 — Benchmark vs. tuned baseline (arXiv:2607.12550)
Ahead of the baseline (−0.1%) but not statistically significant (p=0.9113, wins 3 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.5321
- Idea mean
- 0.5316
- Effect (Δ)
- -0.0005 (−0.1%; negative = idea better)
- Wins
- 3 / 8 paired seeds
- p-value
- 0.9113 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±2.0%
- Mechanism
- Confirmed ✓
- Practical effect
- No effect
- Baseline tuning
- swept over 3 configs
This is a small sequence-forecast proxy rather than a full transformer KV-cache admission/reconstruction implementation: the bench model exposes scalar input sequences and does not expose internal K/V tensors. It does not test actual layer/head-specific Tucker factors, global byte-budget allocation, FP16 KV memory, perplexity, retrieval accuracy, decode throughput, decompression bandwidth, or long-context inference.
How to run: python3 run_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
Implemented a matched sequence-track training proxy using the unchanged transformer architecture with token/feature low-rank reconstruction and residual uniform quantization, plus the required baseline sweep and 8 paired seeds. The NN-scale mechanism signature was confirmed: rank distortion decreased monotonically and two-bit quantization MSE ratios were 28.74, 17.24, and 15.68 versus the predicted 16x. However, the best idea setting only changed mean MSE by -0.00053 with permutation p=0.91125, so there was no significant improvement.
- Baseline
- Best lr=0.001; full 8-seed mean MSE 0.532116, std 0.039885.
- Idea
- Best lr=0.001, rank=16, bits=4; full 8-seed mean MSE 0.531586, std 0.044250; paired delta -0.000530, p=0.91125, 3/8 wins.
Benchmark result
Ahead of the baseline (−0.1%) but not statistically significant (p=0.9113, wins 3 / 8 paired seeds).
- Benchmark
- Sequence modeling
- Model
- transformer_tiny
- Paired seeds
- 8
- Baseline mean
- 0.5321
- Idea mean
- 0.5316
- Effect (Δ)
- -0.0005 (−0.1%; negative = idea better)
- Wins
- 3 / 8 paired seeds
- p-value
- 0.9113 (permutation test, 20 000 shuffles)
- Smallest detectable effect
- ±2.0%
- Mechanism
- Confirmed ✓
- Practical effect
- No effect
- Baseline tuning
- swept over 3 configs
This is a small sequence-forecast proxy rather than a full transformer KV-cache admission/reconstruction implementation: the bench model exposes scalar input sequences and does not expose internal K/V tensors. It does not test actual layer/head-specific Tucker factors, global byte-budget allocation, FP16 KV memory, perplexity, retrieval accuracy, decode throughput, decompression bandwidth, or long-context inference.
How to run: python3 run_bench.py
Artifacts
- bench_report.json 4.1 KB View Raw
- joint_tucker_kv.py 5.9 KB View Raw
- report.md 1.6 KB View
- report_bench_2026-08-30T174447.md 5.1 KB View
- results.json 2.5 KB View Raw
- run_bench.py 4.5 KB View Raw