tq-search

High-performance PyTorch Vector Quantization Engine with Lloyd-Max codebooks and QJL sign-bit residual correction.

Overview

tq-search is a state-of-the-art vector compression and similarity search library built entirely in PyTorch. It implements the optimal quantization algorithms inspired by TurboQuant (Zandieh et al., ICLR 2026 / arXiv:2504.19874) and QJL (arXiv:2406.03482).

Unlike standard vector databases that require decompressing vector databases in RAM or VRAM to perform searches, tq-search performs asymmetric inner product similarity scoring directly on compressed 5-bit representations, reducing GPU memory usage by up to 4x–6x while keeping search results highly accurate.

Key Features

• GPU-Accelerated: Fully integrated with PyTorch, running seamlessly on CUDA and Apple Silicon MPS.
• Batched Compression & Search: Built for high-throughput similarity search, handling millions of candidate vectors concurrently.
• Asymmetric Scoring: Computes dot products directly on the compressed codebooks without decompressing residuals in memory.
• Optimal Distortion: Employs mathematically optimal Lloyd-Max codebooks fit to rotated coordinate distributions.

Installation

Install the library directly from PyPI (once published):

pip install tq-search

Or install it locally in editable mode for development:

git clone https://github.com/barateza/tq-search.git
cd tq-search
pip install -e .

Quickstart

Here is how to compress a database of embeddings and search them using tq-search:

import torch
from tq_search import TurboQuantProd

# 1. Initialize the Quantizer
# We compress 1024-dimensional embeddings into 3-bit MSE + 1-bit QJL (4 bits total)
dim = 1024
bits = 4
quantizer = TurboQuantProd(d=dim, bits=bits, device="cuda")

# 2. Compress your database vectors (Shape: [N, dim])
database_vectors = torch.randn(10000, dim, device="cuda")
# L2-normalize vectors (recommended for maximum cosine similarity accuracy)
database_vectors = database_vectors / torch.norm(database_vectors, dim=-1, keepdim=True)

compressed = quantizer.quantize(database_vectors)
# 'compressed' is a lightweight dict containing:
# - mse_indices: [N, dim] (Stage 1 indices)
# - qjl_signs: [N, dim] (Stage 2 signs)
# - residual_norm: [N] (Dynamic residual scaling factor)

# 3. Perform Asymmetric Query Search (Query remains at full precision)
query = torch.randn(1, dim, device="cuda")
query = query / torch.norm(query, dim=-1, keepdim=True)

# Directly calculate dot products across all 10,000 compressed vectors in VRAM!
scores = quantizer.inner_product(query, compressed)
print("Inner Product Scores:", scores)

Research Credits

This library is a high-performance search adaptation based on the following research papers:

1. TurboQuant: "TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate" (Zandieh et al., ICLR 2026) arXiv:2504.19874
2. QJL Residuals: "QJL: 1-Bit Quantized JL Transform for KV Cache Quantization" (arXiv:2406.03482) arXiv:2406.03482

License

MIT License. Feel free to use, modify, and distribute.