Summary

Implement QJL (Quantized Johnson-Lindenstrauss) — the 1-bit quantizer that provides unbiased inner product estimation. This is the second stage of TurboQuant, applied to the residual from PolarQuant.

Paper Reference

• QJL paper: arXiv 2406.03482
• TurboQuant paper: arXiv 2504.19874, Section 1.3 (Algorithm 2, Stage 2)
• Key property: unbiased inner product estimator at 1-bit with zero memory overhead

Algorithm

Setup:
  Generate random projection matrix S ∈ R^(d×d), S_ij ~ N(0,1)

Quantize(r):
  1. signs ← sign(S @ r)                    # {+1, -1}^d
  2. γ ← ||r||_2                             # residual norm (stored as metadata)
  3. Return (signs, γ)

Dequantize(signs, γ):
  1. r̃ ← √(π/2) / d · γ · S^T @ signs
  2. Return r̃

Requirements

QJL class

• __init__(d, seed) — generates random projection matrix S
• quantize(r) — single (d,) or batch (batch, d) → (signs, norms)
  • signs: int8 array of {+1, -1}
  • norms: float64 scalar or array
• dequantize(signs, norms) — reconstruct approximate residual
• QJL_CONST = √(π/2) — module-level constant

Critical implementation details

• S is a full d×d matrix of N(0,1) entries — O(d²) memory
  • For d=128 (typical head_dim): 128² × 8 bytes = 131 KB — acceptable
  • For d=3072: 72 MB — may need structured alternative later (separate issue)
• sign(0) → +1 (convention, measure-zero event for continuous inputs)
• When norm = 0 (zero residual), dequantize must return zero vector

Mathematical verification needed

• Unbiasedness: E[⟨y, Q⁻¹(Q(x))⟩] = ⟨y, x⟩ (paper Theorem 2)
  • This holds for SINGLE-SIDE quantization (only x quantized, y exact)
  • When BOTH sides are quantized, it's no longer exactly unbiased
• Dequantization constant √(π/2) comes from the expected absolute value of N(0,1)

Tests Required (write FIRST)

1. Signs are binary: all values exactly +1 or -1, dtype int8
2. Zero vector: norm = 0, dequantize returns zero vector
3. Deterministic: same seed + input → same output
4. Batch matches single: quantize(batch)[i] == quantize(single)
5. Norm preservation (approximate): ||dequantize(quantize(x))|| / ||x|| averages near 1.0 over many samples
6. Unbiasedness (single-side): E[⟨y, dequant(quant(x))⟩] ≈ ⟨y, x⟩ for exact y
   • Average over 500+ random (x, y) pairs
   • Mean error should be within 3 standard errors of zero
7. Dequantized scale: √(π/2) / d · γ · S^T @ signs — verify the constant is correct by checking that the estimator is unbiased
8. Large dimension: works at d = 512 without numerical issues

Acceptance Criteria

• Tests written and reviewed with codex BEFORE implementation
• Unbiasedness verified empirically (single-side)
• Zero-input edge case handled correctly
• codex-review on implementation
• Roast review on implementation
• Coverage >95% for qjl.py