remex

Retrieval-validated embedding compression. 2-16x smaller vectors with measured recall.

Based on the rotation + Lloyd-Max scalar quantization insight from TurboQuant (Zandieh et al., ICLR 2026), focused on the use case that matters most to most AI Engineers: embedding storage and retrieval for RAG systems.

Quick start

from remex import Quantizer

# Compress embeddings — no training data needed
pq = Quantizer(d=384, bits=4)            # d = your embedding dimension
compressed = pq.encode(embeddings)       # (n, 384) float32 → compressed
indices, scores = pq.search(compressed, query, k=10)

# Save/load (bit-packed on disk)
compressed.save("index.npz")
from remex import CompressedVectors
loaded = CompressedVectors.load("index.npz")

The quantizer is fully determined by (d, bits, seed, rotation) — no training, no fitting, no index to ship. rotation defaults to "haar" and is part of the encoding exactly as seed is; every container records it, and decoding against the wrong one raises rather than returning wrong-but-plausible vectors.

How it works

Three steps, each with a clear purpose:

1. Random rotation — A fixed orthogonal matrix (Haar-distributed via QR decomposition) transforms any embedding distribution so that coordinates become approximately i.i.d. N(0, 1/d). This is the key insight from TurboQuant: it makes quantization data-oblivious, meaning no training data is required.

2. Lloyd-Max scalar quantization — Each coordinate is independently quantized using optimal boundaries for the N(0, 1/d) distribution. The codebook is computed from the theoretical Gaussian CDF, not from data. This produces the minimum mean-squared-error scalar quantizer for Gaussian inputs.

3. Bit-packing — Indices are stored at their actual bit width (not wasteful uint8), giving honest compression ratios. A 4-bit codebook uses 4 bits per coordinate on disk.

Norms are stored separately as float32, preserving inner-product ranking up to quantization error.

Why not QJL? TurboQuant includes a QJL (quantized Johnson-Lindenstrauss) residual correction stage for unbiased inner product estimation. We omit it because QJL adds variance that hurts retrieval — when only ranking order matters (not absolute scores), the MSE-optimal rotation + Lloyd-Max stage empirically dominates.

Matryoshka bit precision

An n-bit quantized index's top k bits are a valid k-bit code. remex exploits this: encode once at full bit-width, search at any lower precision by right-shifting indices. Centroid tables are precomputed for all bit levels.

This enables two-stage coarse-to-fine retrieval from a single encoded representation:

pq = Quantizer(d=384, bits=8)
compressed = pq.encode(corpus)

# Two-stage: coarse ADC (Asymmetric Distance Computation) scan at reduced bits,
# then full-precision rerank
indices, scores = pq.search_twostage(
    compressed, query, k=10,
    candidates=200,          # coarse pass returns 200 candidates
    coarse_precision=4,      # coarse scan at 4-bit (default: bits-2)
)

The nesting incurs a small penalty vs independently optimized codebooks: ~1.2% at 4-bit, up to ~10% at 2-bit. In practice this matters little for the coarse stage, which only needs to identify the right neighborhood.

Benchmarks

Recall vs bit level (synthetic, d=384, 10k corpus, 200 queries)

Real embeddings (all-MiniLM-L6-v2, d=384, 10k corpus, 500 queries)

Scaling with corpus size (synthetic, 4-bit)

Full benchmark details and distribution sensitivity analysis in bench/RESULTS.md.

When to use remex / when not to

Use remex when

• You want zero training. The quantizer is deterministic and portable — just (d, bits, seed). No codebook to train, no index to ship, no retraining when your corpus changes.
• You need fast encode. Encoding is ~20μs/vector (rotation + searchsorted). Adding new vectors never requires retraining.
• 8-bit caching is enough. At 8-bit (4x compression), R@10 = 0.974 on real embeddings. Near-lossless and much cheaper than float32.
• You want coarse retrieval + reranking. 4-bit R@10=0.707 is enough for a first pass if you rerank the top candidates with a cross-encoder or full-precision search.

Do not use remex when

• You need high recall at aggressive compression on real data. At 4-bit, FAISS PQ (m=96) achieves R@10=0.816 vs remex's 0.707 on real embeddings. Data-adaptive methods exploit structure that data-oblivious methods cannot.
• Your embeddings form very tight clusters. When cluster spread σ < 0.05, 4-bit R@10 drops to 0.53 (from 0.85 at normal spread). Quantization errors flip rankings among near-identical vectors. 8-bit is much more robust (R@10 stays above 0.95).
• You need sublinear search. remex is brute-force only. For >100k vectors, consider FAISS IVF, HNSW, or similar ANN indices. remex's compact encoding can feed into an external ANN index.

Distribution sensitivity (10k corpus, 4-bit, varying cluster tightness)

Detection: If your 4-bit R@10 is significantly below 0.80 on a held-out set, your embeddings likely have tight clusters. Use 8-bit, or switch to a data-adaptive method.

Compression ratios

Honest packed sizes (bit-packed on disk, d=384):

Float32 baseline: 1,536 bytes/vector (15.36 MB per 10k vectors).

In-memory, indices are stored as uint8 for fast search. The PackedVectors class keeps them bit-packed in memory too, using 2-4x less RAM for sub-byte widths.

API reference

Quantizer(d, bits=4, seed=42)

Main quantizer class (formerly PolarQuantizer, which remains available as a deprecated alias).

• d — Vector dimension (must match your embeddings).
• bits — Bits per coordinate: 1-4 or 8. Sweet spot is 3-4. Use 8 for near-lossless.
• seed — Random seed for the rotation matrix. Same seed = same quantizer.

Methods

encode(X) — Quantize (n, d) float32 array. Returns CompressedVectors.

decode(compressed, precision=None) — Reconstruct (n, d) float32 from compressed. Optional precision (1 to bits) for Matryoshka decode.

search(compressed, query, k=10, precision=None) — Find k nearest neighbors by approximate inner product. Caches a dequantized float32 matrix for fast repeated queries. Returns (indices, scores).

search_batch(compressed, queries, k=10, precision=None) — Batch version of search() using matrix multiplication for better throughput. Returns (indices, scores) where both are (n_queries, k).

search_adc(compressed, query, k=10, precision=None, chunk_size=4096) — Memory-efficient search via ADC (Asymmetric Distance Computation) lookup-table scoring. No float32 cache — peak memory is chunk_size * d * 4 bytes (~6 MB). Slower per-query but uses ~5x less RAM. Returns (indices, scores).

search_twostage(compressed, query, k=10, candidates=500, coarse_precision=None) — Two-stage Matryoshka retrieval: ADC coarse scan (no cache) then full-precision rerank on candidates only. Memory-efficient: only the small candidate set is dequantized. Returns (indices, scores).

mse(X, precision=None) — Mean per-vector reconstruction error (L2 squared).

CompressedVectors

Container for quantized data. Created by Quantizer.encode(). Stores indices as uint8 in memory for fast search/decode.

Properties

• n — Number of vectors.
• nbytes — Bit-packed size in bytes (honest compression).
• nbytes_unpacked — In-memory size (uint8 indices + float32 norms).
• compression_ratio — (n * d * 4) / nbytes.
• resident_bytes — Actual RAM including any active caches.

Methods

• save(path) / load(path) — Save/load to .npz with bit-packed indices.
• save_arrow(path) / load_arrow(path) — Save/load to Arrow IPC (Feather v2) format. Requires pyarrow.
• subset(idx) — Return a new CompressedVectors with only the given row indices.
• drop_cache() — Free the dequantized float32 cache to reclaim memory.

PackedVectors

Memory-efficient packed storage. Keeps indices bit-packed in memory, unpacking on demand. Uses 2-4x less RAM than CompressedVectors for sub-byte widths.

from remex import PackedVectors

packed = PackedVectors.from_compressed(compressed)  # pack in memory
packed = PackedVectors.from_rows(rows, norms, d=384, bits=4)  # from DB rows

# ADC and two-stage search work directly on PackedVectors
indices, scores = pq.search_adc(packed, query, k=10)
indices, scores = pq.search_twostage(packed, query, k=10)

# Matryoshka precision reduction
packed_2bit = packed.at_precision(2)

# Convert back if needed
compressed = packed.to_compressed()

Cached search() is not supported on PackedVectors — use search_adc() or search_twostage(), or convert with to_compressed().

GPUSearcher (optional)

GPU-accelerated search wrapper. Requires CuPy or PyTorch with CUDA. Falls back to NumPy.

from remex.gpu import GPUSearcher

searcher = GPUSearcher(pq, compressed)
indices, scores = searcher.search(query, k=10)
indices, scores = searcher.search_adc(query, k=10)
indices, scores = searcher.search_twostage(query, k=10, candidates=200)

IVFCoarseIndex (sublinear coarse-tier scan)

Inverted-file index over the coarse Matryoshka tier. Lets you visit only nprobe of 2**n_bits cells per query, replacing the bandwidth-bound flat coarse scan in two-stage retrieval. Stays data-oblivious — no k-means, no training, no fitting.

from remex import IVFCoarseIndex, Quantizer

pq = Quantizer(d=768, bits=8, seed=42)
compressed = pq.encode(corpus)              # CompressedVectors or PackedVectors

# Mode 1: random-hyperplane LSH (SimHash). Pure data-oblivious — works
# on any embedding distribution. Determined by (d, n_bits, seed).
ivf = IVFCoarseIndex(pq, compressed, n_bits=12, mode="lsh", seed=0)

# Mode 2: sign of the first n_bits post-rotation coords. Free given
# the existing rotation (these bits are already MSBs of the encoded
# indices). Cell balance depends on rotated coords being ~i.i.d.
# Gaussian, which is checked by bench/specter2_eval.py.
ivf = IVFCoarseIndex(pq, compressed, n_bits=12, mode="rotated_prefix")

# Stage-1 only — top-K candidates from the visited cells, ADC scored
indices, scores = ivf.search_coarse(query, k=500, nprobe=8, precision=1)

# End-to-end: IVF coarse + full-precision rerank
indices, scores = ivf.search_twostage(
    query, k=10, candidates=500, nprobe=8, coarse_precision=1
)

Multi-probe is by Hamming distance from the query's hash code: the nprobe cells with the lowest Hamming distance to q_hash are visited (ties broken by cell ID). Setting nprobe = 2**n_bits recovers a flat scan; the index is exact in that limit and tests verify byte-identical agreement with Quantizer.search_adc / Quantizer.search_twostage.

When IVF wins, when flat-scan wins

IVF is for the regime where stage-1 latency is the bottleneck (≥ tens of millions of vectors). The trade-off is recall vs latency:

Flat-scan wins when:

• Corpus < ~10M vectors. Stage-1 is already < 50 ms; the IVF index overhead and per-query hash cost don't pay back.
• Recall@K must equal flat-scan exactly. IVF is approximate by construction — vectors in unvisited cells are missed.
• Embeddings are deeply mixed and queries are uniformly distributed in angle, so cells don't capture meaningful neighborhoods.

Bridge-edge preservation (cross-FoS / cross-partition recall) is benchmarked explicitly in bench/specter2_eval.py — running broad + narrow SPECTER2 partitions concatenated and reporting how many of the flat-scan top-K cross-partition hits the IVF top-K preserves at each nprobe. Both hash modes are content-based (hyperplane signs on the rotated representation), so they don't partition by FoS — but at very low nprobe cross-partition hits drop simply because pool size shrinks.

Memory cost (excluding the corpus)

For 100M vectors at n_bits=12: ~960 MB index overhead vs ~9.6 GB 1-bit coarse memory — about 10% surcharge for ~5–20× stage-1 speedup at moderate nprobe.

Memory profiles (100k vectors, d=384, 8-bit)

Choose search() when latency matters and RAM is available. Choose search_adc() or search_twostage() when memory is constrained (serverless, edge, or very large corpora).

Low-level utilities

from remex import pack, unpack, packed_nbytes
from remex import lloyd_max_codebook, nested_codebooks

• pack(indices, bits) / unpack(packed, bits, n_values) — Bit-pack/unpack uint8 arrays.
• packed_nbytes(n_values, d, bits) — Compute packed byte count.
• lloyd_max_codebook(d, bits) — Generate optimal boundaries and centroids for N(0, 1/d).
• nested_codebooks(d, max_bits) — Build Matryoshka centroid tables for all bit levels 1..max_bits.

vs TurboQuant

TurboQuant (Zandieh et al., ICLR 2026) adds QJL (quantized Johnson-Lindenstrauss) residual correction for unbiased inner product estimates. This is important for KV cache attention, where unbiased estimation matters. For retrieval (ranking by approximate inner product), the QJL variance hurts more than the debiasing helps. remex implements only the MSE-optimal rotation + Lloyd-Max stage, which empirically dominates for nearest-neighbor search.

vs FAISS Product Quantization

Use FAISS when: You have a stable, large corpus, need sublinear search, and can afford training time.

Use remex when: You want zero training, fast encode, frequently changing corpora, or near-lossless 8-bit caching (R@10=0.974 at 4x compression).

vs scalar quantization (naive rounding)

Without the rotation step, scalar quantization on raw embeddings is catastrophically bad — embeddings are highly anisotropic (variance ratios of 10^7x across dimensions). The random rotation spreads information uniformly across coordinates, making scalar quantization viable.

At 3-bit, remex achieves 72-80% R@10 vs ~40% for naive scalar quantization on the same data.

Installation

pip install remex                       # from PyPI (when published)
pip install -e ".[dev]"                 # development: + pytest, pytest-cov
pip install -e ".[bench]"              # benchmarking: + faiss-cpu, sentence-transformers

Testing

pytest                                  # 126 tests (~6 min)
pytest tests/test_polar_embed.py -v     # core tests
pytest tests/test_matryoshka.py -v      # Matryoshka/nested codebook tests
pytest tests/test_adc_gpu.py -v         # ADC and GPU searcher tests
pytest tests/test_packed_vectors.py -v  # PackedVectors tests

Mojo port (polarquant)

A standalone Mojo CLI binary lives in remex/mojo/. It mirrors the encode + ADC search path with no Python runtime dependency, reading .npy corpus files directly and writing a small binary .pq container that the Python library can load via remex.load_pq() (and vice versa via remex.save_pq()).

cd remex/mojo
mojo build -I . polarquant.mojo -o polarquant
./polarquant encode corpus.npy --bits 4 --seed 42 -o corpus.pq
./polarquant search corpus.pq query.npy --k 10 --seed 42

For bit-identical encoding to Python (matching rotations and codebook), use --params P.bin after dumping with remex.save_params(quantizer, P). See remex/mojo/README.md for build, test, and benchmark instructions.

References

• Zandieh et al. (2025). TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate. ICLR 2026. arXiv:2504.19874
• Jégou, Douze & Schmid (2011). Product Quantization for Nearest Neighbor Search. IEEE TPAMI 33(1):117–128. IEEE Xplore — introduces Product Quantization, ADC (Asymmetric Distance Computation), and SDC for approximate nearest neighbor search.
• Kusupati et al. (2022). Matryoshka Representation Learning. NeurIPS 2022. arXiv:2205.13147 — the nested/coarse-to-fine representation idea that inspires remex's bit-level nesting.
• Mezzadri (2007). How to Generate Random Matrices from the Classical Compact Groups. Notices of the AMS 54(5):592–604. arXiv:math-ph/0609050 — the QR-of-Gaussian method for Haar-distributed orthogonal matrices used in remex/rotation.py.
• Lloyd (1982). Least Squares Quantization in PCM. IEEE Trans. Information Theory 28(2):129–137. IEEE Xplore — optimal scalar quantization (Lloyd-Max algorithm) for minimum MSE.

License

MIT