From 1ec50dfa98eb0feee0dadbbf5a86c03ccaf55edb Mon Sep 17 00:00:00 2001 From: ajianaz Date: Fri, 28 Aug 2026 17:01:38 +0700 Subject: [PATCH 1/4] =?UTF-8?q?wip(core):=20residual=20quantization=20?= =?UTF-8?q?=E2=80=94=20encode=20+=20format=20v1.4=20+=20scoring=20(#23)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit WIP — continued tomorrow. Current state: DONE and green (54 unit tests, native 6/6 suites x2, x86_64 50/50, clippy -D warnings 0/0, fmt clean, cora review clean): - with_residual(dim, seed): second-pass Lloyd codes on the residual (rotated value minus code0 centroid) scaled by per-vector RMS - score = raw0*scale0 + raw1*scale2, both terms divided by the reconstructed norm sqrt(sum_sq + rms^2*padded); NEON4/AVX24 batch paths combine both planes with identical expression order - format v1.4 (260): residual f16 scales + codes blocks after the code0 block, before the keyed table; readers accept v1..v1.4, plain writers stay v1.3 - composition verified: keyed + cascade + working_dim all work on residual indexes; v1.4 round trip incl. key table - reconstruction MSE test: residual codes cut MSE 4.5x on the noisy clustered set (0.0147 -> 0.0033) OPEN (tomorrow): recall on the adversarial clustered set is slightly BELOW plain (0.58 vs 0.66 at k=10) despite the much better MSE. Candidates to investigate: estimator denominators (reconstructed vs true norm), f16 rounding of scales2, per-vector rms creating score- scale variance, and validating on the realistic-spread set before concluding. recall test currently asserts a sanity floor only. Signed-off-by: ajianaz --- crates/vecq-core/src/format.rs | 81 ++++++-- crates/vecq-core/src/store.rs | 349 +++++++++++++++++++++++++++++++-- 2 files changed, 399 insertions(+), 31 deletions(-) diff --git a/crates/vecq-core/src/format.rs b/crates/vecq-core/src/format.rs index 59ce3cb..52d6a28 100644 --- a/crates/vecq-core/src/format.rs +++ b/crates/vecq-core/src/format.rs @@ -20,10 +20,15 @@ //! [keyed_entries u32][entries: slot u32 + key u64 each, slots in order] //! ``` //! -//! Readers accept v1, v1.1, v1.2 and v1.3; writers emit 1.3. The seed is -//! stored in the header so the random sign diagonal can be regenerated -//! identically on any platform: identical file -> identical query results, -//! bit for bit. +//! Version 1.4 (stored as 260): a residual-mode index (second-pass codes, +//! issue #23). After the codes block, two extra blocks appear — f16 residual +//! scales (`count` entries) and residual codes (`count * padded/2` bytes) — +//! before the keyed-slot table. +//! +//! Readers accept v1 through v1.4; writers emit 1.3 (plain) or 1.4 +//! (residual). The seed is stored in the header so the random sign diagonal +//! can be regenerated identically on any platform: identical file -> +//! identical query results, bit for bit. use crate::store::VecqIndex; @@ -32,6 +37,7 @@ const V1: u16 = 1; const V1_1: u16 = 257; pub const V1_2: u16 = 258; const V1_3: u16 = 259; +const V1_4: u16 = 260; #[derive(Debug)] pub enum Error { @@ -133,13 +139,14 @@ fn f16_bits_to_f32(h: u16) -> f32 { } impl VecqIndex { - /// Serialize the index to bytes (format version 1.3, f16 scales). + /// Serialize the index to bytes. Plain indexes emit format version 1.3; + /// residual indexes emit 1.4 (extra residual scale + code blocks). /// /// Tombstoned slots are skipped: the output always holds the live vectors /// in slot order, so a round-trip through bytes has the same effect as /// [`VecqIndex::compact`] on disk without disturbing in-memory slot - /// indices. Keys of live keyed slots are stored in the v1.3 keyed-slot - /// table and are fully restored by [`VecqIndex::from_bytes`]. + /// indices. Keys of live keyed slots are stored in the keyed-slot table + /// and are fully restored by [`VecqIndex::from_bytes`]. pub fn to_bytes(&self) -> Vec { let bpv = self.padded_dim() / 2; let reserved: u16 = if self.working_dim() == self.dim() { @@ -147,9 +154,16 @@ impl VecqIndex { } else { self.working_dim() as u16 }; - let mut out = Vec::with_capacity(24 + self.live_slots() * bpv + self.live_slots() * 2); + let version = if self.is_residual() { V1_4 } else { V1_3 }; + let extra = if self.is_residual() { + self.live_slots() * (2 + bpv) + } else { + 0 + }; + let mut out = + Vec::with_capacity(24 + self.live_slots() * bpv + self.live_slots() * 2 + extra); out.extend_from_slice(&MAGIC.to_le_bytes()); - out.extend_from_slice(&V1_3.to_le_bytes()); + out.extend_from_slice(&version.to_le_bytes()); out.extend_from_slice(&reserved.to_le_bytes()); out.extend_from_slice(&(self.dim() as u32).to_le_bytes()); out.extend_from_slice(&self.seed().to_le_bytes()); @@ -166,6 +180,20 @@ impl VecqIndex { } out.extend_from_slice(self.slot_codes(slot, bpv)); } + if self.is_residual() { + for slot in 0..self.slots() { + if !self.slot_alive(slot) { + continue; + } + out.extend_from_slice(&f32_to_f16_bits(self.slot_scale2(slot)).to_le_bytes()); + } + for slot in 0..self.slots() { + if !self.slot_alive(slot) { + continue; + } + out.extend_from_slice(self.slot_codes2(slot, bpv)); + } + } // Keyed-slot table (v1.3): restores the keyed API across reloads. // Slot ids are DENSE positions among the serialized (alive) slots — // the reader's slot space — not the in-memory slot indices, which @@ -195,16 +223,17 @@ impl VecqIndex { return Err(Error::NotAStableFile); } let version = rd_u16(&bytes[4..6]); - if version != V1 && version != V1_1 && version != V1_2 && version != V1_3 { + if version != V1 && version != V1_1 && version != V1_2 && version != V1_3 && version != V1_4 + { return Err(Error::UnsupportedVersion(version)); } let dim = rd_u32(&bytes[8..12]) as usize; let seed = u64::from_le_bytes(bytes[12..20].try_into().unwrap()); let count = rd_u32(&bytes[20..24]) as usize; - // v1.2 / v1.3 carry working_dim in the reserved field (0 = full dim). + // v1.2+ carry working_dim in the reserved field (0 = full dim). let working_dim = match version { - V1_2 | V1_3 => match rd_u16(&bytes[6..8]) as usize { + V1_2 | V1_3 | V1_4 => match rd_u16(&bytes[6..8]) as usize { 0 => dim, w if w <= dim => w, w => { @@ -242,18 +271,38 @@ impl VecqIndex { index.scales = scales; index.n = count; index.init_dense(count); - if version == V1_3 { + if version == V1_4 { + // Residual blocks: f16 scales2, then codes2. + index.residual = true; + if bytes.len() < code_end + count * (2 + codes_bytes) { + return Err(Error::Truncated); + } + let mut off2 = code_end; + for _ in 0..count { + index + .scales2 + .push(f16_bits_to_f32(rd_u16(&bytes[off2..off2 + 2]))); + off2 += 2; + } + index.codes2 = bytes[off2..off2 + count * codes_bytes].to_vec(); + } + if version == V1_3 || version == V1_4 { // Keyed-slot table: [entries u32][slot u32 + key u64 each]. if bytes.len() < code_end + 4 { return Err(Error::Truncated); } - let entries = rd_u32(&bytes[code_end..code_end + 4]) as usize; - let table_end = code_end + 4 + entries * 12; + let key_base = if version == V1_4 { + code_end + count * (2 + codes_bytes) + } else { + code_end + }; + let entries = rd_u32(&bytes[key_base..key_base + 4]) as usize; + let table_end = key_base + 4 + entries * 12; if bytes.len() < table_end { return Err(Error::Truncated); } let mut key_table: Vec> = vec![None; count]; - let mut e = code_end + 4; + let mut e = key_base + 4; for _ in 0..entries { let slot = rd_u32(&bytes[e..e + 4]) as usize; let key = u64::from_le_bytes(bytes[e + 4..e + 12].try_into().unwrap()); diff --git a/crates/vecq-core/src/store.rs b/crates/vecq-core/src/store.rs index 7278727..7ac3ea1 100644 --- a/crates/vecq-core/src/store.rs +++ b/crates/vecq-core/src/store.rs @@ -46,6 +46,13 @@ pub struct VecqIndex { signature: Option>, alive: Vec, // slot -> not tombstoned live: usize, // number of non-tombstoned slots + // Residual quantization (issue #23): second-pass codes quantizing the + // residual (rotated value minus its code0 centroid) re-scaled by a + // per-vector RMS, plus that RMS as the correction scale. Empty in plain + // mode. + pub(crate) residual: bool, + pub(crate) codes2: Vec, // n * padded/2 residual nibbles + pub(crate) scales2: Vec, // per-vector residual RMS } impl VecqIndex { @@ -93,9 +100,30 @@ impl VecqIndex { alive: Vec::new(), live: 0, signature: None, + residual: false, + codes2: Vec::new(), + scales2: Vec::new(), } } + /// Create an empty index with second-pass residual codes (issue #23). + /// + /// Doubles the code storage (~2x padded/2 bytes + one extra f16 scale per + /// vector) in exchange for a finer distance estimate: the residual left + /// by the first Lloyd pass is itself Lloyd-quantized and added to the + /// score. Recall improves most on noise-dominated data; scan cost + /// roughly doubles. Composable with the keyed and cascade layers. + pub fn with_residual(dim: usize, seed: u64) -> Self { + let mut idx = Self::with_working_dim(dim, dim, seed); + idx.residual = true; + idx + } + + /// Whether this index carries second-pass residual codes. + pub fn is_residual(&self) -> bool { + self.residual + } + /// Number of live (searchable) vectors. pub fn len(&self) -> usize { self.live @@ -156,6 +184,14 @@ impl VecqIndex { &self.codes[slot * bpv..(slot + 1) * bpv] } + pub(crate) fn slot_scale2(&self, slot: usize) -> f32 { + self.scales2[slot] + } + + pub(crate) fn slot_codes2(&self, slot: usize, bpv: usize) -> &[u8] { + &self.codes2[slot * bpv..(slot + 1) * bpv] + } + /// Mark the index as holding `count` dense (all-live, keyless) slots; /// used after loading from the file format. pub(crate) fn init_dense(&mut self, count: usize) { @@ -208,8 +244,11 @@ impl VecqIndex { Some(slot) => { // Replace in place: codes change, cascade signatures go stale. self.signature = None; - let scale = self.encode_into(slot * (self.padded / 2), v); - self.scales[slot] = scale; + let (scale0, scale2) = self.encode_into(slot * (self.padded / 2), v); + self.scales[slot] = scale0; + if let Some(s2) = scale2 { + self.scales2[slot] = s2; + } slot } None => { @@ -552,8 +591,11 @@ impl VecqIndex { fn append_slot(&mut self, v: &[f32], key: Option) -> usize { self.signature = None; // codes change: cascade signatures go stale let slot = self.n; - let scale = self.encode_into(slot * (self.padded / 2), v); - self.scales.push(scale); + let (scale0, scale2) = self.encode_into(slot * (self.padded / 2), v); + self.scales.push(scale0); + if let Some(s2) = scale2 { + self.scales2.push(s2); + } self.keys.push(key); self.alive.push(true); self.n += 1; @@ -563,7 +605,7 @@ impl VecqIndex { /// Quantize `v` into the code bytes starting at `base` (extending /// `codes` when appending); returns the unit-norm correction scale. - fn encode_into(&mut self, base: usize, v: &[f32]) -> f32 { + fn encode_into(&mut self, base: usize, v: &[f32]) -> (f32, Option) { assert_eq!(v.len(), self.dim, "vector dim mismatch"); // Normalize over the working dims only (Matryoshka truncation happens // before rotation — see with_working_dim). @@ -584,6 +626,7 @@ impl VecqIndex { self.codes.resize(base + bytes_per_vec, 0); } let mut sum_sq = 0f32; + let mut residual_buf = Vec::new(); for (i, &x) in rotated.iter().enumerate() { let code = lloyd::quantize_4bit(x); let b = base + i / 2; @@ -593,12 +636,48 @@ impl VecqIndex { (self.codes[b] & 0x0F) | (code << 4) }; self.codes[b] = byte; - sum_sq += lloyd::dequantize_4bit(code).powi(2); + let d = lloyd::dequantize_4bit(code); + sum_sq += d.powi(2); + if self.residual { + if residual_buf.is_empty() { + residual_buf.resize(self.padded, 0.0); + } + residual_buf[i] = x - d; + } } // Scale so that the stored vector is unit-norm: dequantized vector q // has norm sqrt(sum_sq); asymmetric scoring multiplies by 1/sqrt(sum_sq). - 1.0 / sum_sq.sqrt() + let scale0 = 1.0 / sum_sq.sqrt(); + if !self.residual { + return (scale0, None); + } + + // Second pass (issue #23): the residual left by the first Lloyd pass + // is roughly Gaussian after scaling by its own RMS, so the same + // N(0,1) codebook applies. Scoring adds raw1 * rms to the estimate. + let rms: f32 = (residual_buf.iter().map(|x| x * x).sum::() / self.padded as f32) + .max(1e-12) + .sqrt(); + if self.codes2.len() < base + bytes_per_vec { + self.codes2.resize(base + bytes_per_vec, 0); + } + for (i, &r) in residual_buf.iter().enumerate() { + let code = lloyd::quantize_4bit(r / rms); + let b = base + i / 2; + let byte = if i % 2 == 0 { + (self.codes2[b] & 0xF0) | code + } else { + (self.codes2[b] & 0x0F) | (code << 4) + }; + self.codes2[b] = byte; + } + // Term coefficients: both terms estimate the cosine of the full + // reconstruction x̂ = x̂0 + rms·d1, so both divide by its norm + // sqrt(sum_sq + rms²·padded). Term1's numerator is rms·raw1 (raw1 = + // q·dequant(c1) estimates q·(r/rms)). + let denom = (sum_sq + rms * rms * self.padded as f32).sqrt(); + ((scale0 * sum_sq.sqrt()) / denom, Some(rms / denom)) } /// Prepare an f32 query in rotated space (call once per query). @@ -649,25 +728,38 @@ impl VecqIndex { let base = idx * (self.padded / 2); let codes = &self.codes[base..base + self.padded / 2]; let q = &pq.rotated[..self.padded]; + let raw0 = self.score_raw(codes, q, &pq.lut); + if !self.residual { + return raw0 * self.scales[idx]; + } + // Residual term: same kernel, same association order, added after the + // first-pass term (bit-identical across all paths). + let codes1 = &self.codes2[base..base + self.padded / 2]; + let raw1 = self.score_raw(codes1, q, &pq.lut); + raw0 * self.scales[idx] + raw1 * self.scales2[idx] + } + + /// Score one vector's code bytes against a prepared query, dispatching to + /// the best available kernel for the target. + #[inline] + fn score_raw(&self, codes: &[u8], q: &[f32], lut: &[f32; 16]) -> f32 { #[cfg(target_arch = "aarch64")] { // NEON is baseline on aarch64. - let raw = unsafe { neon::score_neon(codes, q, &pq.lut) }; - raw * self.scales[idx] + unsafe { neon::score_neon(codes, q, lut) } } #[cfg(target_arch = "x86_64")] { if avx2::available() { // SAFETY: feature availability checked immediately above. - let raw = unsafe { avx2::score_avx2(codes, q, &pq.lut) }; - raw * self.scales[idx] + unsafe { avx2::score_avx2(codes, q, lut) } } else { - score_scalar(codes, q, &pq.lut) * self.scales[idx] + score_scalar(codes, q, lut) } } #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))] { - score_scalar(codes, q, &pq.lut) * self.scales[idx] + score_scalar(codes, q, lut) } } @@ -726,6 +818,15 @@ impl VecqIndex { #[cfg(target_arch = "x86_64")] let use_avx2 = avx2::available(); let mut idx = 0; + // Residual mode combines both code planes per slot; plain mode scales + // only the first pass. The expression order is identical across all + // kernel paths (bit-identity requirement). + let combine = |r0: f32, r1: Option, si: usize| -> f32 { + match r1 { + Some(r1) => r0 * self.scales[si] + r1 * self.scales2[si], + None => r0 * self.scales[si], + } + }; #[cfg(target_arch = "aarch64")] { // Batch 4 vectors per pass: shared q loads + LUT setup. Tombstoned @@ -734,10 +835,17 @@ impl VecqIndex { while idx + 4 <= self.n { let codes4 = &self.codes[idx * bpv..(idx + 4) * bpv]; let raw = unsafe { neon::score_neon4(codes4, q_rot, &pq.lut) }; + let raw1 = if self.residual { + Some(unsafe { + neon::score_neon4(&self.codes2[idx * bpv..(idx + 4) * bpv], q_rot, &pq.lut) + }) + } else { + None + }; for (v, &r) in raw.iter().enumerate() { let si = idx + v; if self.alive[si] { - consider(r * self.scales[si], si, &mut heap); + consider(combine(r, raw1.map(|a| a[v]), si), si, &mut heap); } } idx += 4; @@ -754,10 +862,21 @@ impl VecqIndex { // SAFETY: AVX2 availability checked via `use_avx2`. let raw = unsafe { avx2::score_avx24(codes4, &pq.rotated[..self.padded], &pq.lut) }; + let raw1 = if self.residual { + Some(unsafe { + avx2::score_avx24( + &self.codes2[idx * bpv..(idx + 4) * bpv], + &pq.rotated[..self.padded], + &pq.lut, + ) + }) + } else { + None + }; for (v, &r) in raw.iter().enumerate() { let si = idx + v; if self.alive[si] { - consider(r * self.scales[si], si, &mut heap); + consider(combine(r, raw1.map(|a| a[v]), si), si, &mut heap); } } idx += 4; @@ -2044,4 +2163,204 @@ mod tests { let q = &base[2]; assert_eq!(idx.search_cascade(q, 5, 50), idx.search(q, 5)); } + + // -- residual quantization (second-pass codes) --------------------------- + + #[test] + fn residual_improves_recall_over_plain_on_noisy_data() { + // The adversarial clustered set is exactly where a second-pass + // residual code should pay off: plain 4-bit codes quantize the + // noise-dominated rotated dims coarsely. + let dim = 128; + let base = clustered(600, dim, 30, 19, 0.5); + let mut plain = VecqIndex::new(dim, 42); + let mut resid = VecqIndex::with_residual(dim, 42); + for v in &base { + plain.add(v); + resid.add(v); + } + let queries = &base[590..]; + let recall = |idx: &VecqIndex| -> f32 { + let mut sum = 0f32; + for q in queries { + let truth = exact_top(&base, q, 10); + let got: Vec = idx.search(q, 10).into_iter().map(|(i, _)| i).collect(); + sum += truth.iter().filter(|t| got.contains(t)).count() as f32 / 10.0; + } + sum / queries.len() as f32 + }; + let r_plain = recall(&plain); + let r_resid = recall(&resid); + // TODO(#23): reconstruction MSE improves 4.5x with residual codes, + // yet recall on this adversarial set is slightly BELOW plain + // (0.58 vs 0.66) — the two-term score estimator needs revisiting + // (candidates: true-norm vs reconstructed-norm denominators, + // f16 scale2 rounding, per-vector rms score-scale variance). + // Sanity floor until the estimator is settled. + assert!( + r_resid >= 0.5, + "residual recall {r_resid} (plain {r_plain})" + ); + } + + #[test] + fn residual_search_sorted_and_deterministic() { + let dim = 64; + let mut idx = VecqIndex::with_residual(dim, 7); + for i in 0..40 { + idx.add(&rand_unit(dim, i * 7 + 1)); + } + let q = rand_unit(dim, 555); + let r1 = idx.search(&q, 10); + let r2 = idx.search(&q, 10); + assert_eq!(r1, r2); + assert_eq!(r1.len(), 10); + for w in r1.windows(2) { + assert!(w[0].1 >= w[1].1); + } + assert!(idx.is_residual()); + } + + #[test] + fn residual_round_trips_through_file() { + let dim = 64; + let mut idx = VecqIndex::with_residual(dim, 21); + for i in 0..12 { + idx.add(&rand_unit(dim, i + 400)); + } + let q = rand_unit(dim, 888); + let expected = idx.search(&q, 5); + let bytes = idx.to_bytes(); + assert_eq!(u16::from_le_bytes([bytes[4], bytes[5]]), 260, "v1.4 = 260"); + let back = VecqIndex::from_bytes(&bytes).expect("parse v1.4"); + assert!(back.is_residual()); + let reloaded = back.search(&q, 5); + assert_eq!(reloaded[0].0, expected[0].0); + for ((_, f0), (_, f1)) in reloaded.iter().zip(expected.iter()) { + assert!((f0 - f1).abs() < 1e-3, "{f0} vs {f1}"); + } + // Keyless residual re-serialize stays stable. + assert_eq!(bytes, back.to_bytes()); + } + + #[test] + fn residual_keyed_and_cascade_compose() { + let dim = 64; + let mut idx = VecqIndex::with_residual(dim, 31); + idx.add_keyed(5, &rand_unit(dim, 101)); + idx.add_keyed(6, &rand_unit(dim, 202)); + idx.enable_cascade(); + let q = rand_unit(dim, 101); + let hits = idx.search_keyed(&q, 5); + assert_eq!(hits[0].0, 5); + let casc = idx.search_cascade(&q, 2, 2); + assert_eq!(casc, idx.search(&q, 2)); + let bytes = idx.to_bytes(); + let back = VecqIndex::from_bytes(&bytes).unwrap(); + assert!(back.contains_key(5) && back.contains_key(6)); + assert_eq!(back.search_keyed(&q, 5)[0].0, 5); + } + + #[test] + fn plain_index_stays_v13() { + let mut idx = VecqIndex::new(64, 3); + idx.add(&rand_unit(64, 1)); + assert!(!idx.is_residual()); + let bytes = idx.to_bytes(); + assert_eq!(u16::from_le_bytes([bytes[4], bytes[5]]), 259, "v1.3 = 259"); + } +} + +#[cfg(test)] +mod residual_tests { + use super::*; + + #[test] + fn residual_reconstruction_mse_halved() { + fn next(x: &mut u64) -> f32 { + *x ^= *x << 13; + *x ^= *x >> 7; + *x ^= *x << 17; + *x as f32 / u32::MAX as f32 - 0.5 + } + let dim = 128; + let mut x = 19 | 1; + let mut centroids = Vec::new(); + for _ in 0..10 { + let mut v: Vec = (0..dim).map(|_| next(&mut x)).collect(); + let n: f32 = v.iter().map(|a| a * a).sum::().sqrt(); + v.iter_mut().for_each(|a| *a /= n); + centroids.push(v); + } + let base: Vec> = (0..50) + .map(|i| { + let c = ¢roids[i % 10]; + let mut v: Vec = c.iter().map(|&a| a + 0.5 * next(&mut x)).collect(); + let n: f32 = v.iter().map(|a| a * a).sum::().sqrt(); + v.iter_mut().for_each(|a| *a /= n); + v + }) + .collect(); + let mut plain = VecqIndex::new(dim, 42); + let mut resid = VecqIndex::with_residual(dim, 42); + for v in &base { + plain.add(v); + resid.add(v); + } + // Reconstruct: decode code0 (+code1·scale2) per slot, compare to the + // rotated original (recompute rotation locally). + let bpv = plain.padded() / 2; + let mut mse0 = 0f32; + let mut mse1 = 0f32; + for (slot, v) in base.iter().enumerate() { + let unit: Vec = { + let n: f32 = v.iter().map(|a| a * a).sum::().sqrt(); + v.iter().map(|a| a / n).collect() + }; + let mut rot = Vec::new(); + resid.transform.apply(&unit, &mut rot); + let deq = |codes: &[u8], scale: f32| -> Vec { + let mut acc = vec![0f32; rot.len()]; + for (i, &b) in codes.iter().enumerate() { + acc[2 * i] += lloyd::dequantize_4bit(b & 0x0F) * scale; + acc[2 * i + 1] += lloyd::dequantize_4bit(b >> 4) * scale; + } + acc + }; + // Reconstructions are unit-norm-scaled; the true rotated vector + // has norm sqrt(padded), so scale back up before comparing. + let sp = (plain.padded() as f32).sqrt(); + let x0 = deq( + &plain.codes[slot * bpv..(slot + 1) * bpv], + plain.scales[slot] * sp, + ); + let x1 = deq( + &resid.codes[slot * bpv..(slot + 1) * bpv], + resid.scales[slot] * sp, + ); + let x2 = deq( + &resid.codes2[slot * bpv..(slot + 1) * bpv], + resid.scales2[slot] * sp, + ); + let err = |xh: &[f32]| -> f32 { + xh.iter() + .zip(rot.iter()) + .map(|(a, b)| (a - b) * (a - b)) + .sum::() + / rot.len() as f32 + }; + mse0 += err(&x0); + // residual reconstruction = x1 + x2 + let xtot: Vec = x1.iter().zip(x2.iter()).map(|(a, b)| a + b).collect(); + mse1 += err(&xtot); + } + let m0 = mse0 / base.len() as f32; + let m1 = mse1 / base.len() as f32; + println!("mse plain={:.5} residual={:.5}", m0, m1); + // The whole point of the second pass: reconstruction must improve. + assert!( + m1 < m0 * 0.5, + "residual reconstruction must at least halve plain MSE: {m0} vs {m1}" + ); + } } From 28f29a940caae7d84a866af44a7126e359e89762 Mon Sep 17 00:00:00 2001 From: ajianaz Date: Sat, 29 Aug 2026 17:21:48 +0700 Subject: [PATCH 2/4] =?UTF-8?q?fix(core):=20residual=20scoring=20=E2=80=94?= =?UTF-8?q?=20exact=20reconstruction-norm=20denominator=20(#23)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause of the recall drop (0.58 vs plain 0.66 despite 4.5x better MSE): the two-term score divided by an APPROXIMATE reconstruction norm sqrt(sum_sq + rms^2*padded), but the actual scored reconstruction is x = d0 + rms*d1 whose squared norm includes the cross term 2*rms* (nonzero — the second-pass codes inherit the correlated quantization-error pattern, not an independent Gaussian) and the actual quantized energy . The per-vector approximation error injected score variance (+40% sd, bias identical) that flipped top-10 rankings. Fix: compute the exact ||x||^2 at encode time from the stored code pairs and fold it into the stored scale coefficients. No format change; scoring kernels and bit-identity guarantees untouched. Self-consistency: score(v, slot of v) is now <= 1.0 everywhere (the old denominator produced cosine > 1.0 on some vectors). Measured (adversarial clustered set, n=600/30 clusters, spread 0.5): recall@10 0.58 -> 0.87 (plain 0.66); score-error sd 0.0087 -> 0.0028. Test upgraded from sanity floor to: residual beats plain by >= 0.05 and recall >= 0.8. Local: fmt clean, clippy -D warnings 0/0, 54/54 native (aarch64 NEON). x86_64 validated by CI (the changed encode path is scalar). Signed-off-by: ajianaz --- crates/vecq-core/src/store.rs | 57 ++++++++++++++++++++++++++++------- 1 file changed, 46 insertions(+), 11 deletions(-) diff --git a/crates/vecq-core/src/store.rs b/crates/vecq-core/src/store.rs index 7ac3ea1..6f18a31 100644 --- a/crates/vecq-core/src/store.rs +++ b/crates/vecq-core/src/store.rs @@ -673,10 +673,37 @@ impl VecqIndex { self.codes2[b] = byte; } // Term coefficients: both terms estimate the cosine of the full - // reconstruction x̂ = x̂0 + rms·d1, so both divide by its norm - // sqrt(sum_sq + rms²·padded). Term1's numerator is rms·raw1 (raw1 = - // q·dequant(c1) estimates q·(r/rms)). - let denom = (sum_sq + rms * rms * self.padded as f32).sqrt(); + // reconstruction x̂ = x̂0 + rms·d1, so both divide by ‖x̂‖. Term1's + // numerator is rms·raw1 (raw1 = q·dequant(c1) estimates q·(r/rms)). + // + // #23 (estimator fix): ‖x̂‖² must be the EXACT squared norm of the + // vector we actually score against — including the cross term + // 2·rms·⟨d0, d1⟩ (nonzero: the residual d1 correlates with d0's + // quantization error pattern, not with an independent Gaussian) and + // the actual quantized second-pass energy ⟨d1, d1⟩ (not the raw + // residual rms²·padded). Using sqrt(sum_sq + rms²·padded) instead + // injects per-vector score distortion that RAISES score variance + // enough to flip top-10 rankings: measured recall 0.58 vs plain + // 0.66 on the adversarial clustered set despite 4.5x better MSE + // (bias identical, sd +40%). The cross term and ⟨d1,d1⟩ are folded + // into the stored coefficients at encode time — no format change, + // scoring kernels untouched. + let mut cross = 0f32; // ⟨d0, d1⟩ + let mut norm1_sq = 0f32; // ⟨d1, d1⟩ + for i in 0..self.padded { + let b = base + i / 2; + let (c0, c1) = if i % 2 == 0 { + (self.codes[b] & 0x0F, self.codes2[b] & 0x0F) + } else { + (self.codes[b] >> 4, self.codes2[b] >> 4) + }; + let d0 = lloyd::dequantize_4bit(c0); + let d1 = lloyd::dequantize_4bit(c1); + cross += d0 * d1; + norm1_sq += d1 * d1; + } + let norm_sq = sum_sq + 2.0 * rms * cross + rms * rms * norm1_sq; + let denom = norm_sq.max(1e-12).sqrt(); ((scale0 * sum_sq.sqrt()) / denom, Some(rms / denom)) } @@ -2191,14 +2218,22 @@ mod tests { }; let r_plain = recall(&plain); let r_resid = recall(&resid); - // TODO(#23): reconstruction MSE improves 4.5x with residual codes, - // yet recall on this adversarial set is slightly BELOW plain - // (0.58 vs 0.66) — the two-term score estimator needs revisiting - // (candidates: true-norm vs reconstructed-norm denominators, - // f16 scale2 rounding, per-vector rms score-scale variance). - // Sanity floor until the estimator is settled. + // Residual must clearly BEAT plain on this set: same 514 B/vector + // budget, finer reconstruction. (Estimator note: the two-term score + // must divide by the EXACT reconstruction norm including the + // 2·rms·⟨d0,d1⟩ cross term — see encode_into. With the approximate + // sqrt(sum_sq + rms²·padded) denominator, score variance rose ~40% + // and recall DROPPED to 0.58 despite 4.5x better MSE.) + assert!( + r_resid > r_plain, + "residual must beat plain: resid {r_resid} vs plain {r_plain}" + ); + assert!( + r_resid - r_plain >= 0.05, + "residual gain must be substantial: resid {r_resid} vs plain {r_plain}" + ); assert!( - r_resid >= 0.5, + r_resid >= 0.8, "residual recall {r_resid} (plain {r_plain})" ); } From aad02e1933e1e84965f6fd8cb6453be01f572fde Mon Sep 17 00:00:00 2001 From: ajianaz Date: Sat, 29 Aug 2026 17:48:21 +0700 Subject: [PATCH 3/4] test+docs: real-dataset residual validation + benchmark doc (#23) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add real_residual_validation.rs (ignored by default, requires the EmbeddingGemma dataset from gen_embeddings.py): on the real dataset the fixed residual mode reaches recall@1 0.990 / recall@10 0.984 at ~3x compression with a 1.43x scan cost — vs plain 0.910/0.958 at 5.98x, and usearch HNSW 0.995 at 1x. Plain-path numbers re-verified unchanged: real / vs_usearch / vs_quantizers recall bit-identical to the baseline tables (timings on this Oracle ARM host are ~2x the Apple Silicon baseline uniformly across ALL engines, incl. competitors; ordering preserved). Document the estimator fix, the measured trade-off table, and honest labeling (residual is an opt-in recall mode, plain stays default) in docs/BENCHMARK.md under a new 'Residual quantization' section. Signed-off-by: ajianaz --- .../tests/real_residual_validation.rs | 81 +++++++++++++++++++ docs/BENCHMARK.md | 43 ++++++++++ 2 files changed, 124 insertions(+) create mode 100644 crates/vecq-core/tests/real_residual_validation.rs diff --git a/crates/vecq-core/tests/real_residual_validation.rs b/crates/vecq-core/tests/real_residual_validation.rs new file mode 100644 index 0000000..1e2e55b --- /dev/null +++ b/crates/vecq-core/tests/real_residual_validation.rs @@ -0,0 +1,81 @@ +//! Residual vs plain on the REAL EmbeddingGemma dataset (n=2000 + 100 +//! queries, dim=768, generated by cto/scripts/gen_embeddings.py into +//! /tmp/vecq-bench). Ignored by default: the dataset is not reproducible +//! in CI. Run locally in release before merging estimator/format changes: +//! cargo test -p vecq-core --release --test real_residual_validation -- --ignored --nocapture + +use vecq_core::VecqIndex; + +fn read_f32s(path: &str) -> Vec { + let bytes = std::fs::read(path).expect(path); + bytes + .chunks_exact(4) + .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .collect() +} + +fn cosine(a: &[f32], b: &[f32]) -> f32 { + let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum(); + dot / 1.0 // both already unit-norm (EmbeddingGemma output normalized) +} + +#[test] +#[ignore = "requires the real EmbeddingGemma dataset in /tmp/vecq-bench (gen_embeddings.py)"] +fn real_dataset_residual_vs_plain() { + let (n, nq, dim) = (2000usize, 100usize, 768usize); + let base = read_f32s("/tmp/vecq-bench/base.f32"); + let queries = read_f32s("/tmp/vecq-bench/queries.f32"); + assert_eq!(base.len(), n * dim); + assert_eq!(queries.len(), nq * dim); + let base_v: Vec<&[f32]> = base.chunks(dim).collect(); + let queries_v: Vec<&[f32]> = queries.chunks(dim).collect(); + + let exact_top = |q: &[f32], k: usize| -> Vec { + let mut s: Vec<(usize, f32)> = base_v + .iter() + .enumerate() + .map(|(i, v)| (i, cosine(q, v))) + .collect(); + s.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + s.into_iter().take(k).map(|(i, _)| i).collect() + }; + + let mut plain = VecqIndex::new(dim, 42); + let mut resid = VecqIndex::with_residual(dim, 42); + let t0 = std::time::Instant::now(); + for v in &base_v { + plain.add(v); + resid.add(v); + } + let build = t0.elapsed(); + + for (label, idx) in [("plain", &plain), ("resid", &resid)] { + let t1 = std::time::Instant::now(); + let (mut r1, mut r10) = (0f32, 0f32); + let mut max_score = 0f32; + for q in &queries_v { + let truth = exact_top(q, 10); + let got = idx.search(q, 10); + for (_, s) in &got { + if *s > max_score { + max_score = *s; + } + } + r1 += if got[0].0 == truth[0] { 1.0 } else { 0.0 }; + r10 += truth + .iter() + .filter(|t| got.iter().any(|(i, _)| i == *t)) + .count() as f32 + / 10.0; + } + let dt = t1.elapsed(); + println!( + "{label}: recall@1 {:.3} recall@10 {:.3} | {:.3} ms/q | max score {:.6} (<=1 sanity) | build +search {:?}", + r1 / nq as f32, + r10 / nq as f32, + dt.as_secs_f64() * 1000.0 / nq as f64, + max_score, + build + ); + } +} diff --git a/docs/BENCHMARK.md b/docs/BENCHMARK.md index c9d58d7..ebfb5f3 100644 --- a/docs/BENCHMARK.md +++ b/docs/BENCHMARK.md @@ -121,6 +121,49 @@ therefore be evaluated as a **recall lift at the same 514 B budget** (e.g. the binary-signature cascade (#22) remains the scan-speed lever for larger n. Verdict: proceed with both, benchmarked against this baseline. +## Residual quantization (issue #23) + +Opt-in mode (`VecqIndex::with_residual`): a second Lloyd-Max pass codes the +residual left by the first pass; scoring adds the second term. Format v1.4 +appends the second code block per vector — readers accept v1–v3. + +**Estimator fix (required for the mode to help at all):** the first +implementation divided the two-term score by the approximate norm +`sqrt(sum_sq + rms²·padded)`. The actually-scored reconstruction is +`x̂ = d0 + rms·d1`, whose exact norm includes the cross term `2·rms·⟨d0,d1⟩` +(nonzero: the second-pass codes inherit the correlated quantization-error +pattern, not an independent Gaussian) and the true quantized energy `⟨d1,d1⟩`. +The approximation injected per-vector score distortion that *raised* score +variance (+40% sd, bias unchanged) and flipped top-10 rankings: recall on the +adversarial clustered set **dropped** to 0.58 vs plain 0.66 despite 4.5x +better reconstruction MSE. Fix: compute the exact ‖x̂‖² at encode time from +the stored code pairs and fold it into the stored scale coefficients — no +format change, scoring kernels untouched. Same fix also restored the +self-consistency invariant score(v, v) ≤ 1.0 (the old denominator produced +cosine > 1.0 on some vectors). + +Measured after the fix, real EmbeddingGemma dataset (n=2000 + 100 queries, +dim 768, same corpus as the table above), aarch64 release: + +| mode | recall@1 | recall@10 | bytes/vec | compression | scan cost | +|---|---|---|---|---|---| +| plain 4-bit | 0.910 | 0.958 | 514 | 5.98x | 1.00x | +| plain + residual | **0.990** | **0.984** | 1.028 | ~3.0x | 1.43x | +| usearch f32 HNSW | — | 0.995 | 3,072 | 1x | — | + +Honest labeling: +- Residual roughly doubles the storage (second 4-bit block + second f16 + scale) — it is a recall mode, not a free lunch. Plain stays the default. +- The 1.43x scan cost is the second accumulate term; the cascade (#22) + remains the throughput lever and composes orthogonally. +- The real-dataset comparison lives in + `crates/vecq-core/tests/real_residual_validation.rs` (ignored by default; + requires the dataset from `cto/scripts/gen_embeddings.py`). Re-run it in + release before merging any estimator/format change. +- Plain-path numbers are unchanged by the fix by construction (single-term + path untouched); re-verified: `real`/`vs_usearch`/`vs_quantizers` recall + bit-identical to the baseline tables above. + ## Conclusion The spike validates the technique: training-free 4-bit RHDH + Lloyd-Max From 787768870112b6b787769a9f22590937706b8cc5 Mon Sep 17 00:00:00 2001 From: ajianaz Date: Sat, 29 Aug 2026 17:53:31 +0700 Subject: [PATCH 4/4] =?UTF-8?q?fix(test):=20use=20as=5Fchunks::<4>=20?= =?UTF-8?q?=E2=80=94=20satisfy=20clippy=201.98=20chunks=5Fexact=5Fto=5Fas?= =?UTF-8?q?=5Fchunks=20(#23)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI clippy (1.98) flags chunks_exact with a constant chunk size; local clippy (1.96) does not know this lint yet. as_chunks is stable since 1.83, no MSRV is declared, so the portable form is used. Signed-off-by: ajianaz --- crates/vecq-core/tests/real_residual_validation.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/crates/vecq-core/tests/real_residual_validation.rs b/crates/vecq-core/tests/real_residual_validation.rs index 1e2e55b..24c6dde 100644 --- a/crates/vecq-core/tests/real_residual_validation.rs +++ b/crates/vecq-core/tests/real_residual_validation.rs @@ -8,10 +8,9 @@ use vecq_core::VecqIndex; fn read_f32s(path: &str) -> Vec { let bytes = std::fs::read(path).expect(path); - bytes - .chunks_exact(4) - .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) - .collect() + let (chunks, rest) = bytes.as_chunks::<4>(); + assert!(rest.is_empty(), "{path}: size not a multiple of 4"); + chunks.iter().map(|b| f32::from_le_bytes(*b)).collect() } fn cosine(a: &[f32], b: &[f32]) -> f32 {