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.

[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.

[Mark Hildebrand (hildebrandmw)]

Thanks!