mirror of
https://github.com/TheTom/turboquant_plus.git
synced 2026-10-19 01:30:50 -05:00
Yash Raj Pandey
0cb20bca02
* 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.