TurboQuantRAG

Measuring how vector quantization error propagates through a Retrieval-Augmented Generation pipeline.

This project integrates TurboQuant — a theoretically grounded online vector quantization algorithm — into a full RAG pipeline and measures the compression-quality tradeoff at every stage: from raw MSE distortion, through retrieval recall, to downstream QA exact match. It provides the first empirical Pareto frontier of compression ratio vs QA accuracy for TurboQuant on MS-MARCO + NQ-open.

Course project — AMS 691, Spring 2026.

---

Motivation

The TurboQuant paper benchmarks quantization quality using ANN recall on embedding datasets. What it does not measure is whether those quantization-level improvements actually matter for the task that embeddings are built for: answering questions. This project fills that gap by running a controlled experiment where quantization error is the only variable and the downstream metric is QA exact match.

---

Methods Compared

Bit widths swept: b ∈ {1, 2, 3, 4} for TurboQuant variants.

---

Pipeline

MS-MARCO 500k passages
        │
        ▼
  all-MiniLM-L6-v2 embeddings (dim=384, L2-normalised)
        │
        ├──► TurboQuant-MSE encode   ┐
        ├──► TurboQuant-IP encode    ├── RetrievalIndex (unified interface)
        ├──► FAISS-PQ encode         │
        └──► FP32 store              ┘
                    │
                    ▼
         NQ-open validation queries (1 000 questions)
                    │
                    ▼
         Retrieve top-5 passages (Recall@{1,5,10}, MRR@10)
                    │
                    ▼
         Flan-T5-base reader → answer generation
                    │
                    ▼
         Exact Match, Token F1

---

Key Results

All experiments run on an NVIDIA L4 GPU (24 GB VRAM, GCP us-central1-a).

Pareto Frontier: Compression vs QA Accuracy

TurboQuant-MSE 4-bit achieves statistically equivalent QA accuracy to the FP32 oracle (overlapping 95% CIs) at 8× compression and 8.7× lower query latency.

MSE Distortion (confirms quantization theory)

Distortion decreases ~3.5× per bit — consistent with the 6 dB/bit bound from quantization theory, validating the Lloyd-Max codebook.

Efficiency

Note: TurboQuant ran on GPU; FAISS-PQ on CPU (faiss-gpu unavailable for Python 3.12). Latency comparison is indicative, not directly apples-to-apples.

---

Plots

All plots are in experiments/plots/:

---

Project Structure

TurboQuantRAG/
├── configs/
│   └── experiment.yaml          # all hyperparameters
├── src/
│   ├── quantization/
│   │   ├── codebooks.py         # Lloyd-Max codebook (Beta distribution, 200 iters, 2M samples)
│   │   ├── turboquant.py        # TurboQuantMSE and TurboQuantIP implementations
│   │   └── faiss_pq.py          # FAISS-PQ wrapper and FlatFP32Index oracle
│   ├── retrieval/
│   │   ├── index.py             # RetrievalIndex unified abstraction (add/search/profile)
│   │   ├── embedder.py          # all-MiniLM-L6-v2 batch embedder
│   │   └── metrics.py           # Recall@k, MRR@k
│   ├── data/
│   │   ├── msmarco.py           # MS-MARCO v2.1 loader + QA pair loader
│   │   └── nq_open.py           # NQ-open validation loader
│   ├── qa/
│   │   ├── reader.py            # Flan-T5-base seq2seq reader
│   │   └── metrics.py           # Exact Match, Token F1 (SQuAD-style normalisation)
│   └── pipeline.py              # RAGPipeline end-to-end class
├── scripts/
│   ├── build_index.py           # Step 1: embed 500k passages, save .npy
│   ├── validate_distortion.py   # Step 2: compute MSE/IP distortion for all methods × bits
│   ├── run_retrieval.py         # Step 3: retrieval sweep, saves retrieval_results.json
│   └── run_qa.py                # Step 4: QA sweep, saves qa_results.json
├── notebooks/
│   └── analysis.ipynb           # all plots + summary table (run after experiments)
├── experiments/
│   ├── results/                 # JSON result files + logs
│   └── plots/                   # generated PDFs and CSVs
├── setup_l4.sh                  # one-shot GPU instance setup script
└── requirements.txt

---

Setup

Local (CPU, for analysis only)

git clone https://github.com/<your-username>/TurboQuantRAG.git
cd TurboQuantRAG
python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

GPU instance (recommended for experiments)

The experiments require a GPU with ≥16 GB VRAM. Tested on NVIDIA L4 (24 GB).

# On the GPU instance
bash setup_l4.sh

setup_l4.sh installs Python 3.12, creates a venv, and installs all dependencies. If faiss-gpu is unavailable for your Python version, install faiss-cpu instead and set device: cpu in configs/experiment.yaml.

---

Running Experiments

Run the four steps in order. Each step caches its output so reruns are fast.

source .venv/bin/activate

# Step 1 — embed 500k passages (~25 min on L4)
python scripts/build_index.py --config configs/experiment.yaml

# Step 2 — distortion validation (~2 min)
python scripts/validate_distortion.py --config configs/experiment.yaml

# Step 3 — retrieval sweep, 5000 NQ-open queries (~5 min on L4)
python scripts/run_retrieval.py --config configs/experiment.yaml --n-queries 5000

# Step 4 — QA sweep, 1000 NQ-open questions (~25 min on L4)
python scripts/run_qa.py --config configs/experiment.yaml --n-questions 1000

Results are written to experiments/results/. Then generate all plots:

jupyter nbconvert --to notebook --execute notebooks/analysis.ipynb \
  --output notebooks/analysis_executed.ipynb

---

Implementation Notes

TurboQuant-MSE (Algorithm 1)

1. Draw a random orthogonal matrix R ∈ ℝ^{d×d} (QR decomposition of Gaussian noise, seeded).
2. For each embedding x: normalize → rotate (Rx) → scalar-quantize each coordinate using a Lloyd-Max codebook trained on Beta((d−1)/2, (d−1)/2).
3. At query time: rotate query → asymmetric dot product with stored codes via batched GPU matmul (N,d) @ (d,B) → (N,B).

TurboQuant-IP (Algorithm 2)

Stage 1 — QJL sketch: 1-bit projection sign(Wx) where W ∈ ℝ^{m×d} is a random Gaussian matrix. Hamming distance used as cosine estimator via the matmul-Hamming identity:

agreements = qjl_signed @ q_signed.T   # (N, B), avoids (N, m, B) OOM tensor
cos_est    = agreements / m

Stage 2 — residual MSE quantization of x − cos_est · x̂.

Batched linear scan

All quantized retrieval uses vectorised GPU matmul rather than per-query loops:

• Decode all N stored vectors once → (N, d) on GPU
• Batch queries into blocks of 256 → (N, d) @ (d, B) → (N, B)
• torch.topk along dim=0 → top-k indices per query

This reduced 10-combination retrieval from ~80 minutes to ~1 minute on the L4.

---

Known Limitations

1. Low absolute Recall@10 (~5%). The 500k indexed passages come from the MS-MARCO training split. MS-MARCO validation QA pairs reference passages largely outside this subset. All relative comparisons between methods remain valid; the absolute numbers reflect corpus coverage, not algorithm quality.

2. TurboQuant-IP bits=1 and bits=2 produce identical results. The QJL first stage dominates the score; the 1-bit vs 2-bit residual codebook makes no measurable difference. This is likely caused by codebook collapse at very low bit widths. Treat IP as a single configuration for b≤2.

3. Memory bytes constant across TurboQuant bit widths. The memory_bytes() method returns a fixed value regardless of b. True memory scales linearly with bits; the efficiency plot for TurboQuant memory is inaccurate.

4. Latency comparison is not apples-to-apples. TurboQuant search runs on GPU; FAISS-PQ runs on CPU (faiss-gpu has no Python 3.12 wheel). The 8.7× latency advantage of TurboQuant includes GPU acceleration.

---

Dependencies

transformers 5.x has a known incompatibility (NameError: name 'nn' in accelerate integration). Pin to 4.46.3.

---

Citation

If you use this code or results, please cite the original TurboQuant paper:

@article{turboquant2024,
  title   = {TurboQuant: Online Vector Quantization for High-Quality Embedding Compression},
  year    = {2024},
  url     = {https://arxiv.org/abs/2409.09913}
}

And the datasets:

@inproceedings{msmarco,
  title     = {MS MARCO: A Human Generated MAchine Reading COmprehension Dataset},
  author    = {Bajaj, Payal and others},
  booktitle = {NeurIPS},
  year      = {2016}
}

@article{nqopen,
  title   = {Natural Questions: a Benchmark for Question Answering Research},
  author  = {Kwiatkowski, Tom and others},
  journal = {TACL},
  year    = {2019}
}

---

Author

Sanskruti Deshmukh — AMS 691, Spring 2026