Add RaBitQStats for tracking two-stage search filtering effectiveness (#4723)

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

Adds instrumentation to track the effectiveness of the two-stage filtering strategy in multi-bit RaBitQ indexes. This enables developers and users to measure how much computation is saved by the lower-bound filtering before full distance computation.

Test results demonstrate the filtering is highly effective, with all index types (IndexRaBitQ, IndexIVFRaBitQ, IndexRaBitQFastScan, IndexIVFRaBitQFastScan) achieving 98-99.5% skip rates for both 2-bit and 4-bit configurations. This means only 0.5-2% of candidates require expensive full multi-bit distance computation.

Reviewed By: mdouze

Differential Revision: D89073817

fbshipit-source-id: 5c90c211975c26af66a6989ade6ab4c5c669c153

faiss/CMakeLists.txt

@@ -67,6 +67,7 @@ set(FAISS_SRC
   impl/AdditiveQuantizer.cpp
   impl/RaBitQuantizer.cpp
   impl/RaBitQuantizerMultiBit.cpp
+  impl/RaBitQStats.cpp
   impl/RaBitQUtils.cpp
   impl/ResidualQuantizer.cpp
   impl/LocalSearchQuantizer.cpp
@@ -195,6 +196,7 @@ set(FAISS_HEADERS
   impl/Quantizer.h
   impl/RaBitQuantizer.h
   impl/RaBitQuantizerMultiBit.h
+  impl/RaBitQStats.h
   impl/RaBitQUtils.h
   impl/ResidualQuantizer.h
   impl/ResultHandler.h

faiss/IndexIVFRaBitQ.cpp

@@ -216,6 +216,12 @@ struct RaBitInvertedListScanner : InvertedListScanner {
         // Multi-bit: Two-stage search with adaptive filtering
         size_t nup = 0;
 
+        // Stats tracking for multi-bit two-stage search
+        // n_1bit_evaluations: candidates evaluated using 1-bit lower bound
+        // n_multibit_evaluations: candidates requiring full multi-bit distance
+        size_t local_1bit_evaluations = 0;
+        size_t local_multibit_evaluations = 0;
+
         for (size_t j = 0; j < list_size; j++) {
             if (sel != nullptr) {
                 int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
@@ -225,6 +231,8 @@ struct RaBitInvertedListScanner : InvertedListScanner {
                 }
             }
 
+            local_1bit_evaluations++;
+
             // Stage 1: Compute lower bound using 1-bit codes
             float lower_bound = rabitq_dc->lower_bound_distance(codes);
 
@@ -233,10 +241,11 @@ struct RaBitInvertedListScanner : InvertedListScanner {
             // IP (max-heap): filter if lower_bound > simi[0]
             // Note: Using simi[0] directly (not cached) enables more aggressive
             // filtering as the heap is updated with better candidates
-            bool is_promising = keep_max ? (lower_bound > simi[0])
-                                         : (lower_bound < simi[0]);
+            bool should_refine = keep_max ? (lower_bound > simi[0])
+                                          : (lower_bound < simi[0]);
 
-            if (is_promising) {
+            if (should_refine) {
+                local_multibit_evaluations++;
                 // Lower bound is promising, compute full distance
                 float dis = distance_to_code(codes);
 
@@ -256,6 +265,13 @@ struct RaBitInvertedListScanner : InvertedListScanner {
             }
             codes += code_size;
         }
+
+        // Update global stats atomically
+#pragma omp atomic
+        rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
+#pragma omp atomic
+        rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
+
         return nup;
     }
 

faiss/IndexIVFRaBitQ.h

@@ -13,6 +13,7 @@
 #include <faiss/Index.h>
 #include <faiss/IndexIVF.h>
 
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQuantizer.h>
 
 namespace faiss {

faiss/IndexIVFRaBitQFastScan.cpp

@@ -630,6 +630,13 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
     }
 
     size_t max_positions = std::min<size_t>(32, this->ntotal - idx_base);
+
+    // Stats tracking for two-stage search
+    // n_1bit_evaluations: candidates evaluated using 1-bit lower bound
+    // n_multibit_evaluations: candidates requiring full multi-bit distance
+    size_t local_1bit_evaluations = 0;
+    size_t local_multibit_evaluations = 0;
+
     // Process each candidate vector in the SIMD batch
     for (size_t j = 0; j < max_positions; j++) {
         const int64_t result_id = this->adjust_id(b, j);
@@ -646,6 +653,9 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
                 index->flat_storage.data() + result_id * storage_size;
 
         if (is_multibit) {
+            // Track candidates actually considered for two-stage filtering
+            local_1bit_evaluations++;
+
             // Multi-bit: use FactorsData and two-stage search
             const FactorsData& full_factors =
                     *reinterpret_cast<const FactorsData*>(base_ptr);
@@ -671,6 +681,8 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
                     : (lower_bound < heap_dis[0]); // L2: keep if better
 
             if (should_refine) {
+                local_multibit_evaluations++;
+
                 // Compute local_offset: position within current inverted list
                 size_t local_offset = this->j0 + b * 32 + j;
 
@@ -685,7 +697,6 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
                 }
             }
         } else {
-            // 1-bit: direct distance computation and heap update
             const auto& db_factors =
                     *reinterpret_cast<const BaseFactorsData*>(base_ptr);
 
@@ -705,6 +716,12 @@ void IndexIVFRaBitQFastScan::IVFRaBitQHeapHandler<C>::handle(
             }
         }
     }
+
+    // Update global stats atomically
+#pragma omp atomic
+    rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
+#pragma omp atomic
+    rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
 }
 
 template <class C>

faiss/IndexIVFRaBitQFastScan.h

@@ -12,6 +12,7 @@
 #include <faiss/IndexIVFFastScan.h>
 #include <faiss/IndexIVFRaBitQ.h>
 #include <faiss/IndexRaBitQFastScan.h>
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/RaBitQuantizer.h>
 #include <faiss/impl/simd_result_handlers.h>

faiss/IndexRaBitQ.cpp

@@ -103,8 +103,15 @@ struct Run_search_with_dc_res {
                 resi.begin(q);
                 dc_base->set_query(xq + d * q);
 
+                // Stats tracking for multi-bit two-stage search only
+                // n_1bit_evaluations: candidates evaluated using 1-bit lower
+                // bound n_multibit_evaluations: candidates requiring full
+                // multi-bit distance
+                size_t local_1bit_evaluations = 0;
+                size_t local_multibit_evaluations = 0;
+
                 if (ex_bits == 0) {
-                    // 1-bit: Standard single-stage search
+                    // 1-bit: Standard single-stage search (no stats tracking)
                     for (size_t i = 0; i < ntotal; i++) {
                         if (res.is_in_selection(i)) {
                             float dis = (*dc_base)(i);
@@ -132,6 +139,8 @@ struct Run_search_with_dc_res {
                             const uint8_t* code =
                                     index->codes.data() + i * index->code_size;
 
+                            local_1bit_evaluations++;
+
                             // Stage 1: Compute 1-bit lower bound
                             float lower_bound = dc->lower_bound_distance(code);
 
@@ -141,11 +150,12 @@ struct Run_search_with_dc_res {
                             // lower_bound > resi.threshold Note: Using
                             // resi.threshold directly (not cached) enables more
                             // aggressive filtering as the heap is updated
-                            bool should_compute_full = is_similarity
+                            bool should_refine = is_similarity
                                     ? (lower_bound > resi.threshold)
                                     : (lower_bound < resi.threshold);
 
-                            if (should_compute_full) {
+                            if (should_refine) {
+                                local_multibit_evaluations++;
                                 // Compute full multi-bit distance
                                 float dist_full =
                                         dc->distance_to_code_full(code);
@@ -155,6 +165,13 @@ struct Run_search_with_dc_res {
                     }
                 }
 
+                // Update global stats atomically
+#pragma omp atomic
+                rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
+#pragma omp atomic
+                rabitq_stats.n_multibit_evaluations +=
+                        local_multibit_evaluations;
+
                 resi.end();
             }
         }

faiss/IndexRaBitQ.h

@@ -8,6 +8,7 @@
 #pragma once
 
 #include <faiss/IndexFlatCodes.h>
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQuantizer.h>
 
 namespace faiss {

faiss/IndexRaBitQFastScan.cpp

@@ -541,6 +541,12 @@ void RaBitQHeapHandler<C, with_id_map>::handle(
     // Get storage size once
     const size_t storage_size = rabitq_index->compute_per_vector_storage_size();
 
+    // Stats tracking for multi-bit two-stage search only
+    // n_1bit_evaluations: candidates evaluated using 1-bit lower bound
+    // n_multibit_evaluations: candidates requiring full multi-bit distance
+    size_t local_1bit_evaluations = 0;
+    size_t local_multibit_evaluations = 0;
+
     // Process distances in batch
     for (size_t i = 0; i < max_vectors; i++) {
         const size_t db_idx = base_db_idx + i;
@@ -553,6 +559,9 @@ void RaBitQHeapHandler<C, with_id_map>::handle(
                 rabitq_index->flat_storage.data() + db_idx * storage_size;
 
         if (is_multi_bit) {
+            // Track candidates actually considered for two-stage filtering
+            local_1bit_evaluations++;
+
             const FactorsData& full_factors =
                     *reinterpret_cast<const FactorsData*>(base_ptr);
 
@@ -574,6 +583,7 @@ void RaBitQHeapHandler<C, with_id_map>::handle(
                     : (lower_bound < heap_dis[0]); // L2: keep if better
 
             if (should_refine) {
+                local_multibit_evaluations++;
                 float dist_full = compute_full_multibit_distance(db_idx, q);
 
                 if (Cfloat::cmp(heap_dis[0], dist_full)) {
@@ -602,6 +612,12 @@ void RaBitQHeapHandler<C, with_id_map>::handle(
             }
         }
     }
+
+    // Update global stats atomically
+#pragma omp atomic
+    rabitq_stats.n_1bit_evaluations += local_1bit_evaluations;
+#pragma omp atomic
+    rabitq_stats.n_multibit_evaluations += local_multibit_evaluations;
 }
 
 template <class C, bool with_id_map>

faiss/IndexRaBitQFastScan.h

@@ -11,6 +11,7 @@
 
 #include <faiss/IndexFastScan.h>
 #include <faiss/IndexRaBitQ.h>
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/impl/RaBitQuantizer.h>
 #include <faiss/impl/simd_result_handlers.h>

faiss/impl/RaBitQStats.cpp

@@ -0,0 +1,29 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <faiss/impl/RaBitQStats.h>
+
+namespace faiss {
+
+// NOLINTNEXTLINE(facebook-avoid-non-const-global-variables)
+RaBitQStats rabitq_stats;
+
+void RaBitQStats::reset() {
+    n_1bit_evaluations = 0;
+    n_multibit_evaluations = 0;
+}
+
+double RaBitQStats::skip_percentage() const {
+    const size_t copy_n_1bit_evaluations = n_1bit_evaluations;
+    const size_t copy_n_multibit_evaluations = n_multibit_evaluations;
+    return copy_n_1bit_evaluations > 0
+            ? 100.0 * (copy_n_1bit_evaluations - copy_n_multibit_evaluations) /
+                    copy_n_1bit_evaluations
+            : 0.0;
+}
+
+} // namespace faiss

faiss/impl/RaBitQStats.h

@@ -0,0 +1,56 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <faiss/impl/platform_macros.h>
+#include <cstddef>
+
+namespace faiss {
+
+/// Statistics for RaBitQ multi-bit two-stage search.
+///
+/// These stats are ONLY collected for multi-bit mode (nb_bits > 1).
+/// In 1-bit mode, there is no two-stage filtering - all candidates are
+/// evaluated with a single distance computation, so there is nothing
+/// meaningful to track. For 1-bit mode, both counters remain 0.
+///
+/// Multi-bit mode uses a two-stage search:
+///   Stage 1: Compute 1-bit lower bound distance for all candidates
+///   Stage 2: Compute full multi-bit distance only for promising candidates
+///
+/// The skip_percentage() metric measures filtering effectiveness:
+/// how many candidates were filtered out by the 1-bit lower bound
+/// without needing the more expensive multi-bit distance computation.
+///
+/// WARNING: Statistics are not robust to internal threading nor to
+/// concurrent RaBitQ searches. Use these values in a single-threaded
+/// context to accurately gauge RaBitQ's filtering effectiveness.
+/// Call reset() before search, then read stats after search completes.
+struct RaBitQStats {
+    /// Number of candidates evaluated using 1-bit (lower bound) distance.
+    /// This is the first stage of two-stage search in multi-bit mode.
+    /// Always 0 in 1-bit mode (stats not tracked).
+    size_t n_1bit_evaluations = 0;
+
+    /// Number of candidates that passed 1-bit filtering and required
+    /// full multi-bit distance computation (second stage).
+    /// Always 0 in 1-bit mode (stats not tracked).
+    size_t n_multibit_evaluations = 0;
+
+    void reset();
+
+    /// Compute percentage of candidates skipped (filtered out by 1-bit stage).
+    /// Returns 0 if no candidates were evaluated (including 1-bit mode).
+    double skip_percentage() const;
+};
+
+/// Global stats for RaBitQ indexes
+// NOLINTNEXTLINE(facebook-avoid-non-const-global-variables)
+FAISS_API extern RaBitQStats rabitq_stats;
+
+} // namespace faiss

faiss/python/swigfaiss.swig

@@ -195,6 +195,7 @@ typedef uint64_t size_t;
 #include <faiss/IndexNeuralNetCodec.h>
 
 #include <faiss/impl/RaBitQuantizer.h>
+#include <faiss/impl/RaBitQStats.h>
 #include <faiss/impl/RaBitQUtils.h>
 #include <faiss/IndexRaBitQ.h>
 #include <faiss/IndexRaBitQFastScan.h>
@@ -631,6 +632,7 @@ struct faiss::simd16uint16 {};
 %include <faiss/IndexNeuralNetCodec.h>
 
 %include <faiss/impl/RaBitQuantizer.h>
+%include <faiss/impl/RaBitQStats.h>
 %include <faiss/impl/RaBitQUtils.h>
 %include <faiss/IndexRaBitQ.h>
 %include <faiss/IndexRaBitQFastScan.h>

tests/test_rabitq.py

@@ -2656,5 +2656,98 @@ class TestMultiBitRaBitQIndexFactory(unittest.TestCase):
             self.assertTrue(np.all(I_ivf < ds.nb))
 
 
+class TestRaBitQStats(unittest.TestCase):
+    """Test RaBitQStats tracking for multi-bit two-stage search."""
+
+    INDEX_TYPES = [
+        "IndexRaBitQ",
+        "IndexIVFRaBitQ",
+        "IndexRaBitQFastScan",
+        "IndexIVFRaBitQFastScan",
+    ]
+
+    @classmethod
+    def setUpClass(cls):
+        cls.stats_available = hasattr(faiss, 'cvar') and hasattr(
+            faiss.cvar, 'rabitq_stats'
+        )
+        if cls.stats_available:
+            cls.rabitq_stats = faiss.cvar.rabitq_stats
+
+    def test_stats_reset_and_skip_percentage(self):
+        """Test that stats can be reset and skip_percentage works."""
+        if not self.stats_available:
+            self.skipTest("rabitq_stats not available in Python bindings")
+        self.rabitq_stats.reset()
+        self.assertEqual(self.rabitq_stats.n_1bit_evaluations, 0)
+        self.assertEqual(self.rabitq_stats.n_multibit_evaluations, 0)
+        self.assertEqual(self.rabitq_stats.skip_percentage(), 0.0)
+
+    def test_stats_collected_multibit_all_index_types(self):
+        """Test that stats are collected for all multi-bit index types."""
+        if not self.stats_available:
+            self.skipTest("rabitq_stats not available in Python bindings")
+        ds = datasets.SyntheticDataset(384, 50000, 50000, 10)
+        nlist = 16
+
+        for index_type in self.INDEX_TYPES:
+            for nb_bits in [2, 4]:
+                with self.subTest(index_type=index_type, nb_bits=nb_bits):
+                    self.rabitq_stats.reset()
+
+                    if index_type == "IndexRaBitQ":
+                        index = faiss.IndexRaBitQ(
+                            ds.d, faiss.METRIC_L2, nb_bits
+                        )
+                    elif index_type == "IndexIVFRaBitQ":
+                        quantizer = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
+                        index = faiss.IndexIVFRaBitQ(
+                            quantizer, ds.d, nlist, faiss.METRIC_L2,
+                            True, nb_bits
+                        )
+                        index.nprobe = 4
+                    elif index_type == "IndexRaBitQFastScan":
+                        index = faiss.IndexRaBitQFastScan(
+                            ds.d, faiss.METRIC_L2, 32, nb_bits
+                        )
+                    elif index_type == "IndexIVFRaBitQFastScan":
+                        quantizer = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
+                        index = faiss.IndexIVFRaBitQFastScan(
+                            quantizer, ds.d, nlist, faiss.METRIC_L2,
+                            32, True, nb_bits
+                        )
+                        index.nprobe = 4
+                    else:
+                        raise ValueError(f"Unknown index type: {index_type}")
+
+                    index.train(ds.get_train())
+                    index.add(ds.get_database())
+                    index.search(ds.get_queries(), 10)
+
+                    self.assertGreater(
+                        self.rabitq_stats.n_1bit_evaluations, 0
+                    )
+                    self.assertGreater(
+                        self.rabitq_stats.n_multibit_evaluations, 0
+                    )
+                    # For multi-bit, filtering should skip some candidates
+                    self.assertLess(
+                        self.rabitq_stats.n_multibit_evaluations,
+                        self.rabitq_stats.n_1bit_evaluations,
+                    )
+                    skip_pct = self.rabitq_stats.skip_percentage()
+                    self.assertGreater(skip_pct, 0.0)
+                    self.assertLessEqual(skip_pct, 100.0)
+
+                    n_1bit = self.rabitq_stats.n_1bit_evaluations
+                    n_multibit = self.rabitq_stats.n_multibit_evaluations
+                    print(
+                        f"{index_type} nb_bits={nb_bits}: "
+                        f"n_1bit={n_1bit}, "
+                        f"n_multibit={n_multibit}, "
+                        f"skip={skip_pct:.1f}%"
+                    )
+
+
 if __name__ == "__main__":
     unittest.main()