fix(qjl): use orthogonal projection and sqrt(d) scale factor

[devYRPauli]

What

The QJL stage in turboquant/qjl.py used a Gaussian random matrix S ~ N(0,1)^(d×d) for sign(S · x) with dequantization scale sqrt(π/2) / d. Both that pairing and the corrected orthogonal-S / sqrt(π/2) / sqrt(d) pairing are unbiased on the inner product ⟨x̂, y⟩ — the QJL stage isn't "off." The actual mechanism is variance, not magnitude. On real LLM KV vectors with head_dim=128:

1. The Gaussian projection introduces variance of order ||r||² / d per reconstructed dimension. Attention accumulation over thousands of tokens drives this into word-loop degeneration ("genetic genetic genetic...") the moment QJL is enabled.
2. Both estimators produce the same expected reconstruction norm ≈ √(π/2)·||x||. The sqrt(π/2) / d scale is the unbiased form for Gaussian S; sqrt(π/2) / sqrt(d) is the unbiased form for orthogonal S. The fix isn't a magnitude correction — it's swapping the projection family and matching the scale to it so variance collapses.

Fix

• Generate S as a sign-corrected random orthogonal matrix via QR. Orthogonal projections preserve norms and inner products exactly; reconstruction variance collapses sharply.
• Use the matching unbiased scale sqrt(π/2) / sqrt(d) so the orthogonal-S estimator stays unbiased on ⟨x̂, y⟩.
• Add a shrinkage parameter (default 1.0 = classical paper-faithful unbiased QJL). MMSE-optimal value is 2/np.pi ≈ 0.6366 — derivation in the docstring (from E[||x̂||²] = (π/2)·||x||² and E[⟨x̂, x⟩] = ||x||²).
• Enforce the orthogonality contract in __init__: ||S Sᵀ − I||_F < 1e-10. New TestQJLProjection.test_projection_matrix_is_orthogonal at d ∈ {64, 128, 256, 512} with the tighter < 1e-12 tolerance.

Empirical impact (M1 Pro 16 GB, Qwen 4B, K5/V4 hybrid)

• Needle-in-haystack retrieval at 16K tokens: 0% → 100% with these QJL fixes combined with K5/V4 Hybrid KV cache.
• Verified orthogonality of S to machine precision: ||S Sᵀ − I||_F ≈ 5e-15 to 4e-14 across d ∈ {64, 128, 256, 512}.
• At default shrinkage=1.0, reconstruction norm ratio averages 1.253 = √(π/2) at every tested d, exactly matching the closed-form E[||x̂||²] = (π/2)·||x||².

Reproducer

Full benchmarks, logs and the post-mortem write-up: https://github.com/devYRPauli/turboquant-m1pro-evaluation

Note: the parallel norm_correction work landed in PolarQuant on 86bcbbe — this PR is independent of that and lands on the same main.

[junleen]

您好，这是来自QQ邮箱的自动回复邮件。我将尽快阅读您的邮件，谢谢！

[TheTom]

Thanks. Math fix is right, and unbiased reconstruction is a real win over what was there. Couple of asks before merge.

1. Damping default.

Closed form: for unbiased orthogonal QJL with the √(d) scale, E[||x̂||²] = (π/2)·||x||² exactly, so MMSE-optimal shrinkage is α* = 2/π ≈ 0.6366. Empirical sweep at d ∈ {64,128,256,512} matches to grid resolution. damping=0.7 is ~2% off optimum.

Could you:

• Default the new kwarg to 1.0 (unbiased, paper-faithful, backward-compatible for TurboQuant.dequantize).
• Rename damping to shrinkage. Document 2/np.pi as the MMSE-optimal value in the docstring with the one-line derivation.

Callers who don't pass an explicit kwarg should keep getting classical unbiased QJL.

2. Framing nit on the PR body.

Original code wasn't ~11× too small / disabled. It was ~1.3× too large in norm and high-variance. Both estimators are unbiased on ⟨x̂, y⟩; the gap is variance, not magnitude. Worth a one-line correction in the description. The new framing is the actual mechanism, and it matches the "QJL eliminates bias but explodes variance" finding from turbo4-resurrection.md, now quantified.

3. Tighten tests in this PR.

Existing test_qjl.py passed the destructive code. Please add:

• assert np.linalg.norm(qjl.S @ qjl.S.T - np.eye(d), 'fro') < 1e-12 in the constructor or as a test (orthogonality contract).
• Tighten test_dequantized_has_correct_scale from [0.5, 2.0] to [0.95, 1.05] for the unbiased classical recovery.

4. Heads-up, not a blocker.

I'll update the README §QJL to flag this stage as reference-only. Production TheTom/llama-cpp-turboquant drops QJL on both K and V (recommended config --cache-type-k q8_0 --cache-type-v turbo3), per the 5-group consensus in turbo4-resurrection.md. Your K5/V4 result is the first interesting K-side QJL data I've seen. Would be great to see it pushed to 64K+ context where the variance mechanism historically manifests (buun's CUDA turbo4 degraded -0.28% at 2K to +3.69% at 64K with unit damping + broken Gaussian S; your math fix + shrinkage might mitigate). Separate issue, not a blocker for this fix.

[devYRPauli]

Thanks. Closed-form derivation is cleaner; using that. Working through:

1. Rename + default 1.0 — yes. damping → shrinkage, default 1.0, docstring documents 2/np.pi as MMSE-optimal with the one-line derivation. TurboQuant.dequantize callers without the kwarg get classical unbiased QJL.

2. Framing fix on PR body — correct, mine was wrong. Both estimators are unbiased on ⟨x̂, y⟩; the gap is variance, not magnitude. Rewriting the description: original was high-variance, new estimator collapses variance, which is what lets the two-stage design deliver the paper's ~64% MSE reduction.

3. Tests — adding the orthogonality contract (||S Sᵀ − I||_F < 1e-12) in __init__ and as a separate test.

Quick math check on test_dequantized_has_correct_scale: under the new shrinkage=1.0, E[||x̂||/||x||] ≈ √(π/2) ≈ 1.253, so [0.95, 1.05] would fail. Two options that match your intent:

• Tighten to [1.20, 1.30] (match classical-QJL norm ratio), or
• Repurpose to the unbiased-contract on the inner product itself: |E[⟨x̂, y⟩] − ⟨x, y⟩| < 3·SE over many trials. This is what the paper actually proves; test_inner_product_unbiased_single_side already covers d=256 — could extend to d ∈ {64, 128, 256, 512} and replace the norm test, or keep both.

Which do you prefer?

4. 64K+ follow-up — noted, separate issue. The variance-accumulation mechanism behind buun's −0.28% at 2K → +3.69% at 64K is exactly what the orthogonal projection should suppress; would be useful to confirm. Setting it up on remote hardware (M1 Pro 16 GB constrained me to 16K).

Pushing the rename + default 1.0 + docstring + orthogonality test now, and updating the PR description. Will tag once code is up so you can pick the test option.

[devYRPauli]

Code pushed in fd18dbe (rename + default 1.0 + docstring with 2/π derivation + orthogonality contract ||S Sᵀ − I||_F < 1e-12 at d ∈ {64, 128, 256, 512}). PR description rewritten with the variance framing.

test_dequantized_has_correct_scale is the one open item — at default shrinkage=1.0 it averages 1.253 = √(π/2) exactly, so let me know your call between [1.20, 1.30] (norm-ratio tightening) and pivoting to the inner-product unbiasedness invariant. Either is a one-line push.

[devYRPauli]

Follow-up clarification after re-reading the paper (arXiv 2504.19874, Definition 1 and Algorithm 2).

The commit message here frames the scale change as fixing a formula that was "~11x too small" in the stock code. That is not quite right, and I want to correct the record.

The paper defines the QJL projection S as a d x d matrix with i.i.d. N(0,1) entries, and the dequantization scale as sqrt(pi/2)/d. That pair is self-consistent and unbiased: the Gaussian rows have norm ~sqrt(d), so S^T z comes out about sqrt(d) larger than it would with orthonormal rows, and the 1/d scale absorbs exactly that factor. The stock implementation followed the paper faithfully and was not carrying a scale bug.

What this PR actually does is a coupled substitution:

1. Replace the Gaussian S with a random orthogonal matrix (QR). This is a variance reduction, not a bias correction, and it is the change that matters empirically.
2. Because the matrix rows are now unit norm, the matching unbiased scale becomes sqrt(pi/2)/sqrt(d). Keeping the old 1/d scale with the orthogonal matrix is what would be ~11x too small at d=128.

So the sqrt(d) factor only exists relative to the orthogonal matrix introduced in the same change. The two edits are one substitution and cannot be separated.

The reason the change is worth making is empirical, not a formula fix: the paper-faithful Gaussian QJL degenerates into word loops at head_dim 128 on real LLM KV vectors, even at very short context, while the orthogonal variant (with the matched scale and the shrinkage factor added later) is stable.

I re-ran a small ablation on M1 Pro (Qwen2.5-3B, 2K context) to confirm which ingredient is load-bearing:

Projection	Scale	Shrinkage	Generation
gaussian (paper)	1/d (paper)	1.0	word-loop degeneration
gaussian	1/sqrt(d)	1.0	collapse
orthogonal	1/d	1.0	word-loop degeneration
orthogonal	1/sqrt(d)	1.0	semi-coherent
orthogonal	1/sqrt(d)	0.7	coherent
gaussian (paper)	1/d (paper)	0.7	word-loop degeneration

So orthogonal + matched scale is what escapes degeneration (neither half works alone), and the shrinkage further stabilizes it. Data and method: https://github.com/devYRPauli/turboquant-m1pro-evaluation (reports/reproduction-and-ablation-2026-07-03.md).

None of this changes the merged code, which is correct. It only corrects the "stock scale was 11x too small" description so the history is accurate. The shrinkage=1.0 default (classical unbiased) and 2/pi MMSE-optimal note in the docstring are both right as merged.