Fix SQtqmse codebook scale: fold 1/sqrt(d) into Lloyd-Max table - #5517
Closed
mnorris11 wants to merge 1 commit into
Closed
Fix SQtqmse codebook scale: fold 1/sqrt(d) into Lloyd-Max table#5517mnorris11 wants to merge 1 commit into
mnorris11 wants to merge 1 commit into
Conversation
Summary: `SQtqmse` reconstructions have inflated norms and badly degraded recall since D102408184. Reported as a 1.14.2 -> 1.14.3 regression: at `d=64`, median decoded norm of a unit-norm input went `0.9965` -> `1.1503` (+15%), while `SQtqmse8` looked "essentially unchanged". Fixes facebookresearch#5317 Root cause - D102408184 replaced the `tqmse` codebook construction with hardcoded Lloyd-Max tables: ``` - scalar_quantizer::train_TurboQuantMSE(d, 4, trained); + populate_lloyd_max_trained(4, trained); ``` The old codebook was built for one coordinate of a unit-norm vector in R^d, so its centroids scaled as `1/sqrt(d)`. The new tables are optimal for `N(0, 1)` and do not depend on `d` at all. That is correct for `_eden` and `_tq`, which rescale each vector to unit variance before lookup. Plain `tqmse` does not: it encodes raw unit-norm vectors, whose components are ~`1/sqrt(d)` (~`0.125` at `d=64`), against a codebook ~50x too wide. Nearly everything falls in the innermost cells, so the codebook collapses to a few levels and every component decodes too large: | d | bits | distinct codes used (before D102408184 -> after) | | --- | --- | --- | | 64 | 4 | 16/16 -> **4/16** | | 768 | 4 | 16/16 -> **2/16** | | 768 | 8 | 256/256 -> **32/256** | Gets worse as `d` grows. Fix - Give `populate_lloyd_max_trained` a `scale` argument that multiplies the table. The five `tqmse` cases pass `1/sqrt(d)` (the standard deviation of a unit-norm vector's components) so the codebook once again (like 1.14.2) matches the data it encodes. **No change for `_eden` / `_tq`.** They pass no `scale`, so it defaults to `1` and every table entry is bit-identical to today. Scaling `trained` rather than the encode path is what keeps this a one-line fix per callsite: every SIMD specialization and distance computer reads `this->centroids`, a pointer into `trained`. Notes - - NOT A REVERT TO 1.14.2, but that is fine. 1.14.2 trained on the exact unit-sphere marginal; this scales a Gaussian approximation of it, so centroids differ by up to ~20% at 8 bits. The approximation is not worse: max component error is lower than 1.14.2 and recall is restored (see test plan). The only consequence is that `tqmse` codes are not comparable across 1.14.2 / 1.14.3 / this version -- harmless, because the codebook ships with the index (below). - SERIALIZATION STILL COMPATIBLE: Existing 1.14.3-written `tqmse` indexes stay readable and self-consistent: `read_ScalarQuantizer` loads `trained` verbatim, so decode matches how they were encoded, and `trained` length is unchanged so size validation still passes. They are degraded, not corrupt, and need re-indexing to benefit. No version guard added. Differential Revision: D115440646
Contributor
|
@mnorris11 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115440646. |
Contributor
|
This pull request has been merged in e2f9cca. |
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
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.
Summary:
SQtqmsereconstructions have inflated norms and badly degraded recall sinceD102408184. Reported as a 1.14.2 -> 1.14.3 regression: at
d=64, median decodednorm of a unit-norm input went
0.9965->1.1503(+15%), whileSQtqmse8looked "essentially unchanged".
Fixes #5317
Root cause
D102408184 replaced the
tqmsecodebook construction with hardcoded Lloyd-Maxtables:
The old codebook was built for one coordinate of a unit-norm vector in R^d, so
its centroids scaled as
1/sqrt(d). The new tables are optimal forN(0, 1)and do not depend on
dat all.That is correct for
_edenand_tq, which rescale each vector to unitvariance before lookup. Plain
tqmsedoes not: it encodes raw unit-normvectors, whose components are
1/sqrt(d)(0.125atd=64), against acodebook ~50x too wide. Nearly everything falls in the innermost cells, so the
codebook collapses to a few levels and every component decodes too large:
Gets worse as
dgrows.Fix
Give
populate_lloyd_max_trainedascaleargument that multiplies the table.The five
tqmsecases pass1/sqrt(d)(the standard deviation of a unit-normvector's components) so the codebook once again (like 1.14.2) matches the data it encodes.
No change for
_eden/_tq. They pass noscale, so it defaults to1and every table entry is bit-identical to today.
Scaling
trainedrather than the encode path is what keeps this a one-line fixper callsite: every SIMD specialization and distance computer reads
this->centroids, a pointer intotrained.Notes
unit-sphere marginal; this scales a Gaussian approximation of it, so centroids
differ by up to ~20% at 8 bits. The approximation is not worse: max component
error is lower than 1.14.2 and recall is restored (see test plan). The only
consequence is that
tqmsecodes are not comparable across 1.14.2 / 1.14.3 /this version -- harmless, because the codebook ships with the index (below).
tqmseindexes stay readable and self-consistent:read_ScalarQuantizerloadstrainedverbatim, so decode matches how theywere encoded, and
trainedlength is unchanged so size validation stillpasses. They are degraded, not corrupt, and need re-indexing to benefit. No
version guard added.
Differential Revision: D115440646