GitHub - SeanPedersen/TieredQ: Efficient Vector Quantization - TieredQ: Variance-Aware Bit Allocation

TieredQ is a new vector quantization algorithm that dynamically allocates bits per dimension, providing high recall compression ratio. TieredQ centers and rotates vectors into a PCA basis, then greedily allocates 0–8/12 bits per dimension to reduce quantization error. TieredQ provides same recall per bit at about 36–52% less RAM & disk usage (in low bit budgets like 1,2,4 bits per dimension) compared to RaBitQ and TurboQuant. ANN indexes (f.e. IVF / HNSW) can use the codec to save RAM and disk space (reference IVF_TieredQ implementation is provided).

TieredQ core algorithm:

1. PCA
2. Variance-aware greedy bit allocation per dimension (principal component)
3. Scalar quantization per dimension with allocated bit budget

blog post

Recall per bits used:

Final benchmark report

Install

Install TieredQ from PyPI:

uv add tieredq
# or: pip install tieredq

Python

import gensim.downloader as api
from tieredq import TieredQIndex

vectors = api.load("glove-wiki-gigaword-100").vectors  # (400000, 100), float32
d = vectors.shape[1]
index = TieredQIndex(dimensions=d, code_bits=2 * d)  # 2 code bits per dimension
index.add(vectors)  # first add trains (a seeded 32,768-row sample of larger batches); later adds only encode
scores, ids = index.search(vectors[:2], k=10)  # (2, 10) float32, uint32

index.save("index.tq")
index = TieredQIndex.load("index.tq")

gensim is needed only for this example (uv add gensim). Constructor options after code_bits are keyword-only:

• metric="ip" (inner product) is the default; metric="cosine" normalizes rows and queries, and metric="l2" serves Euclidean search with scores of -|q - x|^2. L2 lifts vectors into one extra dimension with a lift whitened for the query distribution, stored in the model, so saved indexes stay self-contained.
• Bit allocation follows the metric: inner-product and cosine models allocate on PCA variance^1.5 at budgets of at least 1 bit/dim (recall-optimal when queries are distributed like the data, +0.5 to +0.8 R@10 points at equal disk on DBpedia OpenAI/Gemini, GloVe neutral to +1.3, issue #9) in both modes; L2 models use variance (MSE), which is optimal for whitened queries. For ip/cosine budgets below 1 bit/dim, the power falls linearly to 1 at 0.5 bits/dim and stays at 1 below that. allocation_power= overrides it on TieredQIndex.
• mode="fast" is the default and stores the native fast search representation. Models with tiers above 4 bits score through approximate byte tables (L2 models keep exact FP32 scoring); exact_fast_scores=True trains Lloyd-Max books and scans them in exact FP32 instead. That is the batch-throughput option: at 8 bits/dim on AVX-512 (100k, ip) batch search ran about 2x faster than the byte tables at equal recall on Gemini-768 and OpenAI-1536 (with its fine tiers, below) and better recall on GloVe-200, while single-query latency was 2-3x higher (4-10x on NEON and AVX2; #31). Use mode="full" for compact packed rows and FP32 scoring of quantized vectors, with slower search. Both modes scan every vector; saving and loading preserve the selected representation. Read the selection from index.mode. The former fast_mode boolean argument/property has been removed; use mode.
• Fast ip/cosine indexes use fine tiers by default (#26, #31): ranks coded with up to max_tier_bits=12 bits at the same bytes, whose top bits are an ordinary fast code at the automatic cap (6 bits, or ceil(bits/dim) when larger). max_scan_bits overrides that cap; fine tiers require a cap of at least 6 bits. The scan reads only those prefixes, then rescores max(24, 1.5·k) candidates exactly. Training uses them whenever the allocation codes a rank wider than the prefix; fine_tiers=False disables them, and index.fast_prefix_bits is the prefix width (0: none). 100k vectors, mean missed top-10 hits over 3 seeds, fine tiers vs none: Gemini-768 66 → 15, OpenAI-1536 44 → 6 at 8 bits/dim; 209 → 171 and 168 → 115 at 4 bits/dim; GloVe-200 162 → 87 at 8 bits/dim and unchanged at 2-4 (no rank goes wide). Batch QPS vs fine tiers off (100k, #31): +5-6% at 8 bits/dim and -4% to -8% at 3-4 bits/dim for k=10; -5% to -7% at 8 and -22% to -27% at 3-4 bits/dim for k=100 (rescoring 150 candidates). On 1M sets (OpenAI-1536, MS MARCO Cohere-1024, 6-8 bits/dim) top-10 misses fell 46-79% and search got 3-24% faster (likely as the scan reads fewer bits); encoding is 11-34% slower (wider codes). With exact_fast_scores=True the prefix is the 8-bit Lloyd-Max code, scanned exactly, and fine bits refine its cells: at 8 bits/dim mean top-10 misses fell from 28 to 15 (Gemini-768) and 18 to 9 (OpenAI-1536) at the same batch speed. The flat TieredQIndex and TieredQIVFIndex serve fine-tier models (opt in with fine_tiers=True on Python IVF; it defaults to False there). IVF scans the probed lists to the same candidate depth and rescores those candidates. Train nested prefix models with TieredQConfig::fine_tiers = false (the Python nested index does). Nested suffix refinement can independently opt into wider codes with fine_tiers=True.
• In C++, use dv::TieredQIndex(codec) for fast mode or dv::TieredQIndex(codec, dv::TieredQIndexMode::full) for full mode. Inspect the enum with index.mode(); the boolean constructor is removed.
• The first add() needs a representative batch of at least two vectors. Inputs must be contiguous float32 NumPy arrays.
• Queries accept one vector or a matrix. Scores estimate the selected metric: inner product, cosine similarity, or negative squared Euclidean distance; larger is better.
• Fast storage size depends on the model and CPU. Inspect bytes_per_vector for the representation stride and storage_bytes for allocated vector data. Fast scans use AVX-512 or AVX2+F16C on x86, and NEON on ARM. On AVX-512, single-vector queries with compatible packed-code models use a range checkpoint after four tiers to skip uncompetitive suffixes while preserving fast-scan IDs and score bits. Bounds are prepared per query and add no stored bits per vector. Other layouts, backends, and multi-vector queries use the full scanner.
• Matrix queries use OpenMP workers. Set OMP_NUM_THREADS before starting Python to control parallelism.
• mode also steers training. Fast models whose tiers all fit 4 bits use an additive 3-bit codebook, so the x86 nibble lookup stores 3-bit tiers at exactly 3 bits. This adds +1.2 to +3.6 R@10 points on GloVe-200 at 2.25–3.9 bits/dim, at equal storage and speed. It is neutral on SIFT, and DBpedia OpenAI/Gemini at 2–4 bits/dim are unchanged. Full mode keeps the Lloyd-Max 3-bit codebook, and either model serves in either mode. See issue #20 results.

See fast-mode measurements for accuracy and speed. Models with tiers wider than 4 bits now retain original codebook precision in fast mode. This improves recall at a throughput cost; existing numerical-level fast indexes need rebuilding to recover that precision. Their fast storage keeps codes at tier width; files saved with the earlier one-byte-per-coordinate layout still load and are repacked. See the high-bit investigation.

IVF and nested search (Python)

from tieredq import TieredQIVFIndex, TieredQNestedIVFIndex, TieredQPruning

d = vectors.shape[1]
queries = vectors[:2]
ivf = TieredQIVFIndex(dimensions=d, code_bits=2 * d, lists=1024)  # lists=0: one list per codec centroid
ivf.build(vectors)  # trains the codec (seeded 32,768-row sample) and the lists, then indexes
scores, ids = ivf.search(queries, k=10, nprobe=32)

# Prefix codes in an in-RAM IVF; candidates rescored with a 4 bits/dim refinement tier on SSD.
index = TieredQNestedIVFIndex.build(
    vectors, prefix_bits=d, total_bits=4 * d, lists=256, fine_tiers=True
)
index.write_tier("tier.rows", vectors)  # or index.attach_tier(index.encode_tier(vectors)) to keep it in RAM
index.attach_tier("tier.rows", reads="direct", queue_depth=32)
scores, ids = index.search(queries, k=10, nprobe=64, candidates=50)
index.save("nested.tq")  # IVF, refinement and tier layout; attach the tier again after load()

• TieredQIVFIndex(dimensions, code_bits, *, lists=0, metric="ip", fine_tiers=False) always uses fast storage. build(rows) replaces its vectors; only the first build or train_lists(sample) trains the codec. Call train_lists before build to fix the partition on a separate sample.
• TieredQNestedIVFIndex.build(vectors, *, prefix_bits, total_bits, lists=0, metric="ip", tier_rows="prefix_and_suffix", sample=None, fine_tiers=False, max_tier_bits=12) trains a nested index. prefix_bits and total_bits are budgets per vector, with the prefix included in total_bits. For inner product/cosine, fine_tiers=True lets suffix training allocate combined prefix/suffix codes wider than 8 bits, up to max_tier_bits (8–12). The IVF prefix stays fixed; candidates are rescored directly with all suffix bits. max_refinement_bits reports the widest realized code; total_bits and tier_bytes_per_vector report realized storage after block rounding. Wider suffix models use refinement format 2; existing format 1 models still load. In a local Gemini-768 benchmark, an 8-bit/dimension total budget improved full-probe recall@10 from 99.44% to 99.76% (57% fewer misses), with median batch time per query changing from 2.480 to 2.501 ms on one M1 thread with a RAM tier. The 4-bit/dimension control was unchanged. Neither IVF constructor exposes the flat index's mode, allocation_power, or exact_fast_scores.
• index.attach_exact(codes_or_path, scales, offsets, finals=16) adds the exact SQ8 final stage (int8 codes in RAM, or a TieredQNestedIVFIndex.write_sq8_rows(path, codes) file with reads="direct", queue_depth=32 on SSD); detach_exact() removes it.
• IVF searches require keyword nprobe; nested searches also require candidates >= k. Results use 2**32 - 1 ids and -inf scores for slots beyond the rows in the probed lists.
• File tiers default to reads="cached", queue_depth=1. Concurrent reads at larger queue depths require a build with liburing; otherwise reads stay serial. PyPI wheels disable liburing.
• One build trains the prefix, the lists and the refinement from the same seeded sample, so they cannot be mismatched. Pass sample= to train on earlier data instead; its rows also fix the lists for later builds (as TieredQIVFIndex.train_lists does).
• tier_rows="suffix" stores a smaller tier and decodes each candidate's prefix from the lists (see the C++ notes).
• Tier rows kept in another store (a database, an object store): ids = index.candidates(queries, 50, nprobe=64), fetch the encode_tier rows of those ids (shape ids.shape + (index.tier_bytes_per_vector,); ids of 2**32 - 1 are padding), then scores, ids = index.rescore(queries, ids, rows, k=10).
• pruning=TieredQPruning(confidence=8) opts IVF and nested scans into approximate pruning (see below).
• Many lists: ivf.set_coarse_training(iterations=10, assignment_supers=256, assignment_probes=16) before build sets the list training and opts into restricted assignment. After build, ivf.build_router() adds the two-level coarse router, and search(..., routing="auto" | "router" | "exact", super_probes=0) selects the coarse ranking (see the C++ notes on many lists).

C++

Building from source requires CMake 3.18+, a C++20 compiler, OpenMP, and BLAS/LAPACK. liburing is optional for concurrent SSD tier reads.

cmake -S . -B build
cmake --build build -j
OMP_NUM_THREADS=2 ctest --test-dir build --output-on-failure
./build/tieredq_example
cmake --install build --prefix /desired/prefix

Include <tieredq/tieredq.h>, use find_package(TieredQ CONFIG REQUIRED), and link TieredQ::TieredQ. FetchContent also works. See integration.md for layouts, threading, ownership and dependency setup.

dv::TieredQ::train(config, sample, rows, seed) trains on a contiguous sample. To train on part of a larger array without copying it, dv::TieredQ::train_sampled(config, data, count, sample_ids, sample_rows, seed) reads the named rows in place (the same model). About 8–16K training rows reach full recall on 1M-row sets; fewer than ~32 rows per centroid lose recall (#19.2 results).

IVF (C++)

dv::TieredQIVFIndex is a fast-mode inverted-file index. With lists = 0 the codec's k-means centroids (at most 256) are the lists, at no extra storage. Any other count trains a separate coarse quantizer, whose centroids are stored.

dv::TieredQIVFIndex index(codec, /*lists=*/1024);
index.train_lists(sample, sample_rows);  // optional: fix the lists on a sample; builds then only assign rows
index.build(rows, count);
auto hits = index.search(query, /*k=*/10, /*nprobe=*/16);
// Opt-in approximate pruning for models with more than four tiers (about 2+ bits/dim):
auto fast = index.search(query, 10, 16, dv::TieredQPruning{/*confidence=*/8, /*prefix_tiers=*/4});

• Many lists (64K+). set_lists(centroids) fixes the lists to a partition trained elsewhere. build(RowLoader, count) pulls rows by id in bounded chunks, so a base held as SQ8 codes needs no fp32 copy. set_coarse_training({iterations, assignment_supers, assignment_probes}) sets the Lloyd iterations of list training (default 20) and opts into restricted assignment (#48). Each row is then assigned among the members of its assignment_probes nearest supers, a k-means of the lists' centroids. Rows whose nearest list lies in a probed super get the exact label, up to floating-point near-ties. Exact assignment runs as FP32 BLAS sgemm, as in Faiss; query ranking keeps the binary16 panels. On Qwen 1024-D (500K rows, 32K lists), 256 supers and 16 probes kept 99.0% of exact labels and ran 8.5× faster. Exact coarse ranking streams every centroid per query: about 9 ms at 64K × 1024. build_router() adds a two-level compressed router (#47). It ranks about lists/128 supers exactly, scans the 2-bit TieredQ codes of the probed supers' member centroids, and re-ranks the first 64 lists exactly (rerank; 256 cost 0.1-0.3 ms per query at 64K lists without recall above 0.945). The probed supers grow with nprobe. Batches route too: at 64K lists, 32-thread batches through the router reached 1.4× the QPS of the exact GEMM ranking. TieredQRouting{mode, super_probes} forces either path or the number of probed supers per search. The router persists with the index (payload version 5).

  dv::TieredQIVFIndex index(codec, /*lists=*/65536);
  index.set_lists(centroids);  // or set_coarse_training({10, 256, 16}) before training lists in build()
  index.build(loader, count);
  index.build_router();        // rebuilt by every later build()
  auto hits = index.search(query, 40, /*nprobe=*/256);
  auto exact = index.search(query, 40, 256, {}, {dv::TieredQRouting::Mode::exact});

• Pruning. By default the scan prunes exactly: hits equal an unpruned scan. A positive confidence bounds the remaining tiers statistically instead, which is faster and approximate. On OpenAI 100k, confidence 8 kept recall unchanged (±0.0001) at 1.7× the QPS at 4 bits/dim. With prefix_tiers = 3 it gave 1.6–2.4× the QPS for 0.001–0.003 recall at 2 and 4 bits/dim. Confidence 4 or less loses clearly more recall. Platforms: the exact FP32 code scan prunes on AVX-512 only. The nibble byte-table scan (default fast models with tiers above 4 bits) prunes on AVX-512, AVX2 and NEON. The scalar fallback validates the setting and scans unpruned. By default (prefix_tiers = 0, automatic) a nibble scan prunes after all but the last 12.5% of its sections (37.5% with a positive confidence), or after the first four tiers where those leave more; an explicit prefix_tiers prunes after that many tiers. Single-core, 100k documents, 256 lists, on Apple M1 (NEON), EPYC Milan (AVX2) and i9-7980XE (AVX-512) with 384-, 768-, 1024- and 1536-d embeddings: exact pruning returned the unpruned hits in every run and was 1.0–1.3× faster than unpruned at nprobe 64–256 (about 1.0× at nprobe 16); confidence 8 was 1.0–1.8× faster with unchanged recall@10. Flat single-query search prunes the same way, 1.0–1.2× faster except Gemini-768 at 2 bits/dim, 5–18% slower. Larger suffix shares lose with the exact bound, which loosens faster than it saves work. To prune through the exact FP32 scan instead, train with TieredQConfig::target_mode = TieredQIndexMode::full and fine_tiers = false (flat Python exact_fast_scores=True, fine_tiers=False; Python IVF does not expose the exact-scoring training option). Measured at 4 bits/dim, before nibble scans were pruned: on OpenAI-1536 100k exact with confidence 8 was 4–11% faster than nibble scoring at equal or slightly higher recall, and on OpenAI 1M 17–21% faster but 0.4 pt R@10 lower; on Gemini-768 nibble scoring was 20–38% faster, at 0.2 pt R@10 lower. At 2 bits/dim nibble scoring was faster on both (#13).

• Bit cap. At about 2 bits/dim or less, training with TieredQConfig::max_scan_bits = 4 keeps every tier on the nibble lookup kernel. That gave 1.7× RaBitQ's IVF QPS at equal recall and memory (OpenAI 100k), for about 0.006 recall at full probe. Don't use it at 4 bits/dim, where it loses recall.

• Sub-1-bit index + rerank. With rerank, the in-RAM IVF code can be far below 1 bit/dim. On OpenAI-1M (256 lists), a 512-bit prefix (627 index bits/vector) with an INT8 rerank tier reached R@10 0.95 / 0.99 at 2.8x / 3.1x the QPS of RaBitQ 1 bit (1,683 bits) + INT8.

• Nested rerank codes. dv::TieredQRefinement (<tieredq/tieredq_refinement.h>) trains a suffix that refines the prefix towards a finer model. The suffix is read by id for the scan's candidates and scored with the scan's query weights; nesting costs about no recall against a flat code of the full width:

  dv::TieredQNestedIVFIndex (<tieredq/tieredq_nested.h>) serves this end to end. The prefix codes stay in the in-RAM IVF lists. Each candidate's tier row is read by id from a TieredQRowSource: TieredQMemoryRows in RAM, TieredQFileRows with one pread per candidate on SSD (TieredQFileReads::direct bypasses the page cache), or your own store. By default a tier row holds the prefix row and the suffix (TieredQTierRows::prefix_and_suffix, one contiguous read). TieredQTierRows::suffix stores only the suffix and decodes each candidate's prefix from the IVF lists (TieredQIVFIndex::document_rows). That tier is smaller (651 vs 784 B/vector at a 1024-bit prefix, d = 1536) but costs about 3.4 us of decoding per candidate and a 4-byte id-to-document map in RAM.

  // Train the prefix on the same sample and seed as the refinement, with fine tiers disabled.
  auto config = dv::TieredQConfig::target_code_bits(d, d);
  config.fine_tiers = false;
  auto codec = dv::TieredQ::train(config, sample, rows, seed);
  auto refinement = dv::TieredQRefinement::train(codec, /*total_code_bits=*/4 * d, sample, rows, seed,
                                               /*fine_tiers=*/true, /*max_tier_bits=*/12);
  dv::TieredQIVFIndex ivf(codec);
  ivf.build(data, n);
  dv::TieredQNestedIVFIndex index(std::move(ivf), std::move(refinement));
  dv::TieredQRowFileWriter writer("tier.rows", index.tier_bytes_per_vector());
  std::vector<uint8_t> buffer(chunk * index.tier_bytes_per_vector());
  for (size_t first = 0; first < n; first += chunk) {        // ids in build order
    const size_t count = std::min(chunk, n - first);
    index.encode_tier(data + first * d, first, count, buffer.data());
    writer.append(buffer.data(), count);
  }
  writer.finish();
  index.attach_tier(std::make_shared<dv::TieredQFileRows>("tier.rows"));
  auto hits = index.search(query, /*k=*/10, /*nprobe=*/32, /*candidates=*/50);

  An exact final stage (attach_exact) rescores the tier's best finals candidates with exact vectors, e.g. an SQ8 store left on SSD (TieredQSQ8Rows over a TieredQFileRows of int8 codes, direct reads, queue depth 32). On Qwen 8.84M × 1024 at 64K lists, a 6-bit fine tier in RAM plus 16 exact SQ8 rows per query from NVMe matched the in-RAM SQ8 re-rank's recall at every nprobe (0.954-0.986) with 7.8 GB of serving RAM instead of 12.2 GB, for about 0.7 ms per query (hybrid study).

  index.attach_exact(std::make_shared<dv::TieredQSQ8Rows>(
      std::make_shared<dv::TieredQFileRows>("sq8.rows", dv::TieredQFileReads::direct, /*queue_depth=*/32),
      scales, offsets), /*finals=*/16);

• Flat nested search. dv::TieredQNestedIndex (<tieredq/tieredq_nested_flat.h>) scans a narrow prefix for every document and rescores the top candidates with an in-RAM suffix tier (a one-list TieredQNestedIVFIndex). On OpenAI-1536 it is the fast flat path from about 3 bits/dim. A 512-bit prefix with a 3 bits/dim total tier reached RaBitQ's 4-bit recall at 2.2x (100K, 50 candidates) and 2.9x (1M, 100 candidates) its single-query QPS, with 0.76-0.79x its stored bits. The plain flat index is at 0.54x there.

  auto config = dv::TieredQConfig::target_code_bits(d, 512);
  config.fine_tiers = false;
  auto prefix = dv::TieredQ::train(config, sample, rows, seed);
  auto refinement = dv::TieredQRefinement::train(prefix, /*total_code_bits=*/3 * d, sample, rows, seed);
  dv::TieredQNestedIndex index(std::move(prefix), std::move(refinement), data, n);
  auto hits = index.search(query, /*k=*/10, /*candidates=*/100);

  serialize() stores the IVF, the refinement and the tier layout; after loading, attach the tier again. For a tier in an external store, run index().search_batch for the candidates and rescore_batch with their fetched tier rows; it needs no attached tier.

  At 1M, 512 / 768 bit prefixes + a 4-bit refinement (about 700 B/vector) beat RaBitQ's own 5/6-bit codes by 1.2-2.4x QPS at R@10 0.95-0.98 with 28-43% fewer rerank bytes. Refinement recall tops out near 0.984; above that use INT8.

• Results against RaBitQ's IVF at equal measured memory are in issue #19, #23 and the issue #6 eval (rerank results).

Codec comparison benchmark

See benchmark/README.md for the one-command CPU comparison against official RaBitQ on SIFT1M, GloVe-200 and OpenAI3-1536, including measured storage overhead, recall, throughput and latency.

Disclaimer: everything but the core idea and the README was created using both Opus 5.5 and GPT 6.1 Sol