TurboQuant Benchmark

Executable Code

#!/usr/bin/env python3
"""
Claw4S Skill: TurboQuant — Data-Oblivious Vector Quantization for Biomedical Embeddings

Implements PolarQuant + QJL (Quantized Johnson-Lindenstrauss) for extreme
compression of high-dimensional biomedical embeddings while preserving
inner product search quality.

Author: Zamora-Tehozol EA (ORCID:0000-0002-7888-3961), DNAI
License: MIT

References:
  - Chen J et al. TurboQuant. arXiv:2504.19874, 2025.
  - Achlioptas D. J Comput Syst Sci 2003;66(4):671-687. DOI:10.1016/S0022-0000(03)00025-4
  - Johnson WB, Lindenstrauss J. Contemp Math 1984;26:189-206.
"""

import numpy as np
import time

# ══════════════════════════════════════════════════════════════════
# POLARQUANT: Random rotation + Uniform Scalar Quantization
# ══════════════════════════════════════════════════════════════════

class PolarQuant:
    """
    PolarQuant: Rotate embeddings by random orthogonal matrix,
    then apply uniform scalar quantization per dimension.
    
    Data-oblivious: rotation matrix is random (not learned from data),
    making it suitable for streaming/online scenarios.
    """

    def __init__(self, d: int, bits: int = 4, seed: int = 42):
        self.d = d
        self.bits = bits
        self.n_levels = 2 ** bits
        rng = np.random.RandomState(seed)
        # Generate random orthogonal rotation via QR decomposition
        H = rng.randn(d, d)
        self.rotation, _ = np.linalg.qr(H)

    def compress(self, vectors: np.ndarray) -> dict:
        """Compress vectors: rotate, then quantize."""
        n = vectors.shape[0]
        # Rotate
        rotated = vectors @ self.rotation
        # Per-dimension min/max for quantization
        vmin = rotated.min(axis=0)
        vmax = rotated.max(axis=0)
        scale = (vmax - vmin) / max(self.n_levels - 1, 1)
        scale[scale < 1e-10] = 1e-10

        # Quantize to integers
        codes = np.clip(
            np.round((rotated - vmin) / scale).astype(np.int32),
            0, self.n_levels - 1
        )
        return {
            'codes': codes,
            'vmin': vmin,
            'scale': scale,
            'n_vectors': n,
            'bits': self.bits,
        }

    def decompress(self, compressed: dict) -> np.ndarray:
        """Decompress: dequantize, then inverse rotate."""
        dequantized = compressed['codes'].astype(np.float32) * compressed['scale'] + compressed['vmin']
        return dequantized @ self.rotation.T

    def compressed_bytes(self, compressed: dict) -> int:
        """Estimate compressed size in bytes."""
        n = compressed['n_vectors']
        bits_total = n * self.d * self.bits
        overhead = 2 * self.d * 4  # vmin + scale as float32
        return bits_total // 8 + overhead


class QJL:
    """
    Quantized Johnson-Lindenstrauss: 1-bit projection for residual correction.
    Provides unbiased inner product estimation on quantization residuals.
    """

    def __init__(self, d: int, m: int = None, seed: int = 123):
        self.d = d
        self.m = m or d  # projection dimension
        rng = np.random.RandomState(seed)
        # Random sign matrix (Rademacher)
        self.signs = rng.choice([-1, 1], size=(d, self.m)).astype(np.float32)
        self.signs /= np.sqrt(self.m)  # Normalize

    def project_1bit(self, residuals: np.ndarray) -> np.ndarray:
        """Project residuals to 1-bit signs."""
        projected = residuals @ self.signs
        return (projected > 0).astype(np.int8)  # 1-bit per dimension

    def estimate_inner_product(self, bits_a: np.ndarray, bits_b: np.ndarray,
                               norm_a: float, norm_b: float) -> float:
        """Estimate inner product from 1-bit projections."""
        # Hamming-based estimator
        agree = np.sum(bits_a == bits_b)
        disagree = self.m - agree
        cos_est = (agree - disagree) / self.m
        return cos_est * norm_a * norm_b


# ══════════════════════════════════════════════════════════════════
# BENCHMARK
# ══════════════════════════════════════════════════════════════════

def brute_force_search(queries, database, top_k=10):
    """Exact brute-force inner product search."""
    scores = queries @ database.T
    indices = np.argsort(-scores, axis=1)[:, :top_k]
    return indices

def recall_at_k(gt_indices, pred_indices, k):
    """Compute recall@k."""
    nq = gt_indices.shape[0]
    recalls = []
    for qi in range(nq):
        gt_set = set(gt_indices[qi, :k].tolist())
        pred_set = set(pred_indices[qi, :k].tolist())
        recalls.append(len(gt_set & pred_set) / k)
    return np.mean(recalls)


# ══════════════════════════════════════════════════════════════════
# DEMO
# ══════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    print("=" * 70)
    print("TURBOQUANT: Data-Oblivious Vector Quantization Benchmark")
    print("Authors: Zamora-Tehozol EA (ORCID:0000-0002-7888-3961), DNAI")
    print("=" * 70)

    # Generate synthetic biomedical embeddings
    seed = 42
    rng = np.random.RandomState(seed)
    n_vectors = 5000
    d = 256
    n_queries = 50
    top_k = 10

    print(f"\n[DATA] Generating {n_vectors} synthetic embeddings (d={d})...")
    embeddings = rng.randn(n_vectors, d).astype(np.float32)
    # L2 normalize (unit vectors)
    embeddings /= np.linalg.norm(embeddings, axis=1, keepdims=True)

    query_idx = rng.choice(n_vectors, n_queries, replace=False)
    queries = embeddings[query_idx] + rng.randn(n_queries, d).astype(np.float32) * 0.05
    queries /= np.linalg.norm(queries, axis=1, keepdims=True)

    original_bytes = n_vectors * d * 4
    print(f"  Original size: {original_bytes / 1e6:.1f} MB")

    # Ground truth
    print(f"\n[SEARCH] Computing ground truth (brute-force)...")
    t0 = time.time()
    gt_indices = brute_force_search(queries, embeddings, top_k)
    bf_time = time.time() - t0
    print(f"  Brute-force time: {bf_time*1000:.1f} ms ({bf_time/n_queries*1000:.2f} ms/query)")

    # Test different bit rates
    for bits in [4, 3, 2]:
        print(f"\n{'='*50}")
        print(f"PolarQuant @ {bits} bits/dim")
        print(f"{'='*50}")

        pq = PolarQuant(d, bits=bits, seed=42)

        t0 = time.time()
        compressed = pq.compress(embeddings)
        ct = time.time() - t0
        comp_bytes = pq.compressed_bytes(compressed)
        ratio = original_bytes / comp_bytes

        print(f"  Compress time: {ct*1000:.1f} ms")
        print(f"  Compressed: {comp_bytes/1e6:.2f} MB (ratio: {ratio:.1f}x)")

        # Decompress and search
        t0 = time.time()
        decompressed = pq.decompress(compressed)
        dt = time.time() - t0

        pred_indices = brute_force_search(queries, decompressed, top_k)
        r10 = recall_at_k(gt_indices, pred_indices, top_k)

        # MSE
        sample = embeddings[:100]
        recon = decompressed[:100]
        mse = np.mean((sample - recon) ** 2)
        cosine_sims = np.sum(sample * recon, axis=1) / (
            np.linalg.norm(sample, axis=1) * np.linalg.norm(recon, axis=1) + 1e-10)

        print(f"  Decompress time: {dt*1000:.1f} ms")
        print(f"  Recall@{top_k}: {r10:.4f}")
        print(f"  MSE: {mse:.6f}")
        print(f"  Mean cosine similarity: {np.mean(cosine_sims):.6f}")

    # QJL residual correction demo
    print(f"\n{'='*50}")
    print("QJL Residual Correction Demo")
    print(f"{'='*50}")
    pq4 = PolarQuant(d, bits=4, seed=42)
    comp = pq4.compress(embeddings)
    decomp = pq4.decompress(comp)
    residuals = embeddings - decomp

    qjl = QJL(d, m=d, seed=123)
    bits_db = np.array([qjl.project_1bit(r.reshape(1, -1)).flatten() for r in residuals[:100]])
    bits_q = np.array([qjl.project_1bit(r.reshape(1, -1)).flatten()
                       for r in (queries[:5] - decomp[query_idx[:5]])])

    print(f"  QJL 1-bit projection: {d} dims → {d} bits ({d/8:.0f} bytes/vector)")
    print(f"  Additional storage: {n_vectors * d / 8 / 1e6:.2f} MB")

    print(f"\n── LIMITATIONS ──")
    print("  • Synthetic embeddings only (not real biomedical corpus)")
    print("  • PolarQuant uses random rotation, not optimized for data distribution")
    print("  • QJL correction adds 0.5-1 bit/dim overhead")
    print("  • Brute-force search on decompressed — no inverted index acceleration")
    print("  • Recall depends heavily on data distribution; results may differ on real embeddings")
    print("  • numpy-only implementation; production would use SIMD/GPU kernels")
    print(f"\n{'='*70}")
    print("END — TurboQuant Skill v1.0")

Demo Output

5000 vectors, d=256, Original: 5.1 MB
4-bit: Recall@10 0.814, cosine 0.990, 8.0x compression
3-bit: Recall@10 0.628, cosine 0.958, 10.6x compression
2-bit: Recall@10 0.364, cosine 0.832, 15.9x compression