RaBitQ implementation #4235

alexanderguzhva · 2025-03-11T12:24:15Z

This is a reference implementation of the https://arxiv.org/pdf/2405.12497

Jianyang Gao, Cheng Long, "RaBitQ: Quantizing High-Dimensional Vectors with a Theoretical Error Bound for Approximate Nearest Neighbor Search".

The goal is to correctly set up the internals using Faiss.

The following comments for the implementation:

The code does not include the computations for the symmetric distance, because it is absent in the original article. This can be added later, though.
The original RaBitQ includes random matrix rotation as a part of it, but I've decided to rely on external faiss::IndexPreTransform and faiss::RandomRotationMatrix facilities.
Certain features required internal changes in faiss::IndexIVF, but I did that as least invasive as possible, without breaking the backward compatibility.
Not sure about naming convensions, maybe certain classes and structures need to be renamed
METRIC_INNER_PRODUCT is supported as well
More unit tests are needed?
I did not bring any hardware-specific optimizations, bcz this is a reference implementation. Certain simdlib facilities may be added later, if needed

Here's how to use IndexRaBitQ

        ds = datasets.SyntheticDataset(...)

        index_rbq = faiss.IndexRaBitQ(ds.d, faiss.METRIC_L2)
        index_rbq.qb = 8

        # wrap with random rotations
        rrot = faiss.RandomRotationMatrix(ds.d, ds.d)
        rrot.init(rrot_seed)

        index_cand = faiss.IndexPreTransform(rrot, index_rbq)
        index_cand.train(ds.get_train())
        index_cand.add(ds.get_database())

Here's how to use IndexIVFRaBitQ

        ds = datasets.SyntheticDataset(...)

        index_flat = faiss.IndexFlat(ds.d, faiss.METRIC_L2)
        index_rbq = faiss.IndexIVFRaBitQ(index_flat, ds.d, nlist, faiss.METRIC_L2)
        index_rbq.qb = 8

        # wrap with random rotations
        rrot = faiss.RandomRotationMatrix(ds.d, ds.d)
        rrot.init(rrot_seed)

        index_cand = faiss.IndexPreTransform(rrot, index_rbq)
        index_cand.train(ds.get_train())
        index_cand.add(ds.get_database())

mdouze

Looks good! Thanks for deferring the SIMD optimization to later.
I left a few comments.

mdouze · 2025-03-11T13:02:38Z

faiss/IndexIVF.h

     */
    virtual InvertedListScanner* get_InvertedListScanner(
            bool store_pairs = false,
            const IDSelector* sel = nullptr) const;

+    /** Get a scanner for this index (store_pairs means ignore labels).


I would prefer to replace get_InvertedListScanner altogether with the version that takes IVFSearchParameters (and no sel, since IDSelector is a field of IVFSearchParameters)

what about the backward compatibility? This is why I introduced get_InvertedListScanner_2. I mean that I can upgrade the method signature Faiss-wide, but what about the external code?

mdouze · 2025-03-11T13:04:23Z

faiss/impl/RaBitQuantizer.cpp

+    FAISS_ASSERT(codes != nullptr);
+    FAISS_ASSERT(x != nullptr);
+
+    if (n == 0) {


move this test before the asserts

mdouze · 2025-03-11T13:05:21Z

faiss/impl/RaBitQuantizer.cpp

+
+struct FactorsData {
+    // ||or - c||
+    float factor_0 = 0;


why not give the fields proper names? Otherwise you might as well use float factors[4]

mdouze · 2025-03-11T13:07:13Z

faiss/impl/RaBitQuantizer.cpp

+        const uint8_t* query_j = rearranged_rotated_qq.data() + j * di_8b;
+
+        // process 64-bit popcounts
+        unsigned long long count = 0;


please use explicity sized integer types (eg uint64_t)

will do. As far as I remember, I've used unsigned long long, bcz there were problems with compilations on MacOS for these __builtin_popcount functions

mdouze · 2025-03-11T13:16:22Z

faiss/impl/RaBitQuantizer.cpp

+    float factor_2 = 0;
+    // ||or||^2
+    float factor_3 = 0;
+};


In my implementation there are only 2 floats per database vector. Do you store additional ones for efficiency?

alexanderguzhva · 2025-03-11T18:36:38Z

@mdouze two more comments after the discussion

First. IndexIVF::get_InvertedListScanner() signature is made into

    virtual InvertedListScanner* get_InvertedListScanner(
            bool store_pairs = false,
            const IDSelector* sel = nullptr,
            const IVFSearchParameters* params = nullptr) const;

because of the following logic that can override sel

void IndexIVF::search_preassigned(...) const {
    ...
    IDSelector* sel = params ? params->sel : nullptr;
    const IDSelectorRange* selr = dynamic_cast<const IDSelectorRange*>(sel);
    if (selr) {
        if (selr->assume_sorted) {
            sel = nullptr; // use special IDSelectorRange processing
        } else {
            selr = nullptr; // use generic processing
        }
    }
    ....
}

Please let me know your thoughts.

Second. RaBitQ uses 3 factors

struct FactorsData {
    float or_minus_c_l2sqr = 0;
    float dp_multiplier = 0;
    // this is needed to support BOTH L2 and IP on the same data
    float or_l2sqr = 0;
};

The third one or_l2sqr is needed to support both L2 and IP, similar to how PQ / SQ can use the same data for different metrics. So, these three numbers per vector are independent from a chosen metric. These three factors can be reduced to two, if we make a decision to make factors dependent from the chosen metric.
Please let me know if you'd like to have 2 or 3 factors per vector. Just double-checking.

mdouze · 2025-03-21T09:12:38Z

Thanks for the generously commented code!
Sorry I forgot to add I'd prefer to have 2 factors (separate metrics for L2 and IP), since the metric is known at index build time. In the long run, it is important to have a compact representation.
I can live with your change to the get_InvertedListScanner.
If you can implement the 2-factor version, I'll accept the diff.

mdouze · 2025-03-21T09:15:08Z

faiss/impl/platform_macros.h

@@ -76,6 +76,7 @@ inline int __builtin_clzll(uint64_t x) {

 #define __builtin_popcount __popcnt


I'd prefer to call it popcount32 and popcount64 to have explicit bit sizes

facebook-github-bot · 2025-03-21T18:36:35Z

@junjieqi has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

alexanderguzhva · 2025-03-22T18:14:36Z

updated to 2-factor version, added more unit tests

facebook-github-bot · 2025-03-24T17:11:44Z

@junjieqi has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

facebook-github-bot · 2025-03-29T19:29:24Z

@junjieqi merged this pull request in 3a49130.

Summary: it seems that 2937f94 was not included in #4235. This PR fixes this problem. junjieqi Pull Request resolved: #4261 Reviewed By: gtwang01 Differential Revision: D72180816 Pulled By: junjieqi fbshipit-source-id: 55e156a3499fda6f8cdbb99ed941a3cbdd721417

facebook-github-bot added the CLA Signed label Mar 11, 2025

mdouze reviewed Mar 11, 2025

View reviewed changes

alexanderguzhva mentioned this pull request Mar 11, 2025

RaBitQ implementation zilliztech/knowhere#1115

Merged

alexanderguzhva force-pushed the rabitq branch from 1eb0fa8 to 68c3c36 Compare March 11, 2025 18:30

alexanderguzhva force-pushed the rabitq branch 2 times, most recently from 346830c to b0a8619 Compare March 15, 2025 00:22

mdouze reviewed Mar 21, 2025

View reviewed changes

alexanderguzhva added 7 commits March 22, 2025 14:07

RaBitQ implementation

d1cb949

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

fix formatting and compilation

4f6a238

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

fix windows compilation

e65c1c3

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

update factors

f149175

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

modify IndexIVF interface

c933c99

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

fix formatting

eda7764

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

refactor to extra 2 fp32 values per vector, add unit tests

8ee43af

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

alexanderguzhva force-pushed the rabitq branch from 0c0f3f3 to 8ee43af Compare March 22, 2025 18:07

fix a serialization problem

2937f94

Signed-off-by: Alexandr Guzhva <alexanderguzhva@gmail.com>

facebook-github-bot closed this in 3a49130 Mar 29, 2025

facebook-github-bot added the Merged label Mar 29, 2025

alexanderguzhva mentioned this pull request Mar 31, 2025

fix a serialization problem in RaBitQ #4261

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

RaBitQ implementation #4235

RaBitQ implementation #4235

Uh oh!

alexanderguzhva commented Mar 11, 2025

Uh oh!

mdouze left a comment

Uh oh!

mdouze Mar 11, 2025

Uh oh!

alexanderguzhva Mar 11, 2025

Uh oh!

mdouze Mar 11, 2025

Uh oh!

mdouze Mar 11, 2025

Uh oh!

mdouze Mar 11, 2025

Uh oh!

alexanderguzhva Mar 11, 2025

Uh oh!

mdouze Mar 11, 2025

Uh oh!

alexanderguzhva commented Mar 11, 2025

Uh oh!

mdouze commented Mar 21, 2025

Uh oh!

mdouze Mar 21, 2025

Uh oh!

facebook-github-bot commented Mar 21, 2025

Uh oh!

alexanderguzhva commented Mar 22, 2025

Uh oh!

facebook-github-bot commented Mar 24, 2025

Uh oh!

facebook-github-bot commented Mar 29, 2025

Uh oh!

Uh oh!

		@@ -76,6 +76,7 @@ inline int __builtin_clzll(uint64_t x) {

		#define __builtin_popcount __popcnt

RaBitQ implementation #4235

RaBitQ implementation #4235

Uh oh!

Conversation

alexanderguzhva commented Mar 11, 2025

Uh oh!

mdouze left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

alexanderguzhva commented Mar 11, 2025

Uh oh!

mdouze commented Mar 21, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

facebook-github-bot commented Mar 21, 2025

Uh oh!

alexanderguzhva commented Mar 22, 2025

Uh oh!

facebook-github-bot commented Mar 24, 2025

Uh oh!

facebook-github-bot commented Mar 29, 2025

Uh oh!

Uh oh!