feat(index): support Hamming distance in streaming coreset k-means - #8610
feat(index): support Hamming distance in streaming coreset k-means#8610a-erofeev wants to merge 5 commits into
Conversation
462553b to
bc8d8a0
Compare
bc8d8a0 to
6a4591d
Compare
|
Added a few tests for Hamming-distance |
|
@Xuanwo Hi! please review |
wjones127
left a comment
There was a problem hiding this comment.
Reviewed this PR first, so can apply some of the feedback to the other PR. I think this seems like a nice direction, but there are a few minor issues I'd like to see fixed. Plus the module structure could be better. Please address those and I'm happy to merge these PRs.
| #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)] | ||
| enum CounterWidth { | ||
| U8, | ||
| U16, | ||
| U32, | ||
| U64, | ||
| } |
There was a problem hiding this comment.
suggestion: this module could be split up further to make it easier to read. Like CounterWidth and CompactCounters could go into a counter submodule.
| #[async_trait] | ||
| impl StreamingCoresetAlgorithm for FloatCoresetAlgorithm<'_> { |
There was a problem hiding this comment.
suggestion: Similarly here, the module is getting very large. Perhaps we can create a module structure:
- ivf
|-- coreset_algo
|-- float_coreset.rs
|-- hamming_coreset.rs
| } else { | ||
| (loss - last_loss).abs() / last_loss.abs() | ||
| }; | ||
| if kmeans_has_converged(loss, last_loss, params.tolerance) { |
There was a problem hiding this comment.
issue (non-blocking): The arguments are swapped relative to the parameter names: kmeans_has_converged(previous_loss, loss, tol) is called as (loss, last_loss, ...), so the tolerance scales by last_loss.abs(). I believe that's deliberate — it preserves the old tolerance * last_loss scale — but test_convergence_scale_is_sign_independent pins the (previous, current) order, so the next reader will "fix" this call site and silently change the convergence scale. Either swap it (and accept scaling by loss) or add a one-line comment saying the scale reference is intentionally the previous loss.
| last_loss, | ||
| ); | ||
| if (loss - last_loss).abs() < params.tolerance * last_loss { | ||
| let relative_loss_diff = if last_loss == 0.0 { |
There was a problem hiding this comment.
nitpick: relative_loss_diff is computed on every iteration but only read inside the if, and it's a fourth restatement of the convergence arithmetic — move it inside the branch so the log stays the only place that needs it.
| ))); | ||
| } | ||
| let dimension = data.value_length() as usize; | ||
| let sample_rate = data.len().div_ceil(local_k).max(1); |
There was a problem hiding this comment.
issue (non-blocking): div_ceil(local_k) panics when local_k is 0, which DotMetricPolicy::local_coreset_k can produce (see ivf.rs:3600) and nothing here rules out for the Hamming path either.
| let sample_rate = data.len().div_ceil(local_k).max(1); | |
| let local_k = local_k.max(1); | |
| let sample_rate = data.len().div_ceil(local_k).max(1); |
|
|
||
| let mut clusters = heap.into_vec(); | ||
| clusters.sort_by_key(|cluster| cluster.id); | ||
| while clusters.len() < target_k { |
There was a problem hiding this comment.
issue (non-blocking): This padding clones a cluster verbatim, producing byte-identical centroids — duplicate IVF partitions that lose every assignment tie to the lower index and stay empty. Worse, max_by_key returns the last maximum, so each clone becomes the next iteration's pick and all padding entries end up identical. Binary data with near-duplicate summaries makes this reachable. At minimum add a test covering target_k greater than the number of splittable summaries and decide whether to error or perturb instead.
| [cluster_id * coreset.dimension..(cluster_id + 1) * coreset.dimension]; | ||
| Ok((cluster_id, coreset.cost(row_index, centroid)?)) | ||
| }) | ||
| .collect::<Result<Vec<_>>>()? |
There was a problem hiding this comment.
suggestion: This allocates a Vec of k tuples for every summary row just to take a min — at k = 16_384 in refine_weighted that's the dominant cost of the 1,670s training run in your table. A plain loop tracking the best (cluster_id, cost) keeps the same tie-break (lowest id wins) with zero allocation; parallelising the row loop with rayon, as the float path gets for free via compute_membership_and_distances, would be the next step.
| } | ||
|
|
||
| let bit_dimension = self.bit_dimension(); | ||
| let split_bit = (0..bit_dimension) |
There was a problem hiding this comment.
suggestion: max_by calls this comparator ~bit_dimension times and each call recomputes both operands' variance over every row — roughly 2× the necessary work. Compute the per-bit variances once into a Vec<f64> in a single pass over rows, then max_by over that.
| let diff = left - right; | ||
| diff * diff | ||
| }) | ||
| .sum::<f32>() as f64 |
There was a problem hiding this comment.
suggestion: Accumulating the squared differences in f32 loses precision and can overflow to infinity on high-dimensional or large-magnitude vectors — which is precisely the condition the new validate_finite_kmeans_distance turns into a hard error, so some of those failures would be self-inflicted. Accumulate in f64.
| .sum::<f32>() as f64 | |
| .map(|squared| squared as f64) | |
| .sum::<f64>() |
| @@ -1872,6 +1890,64 @@ mod tests { | |||
| }); | |||
There was a problem hiding this comment.
praise: These four tests are the right ones — test_dot_membership_preserves_negative_distances_and_radius and scaled_l2_does_not_stop_early each pin a distinct bug the old code had, and they'd both fail on main. Nice.
b2ffb2f to
2336a83
Compare
2336a83 to
8be4aa9
Compare
There was a problem hiding this comment.
✅ Gate recommendation: approve.
The prior Hamming blocker is fixed: distinct sampled representatives now fill coreset modes that collapse during hierarchical and weighted training, while genuinely insufficient diversity still returns an error. The 257-partition regression and full streaming Hamming test set pass with centroid uniqueness preserved before and after raw refinement.
wjones127
left a comment
There was a problem hiding this comment.
Looks good now. Thanks!
Added support of Hamming distance in streaming coreset k-means.
This PR depends on PR #8463
Benchmark report.
For k=1024 i used generated binary vector dataset with 1M 256bit vectors, not random but with clusters simulation (generate 1024 random vectors (prototypes), for every new vector - take protorype, copy it, flip every bit with probability p=0.05)
For k=16384 i used generated binary vector dataset with 100M 256bit vectors (16384 prototypes)
Percentages are relative to LanceStream stream=64, coreset=16, prefetch=1, refine=0
Exact loss was evaluated over 262,144 vectors as:
Host:
In these tests we got only 2 times less memory consumption (in better case), this because of coreset for Hamming distance consumes more memory comparing to Dot/L2
As an option to decrease memory consumption we can decrease streaming_coreset_rate.