Skip to content

Fix SQtqmse codebook scale: fold 1/sqrt(d) into Lloyd-Max table - #5517

Closed
mnorris11 wants to merge 1 commit into
facebookresearch:mainfrom
mnorris11:export-D115440646
Closed

Fix SQtqmse codebook scale: fold 1/sqrt(d) into Lloyd-Max table#5517
mnorris11 wants to merge 1 commit into
facebookresearch:mainfrom
mnorris11:export-D115440646

Conversation

@mnorris11

Copy link
Copy Markdown
Contributor

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 #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

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
@meta-cla meta-cla Bot added the CLA Signed label Aug 12, 2026
@meta-codesync

meta-codesync Bot commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

@mnorris11 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115440646.

@meta-codesync

meta-codesync Bot commented Aug 12, 2026

Copy link
Copy Markdown
Contributor

This pull request has been merged in e2f9cca.

@meta-codesync meta-codesync Bot added the Merged label Aug 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

SQtqmse4 decoded-vector norms shifted +15% from 1.14.2 to 1.14.3, intentional?

1 participant