|
| 1 | +//! PQ and IVF-PQ codebook training must distribute centroids across the |
| 2 | +//! data even when many input vectors are near-duplicates. |
| 3 | +//! |
| 4 | +//! Spec: k-means initialization selects centroids spread across the data |
| 5 | +//! distribution. The current implementation has two compounding bugs: |
| 6 | +//! |
| 7 | +//! 1. `min_dists[i]` is only updated against `centroids[c - 1]` (the |
| 8 | +//! last centroid), not against the full centroid set. Once two |
| 9 | +//! centroids coincide, `min_dists` stops reflecting "distance to the |
| 10 | +//! nearest centroid," so every subsequent deterministic-argmax pick |
| 11 | +//! lands on the same outlier. |
| 12 | +//! 2. The comment says "k-means++" but the selection is deterministic |
| 13 | +//! farthest-point, so outliers dominate rather than being sampled |
| 14 | +//! proportionally to d². |
| 15 | +//! |
| 16 | +//! Effect: on workloads with repeated prefixes/suffixes (templated chat, |
| 17 | +//! shared headers/footers), most of the 256 centroids alias to one or two |
| 18 | +//! points and PQ recall collapses. |
| 19 | +
|
| 20 | +use nodedb_vector::quantize::pq::PqCodec; |
| 21 | + |
| 22 | +/// Training set of 200 vectors: 190 near-duplicates at the origin plus |
| 23 | +/// 10 outliers scattered across a single subspace. A correct k-means++ |
| 24 | +/// spreads centroids across both clusters; the current farthest-point- |
| 25 | +/// with-broken-min-distance-update collapses to ~2 distinct centroids. |
| 26 | +fn clustered_with_duplicates() -> Vec<Vec<f32>> { |
| 27 | + let mut vecs: Vec<Vec<f32>> = Vec::with_capacity(200); |
| 28 | + // Cluster A: 190 near-identical vectors near origin. |
| 29 | + for i in 0..190 { |
| 30 | + let eps = (i as f32) * 1e-5; |
| 31 | + vecs.push(vec![eps, -eps, eps * 0.5, -eps * 0.5]); |
| 32 | + } |
| 33 | + // Cluster B: 10 outliers at distinct coordinates. |
| 34 | + for j in 0..10 { |
| 35 | + let x = 100.0 + (j as f32) * 10.0; |
| 36 | + vecs.push(vec![x, -x, x * 0.5, -x * 0.5]); |
| 37 | + } |
| 38 | + vecs |
| 39 | +} |
| 40 | + |
| 41 | +fn unique_centroid_count(codec: &PqCodec, vectors: &[Vec<f32>]) -> usize { |
| 42 | + let refs: Vec<&[f32]> = vectors.iter().map(|v| v.as_slice()).collect(); |
| 43 | + let codes = codec.encode_batch(&refs); |
| 44 | + let m = codec.m; |
| 45 | + // Per-subspace unique centroid indices used across the batch. |
| 46 | + let mut min_unique = usize::MAX; |
| 47 | + for sub in 0..m { |
| 48 | + let mut seen = std::collections::HashSet::new(); |
| 49 | + for row in 0..vectors.len() { |
| 50 | + seen.insert(codes[row * m + sub]); |
| 51 | + } |
| 52 | + if seen.len() < min_unique { |
| 53 | + min_unique = seen.len(); |
| 54 | + } |
| 55 | + } |
| 56 | + min_unique |
| 57 | +} |
| 58 | + |
| 59 | +#[test] |
| 60 | +fn pq_kmeans_produces_diverse_centroids_on_duplicate_heavy_data() { |
| 61 | + let vecs = clustered_with_duplicates(); |
| 62 | + let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect(); |
| 63 | + let codec = PqCodec::train(&refs, 4, 2, 16, 20); |
| 64 | + |
| 65 | + let unique = unique_centroid_count(&codec, &vecs); |
| 66 | + assert!( |
| 67 | + unique >= 4, |
| 68 | + "k-means collapsed to {unique} unique centroids per subspace on \ |
| 69 | + duplicate-heavy input; a correct k-means++ should pick at least \ |
| 70 | + 4 distinct cluster representatives for k=16" |
| 71 | + ); |
| 72 | +} |
| 73 | + |
| 74 | +#[test] |
| 75 | +fn pq_distance_table_separates_duplicates_from_outliers() { |
| 76 | + // Spec test: after training, the PQ distance from a duplicate-cluster |
| 77 | + // query to a duplicate vector must be meaningfully smaller than the |
| 78 | + // distance to an outlier vector. Under the collapse bug, most |
| 79 | + // codebook entries alias to one point so all distances look similar. |
| 80 | + let vecs = clustered_with_duplicates(); |
| 81 | + let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect(); |
| 82 | + let codec = PqCodec::train(&refs, 4, 2, 16, 20); |
| 83 | + |
| 84 | + let query = [0.0f32, 0.0, 0.0, 0.0]; |
| 85 | + let table = codec.build_distance_table(&query); |
| 86 | + |
| 87 | + let dup_code = codec.encode(&vecs[0]); // duplicate cluster |
| 88 | + let outlier_code = codec.encode(&vecs[195]); // outlier cluster |
| 89 | + |
| 90 | + let dup_dist = codec.asymmetric_distance(&table, &dup_code); |
| 91 | + let outlier_dist = codec.asymmetric_distance(&table, &outlier_code); |
| 92 | + |
| 93 | + assert!( |
| 94 | + outlier_dist > dup_dist * 10.0, |
| 95 | + "PQ failed to distinguish duplicate (d={dup_dist}) from outlier \ |
| 96 | + (d={outlier_dist}) — codebook collapsed and the two codes encode \ |
| 97 | + to near-identical table entries" |
| 98 | + ); |
| 99 | +} |
| 100 | + |
| 101 | +#[cfg(feature = "ivf")] |
| 102 | +#[test] |
| 103 | +fn ivf_pq_training_does_not_collapse_on_duplicate_heavy_data() { |
| 104 | + use nodedb_vector::DistanceMetric; |
| 105 | + use nodedb_vector::{IvfPqIndex, IvfPqParams}; |
| 106 | + |
| 107 | + let vecs = clustered_with_duplicates(); |
| 108 | + let refs: Vec<&[f32]> = vecs.iter().map(|v| v.as_slice()).collect(); |
| 109 | + let mut idx = IvfPqIndex::new( |
| 110 | + 4, |
| 111 | + IvfPqParams { |
| 112 | + n_cells: 8, |
| 113 | + pq_m: 2, |
| 114 | + pq_k: 16, |
| 115 | + nprobe: 4, |
| 116 | + metric: DistanceMetric::L2, |
| 117 | + }, |
| 118 | + ); |
| 119 | + idx.train(&refs); |
| 120 | + for v in &vecs { |
| 121 | + idx.add(v); |
| 122 | + } |
| 123 | + |
| 124 | + // Query at the origin. Correct training assigns near-duplicates to |
| 125 | + // one cell and outliers to another; the nearest result must come |
| 126 | + // from the duplicate cluster (original indices 0..190). |
| 127 | + let results = idx.search(&[0.0, 0.0, 0.0, 0.0], 5); |
| 128 | + assert!(!results.is_empty(), "IVF-PQ returned no results"); |
| 129 | + for r in &results { |
| 130 | + assert!( |
| 131 | + r.id < 190, |
| 132 | + "IVF-PQ k-means collapse: query at origin returned outlier id={} \ |
| 133 | + instead of a near-duplicate cluster member", |
| 134 | + r.id |
| 135 | + ); |
| 136 | + } |
| 137 | +} |
0 commit comments