[BUG]Recall loss while building IVF-PQ indexes with cosine metric #2958

Issue cuvs

Describe the bug
Building a cosine IVF-PQ index with cuVS gives lower recall than Lance's CPU path on the same data. 15M vectors, dim 2048, 3755 lists, 256 PQ sub-vectors, recall@10, nprobes 100: cuVS 0.674 vs Lance CPU 0.731. At nprobes 1 the two are equal, so clustering and probe selection are healthy.

Root cause. In cosine mode ivf_pq_build.cuh trains the coarse k-means with CosineExpanded on a row-normalised training set and then row-normalises the centroids to unit length. The centroids are no longer cluster means. Lance CPU treats cosine as L2 on unit vectors and keeps plain means: when the metric is cosine, it first applies a NormalizeTransformer to row-normalise the vectors, then switches the distance type from cosine to L2.

Steps/Code to reproduce bug
Measured through Lance (cuVS builds the index, Lance searches it), on a 15M-vector unit-norm set (coyo-ve-base), dim 2048, create_index(index_type="IVF_PQ", metric="cosine", num_partitions=3755, num_sub_vectors=256, num_bits=8), recall@10 against exact cosine ground truth, 1000 queries. The GPU path is the Lance cuVS backend https://github.com/luliCloud/lance-cuvs/tree/cuvs-local-methods-ab-pool (backend cuvs_local). The only difference between the "cosine" and "C" rows below is CUVS_IVF_PQ_COSINE_KMEANS_L2 (0 vs 1) in the cuVS build linked under Method C.

A cuVS-only reproduction (cuvsIvfPqBuild + cuvsIvfPqSearch with metric=CosineExpanded vs the same data pre-normalised with L2Expanded) is not included yet.

Expected behavior
Cosine IVF-PQ should reach the same recall as L2 IVF-PQ on the row-normalised vectors, i.e. the coarse centroids should be the cluster means.

Environment details

• Environment location: Docker container (internal image) on bare-metal 8x H200 (sm_90)
• Method of cuVS install: from source, ./build.sh libcuvs, cuVS 26.12.00 (lance-local branch of https://github.com/jamxia155/cuvs, commit fbdc9b0a), CUDA 13.3
• Lance at commit dc08fbfc4 (CPU reference, unchanged between runs)

Additional context

Three ways to get "L2 on unit vectors with plain-mean centroids" (what Lance CPU does)

Compared against the current cuVS cosine mode (cuVS 26.12.00, CUDA 13.3).

• Baseline, cuVS cosine (unchanged): coarse k-means with CosineExpanded, then the centroids are row-normalised to unit length (not cluster means).
• Methods A and B, outside cuVS (lance-cuvs normalises, asks cuVS for L2; no cuVS change):
  • Method A (auto): lance-cuvs sends L2 to cuVS only if the training sample is already unit length. This requires scanning the sample to check the norms, which works for the method but is not practical at scale. On non-unit data it falls back to cosine, since L2 is not applicable there.
  • Method B (normalize): lance-cuvs normalises the training sample and every transform batch with a CUDA kernel, then always runs cuVS in L2, so it also works on non-unit data. An nsys analysis is still needed to see whether the normalisation kernel overlaps with compute and host-to-device transfer or adds to the critical path.
  • Method C, in cuVS (prototype, potential cuVS PR): the user-facing metric stays CosineExpanded. Internally, a cosine IVF-PQ index is built as L2 on unit-length vectors. In build, the coarse k-means fit/predict use L2 on the row-normalised training sample (already normalised in stock cuVS), and the centroid normalisation is skipped, so the centroids stay cluster means. In transform / extend, each batch is cast to float and row-normalised before the L2 assignment.
  • Proposed API. Add a field to cuvsIvfPqIndexParams that says the input vectors are already normalised (name open, e.g. data_normalized), and let cuVS choose the internal path from it, instead of exposing a "use L2 for cosine k-means" switch. If set, cuVS skips the per-batch normalisation; if not set, cuVS normalises internally. The same L2-on-normalised-vectors path with mean centroids is used in both cases, because the recall loss also appears on non-unit input (2M non-unit rows: cosine 0.678 vs 0.744 with the change). The choice must be stored in the index and in the serialisation, because search has to know which centroid convention an index uses (it needs a ||c||^2 term when the centroids are not unit length).
  • Prototype. The patch is opt-in through the env var CUVS_IVF_PQ_COSINE_KMEANS_L2=1, which the patch adds (stock cuVS has no such option); it is only a stand-in for the field above. It touches ivf_pq_build.cuh and ivf_pq_transform.cuh. NVTX ranges ivf_pq::{build,transform}::cosine_kmeans_as_l2 show it is active: present in the C run, absent in the control.
  • Note. Method C has not been validated through cuVS's own search; our recall results come from cuVS-built indexes searched through Lance.

Method B ties the fix to Lance (extra normalisation pass per batch in lance-cuvs); C fixes it for every cuVS user.

Recall@10 vs nprobes

15M unit vectors, dim 2048, 3755 lists, 1000 queries, one build per cell, cuVS 26.12 + CUDA 13.3 (top1nn backend cutlass; the auto backend gives the same picture). The Lance CPU row is from an earlier CPU run; the CPU path does not depend on CUDA or cuVS.

The gap is not an artefact of where probing stops: it is +0.04 at nprobes 5 and 10 and about +0.05 from nprobes 20 on, and every curve flattens (between nprobes 50 and 100: cosine 0.6719 to 0.6726, C 0.7248 to 0.7278, CPU 0.7297 to 0.7336). nprobes above 100 was not measured. Sampling error is about +-0.5 points.

Timing

1B rows (7.8 TB), 32768 lists, create_index wall time, single runs, same GPU:

The three L2 variants take similar time and are 8-10% faster than cosine, so this is the cost of cosine versus L2 assignment, not something specific to one method. For reference, a separate Lance CPU run of the same build (same dataset and parameters, hardware not verified to match) took 17756 s.

Appendix: potential root cause analysis

Measured on a random 100K sample of the 15M unit-norm dataset (dim 2048, 3755 lists, 256 PQ sub-vectors); all three models are trained and evaluated on the same unit vectors. r = x - c is the residual of a vector against its coarse centroid; PQ encodes r.

The PQ codebook is not worse: PQ loses the same fraction (about 30%) in all three setups. But with unit-length centroids the residuals are off-centre and their mean squared length is about 20% larger (0.730 vs 0.608), so the absolute encoding error is about 15% higher (0.218 vs 0.190).