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

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

The QJL stage used a Gaussian random matrix S ~ N(0,1)^(d×d) for
sign(S · x). Gaussian projections introduce variance of order ||r||^2/d
per dimension; for head_dim=128, this is large relative to signal, and
attention accumulation drives generation into word-loop degeneration
once QJL is enabled on real LLM KV vectors.

Fix:
* Replace Gaussian S with a random orthogonal matrix via QR
  decomposition (sign-corrected to give a proper rotation). Orthogonal
  matrices preserve norms and inner products and reduce reconstruction
  variance sharply.
* Correct the dequantization scale from sqrt(pi/2)/d to sqrt(pi/2)/sqrt(d).
  For d=128, the previous formula was ~11x too small, which effectively
  disabled the QJL correction term even when the stage was enabled.
* Add a damping parameter (default 0.7) to dequantize(), validated as
  the stable operating point in M1 Pro 16K-context needle benchmarks
  on Qwen 4B with K5/V4 hybrid quantization.

Empirical impact on M1 Pro (16 GB, Qwen 4B):
* Needle retrieval at 16K tokens: 0% → 100% with these QJL fixes
  combined with K5/V4 Hybrid configuration.
* Verified orthogonality: ||S S^T - I||_F < 1e-14.
* Existing tests pass (e.g., test_dequantized_has_correct_scale
  reports avg norm ratio 0.877 at d=128, well within [0.5, 2.0]).

Refs: turboquant-m1pro-evaluation reproducer and post-mortem at
https://github.com/devYRPauli/turboquant-m1pro-evaluation

* refactor(qjl): address review — rename damping → shrinkage, default 1.0

Per @TheTom's review on PR #93:

* Rename `damping` kwarg to `shrinkage` on QJL.dequantize and
  TurboQuant.dequantize. Default 1.0 — classical paper-faithful
  unbiased estimator. Existing callers without the kwarg get
  classical unbiased QJL (backward-compatible).
* Docstrings document `2/np.pi ≈ 0.6366` as the MMSE-optimal
  shrinkage with the one-line derivation from E[||x̂||²] =
  (π/2)·||x||² and E[⟨x̂, x⟩] = ||x||².
* Add the orthogonality contract:
  - `__init__` asserts `||S Sᵀ − I||_F < 1e-10` (cheap defensive
    check; required for the unbiased estimator).
  - New `TestQJLProjection.test_projection_matrix_is_orthogonal`
    at d ∈ {64, 128, 256, 512} with the tighter `< 1e-12`
    tolerance @TheTom suggested.

Verified: orthogonality error 5e-15 to 4e-14 across all tested d.
Reconstruction norm ratio under default shrinkage=1.0 averages
1.253 = √(π/2) at every d, exactly matching the closed form.

tests/test_qjl.py

@@ -6,6 +6,21 @@ import pytest
 from turboquant.qjl import QJL, QJL_CONST
 
 
+class TestQJLProjection:
+    """Structural invariants of the projection matrix S."""
+
+    @pytest.mark.parametrize("d", [64, 128, 256, 512])
+    def test_projection_matrix_is_orthogonal(self, d):
+        """S must be orthogonal: ||S Sᵀ − I||_F < 1e-12.
+
+        Required for the classical unbiased estimator E[⟨x̂, y⟩] = ⟨x, y⟩
+        and the closed-form E[||x̂||²] = (π/2)·||x||².
+        """
+        qjl = QJL(d=d, seed=42)
+        err = np.linalg.norm(qjl.S @ qjl.S.T - np.eye(d), "fro")
+        assert err < 1e-12, f"||S Sᵀ − I||_F = {err:.2e} exceeds 1e-12"
+
+
 class TestQJLRoundTrip:
     """QJL quantize → dequantize should approximately preserve vectors."""
 

turboquant/qjl.py

@@ -1,12 +1,15 @@
 """QJL: Quantized Johnson-Lindenstrauss Transform.
 
-1-bit quantization via random projection → sign to compress vectors while
-preserving inner products. Note: this implementation stores a full d×d projection
-matrix (O(d²) memory). For large d, a structured/seeded approach would be needed.
+1-bit quantization via random orthogonal projection → sign to compress
+vectors while preserving inner products. Stores a full d×d projection
+matrix (O(d²) memory); for large d a structured/seeded approach would be
+needed.
 
-Key property: unbiased and optimal at 1-bit.
-    Q_qjl(x) = sign(S · x) where S ~ N(0,1)^(d×d)
-    Q_qjl_inv(z) = √(π/2) / d · S^T · z
+Key properties (orthogonal S):
+    Q_qjl(x) = sign(S · x), S orthogonal
+    Q_qjl_inv(z) = √(π/2) / √d · ||x|| · S^T · z
+    E[⟨x̂, y⟩] = ⟨x, y⟩  (unbiased, paper Theorem 2)
+    E[||x̂||²] = (π/2) · ||x||²
 """
 
 import numpy as np
@@ -18,12 +21,18 @@ QJL_CONST = np.sqrt(np.pi / 2)
 class QJL:
     """Quantized Johnson-Lindenstrauss 1-bit quantizer.
 
+    Uses a random orthogonal projection matrix for minimum variance in the
+    inner-product estimation. Orthogonality is enforced in __init__ and is
+    a required invariant for the classical unbiased estimator.
+
     Usage:
         qjl = QJL(d=128, seed=42)
         signs, norm = qjl.quantize(residual)
         r_hat = qjl.dequantize(signs, norm)
     """
 
+    _ORTHO_TOL = 1e-10
+
     def __init__(self, d: int, seed: int = 123):
         """
         Args:
@@ -32,8 +41,19 @@ class QJL:
         """
         self.d = d
         rng = np.random.default_rng(seed)
-        # Random projection matrix S ∈ R^(d×d), entries ~ N(0,1)
-        self.S = rng.standard_normal((d, d))
+        G = rng.standard_normal((d, d))
+        Q, R = np.linalg.qr(G)
+        diag_signs = np.sign(np.diag(R))
+        Q = Q * diag_signs[np.newaxis, :]
+        self.S = Q
+
+        # Orthogonality contract: ||S Sᵀ − I||_F must be ~0.
+        # Required for E[⟨x̂, y⟩] = ⟨x, y⟩ and E[||x̂||²] = (π/2)·||x||².
+        ortho_err = np.linalg.norm(self.S @ self.S.T - np.eye(d), "fro")
+        assert ortho_err < self._ORTHO_TOL, (
+            f"QJL projection matrix not orthogonal: ||S Sᵀ − I||_F = {ortho_err:.2e} "
+            f"(tolerance {self._ORTHO_TOL:.0e})"
+        )
 
     def quantize(self, r: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
         """Quantize residual vector(s) to sign bits.
@@ -50,28 +70,37 @@ class QJL:
         if single:
             r = r[np.newaxis, :]
 
-        # Compute norms before projection
-        norms = np.linalg.norm(r, axis=1)  # (batch,)
-
-        # Project: S @ r.T → (d, batch), transpose to (batch, d)
+        norms = np.linalg.norm(r, axis=1)
         projected = (self.S @ r.T).T
-
-        # Sign quantization: +1 or -1
         signs = np.sign(projected).astype(np.int8)
-        # Handle exact zeros (extremely rare) — map to +1
         signs[signs == 0] = 1
 
-        # Zero vectors: signs are meaningless, norm=0 ensures dequantize returns zero
         if single:
             return signs[0], norms[0]
         return signs, norms
 
-    def dequantize(self, signs: np.ndarray, norms: np.ndarray) -> np.ndarray:
+    def dequantize(
+        self,
+        signs: np.ndarray,
+        norms: np.ndarray,
+        shrinkage: float = 1.0,
+    ) -> np.ndarray:
         """Dequantize sign bits back to approximate residual.
 
         Args:
             signs: Sign bits, shape (d,) or (batch, d).
             norms: Residual norms, scalar or (batch,).
+            shrinkage: Multiplicative factor applied to the classical QJL
+                reconstruction. The default of ``1.0`` is the paper's
+                classical unbiased estimator: E[⟨x̂, y⟩] = ⟨x, y⟩.
+
+                The MMSE-optimal value is ``2/np.pi ≈ 0.6366``. Derivation:
+                with orthogonal S and the √d scale, E[||x̂||²] = (π/2)·||x||²
+                and E[⟨x̂, x⟩] = ||x||², so the α minimising
+                E[||αx̂ − x||²] = α²·(π/2)·||x||² − 2α·||x||² + ||x||² is
+                α* = 2/π. Shrinking by α* gives a biased but lower-MSE
+                estimator that is preferable for downstream MSE-sensitive
+                consumers (e.g. attention-score correction).
 
         Returns:
             Approximate residual, same shape as original.
@@ -83,13 +112,9 @@ class QJL:
         else:
             signs = signs.astype(np.float64)
 
-        # x̃_qjl = √(π/2) / d · γ · S^T @ signs
-        # S^T @ signs.T → (d, batch), transpose to (batch, d)
-        reconstructed = (self.S.T @ signs.T).T  # (batch, d)
-
-        # Scale by √(π/2) / m * norm, where m = number of projections (= d here)
-        # If projection matrix changes from (d,d) to (m,d), update this divisor
-        scale = QJL_CONST / self.d * norms  # (batch,)
-        reconstructed *= scale[:, np.newaxis]
+        # x̂ = √(π/2) / √d · ||x|| · S^T · signs
+        reconstructed = (self.S.T @ signs.T).T
+        scale = QJL_CONST / np.sqrt(self.d) * norms
+        reconstructed *= (scale[:, np.newaxis] * shrinkage)
 
         return reconstructed[0] if single else reconstructed

turboquant/turboquant.py

@@ -85,11 +85,16 @@ class TurboQuant:
             bit_width=self.bit_width,
         )
 
-    def dequantize(self, compressed: CompressedVector) -> np.ndarray:
+    def dequantize(self, compressed: CompressedVector, shrinkage: float = 1.0) -> np.ndarray:
         """Dequantize back to approximate vector.
 
         Args:
             compressed: CompressedVector from quantize().
+            shrinkage: Multiplicative factor applied to the QJL stage.
+                Default ``1.0`` is the classical unbiased estimator
+                (paper-faithful, backward-compatible). MMSE-optimal is
+                ``2/np.pi ≈ 0.6366`` — see ``QJL.dequantize`` for the
+                derivation.
 
         Returns:
             Reconstructed vector(s), same shape as original.
@@ -98,7 +103,9 @@ class TurboQuant:
         x_mse = self.polar_quant.dequantize(compressed.mse_indices, compressed.vector_norms)
 
         # Stage 2: QJL residual reconstruction
-        x_qjl = self.qjl.dequantize(compressed.qjl_signs, compressed.residual_norms)
+        x_qjl = self.qjl.dequantize(
+            compressed.qjl_signs, compressed.residual_norms, shrinkage=shrinkage
+        )
 
         return x_mse + x_qjl