Repository navigation
Conversation
There was a problem hiding this comment.
🟢 Approval recommended
No active review issues are reported for the normalized-cosine spherical quantization changes.
No open issues 🎉
What changed in this PR
Adds CosineNormalized support to spherical quantization for caller-normalized vectors, integrating it with serialization, in-memory reranking, and benchmark dispatch.
Changes:
- Computes centroids directly from supplied vectors and skips input normalization and pre-scaling for the new mode.
- Reuses compensated cosine distance computation and extends metric serialization.
- Adds coverage for normalized-cosine behavior, serialization, and benchmark selection.
| File | Description |
|---|---|
diskann-quantization/src/utils.rs |
Adds unnormalized column means through a shared implementation and extends corner-case tests. |
diskann-quantization/src/spherical/vectors.rs |
Adds the metric variant, conversions, and normalization contract documentation. |
diskann-quantization/src/spherical/quantizer.rs |
Implements normalized-cosine training and preprocessing, with behavioral and serialization tests. |
diskann-quantization/src/spherical/iface.rs |
Routes the new metric to compensated cosine and extends plan tests. |
diskann-quantization/src/flatbuffers/spherical/supported_metric_generated.rs |
Adds generated FlatBuffers support for the new metric value. |
diskann-quantization/schemas/spherical_v001.fbs |
Appends CosineNormalized to the serialized metric enum. |
diskann-inmem/src/repr/internal/quantization/mod.rs |
Maps normalized cosine to cosine reranking and tests the mapping. |
diskann-benchmark/src/index/inmem2.rs |
Enables normalized-cosine spherical benchmarks and tests dispatch matching. |
Review effort: Balanced
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1524 +/- ##
==========================================
- Coverage 92.16% 91.08% -1.09%
==========================================
Files 592 592
Lines 123582 123663 +81
==========================================
- Hits 113899 112635 -1264
- Misses 9683 11028 +1345
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
Mark Hildebrand (hildebrandmw)
left a comment
There was a problem hiding this comment.
I'm generally not a fan of CosineNormalized. Not because I disagree with what it's trying to accomplish, but I disagree with semantics. It promises that data is normalized at far too low of a level, which interacts poorly when non-floating-point data is involved.
For example, this PR makes it possible to accidentally associate CosineNormalized with integer data, which will end pretty badly (also note, that TryFrom conversion should have been removed since this change makes the conversion lossless, which would have found call sites that were affected).
The optional rescaling also falls into this bucket of somewhat disjoint semantics.
I counter propose an approach that ties both of these together. For example, we can have a InputHints struct provided to compress/distance computer creation that could have hints:
MaybeOutOfDistribution: Input is maybe out of distribution. Inner product could recognize this and do rescaling internally (avoiding callers needing to manually rescale).Normalized: Hint that the input data is normalized. Cosine could avoid normalization, and inner-product could avoid computing the norm in the first place.
This also provides a natural place for future hints.
Reference Issues/PRs
What does this implement/fix? Briefly explain your changes.
Adds CosineNormalized support to spherical quantization. Training computes the centroid directly from the supplied vectors, and compression skips input normalization. The new mode ignores pre-scaling and reuses the existing compensated cosine distance computation.
Any other comments?