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

Summary:
Pull Request resolved: #5517

`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.

Reviewed By: alibeklfc

Differential Revision: D115440646

fbshipit-source-id: a559f6a6a50a65dc3ab32ff3f17cda8796b2d4ae

Author: Michael Norris

faiss/impl/ScalarQuantizer.cpp

@@ -7,6 +7,7 @@
 
 // -*- c++ -*-
 
+#include <cmath>
 #include <cstring>
 #include <memory>
 
@@ -425,14 +426,29 @@ const LloydMaxTable kLloydMaxTables[] = {
         {kLloydMaxCentroids8, kLloydMaxBoundaries8}, // 8
 };
 
-void populate_lloyd_max_trained(size_t mse_bits, std::vector<float>& trained) {
+// The tables are Lloyd-Max optimal for N(0, 1) input. Callers whose input has
+// a different standard deviation pass it as `scale` to stretch the table.
+void populate_lloyd_max_trained(
+        size_t mse_bits,
+        std::vector<float>& trained,
+        float scale = 1.0f) {
     FAISS_THROW_IF_NOT(mse_bits >= 1 && mse_bits <= 8);
     FAISS_THROW_IF_NOT(kLloydMaxTables[mse_bits].centroids);
     size_t k = size_t(1) << mse_bits;
     const auto& t = kLloydMaxTables[mse_bits];
     trained.resize(k + (k - 1));
-    std::copy(t.centroids, t.centroids + k, trained.begin());
-    std::copy(t.boundaries, t.boundaries + k - 1, trained.begin() + k);
+    for (size_t i = 0; i < k; i++) {
+        trained[i] = t.centroids[i] * scale;
+    }
+    for (size_t i = 0; i + 1 < k; i++) {
+        trained[k + i] = t.boundaries[i] * scale;
+    }
+}
+
+// Component scale of a unit-norm vector in R^d.
+float unit_norm_component_scale(size_t d) {
+    FAISS_THROW_IF_NOT(d > 0);
+    return 1.0f / std::sqrt(static_cast<float>(d));
 }
 
 } // namespace
@@ -588,19 +604,24 @@ void ScalarQuantizer::train(size_t n, const float* x) {
             populate_lloyd_max_trained(bits, trained);
             break;
         case QT_1bit_tqmse:
-            populate_lloyd_max_trained(1, trained);
+            populate_lloyd_max_trained(
+                    1, trained, unit_norm_component_scale(d));
             break;
         case QT_2bit_tqmse:
-            populate_lloyd_max_trained(2, trained);
+            populate_lloyd_max_trained(
+                    2, trained, unit_norm_component_scale(d));
             break;
         case QT_3bit_tqmse:
-            populate_lloyd_max_trained(3, trained);
+            populate_lloyd_max_trained(
+                    3, trained, unit_norm_component_scale(d));
             break;
         case QT_4bit_tqmse:
-            populate_lloyd_max_trained(4, trained);
+            populate_lloyd_max_trained(
+                    4, trained, unit_norm_component_scale(d));
             break;
         case QT_8bit_tqmse:
-            populate_lloyd_max_trained(8, trained);
+            populate_lloyd_max_trained(
+                    8, trained, unit_norm_component_scale(d));
             break;
         case QT_2bit_tq:
         case QT_3bit_tq:

tests/test_scalar_quantizer_correctness.py

@@ -67,6 +67,43 @@ def test_tqmse_8bit(self):
         faiss.normalize_L2(self.xb)
         self.do_encode_decode(faiss.ScalarQuantizer.QT_8bit_tqmse, 0.1)
 
+    def test_tqmse_preserves_norm(self):
+        """tqmse encodes unit-norm vectors with no per-vector scale factor in
+        the code, so reconstructions must keep unit norm. The Lloyd-Max table
+        is optimal for unit *variance*, so it only lines up with unit-norm
+        input once scaled by 1/sqrt(d) -- without that, every component lands
+        in the innermost cells and the codebook collapses toward a sign
+        quantizer, inflating norms by 15% at d=64 and 258% at d=768."""
+        for qtype, nbits in (
+            (faiss.ScalarQuantizer.QT_1bit_tqmse, 1),
+            (faiss.ScalarQuantizer.QT_2bit_tqmse, 2),
+            (faiss.ScalarQuantizer.QT_3bit_tqmse, 3),
+            (faiss.ScalarQuantizer.QT_4bit_tqmse, 4),
+            (faiss.ScalarQuantizer.QT_8bit_tqmse, 8),
+        ):
+            for d in (32, 64, 768):
+                with self.subTest(nbits=nbits, d=d):
+                    rs = np.random.RandomState(123)
+                    x = rs.randn(200, d).astype("float32")
+                    faiss.normalize_L2(x)
+
+                    sq = faiss.ScalarQuantizer(d, qtype)
+                    sq.train(x)
+                    decoded = sq.decode(sq.compute_codes(x))
+                    norms = np.linalg.norm(decoded, axis=1)
+
+                    # 1-bit cannot resolve magnitude at all (every component
+                    # reconstructs to +/-c), so it gets a looser band.
+                    tol = 0.25 if nbits == 1 else 0.1
+                    self.assertLess(abs(np.median(norms) - 1.0), tol)
+
+                    # Guard the collapse directly: a codebook matched to the
+                    # data reconstructs most of its 2^nbits distinct levels.
+                    # This is what catches 8-bit, whose norm stays near 1 even
+                    # when collapsed (31 of 256 levels used at d=768).
+                    used = len(np.unique(decoded))
+                    self.assertGreaterEqual(used, 2**nbits // 2)
+
     def test_codes_match_none(self):
         """SQ codes are integer; encode dispatch (sq-dispatch.h) must
         produce bit-identical output at every SIMD level. Catches drift in