quantumclaw/clawhdf5
3
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

fix(ann): HNSW neighbour-selection heuristic — recall 0.31 -> 0.98 at 100K

Browse Source 
Neighbours were selected as the plain closest-M, for both a new node's links
and back-link pruning. On clustered data every link of a node inside a tight
cluster then goes to that same cluster, so clusters become islands a search
entering elsewhere can never reach: recall@10 was 0.87 / 0.67 / 0.31 at
1K / 10K / 100K (384-dim) and flat in ef. Uniform random data — all the
existing tests used — does not show it.

Implement the HNSW paper's Algorithm 4 with keepPrunedConnections: accept a
candidate only if it is closer to the node than to every neighbour already
accepted, then fill spare slots with the closest rejected ones. Recall@10 at
ef=64 is now 1.00 / 1.00 / 0.98 and rises with ef; uniform data improves
slightly. Build is ~3.5x slower at 10K (extra distance evaluations), to be
recovered by the distance-kernel work. The needless rayon fan-out over <=33
distances in prune_connections is gone.

Tests: a clustered-data recall test for bulk build and incremental insert
(scores 0.43 with the old selection), and a unit test of the selection rule.
Harness gains --uniform and --ann-only; before/after in BENCHMARKS.md.

Co-Authored-By: Claude Fable 5.1 <[email protected]>
This commit is contained in:
osobh
2026-09-19 07:43:41 -07:00
co-authored by Claude Fable 5.1
parent eb99de1020
commit 65d219c409
3 changed files with 194 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/clawhdf5-ann/src/hnsw.rs
+151 -16
View File
@@ -252,8 +252,9 @@ impl HnswIndex {
metric,
);
// Select up to m closest neighbors
let selected: Vec<usize> = neighbors.iter().take(max_conn).map(|c| c.id).collect();
let scored: Vec<(usize, f32)> =
neighbors.iter().map(|c| (c.id, c.distance)).collect();
let selected = select_neighbors(vectors, &scored, max_conn, metric);
// Add bidirectional connections
graph[layer][i] = selected.clone();
@@ -382,7 +383,8 @@ impl HnswIndex {
self.ef_construction,
self.metric,
);
let selected: Vec<usize> = neighbors.iter().take(max_conn).map(|c| c.id).collect();
let scored: Vec<(usize, f32)> = neighbors.iter().map(|c| (c.id, c.distance)).collect();
let selected = select_neighbors(&self.vectors, &scored, max_conn, self.metric);
self.graph[layer][id] = selected.clone();
for &neighbor in &selected {
self.graph[layer][neighbor].push(id);
@@ -828,7 +830,52 @@ fn search_layer(
result
}
/// Prune connections for a node to keep only the closest `max_conn` neighbors.
/// Choose up to `max_conn` neighbours for a node from `candidates` (sorted by
/// ascending distance to that node) — the HNSW paper's Algorithm 4 with
/// `keepPrunedConnections`.
///
/// Taking the plain `max_conn` closest is what breaks the graph on clustered
/// data: every link of a node inside a tight cluster goes to that same cluster,
/// so clusters become islands that a search entering elsewhere can never
/// reach, however large `ef` is. Instead a candidate is accepted only if it is
/// closer to the node than to every neighbour already accepted, which spreads
/// links across directions and keeps the long edges that join clusters. Any
/// remaining slots are then filled with the closest rejected candidates, so a
/// node is never left under-connected.
fn select_neighbors(
vectors: &[Vec<f32>],
candidates: &[(usize, f32)],
max_conn: usize,
metric: DistanceMetric,
) -> Vec<usize> {
if candidates.len() <= max_conn {
return candidates.iter().map(|&(id, _)| id).collect();
}
let mut selected: Vec<usize> = Vec::with_capacity(max_conn);
let mut rejected: Vec<usize> = Vec::new();
for &(id, dist_to_node) in candidates {
if selected.len() >= max_conn {
break;
}
let diverse = selected
.iter()
.all(|&s| compute_distance(&vectors[id], &vectors[s], metric) > dist_to_node);
if diverse {
selected.push(id);
} else {
rejected.push(id);
}
}
for id in rejected {
if selected.len() >= max_conn {
break;
}
selected.push(id);
}
selected
}
/// Trim `node`'s neighbour list back to `max_conn` with [`select_neighbors`].
fn prune_connections(
vectors: &[Vec<f32>],
neighbors: &mut Vec<usize>,
@@ -839,22 +886,12 @@ fn prune_connections(
if neighbors.len() <= max_conn {
return;
}
#[cfg(feature = "parallel")]
let mut scored: Vec<(usize, f32)> = {
use rayon::prelude::*;
neighbors
.par_iter()
.map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric)))
.collect()
};
#[cfg(not(feature = "parallel"))]
let mut scored: Vec<(usize, f32)> = neighbors
.iter()
.map(|&n| (n, compute_distance(&vectors[node], &vectors[n], metric)))
.collect();
scored.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
scored.truncate(max_conn);
*neighbors = scored.into_iter().map(|(id, _)| id).collect();
scored.sort_by(|a, b| a.1.total_cmp(&b.1).then(a.0.cmp(&b.0)));
*neighbors = select_neighbors(vectors, &scored, max_conn, metric);
}
// ---------------------------------------------------------------------------
@@ -1043,6 +1080,104 @@ fn get_attr_string(attrs: &[(String, AttrValue)], name: &str) -> Result<String,
mod tests {
use super::*;
/// Tight, well-separated clusters — the shape real embeddings have, and
/// the case plain closest-M neighbour selection fails on: each cluster
/// becomes an island, so recall is capped no matter how large `ef` is.
fn clustered(n: usize, dim: usize, clusters: usize, seed: u64) -> Vec<Vec<f32>> {
let mut state = seed;
let mut next = move || {
state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
let mut z = state;
z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
((z ^ (z >> 31)) >> 40) as f32 / (1u64 << 24) as f32 - 0.5
};
let centres: Vec<Vec<f32>> = (0..clusters)
.map(|_| (0..dim).map(|_| next() * 10.0).collect())
.collect();
(0..n)
.map(|i| {
centres[i % clusters]
.iter()
.map(|c| c + next() * 0.5)
.collect()
})
.collect()
}
fn recall_at_10(
index: &HnswIndex,
vectors: &[Vec<f32>],
queries: &[Vec<f32>],
ef: usize,
) -> f64 {
let mut hits = 0;
for q in queries {
let mut exact: Vec<(usize, f32)> = vectors
.iter()
.enumerate()
.map(|(i, v)| (i, compute_distance(q, v, DistanceMetric::L2)))
.collect();
exact.sort_by(|a, b| a.1.total_cmp(&b.1));
let want: Vec<usize> = exact[..10].iter().map(|e| e.0).collect();
hits += index
.search(q, 10, ef)
.iter()
.filter(|(id, _)| want.contains(id))
.count();
}
hits as f64 / (10 * queries.len()) as f64
}
#[test]
fn clustered_data_keeps_high_recall() {
// Data and queries come from the same clusters: one draw, split.
let mut vectors = clustered(3060, 24, 30, 1);
let queries = vectors.split_off(3000);
let built = HnswIndex::build_with_metric(&vectors, 8, 40, DistanceMetric::L2);
let recall = recall_at_10(&built, &vectors, &queries, 64);
assert!(recall >= 0.95, "bulk build recall@10 = {recall}");
// Incremental inserts go through the same neighbour selection.
let mut incremental = HnswIndex::new(8, 40, DistanceMetric::L2);
for v in &vectors {
incremental.insert(v.clone());
}
let recall = recall_at_10(&incremental, &vectors, &queries, 64);
assert!(recall >= 0.95, "incremental recall@10 = {recall}");
}
#[test]
fn select_neighbors_prefers_diverse_directions_and_fills_up() {
// Node at the origin. Three candidates bunched together on the right,
// one on the left. With room for two, plain closest-M would take two
// from the bunch and lose the only link leftwards.
let vectors = vec![
vec![0.0, 0.0], // 0: the node
vec![1.0, 0.0], // 1
vec![1.1, 0.0], // 2
vec![1.2, 0.0], // 3
vec![-2.0, 0.0], // 4
];
let scored: Vec<(usize, f32)> = (1..5)
.map(|i| {
(
i,
compute_distance(&vectors[0], &vectors[i], DistanceMetric::L2),
)
})
.collect();
assert_eq!(
select_neighbors(&vectors, &scored, 2, DistanceMetric::L2),
[1, 4]
);
// Spare capacity is filled with the closest rejected candidates.
assert_eq!(
select_neighbors(&vectors, &scored, 3, DistanceMetric::L2),
[1, 4, 2]
);
}
fn make_random_vectors(n: usize, dim: usize, seed: u64) -> Vec<Vec<f32>> {
let mut vectors = Vec::with_capacity(n);
let mut state = seed;