Remove centroid from most PQ interfaces.

[hildebrandmw]

Remove the "centroid" from most PQ interfaces. This centroid is/was the dataset centroid and served to shift a vector before compressing or populating the distance lookup table. Here are the problems:

1. This kind of shifting only works when computing L2 distances. This kind of shift does not preserve inner products or cosine similarity. So it is a little bit of a foot-gun and was already not being applied when these distances were being computed, leading to a somewhat confusing scenario where a centroid can be provided to the FixedChunkPQTable constructor and then ignored.

2. When L2 distances are being used - application of the centroid can be done implicitly by adding it to the PQ pivots rather than subtracting it from each vector that gets processed. In this way, it always gets applied.

It's possible this has a small effect on training by conditioning the data a little better, but when we've applied this empirically in the past, it has not made a difference.

Carrying around this centroid is making it awkward to cleanly move PQ to diskann-quantization, so this PR takes the first steps towards removing it.

Logical Flow/Review Order

• diskann_providers/src/model/pq/generate_pivot_arguments.rs: translate_to_center is removed from GeneratePivotArguments.

• diskann_providers/src/model/pq/pq_construction.rs: The function generate_pq_pivots gains a legacy_center_data boolean argument to preserve the behavior of diskann-disk. This is necessary because tests in diskann-disk check against hard-coded pivots. Not centering the data results in different PQ codes due to slightly different floating point calculations, breaking these tests.

  Data translation is completely removed from the _membuf APIs.

• diskann_providers/src/model/pq/fixed_chunk_pq_table.rs: Remove centroids altogether from FixedChunkPQTable.

• diskann-providers/src/storage/pq_storage.rs: This is where backwards compatibility comes in. We cannot change the layout of the storage file and we also need to support previously generated files. To that end, load_pq_pivots_bin is updated to check if the loaded centroid is non-zero. If so, it is accumulated into the pivots to restore the original numerical behavior. Several tests are added to ensure this behavior.

  For saving, when no centroid is available, it can be passed to write_pivot_data as an Option::None and will then be inserted as zeros into the saved file.

• diskann-disk/src/search/pq/pq_dataset.rs and diskann-disk/src/search/pq/quantizer_preprocess.rs: Preprocessing can now be greatly simplified. Because the FixedChunkPQTable no longer has centroids, the TransposedTable (offering much faster pre-processing times) can always be used.

The rest of the changes involve updating the construction sites of FixedChunkPQTable and invoking the new saving interface.

Backwards Compatibility

There exist PQ schemas in the wild that were generated with centroids and we still need to correctly process those.

• diskann-disk: PQ schemas previously generated I believe have been saved with PQStorage. Thus, reloading into a FixedChunkPQTable with PQStorage::load_pq_pivots_bin will merge the centroid with the PQ pivots.
• Users of the _membuf API: Backwards compatibility is handled by the caller (the only real user I know of uses this via FFI, and I will take responsibility for updating that). Again, this can be done by simply folding the centroid into the pivot data, or manually doing the centering. Both of which are quite easy to do.

[Copilot]

Pull request overview

This PR removes explicit dataset-centroid handling from most PQ interfaces, shifts compatibility to the storage/load path by folding legacy centroids into pivots, and simplifies downstream consumers so they can treat PQ tables as centroid-free. In the broader codebase, this moves PQ closer to the diskann-quantization abstractions and removes L2-only centering behavior from general-purpose interfaces.

Changes:

• Removes centroid fields/parameters from FixedChunkPQTable, GeneratePivotArguments, and the _membuf PQ construction/compression APIs.
• Updates PQ storage read/write paths to preserve the on-disk layout and recover legacy behavior by folding non-zero stored centroids into pivots on load.
• Simplifies disk PQ consumers to always use TransposedTable, and updates construction/save call sites across providers, disk builder code, tools, and benchmarks.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
diskann-tools/src/utils/build_pq.rs	Updates tool-side PQ pivot generation call to the new generate_pq_pivots signature.
diskann-providers/src/storage/pq_storage.rs	Adds optional-centroid write path, legacy-centroid folding on load, and storage round-trip tests.
diskann-providers/src/model/pq/pq_construction.rs	Removes centroid handling from in-memory PQ APIs and adds legacy-centering toggle to file-based pivot generation.
diskann-providers/src/model/pq/generate_pivot_arguments.rs	Removes centroid-centering state from PQ generation arguments.
diskann-providers/src/model/pq/fixed_chunk_pq_table.rs	Drops centroid storage/processing from the PQ table type and updates related tests/helpers.
diskann-providers/src/model/pq/distance/test_utils.rs	Updates PQ test helpers to construct centroid-free tables.
diskann-providers/src/model/pq/distance/l2.rs	Stops query preprocessing from subtracting centroids before lookup-table generation.
diskann-providers/src/model/graph/provider/async_/memory_quant_vector_provider.rs	Saves PQ pivots without a centroid payload and updates test construction.
diskann-providers/src/model/graph/provider/async_/inmem/provider.rs	Refreshes public doc examples for the new FixedChunkPQTable::new shape.
diskann-providers/src/model/graph/provider/async_/fast_memory_quant_vector_provider.rs	Switches async fast in-memory PQ provider save/test paths to centroid-free tables.
diskann-providers/src/model/graph/provider/async_/experimental/multi_pq_async.rs	Updates experimental multi-PQ compression call sites to the new membuf API.
diskann-providers/src/model/graph/provider/async_/bf_tree/quant_vector_provider.rs	Adjusts BF-tree quant provider tests to the centroid-free table constructor.
diskann-providers/src/model/graph/provider/async_/bf_tree/provider.rs	Saves BF-tree PQ pivots without centroids and updates related tests.
diskann-providers/src/index/diskann_async.rs	Removes centroid buffers from async index PQ training and test setup.
diskann-disk/src/storage/quant/pq/pq_generation.rs	Threads the new legacy-centering flag through disk PQ generation.
diskann-disk/src/storage/quant/pq/pq_dataset.rs	Replaces the PQTable enum with always-transposed storage in PQData.
diskann-disk/src/storage/quant/pq/mod.rs	Removes the now-obsolete PQTable re-export.
diskann-disk/src/storage/quant/mod.rs	Stops re-exporting PQTable from the public quant module.
diskann-disk/src/search/pq/quantizer_preprocess.rs	Simplifies PQ preprocessing to a single TransposedTable path.
diskann-disk/src/build/builder/quantizer.rs	Saves builder-generated PQ pivots without a centroid payload.
diskann-benchmark/src/backend/exhaustive/product.rs	Updates benchmark-side PQ table construction to the centroid-free constructor.

Comments suppressed due to low confidence (4)

diskann-providers/src/model/pq/pq_construction.rs:166

• This public _membuf entry point no longer exposes the centroid output buffer, which breaks existing callers that still use the 0.50.x signature. The PR description calls out caller-side migration, but the crate versioning policy says these packages obey SemVer, so this signature change needs a shim/deprecation path or a major release.

pub fn generate_pq_pivots_from_membuf<T: Copy + Into<f32>>(
    parameters: &GeneratePivotArguments,
    train_data_slice: &[T],
    offsets: &mut [usize],
    full_pivot_data: &mut [f32],
    rng: &mut (impl Rng + ?Sized),
    cancellation_token: &mut bool,
    pool: RayonThreadPoolRef<'_>,

diskann-providers/src/model/pq/pq_construction.rs:166

• The doc comment for this public function still describes the removed make_zero_mean/centroid parameters, so it no longer matches the actual API. Anyone reading the generated docs will be told to pass buffers that this signature no longer accepts.

pub fn generate_pq_pivots_from_membuf<T: Copy + Into<f32>>(
    parameters: &GeneratePivotArguments,
    train_data_slice: &[T],
    offsets: &mut [usize],
    full_pivot_data: &mut [f32],
    rng: &mut (impl Rng + ?Sized),
    cancellation_token: &mut bool,
    pool: RayonThreadPoolRef<'_>,

diskann-providers/src/model/pq/pq_construction.rs:80

• Adding the new legacy_center_data parameter changes the signature of a public, re-exported function. That is another SemVer break for diskann-providers consumers on 0.50.x, so this needs a compatibility wrapper or a major-version plan instead of replacing the existing API in place.

pub fn generate_pq_pivots<Storage, Random>(
    parameters: GeneratePivotArguments,
    legacy_center_data: bool,
    train_data: &mut [f32],
    pq_storage: &PQStorage,
    storage_provider: &Storage,
    random_provider: RandomProvider<Random>,
    pool: RayonThreadPoolRef<'_>,

diskann-providers/src/storage/pq_storage.rs:98

• Changing write_pivot_data from taking &[f32] to Option<&[f32]> is a source-breaking public API change for diskann-providers consumers. Because the workspace is still on a SemVer-governed 0.50.x line, this needs a compatibility layer or a coordinated major-version bump rather than replacing the signature in place.

    pub fn write_pivot_data<Storage>(
        &self,
        full_pivot_data: &[f32],
        centroid: Option<&[f32]>,
        chunk_offsets: &[usize],
        num_centers: usize,
        dim: usize,
        storage_provider: &Storage,
    ) -> ANNResult<()>

---

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

[codecov-commenter]

Codecov Report

❌ Patch coverage is 97.05882% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.50%. Comparing base (d00cac0) to head (c6219c6).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
diskann-providers/src/storage/pq_storage.rs	93.90%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1010      +/-   ##
==========================================
- Coverage   89.51%   89.50%   -0.02%     
==========================================
  Files         461      459       -2     
  Lines       85920    85549     -371     
==========================================
- Hits        76914    76567     -347     
+ Misses       9006     8982      -24     

Flag	Coverage Δ
miri	89.50% <97.05%> (-0.02%)	⬇️
unittests	89.12% <97.05%> (-0.02%)	⬇️

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

Files with missing lines	Coverage Δ
...iskann-benchmark/src/backend/exhaustive/product.rs	100.00% <ø> (ø)
diskann-disk/src/build/builder/quantizer.rs	90.52% <100.00%> (ø)
diskann-disk/src/search/pq/quantizer_preprocess.rs	97.22% <100.00%> (+9.72%)	⬆️
diskann-disk/src/storage/quant/pq/pq_dataset.rs	96.72% <100.00%> (-0.62%)	⬇️
diskann-disk/src/storage/quant/pq/pq_generation.rs	93.22% <100.00%> (ø)
diskann-providers/src/index/diskann_async.rs	95.99% <100.00%> (-0.02%)	⬇️
...aph/provider/async_/experimental/multi_pq_async.rs	96.38% <ø> (-0.02%)	⬇️
...ovider/async_/fast_memory_quant_vector_provider.rs	98.44% <100.00%> (-0.02%)	⬇️
.../src/model/graph/provider/async_/inmem/provider.rs	90.60% <ø> (ø)
...ph/provider/async_/memory_quant_vector_provider.rs	98.26% <100.00%> (-0.01%)	⬇️
... and 7 more

... and 13 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[harsha-simhadri]

ok with this. could you please run on a few medium size datasets and check recall before merging. Thanks.

[arkrishn94]

Thanks for doing this Mark, lgtm. The only comment I had was on supporting the backward compatibility of generate_pq_pivots for diskann-disk tests. If you can create a small issue/sub-issue on this, we can close it in a quick follow up.

[hildebrandmw]

ok with this. could you please run on a few medium size datasets and check recall before merging. Thanks.

Okay, I did a few runs on Sift-1M and OpenAI-100K (3k dim). Since the latter is already zero centered, I did kind a pathological thing and added 128 to each dimension and scaled each dimension by 256. Since the average magnitude of an open-ai vector is around 0.014, there's just not very much signal if left as-is. I figure this is probably a worst case for any kind of reasonable vector.

Below is an (AI generated, human verified) summary of the results:

• Max degree: 32, L-build: 50, Alpha: 1.2
• Start point: Medoid, Backedge ratio: 1
• Build threads: 10, Single insert mode
• Full precision search: disabled, Max FP vecs: 0

Search Configuration (all jobs)

• KNN: 10, Threads: 10
• L-search swept: 10–100 (step 10)

Job	Dataset	Type	Vectors	Dims	PQ Chunks
1	sift-1m	uint8	1M	128	32
2	sift-1m	uint8	1M	128	16
3	openai-v3 base-100k	float32	100K	3072	384
4	openai-v3 base-100k	float32	100K	3072	300
5	openai-v3 base-100k.shifted	float32	100K	3072	384
6	openai-v3 base-100k.shifted	float32	100K	3072	300

JSON

{
  "search_directories": [
    "<snip>"
  ],
  "output_directory": null,
  "jobs": [
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "sift-1m.gt.bin",
            "num_threads": [
              10
            ],
            "queries": "sift-1m.queries.u8.bin",
            "reps": 1,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "sift-1m.u8.bin",
            "data_type": "uint8",
            "distance": "squared_l2",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 10,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 32,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    },
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "sift-1m.gt.bin",
            "num_threads": [
              10
            ],
            "queries": "sift-1m.queries.u8.bin",
            "reps": 1,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "sift-1m.u8.bin",
            "data_type": "uint8",
            "distance": "squared_l2",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 10,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 16,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    },
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "gt-100k.bin",
            "num_threads": [
              10
            ],
            "queries": "query.bin",
            "reps": 1,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "base-100k.bin",
            "data_type": "float32",
            "distance": "squared_l2",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 10,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 384,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    },
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "gt-100k.bin",
            "num_threads": [
              10
            ],
            "queries": "query.bin",
            "reps": 1,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "base-100k.bin",
            "data_type": "float32",
            "distance": "squared_l2",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 10,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 300,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    },
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "gt-100k.bin",
            "num_threads": [
              10
            ],
            "queries": "query-100k.shifted.bin",
            "reps": 1,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "base-100k.shifted.bin",
            "data_type": "float32",
            "distance": "squared_l2",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 10,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 384,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    },
    {
      "type": "graph-index-build-pq",
      "content": {
        "index_operation": {
          "search_phase": {
            "groundtruth": "gt-100k.bin",
            "num_threads": [
              10
            ],
            "queries": "query-100k.shifted.bin",
            "reps": 1,
            "runs": [
              {
                "recall_k": 10,
                "search_l": [
                  10,
                  20,
                  30,
                  40,
                  50,
                  60,
                  70,
                  80,
                  90,
                  100
                ],
                "search_n": 10
              }
            ],
            "search-type": "topk"
          },
          "source": {
            "alpha": 1.2000000476837158,
            "backedge_ratio": 1.0,
            "data": "base-100k.shifted.bin",
            "data_type": "float32",
            "distance": "squared_l2",
            "index-source": "Build",
            "insert_retry": null,
            "l_build": 50,
            "max_degree": 32,
            "multi_insert": null,
            "num_threads": 10,
            "save_path": null,
            "start_point_strategy": "medoid"
          }
        },
        "max_fp_vecs_per_prune": 0,
        "num_pq_chunks": 300,
        "seed": 13076402859301299683,
        "use_fp_for_search": false
      }
    }
  ]
}