Repository navigation
fix(refine): normalize the refined cosine leaf reference - #382
Draft
pavanmanishd wants to merge 1 commit into
Draft
pavanmanishd wants to merge 1 commit into
pavanmanishd wants to merge 1 commit into
Conversation
The page-backed leaf refinement replaces each leaf's encode reference with the mean of the rows routed to it, but wrote that mean without normalizing it for cosine. Every other centroid write path normalizes cosine leaf centroids before writing them: k-means means drift below unit norm. Two consumers depend on a unit-norm reference: - the Phase-A probe re-rank orders the routed clusters by the squared L2 distance between the rotated query and the rotated centroid, which is only a cosine ranking while both vectors are unit norm, so a reference with a smaller norm looks closer regardless of its direction; and - the head reference centers every RaBitQ residual in the leaf, while the tree write and the parent centroid pages both carry unit-norm leaf references for cosine. Measured on nytimes-256-angular (290k x 256d, cosine, maintenance_work_mem = 64MB so the build is subsampled, 1000 queries, parallel build): with refinement forced on, recall@10 at nprobe=150 goes from 0.9413 to 0.9460 and at nprobe=320 with the full rerank pool from 0.9764 to 0.9806, against 0.9440 and 0.9817 for the unrefined build. The kernel takes a cosine flag and both callers pass their metric. New unit tests pin the kernel's behavior: cosine means are written at unit norm with their direction preserved, L2 means are untouched, empty leaves are skipped and a mean that cancels to the origin stays finite.
Member
|
Looks like a duplicate of #353 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Refs #183.
Trying to reproduce the reported collapse on current main with the issue's dataset and settings, it does not reproduce: with refinement forced on, nytimes-256-angular at nprobe 320 with the full pool measures within half a percent of the unrefined build rather than the 0.9549 to 0.6411 the issue reports. What does reproduce is a smaller, systematic loss from refining, and it is the defect the issue suspected: the refine kernel wrote the refined mean without normalizing it for cosine, while every other centroid write path normalizes, so the head reference was not unit norm — which the Phase-A probe re-rank's squared-L2 ordering assumes.
With the normalization, refinement no longer loses to no-refinement at the sample budget measured (0.9413 to 0.9460 at nprobe 150, 0.9764 to 0.9806 at nprobe 320), and the unit tests pin the kernel invariant: the cosine case fails without it.
Note for reviewers: the maintainer's #353 covers the same normalization, writes both copies (posting head and centroid tree) and handles the zero mean better, so this is at most a test-only follow-up to it.