Fix normalized cosine distance consistency for PQ graphs

[juchen-ms (partychen)]

Problem

For Product-PQ with CosineNormalized, graph-search queries used raw squared L2 while graph pruning used cosine distance.

Vamana pruning compares distance_ik / distance_jk. For normalized vectors, cosine distance is half squared L2, so mixing the two approximately doubled this ratio. This made alpha = 1.2 behave like alpha = 0.6, causing over-pruning and poor recall.

Fix

Use raw squared L2 for Product-PQ Hybrid and Quantized pruning, matching the existing PQ graph-search query approximation.

The change is scoped to Product-PQ pruning. Search behavior, public PQ distance APIs, other metrics, and other quantization strategies are unchanged.

Validation

Regression tests cover all Hybrid operand combinations and the Quantized PQ/PQ pruning path.

The final commit was benchmarked on the first 100,000 BigANN SIFT base vectors and first 1,000 queries, converted to unit-normalized f32 as required by CosineNormalized. Exact ground truth was recomputed for this subset. Configuration: 50 PQ chunks, max_fp_vecs_per_prune = 48, max_degree = 64, l_build = 100, alpha = 1.2, and search_l = 100.

Search results

Build	Recall@10	Mean comparisons	Mean hops	Comparisons / hop
Before	0.0455	1,206.963	119.609	10.09
After	0.9999	3,205.906	103.618	30.94

Graph structure

Build	Graph size	Out-edges	Mean degree	P50	P90	P99	Degree <=4	Degree >=33
Before	8.99 MiB	2,256,586	22.565	22	30	40	0.002%	6.161%
After	22.02 MiB	5,672,071	56.720	57	77	83	0.001%	94.166%

Checks:

• cargo test -p diskann-providers --lib
• cargo clippy --workspace --all-targets -- -D warnings

[Copilot]

Pull request overview

This PR makes Metric::CosineNormalized behavior internally consistent across all PQ-involved distance paths by using 0.5 * squared_l2 for query↔PQ, full↔PQ, and PQ↔PQ comparisons. This prevents pruning/search from comparing incompatible distance scales when PQ vectors are involved.

Changes:

• Introduces a shared scaling constant and applies it to FixedChunkPQTable CosineNormalized distances (query↔PQ and PQ↔PQ).
• Adds a scaled lookup-table construction path for CosineNormalized so query-time evaluation remains a table lookup.
• Adds regression tests to ensure cross-computer consistency and correct hybrid (full/quant) dispatch behavior.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
diskann-providers/src/model/pq/fixed_chunk_pq_table.rs	Defines the scale constant and applies scaled squared-L2 for CosineNormalized in direct PQ distance paths.
diskann-providers/src/model/pq/distance/l2.rs	Adds new_scaled to scale precomputed L2 lookup tables for CosineNormalized preprocessing.
diskann-providers/src/model/pq/distance/dynamic.rs	Wires CosineNormalized to the scaled L2 preprocessing and updates QQ dispatch + tests.
diskann-providers/src/model/graph/provider/async_/distances.rs	Adds a regression test ensuring hybrid full/quant and quant/quant CosineNormalized paths use the scaled L2 definition.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.61%. Comparing base (d1d4869) to head (6e0a3fb).
⚠️ Report is 5 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1298      +/-   ##
==========================================
+ Coverage   92.30%   92.61%   +0.30%     
==========================================
  Files         517      522       +5     
  Lines       98520    99531    +1011     
==========================================
+ Hits        90943    92183    +1240     
+ Misses       7577     7348     -229     

Flag	Coverage Δ
miri	92.61% <100.00%> (+0.30%)	⬆️
unittests	92.58% <100.00%> (+0.30%)	⬆️

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

Files with missing lines	Coverage Δ
diskann-providers/src/index/diskann_async.rs	96.65% <100.00%> (-0.10%)	⬇️
...s/src/model/graph/provider/async_/inmem/product.rs	91.07% <100.00%> (+1.07%)	⬆️

... and 36 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (4)

diskann-providers/src/model/pq/distance/dynamic.rs:505

• The relative tolerance here (6.3e-7) is tighter than other SIMD-vs-scalar distance tests in this crate (commonly 1e-6). Relaxing to 1e-6 would make this regression test less likely to be flaky across platforms.

            assert_relative_eq!(
                cosine_normalized.evaluate_similarity(&*code0, &*code1),
                expected,
                max_relative = 6.3e-7,
            );

diskann-providers/src/model/graph/provider/async_/distances.rs:212

• This assertion uses a very tight relative tolerance (1e-7). Using 1e-6 would better match other SIMD-vs-scalar comparisons in diskann-providers and reduce the risk of cross-platform FP flakiness.

            assert_relative_eq!(quant_quant, expected_quant_quant, max_relative = 1.0e-7);

diskann-providers/src/model/graph/provider/async_/distances.rs:203

• This assertion uses a very tight relative tolerance (1e-7). Using 1e-6 would better match other SIMD-vs-scalar comparisons in diskann-providers and reduce the risk of cross-platform FP flakiness.

This issue also appears on line 212 of the same file.

            assert_relative_eq!(full_quant, expected_full_quant, max_relative = 1.0e-7);

diskann-providers/src/model/pq/distance/dynamic.rs:439

• These new floating-point assertions use a tighter relative tolerance (5e-7) than similar SIMD-vs-scalar comparisons elsewhere in this crate (often 1e-6). Consider relaxing to 1e-6 to reduce cross-arch / compiler flakiness.

This issue also appears on line 501 of the same file.

        assert_relative_eq!(query_distance, expected, max_relative = 5.0e-7);
        assert_relative_eq!(random_access_distance, expected, max_relative = 5.0e-7);
        assert_relative_eq!(
            query_distance,
            random_access_distance,

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (1)

diskann-providers/src/model/pq/distance/dynamic.rs:59

• The doc comment claims that for non-normalized operands the 0.5 * squared_l2 approximation differs from normalized cosine distance only by a positive factor and therefore preserves candidate ordering. That relationship is not generally true when norms vary; the difference is not just a constant scale, and ordering can change. Please adjust the comment to avoid stating an incorrect guarantee.

    /// In other words, half the squared L2 distance equals normalized cosine distance when
    /// both operands are normalized, and the two differ by a positive factor otherwise, so
    /// candidate ordering is preserved either way.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)

diskann-providers/src/model/graph/provider/async_/distances.rs:187

• The hybrid regression test uses non-unit full vectors ([1, 0, 0, 2], [2, 0, 0, 1]), but Metric::CosineNormalized is documented/implemented under a unit-norm assumption. Using unit vectors here would make the test’s intent clearer (CosineNormalized == 0.5*squared-L2) and avoid locking in behavior that only matches scaled-L2 for non-normalized inputs.

            let table = FixedChunkPQTable::new(
                4,
                vec![1.0, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0].into(),
                vec![0, 2, 4].into(),
            )

diskann-providers/src/model/graph/provider/async_/distances.rs:141

• HybridComputer::new maps Metric::CosineNormalized full/full comparisons onto Metric::L2 with a 0.5 scale. This contradicts the PR description’s stated scope that full/full comparisons remain on the native CosineNormalized implementation, and it can also change behavior when full vectors are not perfectly unit-norm. Consider keeping the full-precision path on Metric::CosineNormalized (scale 1.0) and relying on the PQ side’s scaled-L2 approximation for compatibility.

This issue also appears on line 183 of the same file.

        pub fn new(quant: pq::distance::DistanceComputer<'a>, dim: Option<usize>) -> Self {
            let (full_metric, full_scale) = match quant.metric() {
                Metric::CosineNormalized => (Metric::L2, pq::COSINE_NORMALIZED_L2_SCALE),
                metric => (metric, 1.0),
            };

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

[juchen-ms (partychen)]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (2)
diskann-providers/src/model/graph/provider/async_/distances.rs:187

• The hybrid regression test uses non-unit full vectors ([1, 0, 0, 2], [2, 0, 0, 1]), but Metric::CosineNormalized is documented/implemented under a unit-norm assumption. Using unit vectors here would make the test’s intent clearer (CosineNormalized == 0.5*squared-L2) and avoid locking in behavior that only matches scaled-L2 for non-normalized inputs.

            let table = FixedChunkPQTable::new(
                4,
                vec![1.0, 0.0, 0.0, 1.0, 2.0, 0.0, 0.0, 2.0].into(),
                vec![0, 2, 4].into(),
            )

diskann-providers/src/model/graph/provider/async_/distances.rs:141

• HybridComputer::new maps Metric::CosineNormalized full/full comparisons onto Metric::L2 with a 0.5 scale. This contradicts the PR description’s stated scope that full/full comparisons remain on the native CosineNormalized implementation, and it can also change behavior when full vectors are not perfectly unit-norm. Consider keeping the full-precision path on Metric::CosineNormalized (scale 1.0) and relying on the PQ side’s scaled-L2 approximation for compatibility.

This issue also appears on line 183 of the same file.

        pub fn new(quant: pq::distance::DistanceComputer<'a>, dim: Option<usize>) -> Self {
            let (full_metric, full_scale) = match quant.metric() {
                Metric::CosineNormalized => (Metric::L2, pq::COSINE_NORMALIZED_L2_SCALE),
                metric => (metric, 1.0),
            };

Kept full/full scaled-L2 intentionally because Hybrid pruning mixes full/full, full/PQ, and PQ/PQ comparisons. Restoring native CosineNormalized only for full/full would reintroduce incompatible distance definitions, especially for u8/i8. The tests now separate and document unit-vector semantics and integer consistency.

[Wei Wu (wuw92)]

Using raw squared L2 consistently may be valid while these values are used only for internal ordering and pruning, However, the resulting values no longer follow the standard CosineNormalized contract if they are returned to or otherwise consumed by a caller as distances.

The underlying issue seems to be that the Product-PQ CosineNormalized paths already have different contracts:

• query/PQ returns raw squared L2;
• full/full uses the native implementation;
• full/PQ and PQ/PQ use cosine over reconstructed PQ vectors.

Overriding CosineNormalized with L2 in the pruning strategy hides the existing contract mismatch and also changes u8/i8 from angular to Euclidean behavior.

Aditya Krishnan (@arkrishn94) any thoughts on these contracts?

[juchen-ms (partychen)]

Agreed on the contract distinction. The remapped computer here is module-private and used only by Vamana pruning; it is not returned as a CosineNormalized distance, and the public PQ distance behavior remains unchanged.

The narrow intent is to make distance_jk use the same approximation as the search-pool distance_ik, since they participate in the same pruning ratio. I agree this does not resolve the broader Product-PQ contract inconsistency, particularly for u8/i8. If those inputs should preserve angular semantics, the query/PQ path and pruning path would need to be addressed together. I’m happy to follow the contract decision here.

[Magdalen Dobson Manohar (magdalendobson)]

Thank you for finding the bug and contributing this! I was wondering if you could add an integration test that would have caught this, as it seems like a fault in our testing that this wasn't found earlier?

[juchen-ms (partychen)]

Magdalen Dobson Manohar (@magdalendobson) Added in f140b00. I extended the existing SIFT build-and-search coverage with normalized CosineNormalized cases for both the Hybrid Product-PQ path and the quant-only Quantized path. These tests build real indexes and validate searches against exact ground truth. I also verified that the new cases fail against the pre-fix code and pass with this change. Thanks for calling this out!

[Magdalen Dobson Manohar (magdalendobson)]

Thanks for adding testing. Looks good to me now. Request that you take a second look at all the quantizer x metric combinations and make sure all of them are tested properly inside the index test, but not blocking.

[Aditya Krishnan (arkrishn94)]

Hmm, I'm a little worried about adding this kind of bespoke hook to solve the problem here for the following reasons -

• It is not clear for a caller of this function (or user of the underlying distance computer created), what the output metric can or cannot be used for. A caller could mistakenly assume that the underlying metric is L2 when the metric used to create the computer is CosineNormalized for e.g and end up doing incorrect stuff.
• Another implementor of PQ-based prune has to know to call this bespoke function before creating the computer. This kind of implicit knowledge is almost certainly likely to lead to bugs.

It seems to me the problem is that the behavior of the VTable and the QueryComputer dispatch to different underlying estimators for CosineNormalized: the latter seems to do cosine and the former seems to do l2?

Is it possible to fix this by dispatching correctly there? Mark Hildebrand (@hildebrandmw) to correct me if I'm wrong here.

[juchen-ms (partychen)]

Good point. However, changing only the VTable dispatch would not fully resolve the inconsistency: Hybrid full/full comparisons bypass VTable and would still use native CosineNormalized, which is half the squared-L2 scale.
I’ll instead remap the metric in FixedChunkPQTable::create, so the entire Product-PQ graph-build path uses L2 consistently without a pruning-specific hook.

[Mark Hildebrand (hildebrandmw)]

Thanks!