RaBitQ: Fix zero residual sign convention (#5553)

Summary:
The RaBitQ 1 bit quantization, code is quantized comparing positive vs negative or zero
https://github.com/facebookresearch/faiss/blob/718002f7cf269f3f3752dde04a4cc3f8eabb0050/faiss/impl/RaBitQuantizer.cpp#L147-L154

However, during ex bit quantization, the signs are flipped under condition negative vs positive or zero.
https://github.com/facebookresearch/faiss/blob/718002f7cf269f3f3752dde04a4cc3f8eabb0050/faiss/impl/RaBitQuantizerMultiBit.cpp#L305-L310

Alternative implementation can be taking the equality in the one bit path.

Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5553

Reviewed By: mnorris11

Differential Revision: D118529759

Pulled By: alibeklfc

fbshipit-source-id: 319f0ec06952c73616acac25c37c63dfa7f9203c

faiss/impl/RaBitQuantizerMultiBit.cpp

@@ -302,9 +302,9 @@ void quantize_ex_bits(
     }
 
     // Step 6: Handle negative dimensions (flip bits)
-    // For negative residuals, flip all bits: code' = ~code & max_code
+    // For negative or zero residuals, flip all bits: code' = ~code & max_code
     for (size_t i = 0; i < d; i++) {
-        if (residual[i] < 0) {
+        if (residual[i] <= 0) {
             tmp_code[i] = (~tmp_code[i]) & max_code;
         }
     }

tests/test_rabitq.py

@@ -8,6 +8,7 @@ import unittest
 import numpy as np
 import faiss
 from faiss.contrib import datasets
+from faiss.contrib.evaluation import check_ref_knn_with_draws
 
 from common_faiss_tests import for_all_simd_levels, NoneSIMDLevel
 
@@ -948,6 +949,41 @@ class TestMultiBitRaBitQ(unittest.TestCase):
                         err_msg=f"nb_bits={nb_bits}",
                     )
 
+    def test_zero_residual_coordinate(self):
+        """Test that an exactly-zero residual encodes like a negative one.
+
+        The sign bit uses `residual > 0`, so a zero residual must flip its
+        ex-bits too, or it decodes to the largest magnitude instead of the
+        smallest.
+        """
+        ds = datasets.SyntheticDataset(64, 2000, 1, 1)
+        zero_dims = np.arange(0, ds.d, 4)
+
+        for metric in [faiss.METRIC_L2, faiss.METRIC_INNER_PRODUCT]:
+            for nb_bits in [2, 4, 8, 9]:
+                with self.subTest(metric=metric, nb_bits=nb_bits):
+                    index = faiss.IndexRaBitQ(ds.d, metric, nb_bits)
+                    index.train(ds.get_train())
+
+                    # Anchor on the center: the residual must be exactly
+                    # zero, not merely small.
+                    center = faiss.vector_to_array(index.center)
+                    x = ds.get_database().copy()
+                    x[0, zero_dims] = center[zero_dims]
+
+                    x_neg = x.copy()
+                    x_neg[0, zero_dims] = np.nextafter(
+                        x[0, zero_dims], np.float32(-np.inf)
+                    )
+
+                    index.add(x)
+                    Dref, Iref = index.search(ds.get_queries(), 1)
+                    index.reset()
+                    index.add(x_neg)
+                    Dnew, Inew = index.search(ds.get_queries(), 1)
+
+                    check_ref_knn_with_draws(Dref, Iref, Dnew, Inew)
+
 
 if __name__ == "__main__":
     unittest.main()