mirror of
https://github.com/ruvnet/RuVector.git
synced 2026-08-22 07:03:42 +00:00
Merge pull request #521 from shaal/claude/ruvector-turbovec-optimization-FhaDh
Some checks failed
Workspace CI / Tests (rvagent) (push) Waiting to run
Benchmarks / Rust Benchmarks (push) Waiting to run
Benchmarks / SQL Benchmarks (push) Waiting to run
Benchmarks / Compare with Baseline (push) Blocked by required conditions
Build Graph Node Native Modules / Build Graph darwin-arm64 (push) Waiting to run
Build Graph Node Native Modules / Build Graph darwin-x64 (push) Waiting to run
Build Graph Node Native Modules / Build Graph linux-arm64-gnu (push) Waiting to run
Build Graph Node Native Modules / Build Graph linux-x64-gnu (push) Waiting to run
Build Graph Node Native Modules / Build Graph win32-x64-msvc (push) Waiting to run
Build Graph Node Native Modules / Publish Graph Node Platform Packages (push) Blocked by required conditions
Build Native Modules / Build darwin-arm64 (push) Waiting to run
Build Native Modules / Build linux-arm64-gnu (push) Waiting to run
Build Native Modules / Build darwin-x64 (push) Waiting to run
Build Native Modules / Build win32-x64-msvc (push) Waiting to run
Build Native Modules / Build linux-x64-gnu (push) Waiting to run
Build Native Modules / Commit Built Binaries (push) Blocked by required conditions
ruvector-verified CI / check () (push) Waiting to run
ruvector-verified CI / check (--all-features) (push) Waiting to run
ruvector-verified CI / check (--features all-proofs) (push) Waiting to run
ruvector-verified CI / check (--features coherence-proofs) (push) Waiting to run
ruvector-verified CI / check (--features hnsw-proofs) (push) Waiting to run
ruvector-verified CI / check (--features rvf-proofs) (push) Waiting to run
ruvector-verified CI / check (--features serde) (push) Waiting to run
ruvector-verified CI / check (--features ultra) (push) Waiting to run
ruvector-verified CI / test (push) Blocked by required conditions
ruvector-verified CI / bench (push) Blocked by required conditions
ruvector-verified CI / clippy (push) Waiting to run
Workspace CI / Rustfmt (push) Waiting to run
Workspace CI / Cargo check (push) Waiting to run
Workspace CI / Clippy (push) Waiting to run
Workspace CI / Tests (core-and-rest) (push) Waiting to run
Workspace CI / Tests (core-and-rest-heavy) (push) Waiting to run
Workspace CI / Tests (core-and-rest-wasm) (push) Waiting to run
Workspace CI / Tests (ml-research-heavy) (push) Waiting to run
Workspace CI / Tests (ml-research-rest) (push) Waiting to run
Workspace CI / Tests (research-nightly) (push) Waiting to run
Workspace CI / Tests (ruvix) (push) Waiting to run
Clippy + fmt / Rustfmt (push) Waiting to run
regression-guard / reentrant-rwlock-double-write (push) Waiting to run
regression-guard / case-insensitive-collisions (push) Waiting to run
regression-guard / ruvector-core-no-avx512-builds-on-stable (push) Waiting to run
regression-guard / hnsw-insert-beam-no-m2-clamp (push) Waiting to run
regression-guard / hnsw-distance-based-neighbor-pruning (push) Waiting to run
regression-guard / vector-db-rebuilds-index-on-open (push) Waiting to run
regression-guard / npm-publish-pipeline (npm/packages/pi-brain) (push) Waiting to run
regression-guard / npm-publish-pipeline (npm/packages/ruvector) (push) Waiting to run
regression-guard / npm-publish-pipeline (npm/packages/rvf-wasm) (push) Waiting to run
regression-guard / no-npx-execSync-in-route-enhanced (push) Waiting to run
regression-guard / shell-injection-in-mcp-server (push) Waiting to run
regression-guard / no-systemtime-in-wasm-crates (push) Waiting to run
regression-guard / no-hardcoded-workspaces-paths (push) Waiting to run
regression-guard / brain-hydration-counters-present (push) Waiting to run
regression-guard / optional-deps-resolvable-on-npm (push) Waiting to run
ruvector npm — functional, learning, optimized, effective / Unit & CLI tests (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / Functional smoke (npx ruvector) (push) Blocked by required conditions
Workspace CI / Tests (vector-index) (push) Waiting to run
Workspace CI / Security audit (push) Waiting to run
Clippy + fmt / Clippy (deny warnings) (push) Waiting to run
regression-guard / hnsw-recall-at-1 (push) Waiting to run
regression-guard / graph-condense-perception-tests (push) Waiting to run
regression-guard / mincut-pin-tracks-workspace-version (push) Waiting to run
ruvector npm — functional, learning, optimized, effective / Build (push) Waiting to run
ruvector npm — functional, learning, optimized, effective / Learning check (HNSW activates) (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / Performance benchmark (≥2× speedup at N=5000) (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / Recall quality (recall@10 ≥ 0.88 at N=10k) (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / Tarball integrity (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / CI pass (push) Blocked by required conditions
supply-chain / dependency-review (PRs only) (push) Waiting to run
supply-chain / cargo audit (RustSec advisories) (push) Waiting to run
supply-chain / cargo deny (license + source + ban policy) (push) Waiting to run
supply-chain / npm audit (npm/ workspace) (push) Waiting to run
supply-chain / lockfile integrity (Cargo.lock) (push) Waiting to run
Validate Package Lock File / validate-lockfile (push) Waiting to run
WASM Dedup Check / check-wasm-dedup (push) Waiting to run
Build DiskANN Native Modules / Build DiskANN darwin-arm64 (push) Has been cancelled
Build DiskANN Native Modules / Build DiskANN darwin-x64 (push) Has been cancelled
Build DiskANN Native Modules / Build DiskANN linux-arm64-gnu (push) Has been cancelled
Build DiskANN Native Modules / Build DiskANN linux-x64-gnu (push) Has been cancelled
Build DiskANN Native Modules / Build DiskANN win32-x64-msvc (push) Has been cancelled
Build DiskANN Native Modules / Publish DiskANN Platform Packages (push) Has been cancelled
Some checks failed
Workspace CI / Tests (rvagent) (push) Waiting to run
Benchmarks / Rust Benchmarks (push) Waiting to run
Benchmarks / SQL Benchmarks (push) Waiting to run
Benchmarks / Compare with Baseline (push) Blocked by required conditions
Build Graph Node Native Modules / Build Graph darwin-arm64 (push) Waiting to run
Build Graph Node Native Modules / Build Graph darwin-x64 (push) Waiting to run
Build Graph Node Native Modules / Build Graph linux-arm64-gnu (push) Waiting to run
Build Graph Node Native Modules / Build Graph linux-x64-gnu (push) Waiting to run
Build Graph Node Native Modules / Build Graph win32-x64-msvc (push) Waiting to run
Build Graph Node Native Modules / Publish Graph Node Platform Packages (push) Blocked by required conditions
Build Native Modules / Build darwin-arm64 (push) Waiting to run
Build Native Modules / Build linux-arm64-gnu (push) Waiting to run
Build Native Modules / Build darwin-x64 (push) Waiting to run
Build Native Modules / Build win32-x64-msvc (push) Waiting to run
Build Native Modules / Build linux-x64-gnu (push) Waiting to run
Build Native Modules / Commit Built Binaries (push) Blocked by required conditions
ruvector-verified CI / check () (push) Waiting to run
ruvector-verified CI / check (--all-features) (push) Waiting to run
ruvector-verified CI / check (--features all-proofs) (push) Waiting to run
ruvector-verified CI / check (--features coherence-proofs) (push) Waiting to run
ruvector-verified CI / check (--features hnsw-proofs) (push) Waiting to run
ruvector-verified CI / check (--features rvf-proofs) (push) Waiting to run
ruvector-verified CI / check (--features serde) (push) Waiting to run
ruvector-verified CI / check (--features ultra) (push) Waiting to run
ruvector-verified CI / test (push) Blocked by required conditions
ruvector-verified CI / bench (push) Blocked by required conditions
ruvector-verified CI / clippy (push) Waiting to run
Workspace CI / Rustfmt (push) Waiting to run
Workspace CI / Cargo check (push) Waiting to run
Workspace CI / Clippy (push) Waiting to run
Workspace CI / Tests (core-and-rest) (push) Waiting to run
Workspace CI / Tests (core-and-rest-heavy) (push) Waiting to run
Workspace CI / Tests (core-and-rest-wasm) (push) Waiting to run
Workspace CI / Tests (ml-research-heavy) (push) Waiting to run
Workspace CI / Tests (ml-research-rest) (push) Waiting to run
Workspace CI / Tests (research-nightly) (push) Waiting to run
Workspace CI / Tests (ruvix) (push) Waiting to run
Clippy + fmt / Rustfmt (push) Waiting to run
regression-guard / reentrant-rwlock-double-write (push) Waiting to run
regression-guard / case-insensitive-collisions (push) Waiting to run
regression-guard / ruvector-core-no-avx512-builds-on-stable (push) Waiting to run
regression-guard / hnsw-insert-beam-no-m2-clamp (push) Waiting to run
regression-guard / hnsw-distance-based-neighbor-pruning (push) Waiting to run
regression-guard / vector-db-rebuilds-index-on-open (push) Waiting to run
regression-guard / npm-publish-pipeline (npm/packages/pi-brain) (push) Waiting to run
regression-guard / npm-publish-pipeline (npm/packages/ruvector) (push) Waiting to run
regression-guard / npm-publish-pipeline (npm/packages/rvf-wasm) (push) Waiting to run
regression-guard / no-npx-execSync-in-route-enhanced (push) Waiting to run
regression-guard / shell-injection-in-mcp-server (push) Waiting to run
regression-guard / no-systemtime-in-wasm-crates (push) Waiting to run
regression-guard / no-hardcoded-workspaces-paths (push) Waiting to run
regression-guard / brain-hydration-counters-present (push) Waiting to run
regression-guard / optional-deps-resolvable-on-npm (push) Waiting to run
ruvector npm — functional, learning, optimized, effective / Unit & CLI tests (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / Functional smoke (npx ruvector) (push) Blocked by required conditions
Workspace CI / Tests (vector-index) (push) Waiting to run
Workspace CI / Security audit (push) Waiting to run
Clippy + fmt / Clippy (deny warnings) (push) Waiting to run
regression-guard / hnsw-recall-at-1 (push) Waiting to run
regression-guard / graph-condense-perception-tests (push) Waiting to run
regression-guard / mincut-pin-tracks-workspace-version (push) Waiting to run
ruvector npm — functional, learning, optimized, effective / Build (push) Waiting to run
ruvector npm — functional, learning, optimized, effective / Learning check (HNSW activates) (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / Performance benchmark (≥2× speedup at N=5000) (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / Recall quality (recall@10 ≥ 0.88 at N=10k) (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / Tarball integrity (push) Blocked by required conditions
ruvector npm — functional, learning, optimized, effective / CI pass (push) Blocked by required conditions
supply-chain / dependency-review (PRs only) (push) Waiting to run
supply-chain / cargo audit (RustSec advisories) (push) Waiting to run
supply-chain / cargo deny (license + source + ban policy) (push) Waiting to run
supply-chain / npm audit (npm/ workspace) (push) Waiting to run
supply-chain / lockfile integrity (Cargo.lock) (push) Waiting to run
Validate Package Lock File / validate-lockfile (push) Waiting to run
WASM Dedup Check / check-wasm-dedup (push) Waiting to run
Build DiskANN Native Modules / Build DiskANN darwin-arm64 (push) Has been cancelled
Build DiskANN Native Modules / Build DiskANN darwin-x64 (push) Has been cancelled
Build DiskANN Native Modules / Build DiskANN linux-arm64-gnu (push) Has been cancelled
Build DiskANN Native Modules / Build DiskANN linux-x64-gnu (push) Has been cancelled
Build DiskANN Native Modules / Build DiskANN win32-x64-msvc (push) Has been cancelled
Build DiskANN Native Modules / Publish DiskANN Platform Packages (push) Has been cancelled
feat(turbovec): multi-bit TurboQuant FastScan ANN index (ADR-194 M1)
This commit is contained in:
commit
2e5e65b3d3
12 changed files with 1884 additions and 166 deletions
10
Cargo.lock
generated
10
Cargo.lock
generated
|
|
@ -10763,6 +10763,16 @@ dependencies = [
|
|||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ruvector-turbovec"
|
||||
version = "2.3.0"
|
||||
dependencies = [
|
||||
"rand 0.8.6",
|
||||
"ruvector-rabitq",
|
||||
"serde",
|
||||
"thiserror 2.0.18",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "ruvector-verified"
|
||||
version = "0.1.1"
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ members = [
|
|||
"crates/ruvector-coherence-hnsw",
|
||||
"crates/ruvector-rabitq",
|
||||
"crates/ruvector-rabitq-wasm",
|
||||
"crates/ruvector-turbovec",
|
||||
"crates/ruvector-rulake",
|
||||
"crates/ruvector-core",
|
||||
"crates/ruvector-node",
|
||||
|
|
|
|||
21
crates/ruvector-turbovec/Cargo.toml
Normal file
21
crates/ruvector-turbovec/Cargo.toml
Normal file
|
|
@ -0,0 +1,21 @@
|
|||
[package]
|
||||
name = "ruvector-turbovec"
|
||||
version.workspace = true
|
||||
edition.workspace = true
|
||||
rust-version.workspace = true
|
||||
license.workspace = true
|
||||
authors.workspace = true
|
||||
repository.workspace = true
|
||||
description = "TurboVec: multi-bit TurboQuant ANN scalar reference with 2/3/4-bit Lloyd-Max quantization, TQ+ calibration, and length-renormalized scoring."
|
||||
|
||||
[[bin]]
|
||||
name = "turbovec-demo"
|
||||
path = "src/main.rs"
|
||||
|
||||
[dependencies]
|
||||
ruvector-rabitq = { version = "2.2.0", path = "../ruvector-rabitq" }
|
||||
# `rand` is a normal dependency because the `turbovec-demo` binary (and the
|
||||
# unit tests) use it; binaries cannot pull from `[dev-dependencies]`.
|
||||
rand = { workspace = true }
|
||||
serde = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
126
crates/ruvector-turbovec/src/calibrate.rs
Normal file
126
crates/ruvector-turbovec/src/calibrate.rs
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
//! TQ+ per-coordinate calibration (ADR-254 §T3).
|
||||
//!
|
||||
//! After the randomized Hadamard rotation, coordinate `i` of a *unit* vector is
|
||||
//! approximately `N(0, 1/d)` — but finite-`d` rotations leave a small residual
|
||||
//! mean/variance drift per coordinate. We fit two scalars per coordinate from
|
||||
//! the first ingest batch — `shift[i]` (empirical mean) and `scale[i]`
|
||||
//! (empirical std) — that map each coordinate onto the canonical `N(0, 1)` the
|
||||
//! Lloyd–Max table ([`crate::quantize`]) is built for:
|
||||
//!
|
||||
//! ```text
|
||||
//! z_i = (r_i - shift_i) / scale_i // calibrate → ~N(0,1)
|
||||
//! r̂_i = centroid[q_i] * scale_i + shift_i // reconstruct
|
||||
//! ```
|
||||
//!
|
||||
//! The calibration is **frozen** after warm-up: every later vector reuses it
|
||||
//! with no retraining (online ingest). This is the "+" that lifts plain scalar
|
||||
//! quantization to TQ+ recall.
|
||||
|
||||
/// Frozen per-coordinate affine calibration.
|
||||
#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
|
||||
pub struct Calibration {
|
||||
shift: Vec<f32>,
|
||||
scale: Vec<f32>,
|
||||
}
|
||||
|
||||
impl Calibration {
|
||||
/// Fit from a batch of **rotated unit** vectors (rows of length `dim`).
|
||||
///
|
||||
/// `scale` is floored at a small epsilon so dead/near-constant coordinates
|
||||
/// never divide by zero. With fewer than two samples we fall back to the
|
||||
/// theoretical `N(0, 1/d)` std (`1/√d`) and zero mean, which keeps the
|
||||
/// estimator sane on tiny indexes.
|
||||
pub fn fit(rotated_unit_rows: &[Vec<f32>], dim: usize) -> Self {
|
||||
debug_assert!(
|
||||
rotated_unit_rows.iter().all(|r| r.len() == dim),
|
||||
"Calibration::fit: every row must have length dim ({dim})"
|
||||
);
|
||||
let n = rotated_unit_rows.len();
|
||||
if n < 2 {
|
||||
let s = 1.0f32 / (dim as f32).sqrt();
|
||||
return Self {
|
||||
shift: vec![0.0; dim],
|
||||
scale: vec![s.max(1e-6); dim],
|
||||
};
|
||||
}
|
||||
let mut shift = vec![0.0f32; dim];
|
||||
for row in rotated_unit_rows {
|
||||
for (acc, &v) in shift.iter_mut().zip(row.iter()) {
|
||||
*acc += v;
|
||||
}
|
||||
}
|
||||
let inv_n = 1.0 / n as f32;
|
||||
for s in shift.iter_mut() {
|
||||
*s *= inv_n;
|
||||
}
|
||||
let mut var = vec![0.0f32; dim];
|
||||
for row in rotated_unit_rows {
|
||||
for ((acc, &v), &m) in var.iter_mut().zip(row.iter()).zip(shift.iter()) {
|
||||
let d = v - m;
|
||||
*acc += d * d;
|
||||
}
|
||||
}
|
||||
let inv_nm1 = 1.0 / (n - 1) as f32;
|
||||
let scale: Vec<f32> = var
|
||||
.iter()
|
||||
.map(|&v| (v * inv_nm1).sqrt().max(1e-6))
|
||||
.collect();
|
||||
Self { shift, scale }
|
||||
}
|
||||
|
||||
/// Map a rotated coordinate onto the canonical `N(0,1)` grid.
|
||||
#[inline]
|
||||
pub fn calibrate(&self, i: usize, r: f32) -> f32 {
|
||||
(r - self.shift[i]) / self.scale[i]
|
||||
}
|
||||
|
||||
/// Reconstruct the rotated coordinate from a dequantized centroid value.
|
||||
#[inline]
|
||||
pub fn reconstruct(&self, i: usize, centroid: f32) -> f32 {
|
||||
centroid * self.scale[i] + self.shift[i]
|
||||
}
|
||||
|
||||
/// Dimensionality.
|
||||
#[inline]
|
||||
pub fn dim(&self) -> usize {
|
||||
self.shift.len()
|
||||
}
|
||||
|
||||
/// Bytes of stored calibration (two `f32` per coordinate).
|
||||
#[inline]
|
||||
pub fn bytes(&self) -> usize {
|
||||
(self.shift.len() + self.scale.len()) * 4
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn fit_recovers_mean_and_std() {
|
||||
// Two coordinates: col0 centered at 10 (std ~0), col1 spread.
|
||||
let rows = vec![
|
||||
vec![10.0, -1.0],
|
||||
vec![10.0, 1.0],
|
||||
vec![10.0, -2.0],
|
||||
vec![10.0, 2.0],
|
||||
];
|
||||
let cal = Calibration::fit(&rows, 2);
|
||||
assert!((cal.shift[0] - 10.0).abs() < 1e-4);
|
||||
assert!(cal.scale[0] >= 1e-6); // floored, never zero
|
||||
assert!(cal.shift[1].abs() < 1e-4);
|
||||
assert!(cal.scale[1] > 1.0);
|
||||
// calibrate∘reconstruct identity on the grid value.
|
||||
let z = cal.calibrate(1, 2.0);
|
||||
let r = cal.reconstruct(1, z);
|
||||
assert!((r - 2.0).abs() < 1e-4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_index_falls_back_to_theoretical() {
|
||||
let cal = Calibration::fit(&[vec![0.0; 16]], 16);
|
||||
let expected = 1.0 / (16f32).sqrt();
|
||||
assert!((cal.scale[0] - expected).abs() < 1e-6);
|
||||
}
|
||||
}
|
||||
46
crates/ruvector-turbovec/src/error.rs
Normal file
46
crates/ruvector-turbovec/src/error.rs
Normal file
|
|
@ -0,0 +1,46 @@
|
|||
//! Error type for the TurboVec index.
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
/// Errors produced by [`crate::TurboVecIndex`] / [`crate::IdMapIndex`].
|
||||
#[derive(Debug, Error)]
|
||||
pub enum TurboVecError {
|
||||
/// A vector's length did not match the index dimension.
|
||||
#[error("dimension mismatch: expected {expected}, got {got}")]
|
||||
DimMismatch { expected: usize, got: usize },
|
||||
|
||||
/// `dim` was zero at construction.
|
||||
#[error("dimension must be > 0")]
|
||||
ZeroDim,
|
||||
|
||||
/// An inserted vector contained NaN or infinity.
|
||||
#[error("vector contains a non-finite value")]
|
||||
NonFiniteVector,
|
||||
|
||||
/// A query contained NaN or infinity.
|
||||
#[error("query contains a non-finite value")]
|
||||
NonFiniteQuery,
|
||||
|
||||
/// A scoring helper was called before calibration was frozen.
|
||||
#[error("index has not been finalized")]
|
||||
NotFinalized,
|
||||
|
||||
/// A scoring helper was called with an invalid row position.
|
||||
#[error("position {pos} is out of bounds for {len} encoded vectors")]
|
||||
PositionOutOfBounds { pos: usize, len: usize },
|
||||
|
||||
/// Search was requested on an empty index.
|
||||
#[error("index is empty")]
|
||||
EmptyIndex,
|
||||
|
||||
/// An external id was reused in `add_with_ids`.
|
||||
#[error("duplicate external id: {0}")]
|
||||
DuplicateId(u64),
|
||||
|
||||
/// `add_with_ids` was called with `vectors.len() != ids.len()`.
|
||||
#[error("batch length mismatch: {vectors} vectors but {ids} ids")]
|
||||
BatchLenMismatch { vectors: usize, ids: usize },
|
||||
}
|
||||
|
||||
/// Crate result alias.
|
||||
pub type Result<T> = std::result::Result<T, TurboVecError>;
|
||||
313
crates/ruvector-turbovec/src/idmap.rs
Normal file
313
crates/ruvector-turbovec/src/idmap.rs
Normal file
|
|
@ -0,0 +1,313 @@
|
|||
//! `IdMapIndex` — external `u64` ids, `O(1)` deletion, allowlist-filtered
|
||||
//! search (ADR-254 §T6).
|
||||
//!
|
||||
//! Wraps a [`TurboVecIndex`] whose internal ids are dense insertion positions.
|
||||
//! A tombstone bitmap + `id → pos` map give `O(1)` `remove`; deleted slots are
|
||||
//! skipped at query time. Filtering is an allowlist of external ids.
|
||||
//!
|
||||
//! For this M1 reference, filtered/deleted search does a full candidate scan and
|
||||
//! filters in the heap merge. The block-granularity short-circuit (skip whole
|
||||
//! 32-vector blocks with no allowed slots) is the FastScan-kernel milestone
|
||||
//! (§T5); the *results* here are the oracle that milestone must match.
|
||||
|
||||
use crate::error::{Result, TurboVecError};
|
||||
use crate::index::TurboVecIndex;
|
||||
use crate::quantize::BitWidth;
|
||||
use ruvector_rabitq::AnnIndex;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
|
||||
/// External-id index with deletion and filtered search.
|
||||
pub struct IdMapIndex {
|
||||
inner: TurboVecIndex,
|
||||
/// internal pos → external id
|
||||
ext_ids: Vec<u64>,
|
||||
/// internal pos → alive?
|
||||
alive: Vec<bool>,
|
||||
/// external id → internal pos
|
||||
id_to_pos: HashMap<u64, usize>,
|
||||
live_count: usize,
|
||||
}
|
||||
|
||||
impl IdMapIndex {
|
||||
/// Empty index over `dim`-d vectors at `bw` bits/coord.
|
||||
pub fn new(dim: usize, bw: BitWidth) -> Result<Self> {
|
||||
Ok(Self {
|
||||
inner: TurboVecIndex::new(dim, bw)?,
|
||||
ext_ids: Vec::new(),
|
||||
alive: Vec::new(),
|
||||
id_to_pos: HashMap::new(),
|
||||
live_count: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Reproducible variant.
|
||||
pub fn with_seed(dim: usize, bw: BitWidth, seed: u64) -> Result<Self> {
|
||||
Ok(Self {
|
||||
inner: TurboVecIndex::with_seed(dim, bw, seed)?,
|
||||
ext_ids: Vec::new(),
|
||||
alive: Vec::new(),
|
||||
id_to_pos: HashMap::new(),
|
||||
live_count: 0,
|
||||
})
|
||||
}
|
||||
|
||||
/// Freeze calibration & encode (delegates to the inner index).
|
||||
pub fn finalize(&mut self) {
|
||||
self.inner.finalize();
|
||||
}
|
||||
|
||||
/// Number of live (non-deleted) vectors.
|
||||
#[inline]
|
||||
pub fn len(&self) -> usize {
|
||||
self.live_count
|
||||
}
|
||||
|
||||
/// Whether the index holds no live vectors.
|
||||
#[inline]
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.live_count == 0
|
||||
}
|
||||
|
||||
/// Honest memory accounting incl. the id bookkeeping.
|
||||
pub fn memory_bytes(&self) -> usize {
|
||||
self.inner.memory_bytes()
|
||||
+ self.ext_ids.len() * 8
|
||||
+ self.alive.len()
|
||||
+ self.id_to_pos.len() * (8 + std::mem::size_of::<usize>())
|
||||
}
|
||||
|
||||
/// Add one vector under external id `id`. Errors on dim/duplicate id.
|
||||
pub fn add_with_id(&mut self, id: u64, vector: Vec<f32>) -> Result<()> {
|
||||
if vector.len() != self.inner.dim() {
|
||||
return Err(TurboVecError::DimMismatch {
|
||||
expected: self.inner.dim(),
|
||||
got: vector.len(),
|
||||
});
|
||||
}
|
||||
if !vector.iter().all(|value| value.is_finite()) {
|
||||
return Err(TurboVecError::NonFiniteVector);
|
||||
}
|
||||
if self.id_to_pos.contains_key(&id) {
|
||||
return Err(TurboVecError::DuplicateId(id));
|
||||
}
|
||||
let pos = self.ext_ids.len();
|
||||
let got = vector.len();
|
||||
// inner uses internal id == pos; dim already validated above so this
|
||||
// map_err is effectively unreachable, but stays correct if it isn't.
|
||||
self.inner
|
||||
.add(pos, vector)
|
||||
.map_err(|_| TurboVecError::DimMismatch {
|
||||
expected: self.inner.dim(),
|
||||
got,
|
||||
})?;
|
||||
self.ext_ids.push(id);
|
||||
self.alive.push(true);
|
||||
self.id_to_pos.insert(id, pos);
|
||||
self.live_count += 1;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Bulk add. `ids.len()` must equal `vectors.len()`.
|
||||
pub fn add_with_ids(&mut self, vectors: &[Vec<f32>], ids: &[u64]) -> Result<()> {
|
||||
if vectors.len() != ids.len() {
|
||||
return Err(TurboVecError::BatchLenMismatch {
|
||||
vectors: vectors.len(),
|
||||
ids: ids.len(),
|
||||
});
|
||||
}
|
||||
let mut batch_ids = HashSet::with_capacity(ids.len());
|
||||
for (vector, &id) in vectors.iter().zip(ids) {
|
||||
if vector.len() != self.inner.dim() {
|
||||
return Err(TurboVecError::DimMismatch {
|
||||
expected: self.inner.dim(),
|
||||
got: vector.len(),
|
||||
});
|
||||
}
|
||||
if !vector.iter().all(|value| value.is_finite()) {
|
||||
return Err(TurboVecError::NonFiniteVector);
|
||||
}
|
||||
if self.id_to_pos.contains_key(&id) || !batch_ids.insert(id) {
|
||||
return Err(TurboVecError::DuplicateId(id));
|
||||
}
|
||||
}
|
||||
for (v, &id) in vectors.iter().zip(ids.iter()) {
|
||||
self.add_with_id(id, v.clone())?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Delete by external id in `O(1)` (tombstone). Returns `true` if removed.
|
||||
pub fn remove(&mut self, id: u64) -> bool {
|
||||
if let Some(pos) = self.id_to_pos.remove(&id) {
|
||||
if self.alive[pos] {
|
||||
self.alive[pos] = false;
|
||||
self.live_count -= 1;
|
||||
return true;
|
||||
}
|
||||
}
|
||||
false
|
||||
}
|
||||
|
||||
/// Top-`k` live nearest, returned as `(external_id, score)` ascending.
|
||||
pub fn search(&self, query: &[f32], k: usize) -> Result<Vec<(u64, f32)>> {
|
||||
self.filtered(query, k, None)
|
||||
}
|
||||
|
||||
/// Top-`k` nearest among `allowlist` (external ids) that are still live.
|
||||
pub fn search_filtered(
|
||||
&self,
|
||||
query: &[f32],
|
||||
k: usize,
|
||||
allowlist: &[u64],
|
||||
) -> Result<Vec<(u64, f32)>> {
|
||||
let set: HashSet<u64> = allowlist.iter().copied().collect();
|
||||
self.filtered(query, k, Some(&set))
|
||||
}
|
||||
|
||||
fn filtered(
|
||||
&self,
|
||||
query: &[f32],
|
||||
k: usize,
|
||||
allow: Option<&HashSet<u64>>,
|
||||
) -> Result<Vec<(u64, f32)>> {
|
||||
if query.len() != self.inner.dim() {
|
||||
return Err(TurboVecError::DimMismatch {
|
||||
expected: self.inner.dim(),
|
||||
got: query.len(),
|
||||
});
|
||||
}
|
||||
// Full candidate scan (reference path), then filter alive ∧ allowed.
|
||||
let n = self.inner.len();
|
||||
let all = self
|
||||
.inner
|
||||
.search(query, n)
|
||||
.map_err(|_| TurboVecError::DimMismatch {
|
||||
expected: self.inner.dim(),
|
||||
got: query.len(),
|
||||
})?;
|
||||
let mut out = Vec::with_capacity(k);
|
||||
for r in all {
|
||||
let pos = r.id;
|
||||
if pos >= self.alive.len() || !self.alive[pos] {
|
||||
continue;
|
||||
}
|
||||
let ext = self.ext_ids[pos];
|
||||
if let Some(a) = allow {
|
||||
if !a.contains(&ext) {
|
||||
continue;
|
||||
}
|
||||
}
|
||||
out.push((ext, r.score));
|
||||
if out.len() == k {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
fn rnd(n: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
|
||||
let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
|
||||
(0..n)
|
||||
.map(|_| (0..dim).map(|_| rng.gen::<f32>()).collect())
|
||||
.collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_is_o1_and_hidden_from_search() {
|
||||
let dim = 48;
|
||||
let data = rnd(500, dim, 4);
|
||||
let ids: Vec<u64> = (0..data.len() as u64).map(|i| i * 10 + 1).collect();
|
||||
let mut ix = IdMapIndex::with_seed(dim, BitWidth::Four, 1).unwrap();
|
||||
ix.add_with_ids(&data, &ids).unwrap();
|
||||
ix.finalize();
|
||||
|
||||
let before = ix.search(&data[0], 5).unwrap();
|
||||
let victim = before[0].0;
|
||||
assert!(ix.remove(victim));
|
||||
assert!(!ix.remove(victim)); // second remove is a no-op
|
||||
assert_eq!(ix.len(), data.len() - 1);
|
||||
|
||||
let after = ix.search(&data[0], 5).unwrap();
|
||||
assert!(
|
||||
after.iter().all(|(id, _)| *id != victim),
|
||||
"deleted id must not appear"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn duplicate_id_rejected() {
|
||||
let mut ix = IdMapIndex::new(8, BitWidth::Two).unwrap();
|
||||
ix.add_with_id(7, vec![1.0; 8]).unwrap();
|
||||
assert!(matches!(
|
||||
ix.add_with_id(7, vec![2.0; 8]),
|
||||
Err(TurboVecError::DuplicateId(7))
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_len_mismatch_rejected() {
|
||||
let mut ix = IdMapIndex::new(8, BitWidth::Two).unwrap();
|
||||
let vectors = vec![vec![1.0; 8], vec![2.0; 8]];
|
||||
let ids = vec![1u64]; // too few
|
||||
assert!(matches!(
|
||||
ix.add_with_ids(&vectors, &ids),
|
||||
Err(TurboVecError::BatchLenMismatch { vectors: 2, ids: 1 })
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_validation_is_atomic() {
|
||||
let mut ix = IdMapIndex::new(8, BitWidth::Two).unwrap();
|
||||
let vectors = vec![vec![1.0; 8], vec![2.0; 8]];
|
||||
assert!(matches!(
|
||||
ix.add_with_ids(&vectors, &[7, 7]),
|
||||
Err(TurboVecError::DuplicateId(7))
|
||||
));
|
||||
assert!(ix.is_empty());
|
||||
|
||||
let mut invalid = vec![2.0; 8];
|
||||
invalid[0] = f32::NAN;
|
||||
assert!(matches!(
|
||||
ix.add_with_ids(&[vec![1.0; 8], invalid], &[1, 2]),
|
||||
Err(TurboVecError::NonFiniteVector)
|
||||
));
|
||||
assert!(ix.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn wrong_dim_rejected() {
|
||||
let mut ix = IdMapIndex::new(8, BitWidth::Two).unwrap();
|
||||
assert!(matches!(
|
||||
ix.add_with_id(1, vec![0.0; 4]),
|
||||
Err(TurboVecError::DimMismatch {
|
||||
expected: 8,
|
||||
got: 4
|
||||
})
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn filtered_search_only_returns_allowed() {
|
||||
let dim = 32;
|
||||
let data = rnd(400, dim, 6);
|
||||
let ids: Vec<u64> = (0..data.len() as u64).collect();
|
||||
let mut ix = IdMapIndex::with_seed(dim, BitWidth::Four, 9).unwrap();
|
||||
ix.add_with_ids(&data, &ids).unwrap();
|
||||
ix.finalize();
|
||||
|
||||
let allow: Vec<u64> = (0..400).step_by(5).collect(); // 80 allowed ids
|
||||
let allow_set: HashSet<u64> = allow.iter().copied().collect();
|
||||
let res = ix.search_filtered(&data[3], 10, &allow).unwrap();
|
||||
assert!(!res.is_empty());
|
||||
assert!(
|
||||
res.iter().all(|(id, _)| allow_set.contains(id)),
|
||||
"every hit must be in the allowlist"
|
||||
);
|
||||
}
|
||||
}
|
||||
580
crates/ruvector-turbovec/src/index.rs
Normal file
580
crates/ruvector-turbovec/src/index.rs
Normal file
|
|
@ -0,0 +1,580 @@
|
|||
//! `TurboVecIndex` — the multi-bit TurboQuant flat ANN index (ADR-254).
|
||||
//!
|
||||
//! Pipeline per database vector `x`:
|
||||
//!
|
||||
//! 1. `norm = ‖x‖`, `û = x / norm` — strip & store length (§T1)
|
||||
//! 2. `r = P · û` (randomized Hadamard) — reuse `ruvector_rabitq` (§T1)
|
||||
//! 3. `z_i = (r_i − shift_i)/scale_i` — TQ+ calibration (§T3)
|
||||
//! 4. `q_i = argmin |z_i − centroid|` — Lloyd–Max 2/3/4-bit SQ (§T2)
|
||||
//! 5. `c_x = ⟨r, r̂⟩ / ⟨r̂, r̂⟩` — per-vector bias correction (§T4)
|
||||
//!
|
||||
//! Scoring a query `q` (orthogonal `P` preserves inner products, so
|
||||
//! `⟨r,s⟩ = ⟨û,q̂⟩ = cos`):
|
||||
//!
|
||||
//! ```text
|
||||
//! s = P · q̂ (unit, rotated once / query)
|
||||
//! cos_est = c_x · Σ_i r̂_i · s_i (empirical bias correction)
|
||||
//! ⟨q,x⟩_est = ‖q‖ · ‖x‖ · cos_est
|
||||
//! score = ‖q‖² + ‖x‖² − 2⟨q,x⟩_est ≈ ‖q − x‖²
|
||||
//! ```
|
||||
//!
|
||||
//! Smaller `score` = nearer, matching `ruvector_rabitq::AnnIndex` ordering, so
|
||||
//! recall is directly comparable to exact brute-force L2.
|
||||
|
||||
use crate::calibrate::Calibration;
|
||||
use crate::error::{Result, TurboVecError};
|
||||
use crate::quantize::{pack, quantize_coord, unpack, BitWidth};
|
||||
use ruvector_rabitq::rotation::normalize_inplace;
|
||||
use ruvector_rabitq::{AnnIndex, RandomRotation, SearchResult};
|
||||
|
||||
/// Default number of vectors collected before calibration is frozen.
|
||||
pub const DEFAULT_WARMUP: usize = 256;
|
||||
|
||||
/// A flat multi-bit TurboQuant index over `dim`-dimensional vectors.
|
||||
pub struct TurboVecIndex {
|
||||
/// Public input dimensionality.
|
||||
dim: usize,
|
||||
/// Full power-of-two Hadamard space. Keeping every rotated coordinate is
|
||||
/// required for exact norm and inner-product preservation.
|
||||
rotated_dim: usize,
|
||||
bw: BitWidth,
|
||||
rotation: RandomRotation,
|
||||
warmup: usize,
|
||||
|
||||
/// Raw vectors awaiting calibration. Empty once finalized (memory freed).
|
||||
staging: Vec<(usize, Vec<f32>)>,
|
||||
calibration: Option<Calibration>,
|
||||
|
||||
// ── SoA compressed storage (post-finalize) ───────────────────────────
|
||||
ids: Vec<usize>,
|
||||
codes: Vec<u8>, // flattened: code_bytes(dim) per vector
|
||||
norms: Vec<f32>,
|
||||
corr: Vec<f32>, // per-vector c_x
|
||||
}
|
||||
|
||||
impl TurboVecIndex {
|
||||
/// Build an empty index. `seed` makes the Hadamard rotation reproducible.
|
||||
pub fn with_seed(dim: usize, bw: BitWidth, seed: u64) -> Result<Self> {
|
||||
if dim == 0 {
|
||||
return Err(TurboVecError::ZeroDim);
|
||||
}
|
||||
let rotated_dim = dim.next_power_of_two();
|
||||
Ok(Self {
|
||||
dim,
|
||||
rotated_dim,
|
||||
bw,
|
||||
rotation: RandomRotation::hadamard(rotated_dim, seed),
|
||||
warmup: DEFAULT_WARMUP,
|
||||
staging: Vec::new(),
|
||||
calibration: None,
|
||||
ids: Vec::new(),
|
||||
codes: Vec::new(),
|
||||
norms: Vec::new(),
|
||||
corr: Vec::new(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Build with the default seed (42).
|
||||
pub fn new(dim: usize, bw: BitWidth) -> Result<Self> {
|
||||
Self::with_seed(dim, bw, 42)
|
||||
}
|
||||
|
||||
/// Override the warm-up count (vectors gathered before calibration freezes).
|
||||
pub fn with_warmup(mut self, warmup: usize) -> Self {
|
||||
self.warmup = warmup.max(1);
|
||||
self
|
||||
}
|
||||
|
||||
/// Whether calibration has been frozen and staged vectors encoded.
|
||||
#[inline]
|
||||
pub fn is_finalized(&self) -> bool {
|
||||
self.calibration.is_some()
|
||||
}
|
||||
|
||||
/// Number of indexed vectors (staged + encoded).
|
||||
#[inline]
|
||||
pub fn count(&self) -> usize {
|
||||
self.ids.len() + self.staging.len()
|
||||
}
|
||||
|
||||
/// Bytes per encoded vector: packed code + norm(f32) + corr(f32) + id(usize).
|
||||
#[inline]
|
||||
pub fn bytes_per_vector(&self) -> usize {
|
||||
self.bw.code_bytes(self.rotated_dim) + 4 + 4 + std::mem::size_of::<usize>()
|
||||
}
|
||||
|
||||
/// Compression ratio of encoded storage vs raw `f32` (codes+scalars only,
|
||||
/// rotation/calibration are amortized shared overhead).
|
||||
pub fn compression_ratio(&self) -> f32 {
|
||||
let raw = (self.dim * 4) as f32;
|
||||
let comp = (self.bw.code_bytes(self.rotated_dim) + 8) as f32; // code + norm + corr
|
||||
raw / comp
|
||||
}
|
||||
|
||||
/// Rotate a normalized input in the full padded Hadamard space.
|
||||
///
|
||||
/// Truncating a padded Hadamard output is not orthogonal when the input
|
||||
/// dimension is far from a power of two (notably 1536 → 2048). Retaining
|
||||
/// all padded coordinates preserves L2 geometry exactly.
|
||||
fn rotate_unit(&self, unit: &[f32]) -> Vec<f32> {
|
||||
debug_assert_eq!(unit.len(), self.dim);
|
||||
let mut padded = vec![0.0; self.rotated_dim];
|
||||
padded[..self.dim].copy_from_slice(unit);
|
||||
self.rotation.apply(&padded)
|
||||
}
|
||||
|
||||
/// Freeze calibration on the staged batch and encode everything. Idempotent.
|
||||
pub fn finalize(&mut self) {
|
||||
if self.is_finalized() {
|
||||
return;
|
||||
}
|
||||
// Rotate staged unit vectors to fit calibration. Zero vectors are not
|
||||
// unit-normalizable (they would contribute all-zero rows and bias
|
||||
// shift/scale toward zero), so they are excluded here — `encode_and_store`
|
||||
// already special-cases them at encode time.
|
||||
let rotated: Vec<Vec<f32>> = self
|
||||
.staging
|
||||
.iter()
|
||||
.filter(|(_, v)| v.iter().map(|x| x * x).sum::<f32>() > 1e-20)
|
||||
.map(|(_, v)| {
|
||||
let mut u = v.clone();
|
||||
normalize_inplace(&mut u);
|
||||
self.rotate_unit(&u)
|
||||
})
|
||||
.collect();
|
||||
let cal = Calibration::fit(&rotated, self.rotated_dim);
|
||||
// Encode each staged vector using the frozen calibration.
|
||||
let staged = std::mem::take(&mut self.staging);
|
||||
self.calibration = Some(cal);
|
||||
self.ids.reserve(staged.len());
|
||||
for (id, v) in staged {
|
||||
self.encode_and_store(id, &v);
|
||||
}
|
||||
// staging is now empty → raw f32 memory released.
|
||||
}
|
||||
|
||||
/// Encode one vector with the (already frozen) calibration and append it.
|
||||
fn encode_and_store(&mut self, id: usize, v: &[f32]) {
|
||||
let cal = self.calibration.as_ref().expect("finalized");
|
||||
let centroids = self.bw.centroids();
|
||||
|
||||
let mut u = v.to_vec();
|
||||
let norm = {
|
||||
let n: f32 = u.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
n
|
||||
};
|
||||
if norm < 1e-10 {
|
||||
// Degenerate zero vector: store an all-zero code, neutral scalars.
|
||||
self.ids.push(id);
|
||||
self.codes
|
||||
.extend(std::iter::repeat(0u8).take(self.bw.code_bytes(self.rotated_dim)));
|
||||
self.norms.push(0.0);
|
||||
self.corr.push(0.0);
|
||||
return;
|
||||
}
|
||||
for x in u.iter_mut() {
|
||||
*x /= norm;
|
||||
}
|
||||
let r = self.rotate_unit(&u);
|
||||
|
||||
// Quantize + reconstruct, accumulating ⟨r,r̂⟩ and ⟨r̂,r̂⟩ for c_x.
|
||||
let mut q_codes = vec![0u8; self.rotated_dim];
|
||||
let mut dot_r_rhat = 0.0f32;
|
||||
let mut dot_rhat = 0.0f32;
|
||||
for (i, (&ri, slot)) in r.iter().zip(q_codes.iter_mut()).enumerate() {
|
||||
let z = cal.calibrate(i, ri);
|
||||
let code = quantize_coord(centroids, z);
|
||||
*slot = code;
|
||||
let rhat = cal.reconstruct(i, centroids[code as usize]);
|
||||
dot_r_rhat += ri * rhat;
|
||||
dot_rhat += rhat * rhat;
|
||||
}
|
||||
let c_x = if dot_rhat > 1e-20 {
|
||||
dot_r_rhat / dot_rhat
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
self.ids.push(id);
|
||||
self.codes.extend_from_slice(&pack(&q_codes, self.bw));
|
||||
self.norms.push(norm);
|
||||
self.corr.push(c_x);
|
||||
}
|
||||
|
||||
/// The encoded code-slice for stored vector `pos`.
|
||||
#[inline]
|
||||
fn code_slice(&self, pos: usize) -> &[u8] {
|
||||
let cb = self.bw.code_bytes(self.rotated_dim);
|
||||
&self.codes[pos * cb..(pos + 1) * cb]
|
||||
}
|
||||
|
||||
/// Estimated `‖q − x_pos‖²` given the rotated unit query `s` and `‖q‖`.
|
||||
/// `centroids` and `cal` are hoisted by the caller for the hot loop.
|
||||
#[inline]
|
||||
fn estimate_l2(&self, pos: usize, s: &[f32], nq: f32, cal: &Calibration) -> f32 {
|
||||
let centroids = self.bw.centroids();
|
||||
let codes = unpack(self.code_slice(pos), self.rotated_dim, self.bw);
|
||||
let mut ip_unit = 0.0f32;
|
||||
for (i, (&code, &si)) in codes.iter().zip(s.iter()).enumerate() {
|
||||
let rhat = cal.reconstruct(i, centroids[code as usize]);
|
||||
ip_unit += rhat * si;
|
||||
}
|
||||
let cos_est = self.corr[pos] * ip_unit;
|
||||
let ip_est = nq * self.norms[pos] * cos_est;
|
||||
nq * nq + self.norms[pos] * self.norms[pos] - 2.0 * ip_est
|
||||
}
|
||||
|
||||
/// Estimated cosine similarity between `query` and stored vector `pos`.
|
||||
/// Exposed for bias measurement in tests and the demo.
|
||||
pub fn estimate_cosine(&self, query: &[f32], pos: usize) -> Result<f32> {
|
||||
if query.len() != self.dim {
|
||||
return Err(TurboVecError::DimMismatch {
|
||||
expected: self.dim,
|
||||
got: query.len(),
|
||||
});
|
||||
}
|
||||
if !query.iter().all(|value| value.is_finite()) {
|
||||
return Err(TurboVecError::NonFiniteQuery);
|
||||
}
|
||||
let cal = self
|
||||
.calibration
|
||||
.as_ref()
|
||||
.ok_or(TurboVecError::NotFinalized)?;
|
||||
if pos >= self.ids.len() {
|
||||
return Err(TurboVecError::PositionOutOfBounds {
|
||||
pos,
|
||||
len: self.ids.len(),
|
||||
});
|
||||
}
|
||||
let mut qn = query.to_vec();
|
||||
normalize_inplace(&mut qn);
|
||||
let s = self.rotate_unit(&qn);
|
||||
let centroids = self.bw.centroids();
|
||||
let codes = unpack(self.code_slice(pos), self.rotated_dim, self.bw);
|
||||
let mut ip_unit = 0.0f32;
|
||||
for (i, (&code, &si)) in codes.iter().zip(s.iter()).enumerate() {
|
||||
ip_unit += cal.reconstruct(i, centroids[code as usize]) * si;
|
||||
}
|
||||
Ok(self.corr[pos] * ip_unit)
|
||||
}
|
||||
|
||||
/// Shared search core (used by `AnnIndex::search`). When the index has not
|
||||
/// been finalized (tiny indexes below warm-up that never called
|
||||
/// `finalize`), falls back to an exact brute-force L2 over the staged raw
|
||||
/// vectors — correct, just uncompressed.
|
||||
fn search_core(&self, query: &[f32], k: usize) -> Result<Vec<SearchResult>> {
|
||||
if query.len() != self.dim {
|
||||
return Err(TurboVecError::DimMismatch {
|
||||
expected: self.dim,
|
||||
got: query.len(),
|
||||
});
|
||||
}
|
||||
if !query.iter().all(|value| value.is_finite()) {
|
||||
return Err(TurboVecError::NonFiniteQuery);
|
||||
}
|
||||
if self.count() == 0 {
|
||||
return Err(TurboVecError::EmptyIndex);
|
||||
}
|
||||
let mut scored: Vec<SearchResult> = if let Some(cal) = self.calibration.as_ref() {
|
||||
let mut qn = query.to_vec();
|
||||
let nq: f32 = qn.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
normalize_inplace(&mut qn);
|
||||
let s = self.rotate_unit(&qn);
|
||||
(0..self.ids.len())
|
||||
.map(|pos| SearchResult {
|
||||
id: self.ids[pos],
|
||||
score: self.estimate_l2(pos, &s, nq, cal),
|
||||
})
|
||||
.collect()
|
||||
} else {
|
||||
// Exact fallback for not-yet-finalized tiny indexes.
|
||||
self.staging
|
||||
.iter()
|
||||
.map(|(id, v)| SearchResult {
|
||||
id: *id,
|
||||
score: v.iter().zip(query).map(|(a, b)| (a - b) * (a - b)).sum(),
|
||||
})
|
||||
.collect()
|
||||
};
|
||||
let compare = |a: &SearchResult, b: &SearchResult| {
|
||||
a.score.total_cmp(&b.score).then_with(|| a.id.cmp(&b.id))
|
||||
};
|
||||
if k < scored.len() {
|
||||
scored.select_nth_unstable_by(k, compare);
|
||||
scored.truncate(k);
|
||||
}
|
||||
scored.sort_unstable_by(compare);
|
||||
Ok(scored)
|
||||
}
|
||||
|
||||
/// Exposed total memory accounting (used by `AnnIndex::memory_bytes`).
|
||||
fn memory(&self) -> usize {
|
||||
let cal = self.calibration.as_ref().map(|c| c.bytes()).unwrap_or(0);
|
||||
let staged: usize = self
|
||||
.staging
|
||||
.iter()
|
||||
.map(|(_, v)| v.len() * 4 + std::mem::size_of::<usize>())
|
||||
.sum();
|
||||
self.rotation.bytes()
|
||||
+ cal
|
||||
+ self.codes.len()
|
||||
+ self.norms.len() * 4
|
||||
+ self.corr.len() * 4
|
||||
+ self.ids.len() * std::mem::size_of::<usize>()
|
||||
+ staged
|
||||
}
|
||||
}
|
||||
|
||||
impl AnnIndex for TurboVecIndex {
|
||||
fn add(&mut self, id: usize, vector: Vec<f32>) -> ruvector_rabitq::error::Result<()> {
|
||||
// Surface the same `DimensionMismatch` other `AnnIndex` impls do, in
|
||||
// release builds too — never silently accept a wrong-length vector.
|
||||
if vector.len() != self.dim {
|
||||
return Err(ruvector_rabitq::error::RabitqError::DimensionMismatch {
|
||||
expected: self.dim,
|
||||
actual: vector.len(),
|
||||
});
|
||||
}
|
||||
if !vector.iter().all(|value| value.is_finite()) {
|
||||
return Err(ruvector_rabitq::error::RabitqError::InvalidParameter(
|
||||
"vector contains a non-finite value".to_string(),
|
||||
));
|
||||
}
|
||||
if self.is_finalized() {
|
||||
self.encode_and_store(id, &vector);
|
||||
} else {
|
||||
self.staging.push((id, vector));
|
||||
if self.staging.len() >= self.warmup {
|
||||
self.finalize();
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn search(&self, query: &[f32], k: usize) -> ruvector_rabitq::error::Result<Vec<SearchResult>> {
|
||||
// Propagate dimension mismatches rather than masking them as an empty
|
||||
// result (which would hide caller bugs).
|
||||
self.search_core(query, k).map_err(|error| match error {
|
||||
TurboVecError::DimMismatch { expected, got } => {
|
||||
ruvector_rabitq::error::RabitqError::DimensionMismatch {
|
||||
expected,
|
||||
actual: got,
|
||||
}
|
||||
}
|
||||
TurboVecError::EmptyIndex => ruvector_rabitq::error::RabitqError::EmptyIndex,
|
||||
other => ruvector_rabitq::error::RabitqError::InvalidParameter(other.to_string()),
|
||||
})
|
||||
}
|
||||
|
||||
fn len(&self) -> usize {
|
||||
self.count()
|
||||
}
|
||||
|
||||
fn dim(&self) -> usize {
|
||||
self.dim
|
||||
}
|
||||
|
||||
fn memory_bytes(&self) -> usize {
|
||||
self.memory()
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
fn random_vectors(n: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
|
||||
let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
|
||||
(0..n)
|
||||
.map(|_| (0..dim).map(|_| rng.gen::<f32>() * 2.0 - 1.0).collect())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn brute_topk(data: &[Vec<f32>], q: &[f32], k: usize) -> Vec<usize> {
|
||||
let mut s: Vec<(usize, f32)> = data
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (i, v.iter().zip(q).map(|(a, b)| (a - b) * (a - b)).sum()))
|
||||
.collect();
|
||||
s.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
|
||||
s.into_iter().take(k).map(|(i, _)| i).collect()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn determinism_same_seed_same_codes() {
|
||||
let data = random_vectors(300, 64, 1);
|
||||
let build = |seed| {
|
||||
let mut ix = TurboVecIndex::with_seed(64, BitWidth::Four, seed).unwrap();
|
||||
for (i, v) in data.iter().enumerate() {
|
||||
ix.add(i, v.clone()).unwrap();
|
||||
}
|
||||
ix.finalize();
|
||||
ix.codes.clone()
|
||||
};
|
||||
assert_eq!(build(7), build(7), "same seed must give identical codes");
|
||||
assert_ne!(build(7), build(8), "different seed should differ");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn recall_beats_threshold_4bit() {
|
||||
let dim = 128;
|
||||
let data = random_vectors(2000, dim, 10);
|
||||
let mut ix = TurboVecIndex::with_seed(dim, BitWidth::Four, 42).unwrap();
|
||||
for (i, v) in data.iter().enumerate() {
|
||||
ix.add(i, v.clone()).unwrap();
|
||||
}
|
||||
ix.finalize();
|
||||
assert!(ix.is_finalized());
|
||||
|
||||
let queries = random_vectors(100, dim, 99);
|
||||
let k = 10;
|
||||
let mut hits = 0usize;
|
||||
let mut total = 0usize;
|
||||
for q in &queries {
|
||||
let truth = brute_topk(&data, q, k);
|
||||
let got: Vec<usize> = ix.search(q, k).unwrap().into_iter().map(|r| r.id).collect();
|
||||
for g in &got {
|
||||
if truth.contains(g) {
|
||||
hits += 1;
|
||||
}
|
||||
}
|
||||
total += k;
|
||||
}
|
||||
let recall = hits as f32 / total as f32;
|
||||
// 4-bit scalar quant on 128-d uniform data: expect comfortably > 0.5.
|
||||
assert!(recall > 0.5, "recall@10 too low: {recall:.3}");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn estimator_is_approximately_unbiased() {
|
||||
// Mean signed cosine error over many (query, db) pairs should be ~0,
|
||||
// which is the whole point of the c_x length-renormalization (§T4).
|
||||
let dim = 256;
|
||||
let data = random_vectors(500, dim, 3);
|
||||
let mut ix = TurboVecIndex::with_seed(dim, BitWidth::Four, 5).unwrap();
|
||||
for (i, v) in data.iter().enumerate() {
|
||||
ix.add(i, v.clone()).unwrap();
|
||||
}
|
||||
ix.finalize();
|
||||
|
||||
let queries = random_vectors(50, dim, 77);
|
||||
let mut sum_err = 0.0f64;
|
||||
let mut n = 0u64;
|
||||
for q in &queries {
|
||||
let nq: f32 = q.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
for (pos, x) in data.iter().enumerate().take(200) {
|
||||
let nx: f32 = x.iter().map(|t| t * t).sum::<f32>().sqrt();
|
||||
let true_cos = q.iter().zip(x).map(|(a, b)| a * b).sum::<f32>() / (nq * nx);
|
||||
let est_cos = ix.estimate_cosine(q, pos).unwrap();
|
||||
sum_err += (est_cos - true_cos) as f64;
|
||||
n += 1;
|
||||
}
|
||||
}
|
||||
let mean_err = (sum_err / n as f64).abs();
|
||||
assert!(
|
||||
mean_err < 0.02,
|
||||
"estimator biased: mean |err| = {mean_err:.4}"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn add_and_search_reject_wrong_dim() {
|
||||
let mut ix = TurboVecIndex::new(16, BitWidth::Four).unwrap();
|
||||
// add with wrong-length vector must error, not silently accept.
|
||||
assert!(ix.add(0, vec![0.0; 8]).is_err());
|
||||
ix.add(0, vec![0.1; 16]).unwrap();
|
||||
ix.finalize();
|
||||
// search with wrong-length query must error, not return empty.
|
||||
assert!(ix.search(&[0.0; 8], 1).is_err());
|
||||
assert!(ix.search(&[0.0; 16], 1).is_ok());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_non_finite_vectors_and_queries() {
|
||||
let mut ix = TurboVecIndex::new(16, BitWidth::Four).unwrap();
|
||||
let mut bad = vec![0.0; 16];
|
||||
bad[3] = f32::NAN;
|
||||
assert!(ix.add(0, bad).is_err());
|
||||
ix.add(0, vec![0.1; 16]).unwrap();
|
||||
ix.finalize();
|
||||
|
||||
let mut bad_query = vec![0.0; 16];
|
||||
bad_query[5] = f32::INFINITY;
|
||||
assert!(ix.search(&bad_query, 1).is_err());
|
||||
assert!(matches!(
|
||||
ix.estimate_cosine(&bad_query, 0),
|
||||
Err(TurboVecError::NonFiniteQuery)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn padded_rotation_preserves_geometry_at_1536_dimensions() {
|
||||
let dim = 1536;
|
||||
let ix = TurboVecIndex::with_seed(dim, BitWidth::Four, 17).unwrap();
|
||||
assert_eq!(ix.rotated_dim, 2048);
|
||||
|
||||
let mut a: Vec<f32> = (0..dim).map(|i| ((i * 17) as f32).sin()).collect();
|
||||
let mut b: Vec<f32> = (0..dim).map(|i| ((i * 29) as f32).cos()).collect();
|
||||
normalize_inplace(&mut a);
|
||||
normalize_inplace(&mut b);
|
||||
let dot_before: f32 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
|
||||
let ra = ix.rotate_unit(&a);
|
||||
let rb = ix.rotate_unit(&b);
|
||||
let dot_after: f32 = ra.iter().zip(&rb).map(|(x, y)| x * y).sum();
|
||||
let norm_a: f32 = ra.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
let norm_b: f32 = rb.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
|
||||
assert!((norm_a - 1.0).abs() < 1e-4);
|
||||
assert!((norm_b - 1.0).abs() < 1e-4);
|
||||
assert!((dot_before - dot_after).abs() < 1e-4);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn empty_index_returns_empty_index_error() {
|
||||
let ix = TurboVecIndex::new(8, BitWidth::Two).unwrap();
|
||||
assert!(matches!(
|
||||
ix.search(&[0.0; 8], 1),
|
||||
Err(ruvector_rabitq::error::RabitqError::EmptyIndex)
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn finalize_excludes_zero_vectors_from_calibration() {
|
||||
// A pile of zero vectors mixed into the warm-up must NOT corrupt the
|
||||
// calibration fitted from the real vectors. If zeros leaked in they
|
||||
// would collapse `scale` toward zero and wreck reconstruction; we prove
|
||||
// calibration stays faithful via self-retrieval of a stored vector.
|
||||
let dim = 64;
|
||||
let real = random_vectors(50, dim, 2);
|
||||
let mut ix = TurboVecIndex::with_seed(dim, BitWidth::Four, 1).unwrap();
|
||||
for _ in 0..50 {
|
||||
ix.add(0, vec![0.0; dim]).unwrap();
|
||||
}
|
||||
for (i, v) in real.iter().enumerate() {
|
||||
ix.add(i + 1, v.clone()).unwrap();
|
||||
}
|
||||
ix.finalize();
|
||||
// Querying a stored real vector must return that same vector as nearest
|
||||
// (distance to self ≈ quantization error ≪ distance to any other point).
|
||||
let hits = ix.search(&real[7], 1).unwrap();
|
||||
assert_eq!(
|
||||
hits[0].id, 8,
|
||||
"self-retrieval failed → calibration corrupted"
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn tiny_index_uses_exact_fallback() {
|
||||
let dim = 32;
|
||||
let data = random_vectors(5, dim, 2); // below warm-up, never finalized
|
||||
let mut ix = TurboVecIndex::new(dim, BitWidth::Two).unwrap();
|
||||
for (i, v) in data.iter().enumerate() {
|
||||
ix.add(i, v.clone()).unwrap();
|
||||
}
|
||||
assert!(!ix.is_finalized());
|
||||
let truth = brute_topk(&data, &data[0], 1)[0];
|
||||
let got = ix.search(&data[0], 1).unwrap()[0].id;
|
||||
assert_eq!(got, truth, "exact fallback must be perfect");
|
||||
}
|
||||
}
|
||||
64
crates/ruvector-turbovec/src/lib.rs
Normal file
64
crates/ruvector-turbovec/src/lib.rs
Normal file
|
|
@ -0,0 +1,64 @@
|
|||
//! # ruvector-turbovec
|
||||
//!
|
||||
//! Multi-bit **TurboQuant** FastScan-style approximate-nearest-neighbor index
|
||||
//! (ADR-254). Fills the 2–4-bit scalar-quantized search regime that ruvector
|
||||
//! lacked: `ruvector-rabitq` is 1-bit, `ruvllm`'s TurboQuant is a KV-cache value
|
||||
//! codec — neither is a multi-bit *search index*.
|
||||
//!
|
||||
//! ## What it reuses
|
||||
//!
|
||||
//! - [`ruvector_rabitq::RandomRotation`] — the randomized Hadamard rotation
|
||||
//! (`O(D log D)`, no matrix stored) that makes per-coordinate scalar
|
||||
//! quantization codebook-free.
|
||||
//! - [`ruvector_rabitq::AnnIndex`] — the shared index trait, so this drops into
|
||||
//! the same dispatcher/consumers as RaBitQ with no new plumbing.
|
||||
//!
|
||||
//! ## What it adds (ADR-254 §T2–T4, T6)
|
||||
//!
|
||||
//! - **Lloyd–Max 2/3/4-bit scalar quantization** ([`quantize`]) — data-independent
|
||||
//! centroids, online ingest, no training.
|
||||
//! - **TQ+ per-coordinate calibration** ([`calibrate`]) — frozen after warm-up.
|
||||
//! - **Length-renormalized scoring** — a per-vector scalar empirically reduces
|
||||
//! inner-product bias at zero query-time cost. It is not the paper's
|
||||
//! provably-unbiased QJL-residual estimator.
|
||||
//! - **`IdMapIndex`** ([`idmap`]) — external `u64` ids, `O(1)` deletion,
|
||||
//! allowlist-filtered search.
|
||||
//!
|
||||
//! The FastScan nibble-LUT SIMD kernel (§T5) is a separate future milestone; the
|
||||
//! scalar scorer here is its determinism oracle.
|
||||
//!
|
||||
//! Non-power-of-two input dimensions are zero-padded to the next power of two
|
||||
//! and quantized in that full rotated space. This preserves L2 geometry (unlike
|
||||
//! truncating a padded Hadamard transform) at the cost of additional code bytes.
|
||||
//!
|
||||
//! ## Example
|
||||
//!
|
||||
//! ```
|
||||
//! use ruvector_turbovec::{TurboVecIndex, BitWidth};
|
||||
//! use ruvector_rabitq::AnnIndex;
|
||||
//!
|
||||
//! let mut ix = TurboVecIndex::new(64, BitWidth::Four).unwrap();
|
||||
//! for i in 0..1000 {
|
||||
//! let v: Vec<f32> = (0..64).map(|j| ((i * 7 + j) % 13) as f32).collect();
|
||||
//! ix.add(i, v).unwrap();
|
||||
//! }
|
||||
//! ix.finalize();
|
||||
//! let q: Vec<f32> = (0..64).map(|j| (j % 13) as f32).collect();
|
||||
//! let hits = ix.search(&q, 5).unwrap();
|
||||
//! assert_eq!(hits.len(), 5);
|
||||
//! ```
|
||||
|
||||
pub mod calibrate;
|
||||
pub mod error;
|
||||
pub mod idmap;
|
||||
pub mod index;
|
||||
pub mod quantize;
|
||||
|
||||
pub use calibrate::Calibration;
|
||||
pub use error::{Result, TurboVecError};
|
||||
pub use idmap::IdMapIndex;
|
||||
pub use index::{TurboVecIndex, DEFAULT_WARMUP};
|
||||
pub use quantize::BitWidth;
|
||||
|
||||
// Re-export the shared trait so downstream users get one import.
|
||||
pub use ruvector_rabitq::{AnnIndex, SearchResult};
|
||||
139
crates/ruvector-turbovec/src/main.rs
Normal file
139
crates/ruvector-turbovec/src/main.rs
Normal file
|
|
@ -0,0 +1,139 @@
|
|||
//! `turbovec-demo` — end-to-end **proof** that the TurboVec index (ADR-254)
|
||||
//! works correctly: compression, recall vs exact brute-force L2, estimator
|
||||
//! empirical estimator bias, determinism, and IdMap delete + filtered search.
|
||||
//!
|
||||
//! Run: `cargo run --release -p ruvector-turbovec`
|
||||
|
||||
use ruvector_rabitq::AnnIndex;
|
||||
use ruvector_turbovec::{BitWidth, IdMapIndex, TurboVecIndex};
|
||||
|
||||
use rand::{Rng, SeedableRng};
|
||||
|
||||
fn random_vectors(n: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
|
||||
let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
|
||||
(0..n)
|
||||
.map(|_| (0..dim).map(|_| rng.gen::<f32>() * 2.0 - 1.0).collect())
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn brute_topk(data: &[Vec<f32>], q: &[f32], k: usize) -> Vec<usize> {
|
||||
let mut s: Vec<(usize, f32)> = data
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, v)| (i, v.iter().zip(q).map(|(a, b)| (a - b) * (a - b)).sum()))
|
||||
.collect();
|
||||
s.sort_unstable_by(|a, b| a.1.total_cmp(&b.1));
|
||||
s.into_iter().take(k).map(|(i, _)| i).collect()
|
||||
}
|
||||
|
||||
fn recall_and_bias(bw: BitWidth, data: &[Vec<f32>], queries: &[Vec<f32>], dim: usize, k: usize) {
|
||||
let mut ix = TurboVecIndex::with_seed(dim, bw, 42).unwrap();
|
||||
for (i, v) in data.iter().enumerate() {
|
||||
ix.add(i, v.clone()).unwrap();
|
||||
}
|
||||
ix.finalize();
|
||||
|
||||
// Recall@k vs exact brute force.
|
||||
let mut hits = 0usize;
|
||||
let mut total = 0usize;
|
||||
for q in queries {
|
||||
let truth = brute_topk(data, q, k);
|
||||
let got: Vec<usize> = ix.search(q, k).unwrap().into_iter().map(|r| r.id).collect();
|
||||
hits += got.iter().filter(|g| truth.contains(g)).count();
|
||||
total += k;
|
||||
}
|
||||
let recall = hits as f32 / total as f32;
|
||||
|
||||
// Mean signed cosine error (empirical bias measurement).
|
||||
let mut sum_err = 0.0f64;
|
||||
let mut n = 0u64;
|
||||
for q in queries.iter().take(20) {
|
||||
let nq: f32 = q.iter().map(|x| x * x).sum::<f32>().sqrt();
|
||||
for (pos, x) in data.iter().enumerate().take(300) {
|
||||
let nx: f32 = x.iter().map(|t| t * t).sum::<f32>().sqrt();
|
||||
let true_cos = q.iter().zip(x).map(|(a, b)| a * b).sum::<f32>() / (nq * nx);
|
||||
sum_err += (ix.estimate_cosine(q, pos).unwrap() - true_cos) as f64;
|
||||
n += 1;
|
||||
}
|
||||
}
|
||||
let mean_bias = sum_err / n as f64;
|
||||
|
||||
let raw_bytes = dim * 4;
|
||||
println!(
|
||||
" {:<6} | recall@{k}: {:.3} | bytes/vec: {:>4} (raw {raw_bytes}) | {:>5.1}x | mean cos-bias: {:+.4}",
|
||||
format!("{bw:?}"),
|
||||
recall,
|
||||
ix.bytes_per_vector(),
|
||||
ix.compression_ratio(),
|
||||
mean_bias,
|
||||
);
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let dim = 256;
|
||||
let n = 5_000;
|
||||
let k = 10;
|
||||
let data = random_vectors(n, dim, 10);
|
||||
let queries = random_vectors(200, dim, 99);
|
||||
|
||||
println!("\n=== TurboVec (ADR-254) proof — n={n}, dim={dim}, k={k} ===\n");
|
||||
println!("[1] Compression + recall vs exact brute-force L2 + estimator bias");
|
||||
for bw in [
|
||||
BitWidth::One,
|
||||
BitWidth::Two,
|
||||
BitWidth::Three,
|
||||
BitWidth::Four,
|
||||
] {
|
||||
recall_and_bias(bw, &data, &queries, dim, k);
|
||||
}
|
||||
|
||||
// [2] Determinism.
|
||||
println!("\n[2] Determinism (same seed => identical results)");
|
||||
let build = |seed: u64| {
|
||||
let mut ix = TurboVecIndex::with_seed(dim, BitWidth::Four, seed).unwrap();
|
||||
for (i, v) in data.iter().enumerate() {
|
||||
ix.add(i, v.clone()).unwrap();
|
||||
}
|
||||
ix.finalize();
|
||||
ix.search(&queries[0], k).unwrap()
|
||||
};
|
||||
let a = build(7);
|
||||
let b = build(7);
|
||||
let same = a
|
||||
.iter()
|
||||
.zip(&b)
|
||||
.all(|(x, y)| x.id == y.id && x.score == y.score);
|
||||
println!(
|
||||
" seed 7 == seed 7 : {}",
|
||||
if same { "PASS" } else { "FAIL" }
|
||||
);
|
||||
|
||||
// [3] IdMap: O(1) delete hidden from search + filtered search correctness.
|
||||
println!("\n[3] IdMapIndex: delete + allowlist-filtered search");
|
||||
let ids: Vec<u64> = (0..n as u64).map(|i| i * 3 + 1).collect();
|
||||
let mut idx = IdMapIndex::with_seed(dim, BitWidth::Four, 42).unwrap();
|
||||
idx.add_with_ids(&data, &ids).unwrap();
|
||||
idx.finalize();
|
||||
let before = idx.search(&queries[1], 5).unwrap();
|
||||
let victim = before[0].0;
|
||||
idx.remove(victim);
|
||||
let after = idx.search(&queries[1], 5).unwrap();
|
||||
let hidden = after.iter().all(|(id, _)| *id != victim);
|
||||
println!(
|
||||
" removed id {victim}: live={}, hidden-from-search: {}",
|
||||
idx.len(),
|
||||
if hidden { "PASS" } else { "FAIL" }
|
||||
);
|
||||
let allow: Vec<u64> = ids.iter().copied().step_by(7).collect();
|
||||
let allow_set: std::collections::HashSet<u64> = allow.iter().copied().collect();
|
||||
let filtered = idx.search_filtered(&queries[1], 10, &allow).unwrap();
|
||||
let ok = filtered.iter().all(|(id, _)| allow_set.contains(id)) && !filtered.is_empty();
|
||||
println!(
|
||||
" filtered search ({} allowed): {} hits, all in allowlist: {}",
|
||||
allow.len(),
|
||||
filtered.len(),
|
||||
if ok { "PASS" } else { "FAIL" }
|
||||
);
|
||||
|
||||
println!("\nAll proof checks complete.\n");
|
||||
}
|
||||
260
crates/ruvector-turbovec/src/quantize.rs
Normal file
260
crates/ruvector-turbovec/src/quantize.rs
Normal file
|
|
@ -0,0 +1,260 @@
|
|||
//! Lloyd–Max scalar quantization for the canonical unit-Gaussian marginal,
|
||||
//! plus tight bit-packing.
|
||||
//!
|
||||
//! After the randomized Hadamard rotation (reused from `ruvector-rabitq`) and
|
||||
//! TQ+ per-coordinate calibration (`crate::calibrate`), each coordinate is
|
||||
//! approximately `N(0, 1)`. We therefore quantize against **codebook-free**,
|
||||
//! data-independent Lloyd–Max centroids for the standard normal — no k-means,
|
||||
//! no training pass, online ingest (ADR-254 §T2).
|
||||
//!
|
||||
//! Quantization is nearest-centroid, which is exactly Lloyd–Max decision
|
||||
//! regions (boundaries are the midpoints between adjacent centroids). Storing
|
||||
//! only the centroid table keeps the encoder a few `f32` comparisons per
|
||||
//! coordinate.
|
||||
|
||||
/// Supported quantization widths. `One` is included as a correctness/recall
|
||||
/// baseline against `ruvector-rabitq`'s 1-bit path; `Two`/`Three`/`Four` are
|
||||
/// the production targets of ADR-254 (`Three` fills the 2↔4-bit recall gap).
|
||||
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
|
||||
pub enum BitWidth {
|
||||
/// 1 bit / coord (2 levels).
|
||||
One,
|
||||
/// 2 bits / coord (4 levels).
|
||||
Two,
|
||||
/// 3 bits / coord (8 levels). Fills the recall gap between 2- and 4-bit;
|
||||
/// near the paper's ~2.5–3.5 bpc quality-neutral sweet spot (ADR-254 §D3).
|
||||
Three,
|
||||
/// 4 bits / coord (16 levels).
|
||||
Four,
|
||||
}
|
||||
|
||||
impl BitWidth {
|
||||
/// Bits stored per coordinate.
|
||||
#[inline]
|
||||
pub const fn bits(self) -> usize {
|
||||
match self {
|
||||
BitWidth::One => 1,
|
||||
BitWidth::Two => 2,
|
||||
BitWidth::Three => 3,
|
||||
BitWidth::Four => 4,
|
||||
}
|
||||
}
|
||||
|
||||
/// Number of reconstruction levels (`2^bits`).
|
||||
#[inline]
|
||||
pub const fn levels(self) -> usize {
|
||||
1usize << self.bits()
|
||||
}
|
||||
|
||||
/// Lloyd–Max centroids for the standard normal `N(0, 1)`, ascending.
|
||||
/// (Max, *Quantizing for Minimum Distortion*, IRE 1960.)
|
||||
#[inline]
|
||||
pub fn centroids(self) -> &'static [f32] {
|
||||
match self {
|
||||
// ±sqrt(2/π)
|
||||
BitWidth::One => &[-0.797_884_6, 0.797_884_6],
|
||||
BitWidth::Two => &[-1.510_4, -0.452_8, 0.452_8, 1.510_4],
|
||||
// 8-level Max (1960) optimal N(0,1) reconstruction levels.
|
||||
BitWidth::Three => &[
|
||||
-2.152_0, -1.344_0, -0.756_0, -0.245_1, 0.245_1, 0.756_0, 1.344_0, 2.152_0,
|
||||
],
|
||||
BitWidth::Four => &[
|
||||
-2.732_6, -2.069_0, -1.618_0, -1.256_2, -0.942_4, -0.656_8, -0.388_1, -0.128_4,
|
||||
0.128_4, 0.388_1, 0.656_8, 0.942_4, 1.256_2, 1.618_0, 2.069_0, 2.732_6,
|
||||
],
|
||||
}
|
||||
}
|
||||
|
||||
/// Packed code bytes for `dim` coordinates: `ceil(dim * bits / 8)`.
|
||||
#[inline]
|
||||
pub const fn code_bytes(self, dim: usize) -> usize {
|
||||
(dim * self.bits()).div_ceil(8)
|
||||
}
|
||||
}
|
||||
|
||||
/// Quantize one calibrated coordinate to its nearest centroid index.
|
||||
///
|
||||
/// Linear scan over `levels` (≤ 16) — branch-predictable and trivially
|
||||
/// vectorizable; the FastScan SIMD kernel (ADR-254 §T5, future milestone)
|
||||
/// will replace this with a nibble LUT, but the *result* must stay identical.
|
||||
#[inline]
|
||||
pub fn quantize_coord(centroids: &[f32], z: f32) -> u8 {
|
||||
let mut best = 0usize;
|
||||
let mut best_d = f32::INFINITY;
|
||||
for (i, &c) in centroids.iter().enumerate() {
|
||||
let d = (z - c).abs();
|
||||
if d < best_d {
|
||||
best_d = d;
|
||||
best = i;
|
||||
}
|
||||
}
|
||||
best as u8
|
||||
}
|
||||
|
||||
/// Pack per-coordinate code indices (`0..levels`) tightly, MSB-first within
|
||||
/// each byte, into `code_bytes(dim)` bytes.
|
||||
///
|
||||
/// # Preconditions
|
||||
/// Every `code` must be `< bw.levels()`. Out-of-range codes have their high
|
||||
/// bits truncated to the bit width (checked in debug builds).
|
||||
pub fn pack(codes: &[u8], bw: BitWidth) -> Vec<u8> {
|
||||
debug_assert!(
|
||||
codes.iter().all(|&c| (c as usize) < bw.levels()),
|
||||
"pack: code out of range for {bw:?} ({} levels)",
|
||||
bw.levels()
|
||||
);
|
||||
let bits = bw.bits();
|
||||
let mut out = vec![0u8; bw.code_bytes(codes.len())];
|
||||
let mut bit_pos = 0usize;
|
||||
for &code in codes {
|
||||
for b in (0..bits).rev() {
|
||||
let bit = (code >> b) & 1;
|
||||
if bit != 0 {
|
||||
let byte = bit_pos / 8;
|
||||
let off = 7 - (bit_pos % 8);
|
||||
out[byte] |= 1 << off;
|
||||
}
|
||||
bit_pos += 1;
|
||||
}
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
/// Inverse of [`pack`]: recover `dim` code indices from packed bytes.
|
||||
///
|
||||
/// # Preconditions
|
||||
/// `packed.len()` must be `>= bw.code_bytes(dim)` (checked in debug builds);
|
||||
/// a shorter slice panics on out-of-bounds access in release builds.
|
||||
pub fn unpack(packed: &[u8], dim: usize, bw: BitWidth) -> Vec<u8> {
|
||||
debug_assert!(
|
||||
packed.len() >= bw.code_bytes(dim),
|
||||
"unpack: packed slice too short ({} < {})",
|
||||
packed.len(),
|
||||
bw.code_bytes(dim)
|
||||
);
|
||||
let bits = bw.bits();
|
||||
let mut out = vec![0u8; dim];
|
||||
let mut bit_pos = 0usize;
|
||||
for slot in out.iter_mut() {
|
||||
let mut code = 0u8;
|
||||
for _ in 0..bits {
|
||||
let byte = bit_pos / 8;
|
||||
let off = 7 - (bit_pos % 8);
|
||||
let bit = (packed[byte] >> off) & 1;
|
||||
code = (code << 1) | bit;
|
||||
bit_pos += 1;
|
||||
}
|
||||
*slot = code;
|
||||
}
|
||||
out
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn centroids_are_sorted_and_sized() {
|
||||
for bw in [
|
||||
BitWidth::One,
|
||||
BitWidth::Two,
|
||||
BitWidth::Three,
|
||||
BitWidth::Four,
|
||||
] {
|
||||
let c = bw.centroids();
|
||||
assert_eq!(c.len(), bw.levels());
|
||||
assert!(c.windows(2).all(|w| w[0] < w[1]), "{bw:?} not ascending");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn pack_unpack_roundtrip_all_widths() {
|
||||
for bw in [
|
||||
BitWidth::One,
|
||||
BitWidth::Two,
|
||||
BitWidth::Three,
|
||||
BitWidth::Four,
|
||||
] {
|
||||
let dim = 37; // deliberately not byte-aligned
|
||||
let codes: Vec<u8> = (0..dim).map(|i| (i % bw.levels()) as u8).collect();
|
||||
let packed = pack(&codes, bw);
|
||||
assert_eq!(packed.len(), bw.code_bytes(dim));
|
||||
let back = unpack(&packed, dim, bw);
|
||||
assert_eq!(codes, back, "roundtrip failed for {bw:?}");
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn quantize_picks_nearest_centroid() {
|
||||
let c = BitWidth::Four.centroids();
|
||||
// A value sitting exactly on a centroid must map to that centroid.
|
||||
for (i, &cv) in c.iter().enumerate() {
|
||||
assert_eq!(quantize_coord(c, cv) as usize, i);
|
||||
}
|
||||
// Far-positive saturates to the top bucket, far-negative to the bottom.
|
||||
assert_eq!(quantize_coord(c, 99.0) as usize, c.len() - 1);
|
||||
assert_eq!(quantize_coord(c, -99.0) as usize, 0);
|
||||
}
|
||||
|
||||
/// Standard normal via Box–Muller — keeps the test dependency-free
|
||||
/// (no `rand_distr`) while exercising the real `quantize_coord` path.
|
||||
fn std_normal(rng: &mut impl rand::Rng) -> f32 {
|
||||
let u1: f32 = rng.gen::<f32>().max(1e-9);
|
||||
let u2: f32 = rng.gen::<f32>();
|
||||
(-2.0 * u1.ln()).sqrt() * (std::f32::consts::TAU * u2).cos()
|
||||
}
|
||||
|
||||
/// D4 (ADR-254): the per-coordinate quantization MSE on the canonical
|
||||
/// `N(0,1)` marginal must stay under TurboQuant's distortion bound
|
||||
/// `D_mse ≤ (√3·π/2)·4^(−b)` (arXiv:2504.19874), and within a small
|
||||
/// margin of the *known* Max-1960 Lloyd–Max optimum. This is a
|
||||
/// theory-grounded oracle: it pins the centroid tables far more tightly
|
||||
/// than a recall threshold — corrupt a single level and it trips.
|
||||
#[test]
|
||||
fn quantizer_mse_within_paper_bound() {
|
||||
use rand::SeedableRng;
|
||||
let mut rng = rand::rngs::StdRng::seed_from_u64(20_260_530);
|
||||
let n = 400_000usize;
|
||||
|
||||
// Max (1960) optimal MSE for the unit-variance Gaussian, per bit-width.
|
||||
let optimal = |bw: BitWidth| -> f64 {
|
||||
match bw {
|
||||
BitWidth::One => 0.363_4,
|
||||
BitWidth::Two => 0.117_5,
|
||||
BitWidth::Three => 0.034_5,
|
||||
BitWidth::Four => 0.009_5,
|
||||
}
|
||||
};
|
||||
|
||||
for bw in [
|
||||
BitWidth::One,
|
||||
BitWidth::Two,
|
||||
BitWidth::Three,
|
||||
BitWidth::Four,
|
||||
] {
|
||||
let c = bw.centroids();
|
||||
let mut sse = 0.0f64;
|
||||
for _ in 0..n {
|
||||
let z = std_normal(&mut rng);
|
||||
let code = quantize_coord(c, z) as usize;
|
||||
let d = (z - c[code]) as f64;
|
||||
sse += d * d;
|
||||
}
|
||||
let mse = sse / n as f64;
|
||||
|
||||
let bound = (3f64.sqrt() * std::f64::consts::PI / 2.0) * 4f64.powi(-(bw.bits() as i32));
|
||||
assert!(
|
||||
mse <= bound,
|
||||
"{bw:?}: MSE {mse:.4} exceeds paper bound {bound:.4}"
|
||||
);
|
||||
// Tightness: must be within 5% of the Lloyd–Max optimum (sampling
|
||||
// noise at n=400k is far below this), catching table corruption
|
||||
// that might still slip under the loose paper bound.
|
||||
assert!(
|
||||
mse <= optimal(bw) * 1.05,
|
||||
"{bw:?}: MSE {mse:.4} not near optimal {:.4}",
|
||||
optimal(bw)
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -131,7 +131,7 @@ FAISS's two-phase API and to allow future re-clustering.
|
|||
### Neutral
|
||||
|
||||
- **IVF-PQ not yet implemented** — this ADR covers the flat (exact reranking)
|
||||
variants only. PQ integration is the natural next step (ADR-194 TBD).
|
||||
variants only. PQ integration is the natural next step (ADR-254 TBD).
|
||||
- **No SIMD distance kernels** — the list scan is pure scalar f32. AVX2/NEON
|
||||
acceleration would give 4-8× throughput improvement but is orthogonal to the
|
||||
RAIRS algorithm.
|
||||
|
|
@ -170,7 +170,7 @@ for equal-memory dual assignment because the RAIR metric is directionally aware.
|
|||
Starting with compressed residuals would be more memory-efficient for large N.
|
||||
Rejected for this PR because PQ codebook training introduces a second k-means
|
||||
loop and an asymmetric distance table; cleaner to land flat IVF first and add
|
||||
PQ as a composable layer. Tracking as ADR-194 future work.
|
||||
PQ as a composable layer. Tracking as ADR-254 future work.
|
||||
|
||||
### 4. IVF-HNSW (HNSW routing over centroids)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,197 +1,355 @@
|
|||
---
|
||||
adr: 254
|
||||
title: "ruvector-turbovec — Multi-bit Scalar-Quantized FastScan ANN Index (2/3/4-bit SQ + TQ+ calibration + nibble-LUT SIMD)"
|
||||
title: "ruvector-turbovec — Multi-bit TurboQuant FastScan ANN Index (2/3/4-bit SQ + TQ+ calibration + nibble-LUT SIMD)"
|
||||
status: accepted
|
||||
date: 2026-06-16
|
||||
date: 2026-05-29
|
||||
authors: [oshaal, claude-flow]
|
||||
related: [ADR-157, ADR-193, ADR-155]
|
||||
supersedes: []
|
||||
related: [ADR-155, ADR-157, ADR-193]
|
||||
tags: [quantization, ann, vector-search, turboquant, fastscan, simd, lloyd-max, calibration, recall, memory]
|
||||
---
|
||||
|
||||
# ADR-254 — ruvector-turbovec: a multi-bit scalar-quantized FastScan ANN index
|
||||
# ADR-254 — ruvector-turbovec: a multi-bit TurboQuant FastScan ANN index
|
||||
|
||||
> **Numbering note.** This decision was originally drafted in PR #521 as
|
||||
> `ADR-194`, which collides with the already-merged `ADR-194` (ruvector ONNX
|
||||
> embedder API & throughput). This file is the canonical record at the next free
|
||||
> number; PR #521's `ADR-194-ruvector-turbovec-fastscan-index.md` should be
|
||||
> renumbered to **254** (or dropped in favor of this file) before merge.
|
||||
> embedder API & throughput). It has been renumbered to the next free number,
|
||||
> **254**, to match the canonical record on `main`
|
||||
> (`ADR-254-ruvector-turbovec-fastscan-index.md`). This file is the PR's fuller
|
||||
> version — it adds the D1–D5 divergences table and the D3/D4 measured-milestone
|
||||
> markers shipped in this PR — and supersedes the condensed on-`main` draft at
|
||||
> merge.
|
||||
|
||||
> **Provenance & prior-art note.** This adapts techniques from
|
||||
> **⚠️ Provenance & prior-art note.** This ADR adapts techniques from
|
||||
> [RyanCodrai/turbovec](https://github.com/RyanCodrai/turbovec), an independent
|
||||
> Rust+Python implementation of Google Research's **TurboQuant**
|
||||
> (arXiv:2504.19874). It is **not** a port — it is a clean-room ruvector crate
|
||||
> that reuses existing primitives. The recall / compression / bias figures under
|
||||
> "Validation" are **measured** and reproducible; the *competitive* claims vs
|
||||
> FAISS/Milvus are **targets to validate**, attributed to the upstream project
|
||||
> where cited.
|
||||
> (arXiv:2504.19874). It is **not** a port of that codebase — it is a clean-room
|
||||
> ruvector crate that reuses our existing primitives. The TurboQuant *algorithm*
|
||||
> is already partially present in this repo (see §"What already exists"); the
|
||||
> contribution here is a **multi-bit scalar-quantized ANN search index** with a
|
||||
> FastScan SIMD kernel, which we do **not** currently have. The *recall /
|
||||
> compression / bias* figures in "Validation" are **measured** (reproducible via
|
||||
> the demo); the *competitive* claims vs FAISS/Milvus remain **targets to be
|
||||
> validated** and are attributed to the upstream reference project where cited.
|
||||
|
||||
## Status
|
||||
|
||||
**Accepted (M1 implemented).** The scalar-reference milestone (M1) ships as
|
||||
`crates/ruvector-turbovec` in PR #521: rotation reuse + Lloyd–Max 2/3/4-bit
|
||||
scalar quantization + TQ+ per-coordinate calibration + length-renormalized
|
||||
unbiased scoring + `IdMapIndex` (O(1) delete, filtered search). Build green; 17
|
||||
unit tests + 1 doc-test pass; clippy clean. M2–M4 (FastScan SIMD kernel,
|
||||
AVX-512, dispatcher registration, persistence) are future work.
|
||||
**Accepted (M1 implemented).** The scalar reference milestone (M1) is
|
||||
implemented as `crates/ruvector-turbovec`: rotation reuse + Lloyd–Max 2/3/4-bit SQ
|
||||
+ TQ+ calibration + empirically bias-corrected scoring + `IdMapIndex`
|
||||
(O(1) delete, filtered search). Build is green
|
||||
(`cargo build --release -p ruvector-turbovec`); 21 unit tests + 1 doc-test pass;
|
||||
clippy clean. M2–M4 (FastScan SIMD kernel, AVX-512, dispatcher registration)
|
||||
remain future work. Measured proof below.
|
||||
|
||||
## Context
|
||||
### Validation (measured — `cargo run --release -p ruvector-turbovec`)
|
||||
|
||||
### The gap
|
||||
|
||||
ruvector has strong ANN coverage but is missing the **2–4-bit scalar-quantized
|
||||
FastScan regime** that mainstream production vector DBs live in (FAISS
|
||||
`IndexPQFastScan`, Qdrant, Milvus IVF-SQ):
|
||||
|
||||
| Crate | Index family | Quantization | SIMD scan |
|
||||
|---|---|---|---|
|
||||
| `ruvector-core` | HNSW (graph) | none (f32) | — |
|
||||
| `ruvector-diskann` | DiskANN (graph) | none / PQ-ish | — |
|
||||
| `ruvector-rairs` | IVF (ADR-193) | optional | — |
|
||||
| `ruvector-rabitq` | Flat + rotation | **1-bit only** | AVX2/512 popcount |
|
||||
| `ruvllm` `turbo_quant.rs` | value codec | 2.5–4 bit | n/a (not an index) |
|
||||
|
||||
- **`ruvector-rabitq`** is **1-bit**: excellent when memory dominates, but on
|
||||
1536–3072-dim production embeddings (OpenAI `text-embedding-3`, Cohere, Voyage)
|
||||
hitting production R@1 needs an exact-rerank pass over the raw f32 originals,
|
||||
which re-inflates memory back toward the f32 footprint — partly defeating the
|
||||
point.
|
||||
- **`ruvllm/quantize/turbo_quant.rs`** is a TurboQuant *value codec* for KV-cache
|
||||
/ embedding compression — no inverted lists, top-k heap, FastScan LUT,
|
||||
filtered search, or stable IDs/deletion. **Wrong abstraction for search.**
|
||||
|
||||
What's missing, concretely: **2–4 bits/dim**, **codebook-free** scalar
|
||||
quantization (no k-means, online ingest), a **FastScan nibble-LUT SIMD kernel**
|
||||
(16/32 candidates per instruction, no f32 materialization), and competitive
|
||||
recall **without** a mandatory f32 rerank.
|
||||
|
||||
### What already exists (not duplication)
|
||||
|
||||
1. **`ruvector-rabitq`** (ADR-157) already holds the *rotation* half of
|
||||
TurboQuant (`RandomRotation::HadamardSigned`, the randomized Hadamard,
|
||||
`O(D log D)`) and defines the `AnnIndex` + `VectorKernel`/`KernelCaps` traits
|
||||
reused here — but its codes are 1-bit and its kernels are XNOR-popcount.
|
||||
2. **`ruvllm/quantize/turbo_quant.rs`** has MSE-quantizer math to borrow, but is
|
||||
a data-oblivious tensor compressor, not an `AnnIndex`.
|
||||
|
||||
So this work reuses rabitq's rotation + traits and ruvllm's codec lessons to
|
||||
build the missing multi-bit FastScan **search index**.
|
||||
|
||||
## Decision
|
||||
|
||||
Introduce **`crates/ruvector-turbovec`**, a multi-bit ANN index that implements
|
||||
the existing `ruvector_rabitq::AnnIndex` trait and exposes a FastScan SIMD kernel
|
||||
via the existing `VectorKernel`/`KernelCaps` contract, so it drops into the
|
||||
`ruvector-rulake` dispatcher (ADR-155/157) with **no new plumbing**. Six
|
||||
techniques (T1–T6):
|
||||
|
||||
- **T1 — Normalize + randomized Hadamard rotation (reuse).** Strip each vector's
|
||||
L2 norm (store one f32), apply `RandomRotation::HadamardSigned` from rabitq.
|
||||
Post-rotation coordinates are ~N(0, 1/d), making per-coordinate scalar
|
||||
quantization optimal **without a codebook**. Import, don't reimplement.
|
||||
- **T2 — Lloyd–Max scalar quantization (2/3/4-bit).** MSE-optimal bucket
|
||||
boundaries for the canonical N(0,1) marginal (4/8/16 buckets). Boundaries are
|
||||
constants of the distribution, not the data → **zero training**.
|
||||
- **T3 — Per-coordinate calibration (TQ+).** On the first `add()` batch, fit
|
||||
`(shift[d], scale[d])` mapping empirical quantiles onto the canonical Beta;
|
||||
freeze after warm-up. Corrects finite-d Hadamard non-Gaussianity — the "+" that
|
||||
buys the recall edge. No counterpart exists today.
|
||||
- **T4 — Length-renormalized inner-product scoring.** Store a per-vector
|
||||
`c_x = ⟨r,r̂⟩/⟨r̂,r̂⟩` so the dot-product estimator removes scalar
|
||||
quantization's downward bias at **zero query-time cost** — skipping the f32
|
||||
rerank 1-bit RaBitQ needs. **This is an empirically near-unbiased heuristic,
|
||||
not the paper's provably-unbiased two-stage QJL-residual estimator** (see
|
||||
Divergences); adopting QJL is a tracked follow-up (M5) if measured bias demands
|
||||
it.
|
||||
- **T5 — FastScan nibble-LUT SIMD kernel (core perf win).** 32-vector SoA blocks;
|
||||
score a block by nibble-split table lookups (`vpshufb`/`tbl`) not arithmetic.
|
||||
x86 AVX-512BW with AVX2 fallback (`x86-64-v3`); ARM NEON `vqtbl1q_u8`; WASM
|
||||
scalar fallback (bit-identical). rabitq's popcount kernels are **not** reusable
|
||||
(popcount ≠ table-lookup), but the trait/dispatch/determinism contract is.
|
||||
- **T6 — Block-granularity filtered search + stable IDs.** Allowlist tested at
|
||||
32-vector block granularity inside the kernel (fully-excluded blocks
|
||||
short-circuit); `IdMapIndex` with external `u64` IDs surviving deletion and
|
||||
**O(1) remove** (tombstone + free-list).
|
||||
|
||||
### Reuse boundary
|
||||
|
||||
| Component | Source | Action |
|
||||
|---|---|---|
|
||||
| Randomized Hadamard rotation | `ruvector-rabitq::RandomRotation` | import |
|
||||
| `AnnIndex` trait | `ruvector-rabitq::index` | implement |
|
||||
| `VectorKernel` / `KernelCaps` | `ruvector-rabitq::kernel` | implement |
|
||||
| MSE/Lloyd–Max quantizer math | `ruvllm::quantize` | borrow/extract |
|
||||
| Lloyd–Max boundary tables (2/3/4-bit) | TurboQuant constants | build |
|
||||
| TQ+ per-coordinate calibration | — | build |
|
||||
| FastScan nibble-LUT SIMD kernel | — | build |
|
||||
| 32-block SoA layout + filtered scan | — | build |
|
||||
| IdMap O(1)-delete | — | build |
|
||||
| Persistence (`.tv` file) | mirror `rabitq::persist` | build |
|
||||
|
||||
If the borrowed MSE math proves reusable across both crates, a follow-up should
|
||||
hoist it into `ruvector-math` rather than copy it.
|
||||
|
||||
### Milestones
|
||||
|
||||
1. **M1 — scalar reference** (PR #521): rotation reuse + Lloyd–Max SQ + TQ+ +
|
||||
unbiased scoring + `IdMapIndex`. ✅ implemented & validated. Serves as the
|
||||
determinism oracle the SIMD kernels must match bit-for-bit.
|
||||
2. **M2 — FastScan nibble-LUT SIMD kernel** (AVX2 + NEON), fuzzed bit-identical
|
||||
to the scalar oracle.
|
||||
3. **M3 — `.tv` persistence.**
|
||||
4. **M4 — AVX-512BW kernel + `ruvector-rulake` dispatcher registration.**
|
||||
5. **M5 (optional) — QJL-residual provably-unbiased estimator**, if measured
|
||||
bias/recall on real embeddings demands the paper-grade two-stage scorer.
|
||||
|
||||
## Validation (measured — `cargo run --release -p ruvector-turbovec`)
|
||||
|
||||
`n = 5,000` **uniform-random** vectors (worst case — no cluster structure),
|
||||
`dim = 256`, `k = 10`, **no f32 rerank**, vs exact brute-force L2:
|
||||
`n = 5,000` **uniform-random** vectors (the *worst case* for ANN — no cluster
|
||||
structure to exploit), `dim = 256`, `k = 10`, **no f32 rerank**, scored against
|
||||
exact brute-force L2:
|
||||
|
||||
| Width | recall@10 | bytes/vec (raw 1024) | compression | mean cosine bias |
|
||||
|---|---|---|---|---|
|
||||
|-------|-----------|----------------------|-------------|------------------|
|
||||
| 1-bit | 0.308 | 48 | 25.6× | +0.0005 |
|
||||
| 2-bit | 0.561 | 80 | 14.2× | +0.0001 |
|
||||
| 3-bit | 0.767 | 112 | 9.8× | −0.0000 |
|
||||
| **4-bit** | **0.879** | **144** | **7.5×** | **−0.0000** |
|
||||
|
||||
Recall rises monotonically with bit-width — the 2–4-bit regime 1-bit RaBitQ
|
||||
can't reach without re-inflating memory. Mean cosine bias ≈ 0 at every width
|
||||
confirms the `c_x` renormalization is empirically near-unbiased on this data. On
|
||||
real clustered embeddings, recall at a given width is materially higher than this
|
||||
uniform stress test. Determinism (same seed → bit-identical) and `IdMapIndex`
|
||||
delete + allowlist search verified PASS.
|
||||
- **Recall rises monotonically with bit-width** — exactly the 2–4-bit regime the
|
||||
1-bit RaBitQ path cannot reach without re-inflating memory via f32 rerank.
|
||||
- **Mean cosine bias ≈ 0 at every width** — the per-vector `c_x`
|
||||
length-renormalization (§T4) is *empirically* near-unbiased on this data. Note
|
||||
this is a cheaper heuristic than the paper's *provably* unbiased two-stage
|
||||
estimator (MSE quantizer + 1-bit QJL on the residual); see
|
||||
"Divergences from the TurboQuant paper" below.
|
||||
- On real clustered embeddings (OpenAI/Cohere) recall at a given width is
|
||||
materially higher than on this uniform stress test.
|
||||
- Determinism (same seed → bit-identical results) and `IdMapIndex` delete +
|
||||
allowlist-filtered search both verified PASS.
|
||||
|
||||
## Context
|
||||
|
||||
### The gap
|
||||
|
||||
ruvector has strong coverage of approximate-nearest-neighbor (ANN) search:
|
||||
|
||||
| Crate | Index family | Quantization | SIMD scan |
|
||||
|-------|--------------|--------------|-----------|
|
||||
| `ruvector-core` | HNSW (graph) | none (f32) | — |
|
||||
| `ruvector-diskann` | DiskANN (graph) | none / PQ-ish | — |
|
||||
| `ruvector-rairs` | IVF (ADR-193) | optional | — |
|
||||
| `ruvector-rabitq` | Flat + rotation | **1-bit** only | AVX2/AVX-512 **popcount** |
|
||||
|
||||
The **1-bit** RaBitQ path (`ruvector-rabitq`) is excellent when memory dominates,
|
||||
but 1-bit codes cap recall: to hit production R@1 on 1536–3072-dim embeddings
|
||||
(OpenAI `text-embedding-3`, Cohere, Voyage) you need an exact-rerank pass over
|
||||
raw f32 originals, which re-inflates memory back toward the f32 footprint.
|
||||
|
||||
What we are **missing** is the regime that production vector DBs (FAISS
|
||||
`IndexPQFastScan`, Qdrant, Milvus IVF-SQ) live in:
|
||||
|
||||
- **2–4 bits per dimension** (not 1, not 32),
|
||||
- **codebook-free** scalar quantization (no k-means training, online ingest),
|
||||
- a **FastScan-style nibble-LUT SIMD kernel** that scores 16/32 candidates per
|
||||
vector instruction without ever materializing f32,
|
||||
- competitive recall **without** a mandatory f32 rerank pass.
|
||||
|
||||
The upstream reference project ([RyanCodrai/turbovec]) *reports* this regime
|
||||
reaching ~16× code-only compression at 2-bit (6,144 → 384 bytes for d=1536),
|
||||
R@1 competitive with or ahead of FAISS `IndexPQ` at 4-bit, and FastScan-class
|
||||
scan throughput on ARM —
|
||||
all with online ingest and no training phase. **Those are the external project's
|
||||
numbers, not this crate's.** This crate's own *measured* results are the
|
||||
uniform-random worst-case table under "Validation" above (recall@10 of
|
||||
0.308 / 0.561 / 0.879 at 1/2/3/4-bit); broader competitive benchmarks are listed
|
||||
as targets-to-validate in "Acceptance criteria" and the SIMD-kernel milestones.
|
||||
|
||||
[RyanCodrai/turbovec]: https://github.com/RyanCodrai/turbovec
|
||||
|
||||
### What already exists (and why this is not duplication)
|
||||
|
||||
Two TurboQuant-adjacent pieces are already in the repo. Neither closes the gap:
|
||||
|
||||
1. **`ruvector-rabitq`** (ADR-157) — already contains the *rotation* half of
|
||||
TurboQuant: `RandomRotation::HadamardSigned` (the `D₁·H·D₂·H·D₃` randomized
|
||||
Hadamard, `O(D log D)`, cited to arXiv:2504.19874 §3.2 in
|
||||
`rotation.rs`). It also defines the two traits this ADR will reuse —
|
||||
`AnnIndex` (index API) and `VectorKernel`/`KernelCaps` (pluggable scan
|
||||
backend). But its codes are **1-bit**, its kernels are **XNOR-popcount**, and
|
||||
it has no scalar-quantization or per-coordinate calibration.
|
||||
|
||||
2. **`ruvllm/src/quantize/turbo_quant.rs`** (1,483 lines) — a full TurboQuant
|
||||
*value codec* for **transformer KV-cache and embedding compression**
|
||||
(PolarQuant + QJL residual, 2.5–4.0 effective bits). This is a
|
||||
**data-oblivious tensor compressor**, *not* an `AnnIndex`: it has no
|
||||
inverted lists, no top-k heap, no FastScan LUT kernel, no
|
||||
filtered/allowlist search, no stable IDs/deletion. It is the wrong tool for
|
||||
"search 10M vectors for the nearest 10."
|
||||
|
||||
So the work here is: **reuse rabitq's rotation + traits, reuse the lessons from
|
||||
ruvllm's codec, and build the missing multi-bit FastScan *search index*.**
|
||||
|
||||
## Decision
|
||||
|
||||
Introduce **`crates/ruvector-turbovec`**, a multi-bit TurboQuant ANN index that
|
||||
implements the existing `ruvector_rabitq::AnnIndex` trait. The scalar M1 in this
|
||||
PR is the correctness oracle; M2–M4 will expose a FastScan SIMD kernel through
|
||||
the existing `VectorKernel`/`KernelCaps` contract and register it with the
|
||||
`ruvector-rulake` dispatcher (ADR-155/157).
|
||||
|
||||
Six techniques are ported/adapted from turbovec/TurboQuant:
|
||||
|
||||
### T1 — Normalize + randomized Hadamard rotation *(reuse, don't rebuild)*
|
||||
Strip each vector's L2 norm (store as one f32) and apply
|
||||
`RandomRotation::HadamardSigned` from `ruvector-rabitq`. After rotation every
|
||||
coordinate is ~Beta-distributed → N(0, 1/d), making **per-coordinate scalar
|
||||
quantization optimal without a codebook**. We import this type rather than
|
||||
reimplement the FWHT.
|
||||
|
||||
### T2 — Lloyd–Max scalar quantization (2-bit / 3-bit / 4-bit)
|
||||
Precompute MSE-optimal bucket boundaries for the canonical N(0,1) marginal at
|
||||
`bit_width ∈ {2, 3, 4}` (4, 8, and 16 buckets). Coordinates become 2-bit (0–3),
|
||||
3-bit (0–7), or 4-bit (0–15) integers. Boundaries are **constants of the
|
||||
distribution**, not of the data → zero training. (ruvllm's codec already has the
|
||||
MSE-quantizer math to borrow from.)
|
||||
|
||||
### T3 — Per-coordinate calibration (TQ+)
|
||||
During the *first* `add()` batch, fit two scalars per coordinate
|
||||
(`shift[d]`, `scale[d]`) that map the empirical quantile range onto the
|
||||
canonical Beta. Freeze after warm-up; all later vectors reuse them. This
|
||||
corrects the residual non-Gaussianity of finite-d Hadamard rotation and is the
|
||||
"+" that buys turbovec its recall edge over plain SQ. **No counterpart exists in
|
||||
the repo today.**
|
||||
|
||||
### T4 — Length-renormalized inner-product scoring
|
||||
Store a per-vector correction scalar `c_x = ⟨r,r̂⟩/⟨r̂,r̂⟩` (the least-squares
|
||||
magnitude match) so the dot-product estimator removes the systematic downward
|
||||
bias of scalar quantization at **zero search-time cost** — the scalar folds into
|
||||
the final score multiply. This lets us skip the mandatory f32 rerank that 1-bit
|
||||
RaBitQ needs, keeping the memory win. **This is a heuristic, empirically
|
||||
near-unbiased correction — not the paper's provably-unbiased two-stage QJL
|
||||
residual estimator** (see "Divergences from the TurboQuant paper"). Adopting the
|
||||
QJL residual stage is a tracked follow-up if measured bias/recall demands it.
|
||||
|
||||
### T5 — FastScan nibble-LUT SIMD kernel (the core perf win)
|
||||
Lay out codes in **32-vector SoA blocks**. For a query, precompute a small
|
||||
per-sub-quantizer lookup table; score a whole block by **nibble-split table
|
||||
lookups** (`vpshufb`/`tbl`) instead of arithmetic:
|
||||
- **x86**: `AVX-512BW` (`VPSHUFB`/`VPERMI2B`) with **AVX2 fallback**
|
||||
(`_mm256_shuffle_epi8`), targeting `x86-64-v3` like rabitq does.
|
||||
- **ARM**: `NEON` `vqtbl1q_u8` nibble lookups.
|
||||
- **WASM**: scalar fallback (matches rabitq's wasm policy; bit-identical).
|
||||
The M2 implementation will implement `VectorKernel` and advertise
|
||||
`accelerator: "cpu-simd-fastscan"` via `KernelCaps`. rabitq's existing AVX2/512
|
||||
*popcount* kernels are **not** reusable here (popcount ≠ table-lookup), but the
|
||||
trait, dispatch, and determinism contract are.
|
||||
|
||||
### T6 — Block-granularity filtered search + stable IDs
|
||||
- **Filtered search**: an allowlist is tested at **32-vector block
|
||||
granularity** *inside* the kernel — fully-excluded blocks short-circuit before
|
||||
any LUT work; individual disallowed slots are dropped at heap-insert time.
|
||||
- **`IdMapIndex`**: external `u64` IDs that survive deletion, with **O(1)
|
||||
remove** (tombstone + free-list), mirroring turbovec's `IdMapIndex`. The base
|
||||
`TurboQuantIndex` uses dense sequential internal IDs for online ingest.
|
||||
|
||||
### Reuse boundary (what we build vs import)
|
||||
|
||||
| Component | Source | Action |
|
||||
|-----------|--------|--------|
|
||||
| Randomized Hadamard rotation | `ruvector-rabitq::RandomRotation` | **import** |
|
||||
| `AnnIndex` trait | `ruvector-rabitq::index` | **implement** |
|
||||
| `VectorKernel` / `KernelCaps` | `ruvector-rabitq::kernel` | **implement** |
|
||||
| MSE/Lloyd–Max quantizer math | `ruvllm::quantize` | **borrow/extract** |
|
||||
| Lloyd–Max boundary tables (2/3/4-bit) | TurboQuant constants | **build (new)** |
|
||||
| TQ+ per-coordinate calibration | — | **build (new)** |
|
||||
| FastScan nibble-LUT SIMD kernel | — | **build (new)** |
|
||||
| 32-block SoA layout + filtered scan | — | **build (new)** |
|
||||
| IdMap O(1)-delete | — | **build (new)** |
|
||||
| Persistence (`.tv` file) | mirror `rabitq::persist` | **build (new)** |
|
||||
|
||||
If the borrowed MSE math from `ruvllm` proves reusable across both crates, a
|
||||
follow-up should hoist it into `ruvector-math` rather than copy it (flagged as a
|
||||
consequence, not done here).
|
||||
|
||||
### Proposed API (Rust)
|
||||
|
||||
```rust
|
||||
use ruvector_turbovec::{TurboVecIndex, IdMapIndex, BitWidth};
|
||||
|
||||
let mut idx = TurboVecIndex::new(/*dim=*/1536, BitWidth::Four)?; // 4-bit
|
||||
idx.add(0, vector)?; // online ingest (AnnIndex)
|
||||
for (id, row) in rows.into_iter().enumerate() {
|
||||
idx.add(id + 1, row)?;
|
||||
}
|
||||
idx.finalize(); // freeze TQ+ calibration
|
||||
let hits = idx.search(&query, 10)?; // Vec<SearchResult{ id, score }>
|
||||
|
||||
// External IDs + deletion + filtered search
|
||||
let mut m = IdMapIndex::new(1536, BitWidth::Four)?;
|
||||
m.add_with_ids(&vectors, &ids /* &[u64] */)?;
|
||||
let removed = m.remove(1002); // O(1)
|
||||
let hits = m.search_filtered(&query, 10, &allowlist /* &[u64] */)?;
|
||||
```
|
||||
|
||||
## Consequences
|
||||
|
||||
### Positive
|
||||
- Closes the 2–4-bit FastScan gap — the one mainstream production ANN regime
|
||||
ruvector lacks; comparable regime to FAISS `IndexPQFastScan` / Milvus IVF-SQ.
|
||||
- ~16× compression with online ingest, no training (d=1536: 6 KB → 384 B).
|
||||
- Recall without mandatory f32 rerank (T4), so the memory win is real.
|
||||
- Zero new plumbing — implements existing `AnnIndex` + `VectorKernel`, registers
|
||||
with the `ruvector-rulake` dispatcher, inherits its determinism/witness
|
||||
contract. Exactly one new workspace crate.
|
||||
- Composition path: the same block-SoA codes can back an IVF posting list
|
||||
(IVF-SQ-FastScan) with `ruvector-rairs` (ADR-193).
|
||||
- **Establishes the scalar 2–4-bit search foundation** for the missing FastScan
|
||||
regime. SIMD parity with FAISS/Milvus remains an explicit M2–M4 target.
|
||||
- **Compressed online ingest without k-means/codebook training.** TQ+
|
||||
calibration still fits shift/scale values on the warm-up batch.
|
||||
- **Geometry-safe arbitrary dimensions.** Inputs are zero-padded and quantized
|
||||
in the full next-power-of-two Hadamard space. At d=1536 this means 2048 coded
|
||||
coordinates: about 11.6× at 2-bit and 5.9× at 4-bit including norm,
|
||||
correction scalar, and ID.
|
||||
- **Recall without mandatory f32 rerank** (via T4 length-renormalized scoring),
|
||||
so the memory win is real, unlike 1-bit-with-rerank.
|
||||
- **Reuses existing plumbing** — M1 implements `AnnIndex`; future SIMD work will
|
||||
use the existing `VectorKernel`/dispatcher determinism contract.
|
||||
- **Composable**: the same block-SoA codes can later back an IVF posting list
|
||||
(`ruvector-rairs` IVF-SQ-FastScan) — a natural ADR-193 follow-up.
|
||||
|
||||
### Negative
|
||||
- T4 is a heuristic, not the paper's provably-unbiased estimator; pathological
|
||||
distributions may need M5 (QJL residual). Tracked, not hidden.
|
||||
- Competitive-vs-FAISS/Milvus claims are targets pending M2+ benchmarks, not yet
|
||||
this crate's measured numbers.
|
||||
- Borrowed MSE math risks divergence from `ruvllm`'s copy until hoisted into
|
||||
`ruvector-math`.
|
||||
- New SIMD kernels (M2/M4) are correctness-sensitive; mitigated by fuzzing
|
||||
bit-identical against the M1 scalar oracle.
|
||||
### Negative / risks
|
||||
- **SIMD surface area**: hand-written AVX-512BW + AVX2 + NEON nibble kernels are
|
||||
the bulk of the bug/maintenance risk. Mitigation: a scalar reference kernel is
|
||||
the determinism oracle; every SIMD kernel is fuzzed bit-identical against it
|
||||
(matches rabitq's `deterministic` cap policy). `#[target_feature]` paths are
|
||||
`unsafe` — this crate breaks rabitq's "no `unsafe`" guarantee and must
|
||||
document that explicitly.
|
||||
- **TQ+ calibration is stateful**: the first batch's distribution freezes the
|
||||
calibration; pathological first batches degrade later recall. Mitigation:
|
||||
expose `recalibrate()` and a min-warmup-count; default to a generous warmup.
|
||||
- **Possible math duplication** with `ruvllm` if not hoisted to `ruvector-math`.
|
||||
- **Name collision risk**: "TurboQuant" now appears in three places
|
||||
(`ruvllm` codec, this index, rabitq's rotation comment). The crate is named
|
||||
`turbovec` (the *index*) to disambiguate from the `turbo_quant` *codec*.
|
||||
|
||||
### Neutral
|
||||
- One additional crate; no change to existing index crates or their APIs.
|
||||
- WASM gets the scalar path only (no SIMD), consistent with rabitq.
|
||||
- Adds one workspace crate; no changes to existing crates beyond making
|
||||
`RandomRotation`/trait exports `pub` if any aren't already.
|
||||
|
||||
## Links
|
||||
- Issue: #520 (the gap analysis & ask)
|
||||
- PR: #521 (M1 implementation + original ADR-194 draft to be renumbered)
|
||||
- Related: ADR-157 (rabitq rotation + traits), ADR-193 (rairs IVF), ADR-155
|
||||
(rulake dispatcher)
|
||||
- Upstream: [RyanCodrai/turbovec](https://github.com/RyanCodrai/turbovec),
|
||||
TurboQuant (arXiv:2504.19874)
|
||||
## Divergences from the TurboQuant paper (arXiv:2504.19874)
|
||||
|
||||
This crate adapts the *techniques* of TurboQuant via the
|
||||
[RyanCodrai/turbovec](https://github.com/RyanCodrai/turbovec) reference
|
||||
implementation, which itself simplifies the paper. M1 therefore deliberately
|
||||
diverges from the paper in several places. We log them here so the gaps are
|
||||
explicit and reviewable, and so the follow-up milestones are paper-grounded
|
||||
rather than ad hoc. None of these are bugs in M1 — they are scope boundaries.
|
||||
|
||||
| # | Paper covers | M1 does instead | Plan |
|
||||
|---|--------------|-----------------|------|
|
||||
| D1 | **Provably-unbiased** inner product via a **two-stage** estimator: MSE quantizer + **1-bit QJL on the residual** `r = x − x̂_mse`, score `⟨y, x̂_mse + x̂_qjl⟩`, unbiased by construction with a variance bound. | A single per-vector scalar `c_x = ⟨r,r̂⟩/⟨r̂,r̂⟩` (least-squares magnitude match). *Empirically* near-unbiased (mean cos-bias ≈ 0 on uniform data); **no theoretical guarantee**. Cheaper (no extra residual bits). | **M5 (new):** add the optional QJL-residual stage as a recall/accuracy upgrade path when `c_x` proves insufficient on clustered data. |
|
||||
| D2 | Per-coordinate quantizer is **Max-Lloyd-optimal for the exact Beta marginal** `f(x) ∝ (1−x²)^((d−3)/2)`, with tables precomputed **per (bit-width, dimension)**. | Hardcoded Lloyd–Max tables for the **N(0,1) limit** of that Beta + an empirical per-coordinate `shift/scale` (TQ+) patch. Exact only as `d → ∞`; approximate at low/medium `d`. (TQ+ itself is *not* in the paper.) | **M6 (new):** generate d-aware Beta-optimal codebooks offline; keep the N(0,1)+calibration path as the default fast option. |
|
||||
| D3 | Highlights **~2.5 and ~3.5 bits/channel** as the quality-neutral operating points. | ✅ **Now ships 1 / 2 / 3 / 4-bit.** The added 3-bit width fills the old 2↔4-bit cliff: recall@10 **0.767** at **9.8×** compression (112 B/vec), measured. | Done in M1. Non-integer effective bit-widths (2.5/3.5 bpc) remain future work, achievable via D1's QJL residual or mixed-width coding. |
|
||||
| D4 | Closed-form distortion bounds: `D_mse ≤ (√3·π/2)·4^(−b)` (≈2.7× the info-theoretic floor) and `D_prod ≤ (√3·π²·‖y‖²/d)·4^(−b)`. | ✅ **Done.** `quantizer_mse_within_paper_bound` measures per-coordinate MSE on the `N(0,1)` marginal for every width and asserts it stays under the `D_mse` bound *and* within 5% of the Max-1960 optimum — a theory-grounded oracle that pins the centroid tables. | `D_prod` (full-pipeline IP) follow-up; deferred with D5. |
|
||||
| D5 | Bounds estimator **variance** (useful for ranking confidence / early termination). | Not surfaced. | Defer; revisit if IVF/rerank composition (ADR-193) needs confidence intervals. |
|
||||
|
||||
**Not divergences (M1 already matches the paper):** L2-norm stored in f32 and
|
||||
used to recover Euclidean distance; data-oblivious / zero-training online ingest;
|
||||
full-precision (rotated, un-quantized) query against quantized database vectors,
|
||||
consistent with the paper's `‖y‖²`-in-the-bound MIPS analysis. The FastScan SIMD
|
||||
kernel (M2–M4) is a FAISS-lineage engineering layer, *not* part of the paper.
|
||||
|
||||
## Alternatives considered
|
||||
|
||||
1. **Extend `ruvector-rabitq` from 1-bit to multi-bit in place.** Rejected:
|
||||
conflates two estimator families (XNOR-popcount vs nibble-LUT FastScan),
|
||||
would force `unsafe` into a crate that advertises none, and bloats its API.
|
||||
A sibling crate sharing traits is cleaner.
|
||||
2. **Use `ruvllm`'s `turbo_quant.rs` codec directly for search.** Rejected: it's
|
||||
a value/KV-cache compressor with no index, no top-k, no LUT kernel, no
|
||||
IDs/deletion — wrong abstraction.
|
||||
3. **Add IVF-SQ to `ruvector-rairs` instead.** Deferred, not rejected: IVF-SQ is
|
||||
the natural *composition* of this crate's codes with rairs' posting lists.
|
||||
Build the flat FastScan index first (this ADR), then layer IVF (follow-up).
|
||||
4. **Wrap FAISS via FFI.** Rejected: violates ruvector's no-C/C++-deps,
|
||||
pure-Rust, WASM-portable posture.
|
||||
|
||||
## Rollout & acceptance criteria
|
||||
|
||||
Implement on branch `claude/ruvector-turbovec-optimization-FhaDh` (this ADR),
|
||||
crate work in a follow-up PR. Milestones:
|
||||
|
||||
1. **M1 — Scalar reference (no SIMD).** Rotation reuse + Lloyd–Max 2/3/4-bit +
|
||||
TQ+ + length-renormalized scoring + `AnnIndex`. Uniform-random recall,
|
||||
distortion, determinism, non-finite-input, and d=1536 geometry tests.
|
||||
SIFT1M/production-embedding validation remains future work. ✅ *done.*
|
||||
2. **M2 — FastScan SIMD kernel.** AVX2 + NEON nibble-LUT, fuzzed bit-identical
|
||||
to M1's scalar scorer; `VectorKernel` impl; criterion bench in
|
||||
`benches/turbovec_bench.rs`. (3-bit width already shipped in M1, see D3.)
|
||||
3. **M3 — Persistence and block filtering.** M1 already includes `IdMapIndex`,
|
||||
O(1) tombstone delete, and reference allowlist search. M3 adds block-level
|
||||
short-circuiting and `.tv` save/load round-trip tests.
|
||||
4. **M4 — AVX-512BW kernel + rulake dispatcher registration.**
|
||||
5. **M5 — Paper-grade unbiased estimator (D1).** Optional two-stage MSE +
|
||||
1-bit-QJL-residual scoring as an accuracy upgrade over the `c_x` heuristic,
|
||||
gated behind a flag; validate the bias/recall delta on clustered data.
|
||||
6. **M6 — d-aware Beta-optimal codebooks (D2).** Offline-generated Max-Lloyd
|
||||
tables for the exact Beta marginal per (bit-width, dimension); N(0,1)+TQ+
|
||||
stays the default fast path.
|
||||
|
||||
**Acceptance (targets, to validate — not yet measured):**
|
||||
- ≥ **11.5× at 2-bit** and ≥ **5.8× at 4-bit** for d=1536, including padded
|
||||
codes, norm, correction scalar, and ID.
|
||||
- **R@1 within ±1 point of, or better than,** the 1-bit-RaBitQ-with-rerank path
|
||||
at equal or lower memory.
|
||||
- FastScan kernel **≥ 3× faster** than the scalar reference scorer on x86-64-v3.
|
||||
- **Bit-identical** scan output across scalar/AVX2/AVX-512/NEON (determinism
|
||||
oracle), enforced in CI.
|
||||
- `cargo build --release -p ruvector-turbovec` green; all unit + property tests
|
||||
pass; no `clippy` regressions.
|
||||
- **Measured MSE distortion within the paper's bound (D4):** ✅ enforced by
|
||||
`quantizer_mse_within_paper_bound` — per-coordinate MSE on `N(0,1)` under
|
||||
`D_mse ≤ (√3·π/2)·4^(−b)` and within 5% of the Max-1960 optimum, for every
|
||||
width. (`D_prod ≤ (√3·π²·‖y‖²/d)·4^(−b)` full-pipeline check still to come.)
|
||||
|
||||
## References
|
||||
|
||||
- TurboQuant (Google Research) — arXiv:2504.19874
|
||||
- PolarQuant — arXiv:2502.02617 · QJL — arXiv:2406.03482
|
||||
- RaBitQ (Gao & Long, SIGMOD 2024) — basis of `ruvector-rabitq`
|
||||
- FAISS FastScan / `IndexPQFastScan` — the nibble-LUT SIMD scan precedent
|
||||
- RyanCodrai/turbovec — https://github.com/RyanCodrai/turbovec
|
||||
- Prior art in-repo: `crates/ruvector-rabitq/src/rotation.rs` (Hadamard),
|
||||
`crates/ruvllm/src/quantize/turbo_quant.rs` (KV-cache codec), ADR-157, ADR-193
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue