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 91.49%. Comparing base (59dd048) to head (6948332).
⚠️ Report is 9 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1298      +/-   ##
==========================================
+ Coverage   90.59%   91.49%   +0.89%     
==========================================
  Files         513      516       +3     
  Lines       99091    98330     -761     
==========================================
+ Hits        89775    89963     +188     
+ Misses       9316     8367     -949     

Flag	Coverage Δ
miri	91.49% <100.00%> (+0.89%)	⬆️
unittests	91.16% <100.00%> (+0.88%)	⬆️

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

Files with missing lines	Coverage Δ
...s/src/model/graph/provider/async_/inmem/product.rs	91.83% <100.00%> (+1.83%)	⬆️

... and 263 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.

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