Skip to content

[diskann-quantization] Add input hints for spherical normalization and query rescaling - #1524

Open
xinyuwen2 wants to merge 2 commits into
mainfrom
wxy/sphericalQ_cosine_normalized
Open

xinyuwen2 wants to merge 2 commits into
mainfrom
wxy/sphericalQ_cosine_normalized

Conversation

@xinyuwen2

@xinyuwen2 xinyuwen2 commented Oct 28, 2026 •

Copy link
Copy Markdown
Contributor
  • Does this PR have a descriptive title that could go in our release notes?
  • Does this PR add any new dependencies?
  • Does this PR modify any existing APIs?
  • Is the change to the API backwards compatible?
  • Should this result in any changes to our documentation, either updating existing docs or adding new ones?

Reference Issues/PRs

What does this implement/fix? Briefly explain your changes.

Adds InputHints to spherical training, compression, and query-computer creation, separating input properties from the distance metric. MaybeOutOfDistribution enables internal rescaling of inner-product queries. Combined with Normalized, it also avoids computing the input norm. Stored-data compression through the quantizer interface disables query rescaling to preserve inner-product rankings.

Any other comments?

@xinyuwen2
xinyuwen2 requested review from a team and a balanced review from Copilot October 28, 2026 05:33

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 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-commenter

Codecov Comments Bot (codecov-commenter) commented Oct 28, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 97.87234% with 11 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.05%. Comparing base (65443a6) to head (bd15a80).

Files with missing lines Patch % Lines
diskann-garnet/src/quantization.rs 84.61% 8 Missing ⚠️
diskann-inmem/src/repr/quantized/spherical.rs 95.12% 2 Missing ⚠️
...-benchmark/src/index/bftree/spherical_streaming.rs 0.00% 1 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1524      +/-   ##
==========================================
+ Coverage   91.03%   91.05%   +0.01%     
==========================================
  Files         602      603       +1     
  Lines      124080   124479     +399     
==========================================
+ Hits       112955   113343     +388     
- Misses      11125    11136      +11     
Flag Coverage Δ
miri 91.05% <97.87%> (+0.01%) ⬆️
unittests 90.76% <97.85%> (+0.02%) ⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines Coverage Δ
diskann-benchmark/src/exhaustive/spherical.rs 93.53% <100.00%> (+0.02%) ⬆️
diskann-benchmark/src/index/bftree/spherical.rs 70.40% <100.00%> (ø)
diskann-benchmark/src/index/inmem/spherical.rs 53.94% <100.00%> (ø)
...-benchmark/src/index/inmem2/quantized/spherical.rs 96.90% <100.00%> (+6.58%) ⬆️
diskann-disk/src/build/builder/quantizer.rs 93.87% <100.00%> (-0.07%) ⬇️
diskann-garnet/src/provider.rs 90.36% <100.00%> (+0.15%) ⬆️
diskann-inmem/integration/index/runner.rs 92.14% <100.00%> (+0.36%) ⬆️
diskann-providers/src/index/diskann_async.rs 96.94% <100.00%> (ø)
...src/model/graph/provider/async_/inmem/spherical.rs 93.76% <100.00%> (ø)
diskann-quantization/src/spherical/hints.rs 100.00% <100.00%> (ø)
... and 7 more

... and 3 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@xinyuwen2
xinyuwen2 force-pushed the wxy/sphericalQ_cosine_normalized branch from cb8d4eb to f26002c Compare October 29, 2026 08:31
@xinyuwen2

Copy link
Copy Markdown
Contributor Author

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.

Thank you very much for pointing this out. I realize I had oversimplified this change. I've rewritten the code.

@xinyuwen2 xinyuwen2 changed the title [diskann-quantization] Support normalized cosine in spherical quantization [diskann-quantization] Add input hints for spherical normalization and query rescaling Oct 30, 2026
Separate normalization and query rescaling hints from the distance metric. Preserve default preprocessing and serialization, propagate hints through in-memory and Garnet adapters, and retain cosine normalization for converted integer inputs.

Cover mixed normalized/raw inputs, integer fallback, query layouts, rescaling, serialization, and benchmark dispatch.
@xinyuwen2
xinyuwen2 force-pushed the wxy/sphericalQ_cosine_normalized branch from 788b1e1 to bd15a80 Compare October 30, 2026 05:43

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants