Expose metric-aware quantizer centroid distance - #1395
cherryrhea wants to merge 2 commits into
Conversation
There was a problem hiding this comment.
🟡 Changes recommended
The public trait change is breaking, and cosine correction and overflow handling remain unresolved.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Adds a metric-aware squared centroid-distance API to spherical quantizers.
Changes:
- Implements distance calculation and validation.
- Re-exports the new error type.
- Forwards the API through
Quantizerimplementations and adds tests.
File summaries
| File | Summary |
|---|---|
diskann-quantization/src/spherical/quantizer.rs |
Implements centroid-distance logic, validation, and tests. |
diskann-quantization/src/spherical/mod.rs |
Re-exports the new error type. |
diskann-quantization/src/spherical/iface.rs |
Adds and forwards the public trait method. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| /// Compute the squared L2 distance from `centroid` to the training centroid. | ||
| fn centroid_squared_distance( | ||
| &self, | ||
| centroid: &[f32], | ||
| ) -> Result<f64, quantizer::CentroidDistanceError>; |
| let correction = match self.metric { | ||
| SupportedMetric::Cosine => { | ||
| let norm = (FastL2Norm).evaluate(centroid); | ||
| if norm == 0.0 { 1.0 } else { 1.0 / norm } | ||
| } |
| if !correction.is_finite() { | ||
| return Err(CentroidDistanceError::NonFiniteValue); | ||
| } |
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
Thanks - a few maintainability requests on my end with respect to testing. It's okay that this is technically an API breaking change.
Also feel free to ignore Copilot's claims about divergence in the cosine correction. There appears to be an existing bug with pre-scales not being correctly rejected when Cosine is used. As long as the implementation follows SphericalQuantizer::preprocess, it's fine.
| }); | ||
| } | ||
|
|
||
| if centroid.iter().any(|value| !value.is_finite()) { |
There was a problem hiding this comment.
Instead of checking the input like this, we can check the output instead. This will be more efficient (no need to loop multiple times over the input) and more correct (overflowed accumulators during the norm computation will be detected.
There was a problem hiding this comment.
Addressed in 2b5ffd2. The separate input scan is removed; the implementation now validates the final computed distance. This avoids an extra pass and catches non-finite results caused by intermediate overflow. I also added a finite-input overflow regression using f32::MAX.
| return Err(CentroidDistanceError::NonFiniteValue); | ||
| } | ||
|
|
||
| let correction = f64::from(correction); |
There was a problem hiding this comment.
What's the motivation for using f64 here? Are there observed accuracy issues with f32?
There was a problem hiding this comment.
Addressed in 2b5ffd2. I did not have measured evidence requiring f64, so the method now follows the existing preprocessing numeric contract and computes/returns f32. The downstream C FFI can widen the result at its ABI boundary where it exposes a double.
| } | ||
|
|
||
| #[test] | ||
| fn centroid_squared_distance_uses_metric_correction() { |
There was a problem hiding this comment.
I'm worried that a few spot-checks for behavior will miss the larger contract requirement that preprocess and centroid_squared_distance should agree.
I'd suggest hooking into the existing test_l2/test_ip/test_cosine functions and checking that centroid_squared_distance is within some small tolerance of Preprocessed.shifted_norm. This will provide more coverage (more tests cases) and ensure that if one path gets updated, the other needs to as well.
There was a problem hiding this comment.
Addressed in 2b5ffd2. I removed the standalone metric spot test and added a shared invariant assertion to the existing test_l2, test_ip, and test_cosine helpers. Every exercised trial now checks centroid_squared_distance(v) against preprocess(v).shifted_norm squared within f32 tolerance, so future preprocessing changes must remain synchronized.
|
cherryrhea please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
Summary
SphericalQuantizer::preprocessQuantizertrait implementationsMotivation
LibVector needs to compare a current data centroid with the centroid persisted in a trained quantizer. Keeping this computation in
diskann-quantizationavoids reconstructing private quantizer state in an external FFI layer and preserves metric-specific preprocessing semantics, especially for cosine.Validation
cargo test -p diskann-quantizationcargo fmt --all -- --checkpassesgit diff --checkpasses