diff --git a/tests/test_qjl.py b/tests/test_qjl.py index 7915196..6b53c0f 100644 --- a/tests/test_qjl.py +++ b/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.""" diff --git a/turboquant/qjl.py b/turboquant/qjl.py index f7a4b39..7daa0e1 100644 --- a/turboquant/qjl.py +++ b/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 diff --git a/turboquant/turboquant.py b/turboquant/turboquant.py index 2a3bc88..f87ab2a 100644 --- a/turboquant/turboquant.py +++ b/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