Replace squared-distance IP with direct dot-product in multi-bit RaBitQ (#4877)

Summary:
Pull Request resolved: #4877

Refactor the multi-bit RaBitQ inner product computation from the
squared-distance-then-convert approach to a direct dot-product formulation.

Before: IP = -0.5 * (||q-c||² + (||r||² - ||x||²) + (-2·||r||/ipnorm)·ex_ip - ||q||²)
After:  IP = <q,c> + <c,r> + (||r||/ipnorm)·ex_ip

Both are mathematically equivalent to <q, x>. The new form is simpler and
makes the code structurally immune to the D95166460 bug class where
the degenerate case (x ≈ centroid) required complex metric-specific
branching that was easy to get wrong.

Changes:
- Per-document factors (compute_ex_factors): IP branch now computes
  f_add_ex = <c, r> and f_rescale_ex = ||r||/ipnorm (positive, no -2 factor)
- Degenerate case simplified to metric-agnostic f_add_ex=0, f_rescale_ex=0
- Core distance function (compute_full_multibit_distance): accepts a single
  qr_base parameter instead of qr_to_c_L2sqr + qr_norm_L2sqr, eliminates
  the -0.5*(dist - qr_norm_L2sqr) IP post-processing
- Added q_dot_c field to QueryFactorsData, computed at all query-setup sites
- L2 path, 1-bit path, and SIMD kernels are completely unchanged
- All four index types covered: IndexRaBitQ, IndexIVFRaBitQ,
  IndexRaBitQFastScan, IndexIVFRaBitQFastScan

Breaking change: serialized multi-bit IP indexes must be re-encoded.
L2 indexes are unaffected.

Reviewed By: ddrcoder, latham-meta

Differential Revision: D95419974

fbshipit-source-id: 3a5a33b5a3065d172a2f57578f3a78bdf562b87c

Author: alibeklfc

faiss/IndexIVFRaBitQFastScan.cpp

@@ -283,10 +283,11 @@ void IndexIVFRaBitQFastScan::compute_residual_LUT(
             rotated_q,
             rotated_qq);
 
-    // Override query norm for inner product if original query is provided
     if (metric_type == MetricType::METRIC_INNER_PRODUCT &&
         original_query != nullptr) {
         query_factors.qr_norm_L2sqr = fvec_norm_L2sqr(original_query, d);
+        query_factors.q_dot_c = query_factors.qr_norm_L2sqr -
+                fvec_inner_product(original_query, residual, d);
     }
 
     const size_t ex_bits = rabitq.nb_bits - 1;
@@ -813,8 +814,9 @@ float IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::
             ex_code,
             ex_fac,
             query_factors.rotated_q.data(),
-            query_factors.qr_to_c_L2sqr,
-            query_factors.qr_norm_L2sqr,
+            (index->metric_type == MetricType::METRIC_INNER_PRODUCT)
+                    ? query_factors.q_dot_c
+                    : query_factors.qr_to_c_L2sqr,
             dim,
             ex_bits,
             index->metric_type);

faiss/IndexRaBitQFastScan.cpp

@@ -751,8 +751,9 @@ float RaBitQHeapHandler<C, with_id_map>::compute_full_multibit_distance(
             ex_code,
             ex_fac,
             query_factors.rotated_q.data(),
-            query_factors.qr_to_c_L2sqr,
-            query_factors.qr_norm_L2sqr,
+            (rabitq_index->metric_type == MetricType::METRIC_INNER_PRODUCT)
+                    ? query_factors.q_dot_c
+                    : query_factors.qr_to_c_L2sqr,
             dim,
             ex_bits,
             rabitq_index->metric_type);

faiss/impl/RaBitQUtils.cpp

@@ -244,8 +244,12 @@ QueryFactorsData compute_query_factors(
 
     // Compute query norm for inner product metric
     query_factors.qr_norm_L2sqr = 0.0f;
+    query_factors.q_dot_c = 0.0f;
     if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
         query_factors.qr_norm_L2sqr = fvec_norm_L2sqr(query, d);
+        if (centroid != nullptr) {
+            query_factors.q_dot_c = fvec_inner_product(query, centroid, d);
+        }
     }
 
     return query_factors;
@@ -307,8 +311,7 @@ float compute_full_multibit_distance(
         const uint8_t* ex_code,
         const ExtraBitsFactors& ex_fac,
         const float* rotated_q,
-        float qr_to_c_L2sqr,
-        float qr_norm_L2sqr,
+        float qr_base,
         size_t d,
         size_t ex_bits,
         MetricType metric_type) {
@@ -317,11 +320,9 @@ float compute_full_multibit_distance(
     float ex_ip = rabitq::multibit::compute_inner_product(
             sign_bits, ex_code, rotated_q, d, ex_bits, cb);
 
-    float dist = qr_to_c_L2sqr + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
+    float dist = qr_base + ex_fac.f_add_ex + ex_fac.f_rescale_ex * ex_ip;
 
-    if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
-        dist = -0.5f * (dist - qr_norm_L2sqr);
-    } else {
+    if (metric_type == MetricType::METRIC_L2) {
         dist = std::max(0.0f, dist);
     }
 

faiss/impl/RaBitQUtils.h

@@ -70,6 +70,7 @@ struct QueryFactorsData {
 
     float qr_to_c_L2sqr = 0;
     float qr_norm_L2sqr = 0;
+    float q_dot_c = 0; // <query, centroid> for IP metric; 0 for L2
 
     float int_dot_scale = 1;
 
@@ -320,8 +321,7 @@ inline int extract_code_inline(
  * @param ex_code         packed ex-bit codes
  * @param ex_fac          ex-bit factors (f_add_ex, f_rescale_ex)
  * @param rotated_q       rotated query vector
- * @param qr_to_c_L2sqr   precomputed ||query_rotated - centroid||^2
- * @param qr_norm_L2sqr   precomputed ||query_rotated||^2 (0 for L2 metric)
+ * @param qr_base         precomputed base term: ||q-c||^2 for L2, <q,c> for IP
  * @param d               dimensionality
  * @param ex_bits         number of extra bits (nb_bits - 1)
  * @param metric_type     distance metric (L2 or Inner Product)
@@ -332,8 +332,7 @@ float compute_full_multibit_distance(
         const uint8_t* ex_code,
         const ExtraBitsFactors& ex_fac,
         const float* rotated_q,
-        float qr_to_c_L2sqr,
-        float qr_norm_L2sqr,
+        float qr_base,
         size_t d,
         size_t ex_bits,
         MetricType metric_type);

faiss/impl/RaBitQuantizer.cpp

@@ -314,13 +314,15 @@ float RaBitQDistanceComputerNotQ::distance_to_code_full(const uint8_t* code) {
             ex_code + (d * ex_bits + 7) / 8);
 
     // Call shared utility directly with rotated_q pointer
+    float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT)
+            ? query_fac.q_dot_c
+            : query_fac.qr_to_c_L2sqr;
     return rabitq_utils::compute_full_multibit_distance(
             binary_data,
             ex_code,
             *ex_fac,
             rotated_q.data(),
-            query_fac.qr_to_c_L2sqr,
-            query_fac.qr_norm_L2sqr,
+            qr_base,
             d,
             ex_bits,
             metric_type);
@@ -366,6 +368,8 @@ void RaBitQDistanceComputerNotQ::set_query(const float* x) {
     if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
         // precompute if needed
         query_fac.qr_norm_L2sqr = fvec_norm_L2sqr(x, d);
+        query_fac.q_dot_c =
+                centroid ? fvec_inner_product(x, centroid, d) : 0.0f;
     }
 }
 
@@ -480,13 +484,15 @@ float RaBitQDistanceComputerQ::distance_to_code_full(const uint8_t* code) {
             ex_code + (d * ex_bits + 7) / 8);
 
     // Call shared utility directly with rotated_q pointer
+    float qr_base = (metric_type == MetricType::METRIC_INNER_PRODUCT)
+            ? query_fac.q_dot_c
+            : query_fac.qr_to_c_L2sqr;
     return rabitq_utils::compute_full_multibit_distance(
             binary_data,
             ex_code,
             *ex_fac,
             rotated_q.data(),
-            query_fac.qr_to_c_L2sqr,
-            query_fac.qr_norm_L2sqr,
+            qr_base,
             d,
             ex_bits,
             metric_type);

faiss/impl/RaBitQuantizerMultiBit.cpp

@@ -216,25 +216,13 @@ void compute_ex_factors(
         ex_factors.f_add_ex = l2_sqr;
         ex_factors.f_rescale_ex = ipnorm_inv * -2.0f * norm;
     } else {
-        // For IP: Need to compute ||x||^2 for the correction term
-        // The distance formula expects f_add_ex = ||x - c||^2 - ||x||^2
-        // to match the 1-bit formula's or_minus_c_l2sqr for IP metric
-        //
-        // Reconstruct ||x||^2 from residual and centroid:
-        // x = residual + centroid, so ||x||^2 = ||residual + centroid||^2
-        float or_L2sqr = 0;
-        if (centroid != nullptr) {
-            for (size_t i = 0; i < d; i++) {
-                float x_val = residual[i] + centroid[i];
-                or_L2sqr += x_val * x_val;
-            }
-        } else {
-            // If no centroid, x = residual
-            or_L2sqr = l2_sqr;
-        }
-
-        ex_factors.f_add_ex = l2_sqr - or_L2sqr;
-        ex_factors.f_rescale_ex = ipnorm_inv * -2.0f * norm;
+        // For IP: direct dot-product formulation
+        // f_add_ex = <c, r> (dot product of centroid and residual)
+        // f_rescale_ex = ||r|| / ipnorm (positive scaling)
+        float c_dot_r =
+                centroid ? fvec_inner_product(residual, centroid, d) : 0.0f;
+        ex_factors.f_add_ex = c_dot_r;
+        ex_factors.f_rescale_ex = ipnorm_inv * norm;
     }
 }
 
@@ -276,12 +264,14 @@ void quantize_ex_bits(
     float norm_sqr = fvec_norm_L2sqr(residual, d);
     float norm = std::sqrt(norm_sqr);
 
-    // Handle degenerate case
+    // Handle degenerate case: residual is (near-)zero, meaning x ≈ centroid.
+    // For both L2 and IP, f_add_ex and f_rescale_ex are trivially zero:
+    //   L2: ||r||² ≈ 0, IP: <c,r> ≈ 0 and ||r||/ipnorm ≈ 0
     if (norm < 1e-10f) {
         size_t code_size = (d * ex_bits + 7) / 8;
         memset(ex_code, 0, code_size);
         ex_factors.f_add_ex = 0.0f;
-        ex_factors.f_rescale_ex = 1.0f;
+        ex_factors.f_rescale_ex = 0.0f;
         return;
     }
 

tests/test_rabitq.py

@@ -965,6 +965,48 @@ def test_factory_end_to_end(self):
             self.assertEqual(D_ivf.shape, (ds.nq, 5))
             self.assertTrue(np.all(I_ivf >= 0))
 
+    def test_degenerate_centroid_distance(self):
+        """Test that a doc identical to its centroid gets correct distances.
+
+        When a vector's residual (x - centroid) has near-zero norm,
+        quantize_ex_bits hits a degenerate early-return path. A previous
+        bug set f_rescale_ex=1 (should be 0) and f_add_ex=0 (should be
+        metric-dependent), causing wildly wrong multi-bit distances.
+        """
+        d = 128
+        nlist = 16
+        rs = np.random.RandomState(42)
+        xt = rs.randn(2000, d).astype("float32")
+        xt /= np.linalg.norm(xt, axis=1, keepdims=True)
+
+        query = rs.randn(1, d).astype("float32")
+        query /= np.linalg.norm(query)
+
+        for metric in [faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT]:
+            for nb_bits in [2, 3, 4]:
+                with self.subTest(metric=metric, nb_bits=nb_bits):
+                    quantizer = faiss.IndexFlat(d, metric)
+                    index = faiss.IndexIVFRaBitQ(
+                        quantizer, d, nlist, metric, True, nb_bits
+                    )
+                    index.train(xt)
+                    index.nprobe = nlist  # exhaustive
+
+                    # Add the centroid itself as the only document
+                    centroid = index.quantizer.reconstruct(0).reshape(1, d)
+                    index.add(centroid)
+
+                    D, I = index.search(query, 1)
+
+                    if metric == faiss.METRIC_INNER_PRODUCT:
+                        true_dist = (query @ centroid.T).item()
+                    else:
+                        true_dist = float(np.sum((query - centroid) ** 2))
+
+                    np.testing.assert_allclose(
+                        D[0, 0], true_dist, atol=0.15,
+                        err_msg=f"nb_bits={nb_bits}")
+
 
 class TestRaBitQStats(unittest.TestCase):
     """Test RaBitQStats tracking for multi-bit two-stage search."""