mirror of
https://github.com/ruvnet/RuVector.git
synced 2026-07-10 01:38:44 +00:00
feat(bet4): M2/M3 — steelman B&B + PCA-8 control + matched-recall sweep
- kernel: search_bnb_skip — the STEELMAN. Centroid-distance order (the effective nprobe ordering) + per-cluster LB-skip (correctness-safe in any order, unlike the LB-order global break). The strongest cluster-level B&B: if it can't beat tuned nprobe, the bound doesn't pay. - pca: minimal power-iteration top-m PCA (no linalg dep) for the low-dim control — projects real arxiv features to 8-d where the bound is tight. - examples/ivf_pruning_sweep: 3 contenders share one index per nclusters (plain nprobe / B&B LB-order / B&B steelman) x 2 regimes (128-d, PCA-8), exact-regime pruning probe, matched-recall@0.95, frozen-gate verdict. RESULT (n=20k & n=50k both): steelman = 1.00x evals vs nprobe in EVERY cell, BOTH regimes. NO-GO. Mechanism is structural, not dimensional: the LB bound only prunes FAR clusters that tuned nprobe already skips, so it's redundant with nprobe's centroid-distance cutoff. Exact-prune fraction scales correctly with dim (0-13% @128-d, 8-87% @PCA-8) => kernel sound; the redundancy is fundamental. LB-ORDER (faithful BET-2 kernel) is strictly WORSE (0.18-0.25x) — LB-ordering probes far large-radius clusters early.
This commit is contained in:
parent
d36f4e043d
commit
762fa976b2
4 changed files with 318 additions and 0 deletions
198
crates/ruvector-bet4-ivf-bench/examples/ivf_pruning_sweep.rs
Normal file
198
crates/ruvector-bet4-ivf-bench/examples/ivf_pruning_sweep.rs
Normal file
|
|
@ -0,0 +1,198 @@
|
|||
//! BET 4 matched-recall sweep (M2/M3): LB-ordered branch-and-bound IVF probing vs the tuned plain
|
||||
//! `IvfFlat` `nprobe` incumbent, on real 128-d arxiv embeddings AND a PCA-8 low-dim control.
|
||||
//!
|
||||
//! Three contenders share one index per `nclusters` (built once): plain `nprobe` (incumbent),
|
||||
//! B&B in **LB-order** (the faithful BET-2 `RegionPruneIvf` kernel), and the **steelman** B&B —
|
||||
//! centroid-distance order + LB-skip (the strongest version: if it can't beat `nprobe`, the bound
|
||||
//! doesn't pay). Reports the exact-regime pruning fraction, matched-recall cost, and checks the
|
||||
//! FROZEN gate (docs/plans/bet4-ivf-pruning/PRE-REGISTRATION.md) on the steelman ratio.
|
||||
//!
|
||||
//! Run: `cargo run --release -p ruvector-bet4-ivf-bench --example ivf_pruning_sweep -- [N]`
|
||||
|
||||
use ruvector_bet4_ivf_bench::data::load_feat_csv;
|
||||
use ruvector_bet4_ivf_bench::kernel::BnBIvf;
|
||||
use ruvector_bet4_ivf_bench::oracle::{brute_force_topk, recall_at_k};
|
||||
use ruvector_bet4_ivf_bench::pca::project_topm;
|
||||
use ruvector_rairs::SearchResult;
|
||||
use std::time::Instant;
|
||||
|
||||
const K: usize = 10;
|
||||
const R_TARGET: f64 = 0.95;
|
||||
const NCLUSTERS: [usize; 3] = [64, 256, 1024];
|
||||
|
||||
fn main() {
|
||||
let args: Vec<String> = std::env::args().collect();
|
||||
let n_req: usize = args.get(1).and_then(|s| s.parse().ok()).unwrap_or(20_000);
|
||||
let data =
|
||||
std::env::var("BET4_DATA").unwrap_or_else(|_| "target/m1-data/node-feat-100k.csv".into());
|
||||
|
||||
let corpus = load_feat_csv(&data, n_req).unwrap_or_else(|e| {
|
||||
eprintln!("failed to load {data}: {e}");
|
||||
std::process::exit(1);
|
||||
});
|
||||
let n = corpus.len();
|
||||
let dim = corpus.first().map(|v| v.len()).unwrap_or(0);
|
||||
println!("# BET4 sweep n={n} dim={dim} k={K} R_target={R_TARGET} data={data}\n");
|
||||
|
||||
run_regime("128-d (real arxiv features)", &corpus);
|
||||
|
||||
println!("\n# Projecting to PCA-8 (low-dim control)…");
|
||||
let t = Instant::now();
|
||||
let corpus8 = project_topm(&corpus, 8, 60);
|
||||
println!("# PCA done in {:?}\n", t.elapsed());
|
||||
run_regime("PCA-8 (low-dim control — bound should be TIGHT, B&B should WIN)", &corpus8);
|
||||
}
|
||||
|
||||
fn run_regime(label: &str, corpus: &[Vec<f32>]) {
|
||||
let n = corpus.len();
|
||||
let dim = corpus[0].len();
|
||||
let nq = 200.min(n);
|
||||
let queries: Vec<usize> = (0..nq).collect();
|
||||
let truth: Vec<Vec<usize>> = queries
|
||||
.iter()
|
||||
.map(|&q| brute_force_topk(corpus, &corpus[q], K))
|
||||
.collect();
|
||||
|
||||
println!("════ REGIME: {label} (dim={dim}) ════");
|
||||
let mut cells: Vec<Cell> = Vec::new();
|
||||
|
||||
for &nc in &NCLUSTERS {
|
||||
let t_build = Instant::now();
|
||||
let idx = BnBIvf::build(corpus, nc, 15, 42);
|
||||
let nc_eff = idx.num_lists();
|
||||
let build = t_build.elapsed();
|
||||
|
||||
// Exact-regime pruning fraction (LB-order full budget).
|
||||
let mut pruned = 0.0;
|
||||
for &q in &queries {
|
||||
let (_r, _e, probed) = idx.search(&corpus[q], K, None);
|
||||
pruned += (nc_eff - probed) as f64 / nc_eff as f64;
|
||||
}
|
||||
let prune_frac = pruned / nq as f64;
|
||||
|
||||
let grid = knob_grid(nc_eff);
|
||||
let plain = matched(&queries, corpus, &truth, &grid, |q, knob| {
|
||||
let (r, ev, _) = idx.search_nprobe(q, K, knob);
|
||||
(ids(&r), ev)
|
||||
});
|
||||
let bnb_lb = matched(&queries, corpus, &truth, &grid, |q, knob| {
|
||||
let (r, ev, _) = idx.search(q, K, Some(knob));
|
||||
(ids(&r), ev)
|
||||
});
|
||||
let bnb_skip = matched(&queries, corpus, &truth, &grid, |q, knob| {
|
||||
let (r, ev, _) = idx.search_bnb_skip(q, K, Some(knob));
|
||||
(ids(&r), ev)
|
||||
});
|
||||
|
||||
let eval_ratio = plain.evals / bnb_skip.evals.max(1.0);
|
||||
let wall_ratio = plain.wall_ns as f64 / bnb_skip.wall_ns.max(1) as f64;
|
||||
|
||||
println!("\n## nclusters={nc_eff} (build {build:?}) exact-regime prune={:.1}%", prune_frac * 100.0);
|
||||
print_row("plain nprobe (incumbent)", &plain);
|
||||
print_row("B&B LB-order (BET-2 kernel)", &bnb_lb);
|
||||
print_row("B&B steelman (cdist+LB-skip)", &bnb_skip);
|
||||
println!(
|
||||
" steelman vs incumbent: eval {eval_ratio:.2}x wall {wall_ratio:.2}x"
|
||||
);
|
||||
|
||||
cells.push(Cell { nc: nc_eff, eval_ratio, wall_ratio, prune_frac });
|
||||
}
|
||||
|
||||
verdict(label, &cells);
|
||||
}
|
||||
|
||||
struct Cell {
|
||||
nc: usize,
|
||||
eval_ratio: f64,
|
||||
wall_ratio: f64,
|
||||
prune_frac: f64,
|
||||
}
|
||||
|
||||
struct Matched {
|
||||
knob: usize,
|
||||
recall: f64,
|
||||
evals: f64,
|
||||
wall_ns: u128,
|
||||
}
|
||||
|
||||
fn print_row(name: &str, m: &Matched) {
|
||||
println!(
|
||||
" {name:<32} knob={:<4} recall={:.4} evals/q={:>8.0} wall/q={:>6}µs",
|
||||
m.knob,
|
||||
m.recall,
|
||||
m.evals,
|
||||
m.wall_ns / 1000
|
||||
);
|
||||
}
|
||||
|
||||
/// First knob (ascending) whose mean recall ≥ `R_TARGET`, with its mean member-evals and wall-time;
|
||||
/// falls back to the largest knob if none reaches target.
|
||||
fn matched<F>(
|
||||
queries: &[usize],
|
||||
corpus: &[Vec<f32>],
|
||||
truth: &[Vec<usize>],
|
||||
grid: &[usize],
|
||||
search: F,
|
||||
) -> Matched
|
||||
where
|
||||
F: Fn(&[f32], usize) -> (Vec<usize>, usize),
|
||||
{
|
||||
let mut last = Matched { knob: 0, recall: 0.0, evals: 0.0, wall_ns: 0 };
|
||||
for &knob in grid {
|
||||
let t = Instant::now();
|
||||
let mut rec = 0.0;
|
||||
let mut ev = 0usize;
|
||||
for (qi, &q) in queries.iter().enumerate() {
|
||||
let (got, e) = search(&corpus[q], knob);
|
||||
ev += e;
|
||||
rec += recall_at_k(&truth[qi], &got, K);
|
||||
}
|
||||
let wall_ns = t.elapsed().as_nanos() / queries.len() as u128;
|
||||
last = Matched {
|
||||
knob,
|
||||
recall: rec / queries.len() as f64,
|
||||
evals: ev as f64 / queries.len() as f64,
|
||||
wall_ns,
|
||||
};
|
||||
if last.recall >= R_TARGET {
|
||||
return last;
|
||||
}
|
||||
}
|
||||
last
|
||||
}
|
||||
|
||||
fn knob_grid(maxv: usize) -> Vec<usize> {
|
||||
let mut g = Vec::new();
|
||||
let mut x = 1usize;
|
||||
while x < maxv {
|
||||
g.push(x);
|
||||
x = ((x as f64) * 1.5).ceil() as usize;
|
||||
}
|
||||
g.push(maxv);
|
||||
g.dedup();
|
||||
g
|
||||
}
|
||||
|
||||
fn ids(res: &[SearchResult]) -> Vec<usize> {
|
||||
res.iter().map(|r| r.id).collect()
|
||||
}
|
||||
|
||||
fn verdict(label: &str, cells: &[Cell]) {
|
||||
let all_win = cells.iter().all(|c| c.eval_ratio >= 2.0 && c.wall_ratio > 1.0);
|
||||
let any_kill = cells.iter().any(|c| c.eval_ratio < 1.5 || c.wall_ratio < 1.0);
|
||||
let v = if all_win {
|
||||
"WIN (≥2× evals AND wall-clock win across all nclusters)"
|
||||
} else if any_kill {
|
||||
"KILL / NO-GO (<1.5× somewhere or wall reversed — bound too loose to pay)"
|
||||
} else {
|
||||
"QUALIFIED (1.5–2×, or mixed)"
|
||||
};
|
||||
println!("\n ── verdict [{label}] ──");
|
||||
for c in cells {
|
||||
println!(
|
||||
" nclusters={:<5} steelman eval={:.2}x wall={:.2}x exact-prune={:.1}%",
|
||||
c.nc, c.eval_ratio, c.wall_ratio, c.prune_frac * 100.0
|
||||
);
|
||||
}
|
||||
println!(" => {v}");
|
||||
}
|
||||
|
|
@ -157,6 +157,52 @@ impl BnBIvf {
|
|||
(finalize(heap), member_evals, probed)
|
||||
}
|
||||
|
||||
/// The **steelman B&B**: visit clusters in centroid-distance order (the effective `nprobe`
|
||||
/// ordering, so τ tightens fast), but **skip** scanning any cluster the lower bound proves
|
||||
/// cannot hold a top-k point (`LB(q,c) ≥ τ`). Unlike [`search`](Self::search)'s global early
|
||||
/// `break`, skipping is correctness-safe in *any* visit order (a skipped cluster genuinely
|
||||
/// cannot contain a closer point); a global break would be unsound here because a later,
|
||||
/// large-radius cluster can have a *smaller* LB than the current one.
|
||||
///
|
||||
/// `max_probe` caps the number of clusters **considered** (the apples-to-apples budget against
|
||||
/// `nprobe`); LB-skips save member scans within that budget. This is the strongest version of
|
||||
/// the bet — if it cannot beat `nprobe`, the bound itself doesn't pay. Returns
|
||||
/// `(top-k, member_evals, clusters_considered)`.
|
||||
pub fn search_bnb_skip(
|
||||
&self,
|
||||
q: &[f32],
|
||||
k: usize,
|
||||
max_probe: Option<usize>,
|
||||
) -> (Vec<SearchResult>, usize, usize) {
|
||||
let nclusters = self.centroids.len();
|
||||
let mut order: Vec<(f32, usize)> = (0..nclusters)
|
||||
.map(|c| (l2(q, &self.centroids[c]), c))
|
||||
.collect();
|
||||
order.sort_by(|a, b| a.0.total_cmp(&b.0));
|
||||
let cap = max_probe.unwrap_or(nclusters).min(nclusters);
|
||||
|
||||
let mut heap: BinaryHeap<Cand> = BinaryHeap::with_capacity(k + 1);
|
||||
let mut member_evals = 0usize;
|
||||
let mut considered = 0usize;
|
||||
for (dc, c) in order {
|
||||
if considered >= cap {
|
||||
break;
|
||||
}
|
||||
considered += 1;
|
||||
if heap.len() == k {
|
||||
let kth = heap.peek().unwrap().dist;
|
||||
if (dc - self.radii[c]).max(0.0) >= kth {
|
||||
continue; // LB-skip: provably cannot improve the top-k
|
||||
}
|
||||
}
|
||||
for (id, v) in &self.lists[c] {
|
||||
member_evals += 1;
|
||||
consider(&mut heap, k, *id, l2(q, v));
|
||||
}
|
||||
}
|
||||
(finalize(heap), member_evals, considered)
|
||||
}
|
||||
|
||||
/// The **plain-IVF incumbent** strategy on this same shared index: visit the `nprobe` nearest
|
||||
/// centroids (by centroid distance) and scan **all** their members — no lower-bound ordering,
|
||||
/// no early termination. This is exactly `ruvector-rairs::IvfFlat::search`'s algorithm
|
||||
|
|
|
|||
|
|
@ -12,5 +12,6 @@
|
|||
pub mod data;
|
||||
pub mod kernel;
|
||||
pub mod oracle;
|
||||
pub mod pca;
|
||||
|
||||
pub use kernel::BnBIvf;
|
||||
|
|
|
|||
73
crates/ruvector-bet4-ivf-bench/src/pca.rs
Normal file
73
crates/ruvector-bet4-ivf-bench/src/pca.rs
Normal file
|
|
@ -0,0 +1,73 @@
|
|||
//! Minimal top-`m` PCA via power iteration + deflation — for BET 4's **low-dimensional control**.
|
||||
//!
|
||||
//! Projecting the real arxiv features onto their top principal components gives the *same data*
|
||||
//! at low intrinsic dimensionality, where the triangle-inequality cluster bound should be tight
|
||||
//! and the B&B kernel is expected to WIN — proving the kernel/harness are sound and isolating
|
||||
//! high-dimensional distance concentration as the cause of any 128-d NO-GO. No linalg dependency.
|
||||
|
||||
/// Project `data` (n × dim) onto its top `m` principal components, returning n × m coordinates.
|
||||
/// Data is mean-centered first; components found by power iteration with deflation (`iters` steps
|
||||
/// each). f64 accumulation for numerical stability.
|
||||
pub fn project_topm(data: &[Vec<f32>], m: usize, iters: usize) -> Vec<Vec<f32>> {
|
||||
let n = data.len();
|
||||
if n == 0 {
|
||||
return Vec::new();
|
||||
}
|
||||
let dim = data[0].len();
|
||||
|
||||
let mut mean = vec![0.0f64; dim];
|
||||
for v in data {
|
||||
for (d, &x) in v.iter().enumerate() {
|
||||
mean[d] += x as f64;
|
||||
}
|
||||
}
|
||||
for x in &mut mean {
|
||||
*x /= n as f64;
|
||||
}
|
||||
let centered: Vec<Vec<f64>> = data
|
||||
.iter()
|
||||
.map(|v| (0..dim).map(|d| v[d] as f64 - mean[d]).collect())
|
||||
.collect();
|
||||
|
||||
let mut comps: Vec<Vec<f64>> = Vec::with_capacity(m.min(dim));
|
||||
for c in 0..m.min(dim) {
|
||||
let mut v = vec![0.0f64; dim];
|
||||
v[c % dim] = 1.0;
|
||||
for _ in 0..iters {
|
||||
// u = Σ_i (x_i · v) x_i — covariance-times-v without forming the covariance matrix.
|
||||
let mut u = vec![0.0f64; dim];
|
||||
for x in ¢ered {
|
||||
let dot: f64 = x.iter().zip(&v).map(|(a, b)| a * b).sum();
|
||||
for (d, &xd) in x.iter().enumerate() {
|
||||
u[d] += dot * xd;
|
||||
}
|
||||
}
|
||||
// Deflate against already-found components (Gram–Schmidt).
|
||||
for prev in &comps {
|
||||
let proj: f64 = u.iter().zip(prev).map(|(a, b)| a * b).sum();
|
||||
for (d, &pd) in prev.iter().enumerate() {
|
||||
u[d] -= proj * pd;
|
||||
}
|
||||
}
|
||||
let norm = u.iter().map(|x| x * x).sum::<f64>().sqrt();
|
||||
if norm < 1e-12 {
|
||||
break;
|
||||
}
|
||||
for x in &mut u {
|
||||
*x /= norm;
|
||||
}
|
||||
v = u;
|
||||
}
|
||||
comps.push(v);
|
||||
}
|
||||
|
||||
centered
|
||||
.iter()
|
||||
.map(|x| {
|
||||
comps
|
||||
.iter()
|
||||
.map(|comp| x.iter().zip(comp).map(|(a, b)| a * b).sum::<f64>() as f32)
|
||||
.collect()
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue