Simplify RaBitQ slightly, improving speed and recall (#4550)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/4550
## Introduction
The [RaBitQ paper](https://arxiv.org/abs/2405.12497) describes a method for implementing vector inner products with bitwise operations, quantizing the query with a uniform scalar quantizer, then performing a **bitwise and** of a single bit position of the query vector with the document vector, then using a `popcount()` to accumulate the inner product. This is done for each query bit position in `[0, B_q)`, and an additional `popcount()` on the document vector alone. This additional `popcount()` can be avoided by rearranging the arithmetic.

The quantization of the query, `q_u`, is specified in equation 18 of the paper, reproduced here:
```lang=txt
q' = q @ P.T
v_l, v_h = min(q'), max(q')
delta = v_h - v_l
u_i[i] = sample(uniform(0,1))
q_u[i] = floor((q'[i] - v_l) / (v_h - v_l) + u[i])
```

The *unsigned dot product* is then estimated as (eq. 22):
```lang=txt
x_b[i] = (o @ P.T)[i] > 0 ? 1 : 0
q_b[i, j] = (q_u[i] >> j) & 1
unsigned_inner_product = sum_j(2**j * sum_i(x_b[i] * q[i, j])
```
Note that `x_b[i] * q[i, j]` is implemented as a bitwise AND.

This is then used to estimate the real inner product (eq. 20):
```lang=txt
inner_product_estimate
  = (2 * delta * unsigned_inner_product
     + 2 * v_l * sum_i(x_b[i])
     - delta * sum(q_u[i])
     - 1) / sqrt(D)
 ```

## Adjustment
The presence of `sum_i(x_b[i])` can be eliminated by redefining the quantization scheme used to construct `q_b[i,j]`. Instead of mapping the query to an unsigned integer, the original value is scaled, then rounded to a **signed, odd integer** in a manner very similar to a paper which builds upon RaBitQ (https://arxiv.org/abs/2409.09913).
```lang=txt
           |-------|-------|-------|-------|-------|-------|-------|-------|
 q'[i]*m: -8   v  -6   v  -4   v  -2   v   0   v  +2   v  +4   v  +6   v  +8
               |       |       |       |       |       |       |       |
 odd_int:     -7      -5      -3      -1      +1      +3      +5      +7
               |       |       |       |       |       |       |       |
               v       v       v       v       v       v       v       v
  q_s[i]:      0       1       2       3       4       5       6       7
```

This has the convenient property that `odd_int = sum(2**j * (2 * q_s[i, j] - 1))`. That is, each component is encoded as `(±4) + (±2) + (±1)`.

With both queries and documents now mapped to the signed, odd integer domain, the signed integer dot product can be computed more directly. All query bits represent `±2^j`, and all document bits represent `±1`, so the maximum possible dot product is `(2^0 + 2^1 + ... + 2^j) * D = (2^(B_q-1) - 1) * D`. Every *mismatched bit* represents a *mismatched sign*, so each such bit mismatch represents a replacement of an inner product term `+2^j` with `-2^j`, a reduction of `2 * 2^j`.

So, the signed dot product can be estimated as:
```lang=txt
signed_inner_product
  = (2 ** B_q - 1) * D
  - 2 * sum_j(2**j * sum_i(xor(q_s[i,j], x_b[i])))
```
Which can then be simply rescaled to to estimate the real inner product:
```
scalar = norm(q) / (norm(q_s) * sqrt(D))
inner_product_estimate = signed_inner_product * scalar
```

## Results
But because queries are expected to be centered around 0 due to random rotation, there is minimal loss in accuracy by defining the scalar quantizer such that the bias is fixed at zero. To the contrary, the centering of the scalar quantizer appears to **correct a systematic bias** in the RaBitQ quantizer which may lead to query components being disproportionately rounded up or down depending on the value of the `bias`. Particularly for smaller values of `B_q`, this imbalance can lead to significant systematic bias, adversely affecting recall. For `B_q = 1` in particular, experiments show that this centering **halves recall loss**.

Costs and recall losses are consistently reduced across all `B_q` values tested, for both a 1024-d and 128-d test corpus. Results for a 1024-d test:
* `B_q = 1`: 60.3% -> 86.16% recall (2.9x loss reduction)
* `B_q = 2`: 87.3% -> 87.5% recall (+0.2 pp)
* `B_q = 3`: 87.9% -> 88.1% recall (+0.2 pp)

## Interfaces / Compatibility
This new mode is gated by a parameter on the distance computer, exposed through `IVFRaBitQSearchParameters::centered` among other APIs. For fairer performance comparisons, the **existing implementation** is also sped up with a specialized SIMD implementations for the doc-side `popcount()`.

Reviewed By: luciang

Differential Revision: D80904214

fbshipit-source-id: 6075deeda13420ff08adac80d7ad5bf49552aff4
This commit is contained in:
Tom Jackson
2025-09-09 19:00:58 -07:00
committed by Facebook GitHub Bot
parent bc3e3a1ef1
commit 6470b8d9d0
9 changed files with 218 additions and 48 deletions
+12
View File
@@ -61,6 +61,15 @@ void bench_rabitq_and_dot_product(benchmark::State& state) {
});
}
void bench_rabitq_xor_dot_product(benchmark::State& state) {
bench_rabitq_generic(
state,
[](const uint8_t* q, const uint8_t* x, size_t size, size_t qb)
-> int64_t {
return rabitq::bitwise_xor_dot_product(q, x, size, qb);
});
}
void bench_rabitq_and_dot_product_with_sum(benchmark::State& state) {
bench_rabitq_generic(
state,
@@ -80,6 +89,9 @@ BENCHMARK(bench_rabitq_sum)->ArgsProduct({{0}, dims})->ArgNames({"qb", "d"});
BENCHMARK(bench_rabitq_and_dot_product)
->ArgsProduct({qbs, dims})
->ArgNames({"qb", "d"});
BENCHMARK(bench_rabitq_xor_dot_product)
->ArgsProduct({qbs, dims})
->ArgNames({"qb", "d"});
BENCHMARK(bench_rabitq_and_dot_product_with_sum)
->ArgsProduct({qbs, dims})
->ArgNames({"qb", "d"});
+12 -5
View File
@@ -155,15 +155,18 @@ struct RaBitInvertedListScanner : InvertedListScanner {
std::unique_ptr<FlatCodesDistanceComputer> dc;
uint8_t qb = 0;
bool centered = false;
explicit RaBitInvertedListScanner(
const IndexIVFRaBitQ& ivf_rabitq_in,
bool store_pairs = false,
const IDSelector* sel = nullptr,
uint8_t qb_in = 0)
uint8_t qb_in = 0,
bool centered = false)
: InvertedListScanner(store_pairs, sel),
ivf_rabitq{ivf_rabitq_in},
qb{qb_in} {
qb{qb_in},
centered(centered) {
keep_max = is_similarity_metric(ivf_rabitq.metric_type);
code_size = ivf_rabitq.code_size;
}
@@ -196,7 +199,7 @@ struct RaBitInvertedListScanner : InvertedListScanner {
// both query_vector and centroid are available!
// set up DistanceComputer
dc.reset(ivf_rabitq.rabitq.get_distance_computer(
qb, reconstructed_centroid.data()));
qb, reconstructed_centroid.data(), centered));
dc->set_query(query_vector.data());
}
@@ -208,12 +211,15 @@ InvertedListScanner* IndexIVFRaBitQ::get_InvertedListScanner(
const IDSelector* sel,
const IVFSearchParameters* search_params_in) const {
uint8_t used_qb = qb;
bool centered = false;
if (auto params = dynamic_cast<const IVFRaBitQSearchParameters*>(
search_params_in)) {
used_qb = params->qb;
centered = params->centered;
}
return new RaBitInvertedListScanner(*this, store_pairs, sel, used_qb);
return new RaBitInvertedListScanner(
*this, store_pairs, sel, used_qb, centered);
}
void IndexIVFRaBitQ::reconstruct_from_offset(
@@ -278,7 +284,8 @@ float IVFRaBitDistanceComputer::operator()(idx_t i) {
float distance = 0;
std::unique_ptr<FlatCodesDistanceComputer> dc(
parent->rabitq.get_distance_computer(parent->qb, centroid.data()));
parent->rabitq.get_distance_computer(
parent->qb, centroid.data(), /*centered=*/false));
dc->set_query(q);
distance = dc->distance_to_code(code);
+1
View File
@@ -19,6 +19,7 @@ namespace faiss {
struct IVFRaBitQSearchParameters : IVFSearchParameters {
uint8_t qb = 0;
bool centered = false;
};
// * by_residual is true, just by design
+13 -10
View File
@@ -55,16 +55,17 @@ void IndexRaBitQ::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
FlatCodesDistanceComputer* IndexRaBitQ::get_FlatCodesDistanceComputer() const {
FlatCodesDistanceComputer* dc =
rabitq.get_distance_computer(qb, center.data());
rabitq.get_distance_computer(qb, center.data(), centered);
dc->code_size = rabitq.code_size;
dc->codes = codes.data();
return dc;
}
FlatCodesDistanceComputer* IndexRaBitQ::get_quantized_distance_computer(
const uint8_t qb) const {
const uint8_t qb,
bool centered) const {
FlatCodesDistanceComputer* dc =
rabitq.get_distance_computer(qb, center.data());
rabitq.get_distance_computer(qb, center.data(), centered);
dc->code_size = rabitq.code_size;
dc->codes = codes.data();
return dc;
@@ -76,6 +77,7 @@ struct Run_search_with_dc_res {
using T = void;
uint8_t qb = 0;
bool centered = false;
template <class BlockResultHandler>
void f(BlockResultHandler& res, const IndexRaBitQ* index, const float* xq) {
@@ -87,7 +89,7 @@ struct Run_search_with_dc_res {
#pragma omp parallel // if (res.nq > 100)
{
std::unique_ptr<FlatCodesDistanceComputer> dc(
index->get_quantized_distance_computer(qb));
index->get_quantized_distance_computer(qb, centered));
SingleResultHandler resi(res);
#pragma omp for
for (int64_t q = 0; q < res.nq; q++) {
@@ -114,14 +116,15 @@ void IndexRaBitQ::search(
float* distances,
idx_t* labels,
const SearchParameters* params_in) const {
uint8_t used_qb = qb;
if (auto params = dynamic_cast<const RaBitQSearchParameters*>(params_in)) {
used_qb = params->qb;
}
const IDSelector* sel = (params_in != nullptr) ? params_in->sel : nullptr;
Run_search_with_dc_res r;
r.qb = used_qb;
if (auto params = dynamic_cast<const RaBitQSearchParameters*>(params_in)) {
r.qb = params->qb;
r.centered = params->centered;
} else {
r.qb = this->qb;
r.centered = this->centered;
}
dispatch_knn_ResultHandler(
n, distances, labels, k, metric_type, sel, r, this, x);
+6 -1
View File
@@ -14,6 +14,7 @@ namespace faiss {
struct RaBitQSearchParameters : SearchParameters {
uint8_t qb = 0;
bool centered = false;
};
struct IndexRaBitQ : IndexFlatCodes {
@@ -26,6 +27,9 @@ struct IndexRaBitQ : IndexFlatCodes {
// use '0' to disable quantization and use raw fp32 values.
uint8_t qb = 0;
// quantize the query with a zero-centered scalar quantizer.
bool centered = false;
IndexRaBitQ();
explicit IndexRaBitQ(idx_t d, MetricType metric = METRIC_L2);
@@ -42,7 +46,8 @@ struct IndexRaBitQ : IndexFlatCodes {
// returns a quantized-to-qb bits DC if qb_in > 0
// returns a default fp32-based DC if qb_in == 0
FlatCodesDistanceComputer* get_quantized_distance_computer(
const uint8_t qb_in) const;
const uint8_t qb_in,
bool centered) const;
// Don't rely on sa_decode(), bcz it is good for IP, but not for L2.
// As a result, use get_FlatCodesDistanceComputer() for the search.
+71 -25
View File
@@ -32,6 +32,8 @@ struct QueryFactorsData {
float qr_to_c_L2sqr = 0;
float qr_norm_L2sqr = 0;
float int_dot_scale = 1;
};
static size_t get_code_size(const size_t d) {
@@ -310,6 +312,7 @@ struct RaBitDistanceComputerQ : RaBitDistanceComputer {
// the number of bits for SQ quantization of the query (qb > 0)
uint8_t qb = 8;
bool centered = false;
// the smallest value divisible by 8 that is not smaller than dim
size_t popcount_aligned_dim = 0;
@@ -333,20 +336,30 @@ float RaBitDistanceComputerQ::distance_to_code(const uint8_t* code) {
const uint8_t* binary_data = code;
const FactorsData* fac = reinterpret_cast<const FactorsData*>(code + size);
// See RaBitDistanceComputerNotQ::distance_to_code() for baseline code.
float dot_qo = rabitq::bitwise_and_dot_product(
rearranged_rotated_qq.data(), binary_data, size, qb);
// It was a willful decision (after the discussion) to not to pre-cache
// the sum of all bits, just in order to reduce the overhead per vector.
auto sum_q = rabitq::popcount(binary_data, size);
// this is ||or - c||^2 - (IP ? ||or||^2 : 0)
float final_dot = 0;
// dot-product itself
final_dot += query_fac.c1 * dot_qo;
// normalizer coefficients
final_dot += query_fac.c2 * sum_q;
// normalizer coefficients
final_dot -= query_fac.c34;
if (centered) {
int64_t int_dot = ((1 << qb) - 1) * d;
int_dot -= 2 *
rabitq::bitwise_xor_dot_product(
rearranged_rotated_qq.data(), binary_data, size, qb);
final_dot += int_dot * query_fac.int_dot_scale;
} else {
// See RaBitDistanceComputerNotQ::distance_to_code() for baseline code.
auto dot_qo = rabitq::bitwise_and_dot_product(
rearranged_rotated_qq.data(), binary_data, size, qb);
// It was a willful decision (after the discussion) to not to pre-cache
// the sum of all bits, just in order to reduce the overhead per vector.
// process 64-bit popcounts
auto sum_q = rabitq::popcount(binary_data, size);
// dot-product itself
final_dot += query_fac.c1 * dot_qo;
// normalizer coefficients
final_dot += query_fac.c2 * sum_q;
// normalizer coefficients
final_dot -= query_fac.c34;
}
// this is ||or - c||^2 - (IP ? ||or||^2 : 0)
const float or_c_l2sqr = fac->or_minus_c_l2sqr;
@@ -370,11 +383,30 @@ float RaBitDistanceComputerQ::distance_to_code(const uint8_t* code) {
}
}
namespace {
// Ideal quantizer radii for quantizers of 1..8 bits, optimized to minimize L2
// reconstruction error.
const float z_max_by_qb[8] = {
0.79688, // qb = 1.
1.49375,
2.05078,
2.50938,
2.91250,
3.26406,
3.59844,
3.91016, // qb = 8.
};
} // namespace
void RaBitDistanceComputerQ::set_query(const float* x) {
FAISS_ASSERT(x != nullptr);
FAISS_ASSERT(
(metric_type == MetricType::METRIC_L2 ||
metric_type == MetricType::METRIC_INNER_PRODUCT));
FAISS_THROW_IF_NOT(qb <= 8);
FAISS_THROW_IF_NOT(qb > 0);
// compute the distance from the query to the centroid
if (centroid != nullptr) {
@@ -398,26 +430,36 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
// quantize the query. compute min and max
float v_min = std::numeric_limits<float>::max();
float v_max = std::numeric_limits<float>::lowest();
for (size_t i = 0; i < d; i++) {
const float v_q = rotated_q[i];
v_min = std::min(v_min, v_q);
v_max = std::max(v_max, v_q);
if (centered) {
float z_max = z_max_by_qb[qb - 1];
float v_radius = z_max * std::sqrt(query_fac.qr_to_c_L2sqr / d);
v_min = -v_radius;
v_max = v_radius;
} else {
for (size_t i = 0; i < d; i++) {
const float v_q = rotated_q[i];
v_min = std::min(v_min, v_q);
v_max = std::max(v_max, v_q);
}
}
const float pow_2_qb = 1 << qb;
const float delta = (v_max - v_min) / (pow_2_qb - 1);
const uint8_t max_code = (1 << qb) - 1;
const float delta = (v_max - v_min) / max_code;
const float inv_delta = 1.0f / delta;
size_t sum_qq = 0;
int64_t sum2_signed_odd_int = 0;
for (int32_t i = 0; i < d; i++) {
const float v_q = rotated_q[i];
// a default non-randomized SQ
const int v_qq = std::round((v_q - v_min) * inv_delta);
rotated_qq[i] = std::min(255, std::max(0, v_qq));
const uint8_t v_qq = std::clamp<float>(
std::round((v_q - v_min) * inv_delta), 0, max_code);
rotated_qq[i] = v_qq;
sum_qq += v_qq;
if (centered) {
int64_t signed_odd_int = int64_t(v_qq) * 2 - max_code;
sum2_signed_odd_int += signed_odd_int * signed_odd_int;
}
}
// rearrange the query vector
@@ -438,6 +480,8 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
query_fac.c1 = 2 * delta * inv_d;
query_fac.c2 = 2 * v_min * inv_d;
query_fac.c34 = inv_d * (delta * sum_qq + d * v_min);
query_fac.int_dot_scale =
std::sqrt(query_fac.qr_to_c_L2sqr / (sum2_signed_odd_int * d));
if (metric_type == MetricType::METRIC_INNER_PRODUCT) {
// precompute if needed
@@ -447,7 +491,8 @@ void RaBitDistanceComputerQ::set_query(const float* x) {
FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
uint8_t qb,
const float* centroid_in) const {
const float* centroid_in,
bool centered) const {
if (qb == 0) {
auto dc = std::make_unique<RaBitDistanceComputerNotQ>();
dc->metric_type = metric_type;
@@ -461,6 +506,7 @@ FlatCodesDistanceComputer* RaBitQuantizer::get_distance_computer(
dc->d = d;
dc->centroid = centroid_in;
dc->qb = qb;
dc->centered = centered;
return dc.release();
}
+2 -1
View File
@@ -72,7 +72,8 @@ struct RaBitQuantizer : Quantizer {
// specify qb > 0 to have SQ qb-bits query
FlatCodesDistanceComputer* get_distance_computer(
uint8_t qb,
const float* centroid_in = nullptr) const;
const float* centroid_in = nullptr,
bool centered = false) const;
};
} // namespace faiss
+83
View File
@@ -293,6 +293,89 @@ inline uint64_t bitwise_and_dot_product(
return sum;
}
/**
* Compute dot product between query and binary data using popcount operations.
*
* @param query Pointer to rearranged rotated query data
* @param data Pointer to binary data
* @param d Dimension
* @param qb Number of quantization bits
* @return Unsigned integer dot product
*/
inline uint64_t bitwise_xor_dot_product(
const uint8_t* query,
const uint8_t* data,
size_t size,
size_t qb) {
uint64_t sum = 0;
size_t offset = 0;
#if defined(__AVX512F__)
// Handle 512-bit chunks.
if (size_t step = 512 / 8; offset + step <= size) {
__m512i sum_512 = _mm512_setzero_si512();
for (; offset + step <= size; offset += step) {
__m512i v_x = _mm512_loadu_si512((const __m512i*)(data + offset));
for (int j = 0; j < qb; j++) {
__m512i v_q = _mm512_loadu_si512(
(const __m512i*)(query + j * size + offset));
__m512i v_xor = _mm512_xor_si512(v_q, v_x);
__m512i v_popcnt = popcount_512(v_xor);
__m512i v_shifted = _mm512_slli_epi64(v_popcnt, j);
sum_512 = _mm512_add_epi64(sum_512, v_shifted);
}
}
sum += _mm512_reduce_add_epi64(sum_512);
}
#endif
#if defined(__AVX2__)
if (size_t step = 256 / 8; offset + step <= size) {
__m256i sum_256 = _mm256_setzero_si256();
for (; offset + step <= size; offset += step) {
__m256i v_x = _mm256_loadu_si256((const __m256i*)(data + offset));
for (int j = 0; j < qb; j++) {
__m256i v_q = _mm256_loadu_si256(
(const __m256i*)(query + j * size + offset));
__m256i v_xor = _mm256_xor_si256(v_q, v_x);
__m256i v_popcnt = popcount_256(v_xor);
__m256i v_shifted = _mm256_slli_epi64(v_popcnt, j);
sum_256 = _mm256_add_epi64(sum_256, v_shifted);
}
}
sum += reduce_add_256(sum_256);
}
#endif
#if defined(__SSE4_1__)
__m128i sum_128 = _mm_setzero_si128();
for (size_t step = 128 / 8; offset + step <= size; offset += step) {
__m128i v_x = _mm_loadu_si128((const __m128i*)(data + offset));
for (int j = 0; j < qb; j++) {
__m128i v_q = _mm_loadu_si128(
(const __m128i*)(query + j * size + offset));
__m128i v_xor = _mm_xor_si128(v_q, v_x);
__m128i v_popcnt = popcount_128(v_xor);
__m128i v_shifted = _mm_slli_epi64(v_popcnt, j);
sum_128 = _mm_add_epi64(sum_128, v_shifted);
}
}
sum += reduce_add_128(sum_128);
#endif
for (size_t step = 64 / 8; offset + step <= size; offset += step) {
const auto yv = *(const uint64_t*)(data + offset);
for (int j = 0; j < qb; j++) {
const auto qv = *(const uint64_t*)(query + j * size + offset);
sum += __builtin_popcountll(qv ^ yv) << j;
}
}
for (; offset < size; ++offset) {
const auto yv = *(data + offset);
for (int j = 0; j < qb; j++) {
const auto qv = *(query + j * size + offset);
sum += __builtin_popcount(qv ^ yv) << j;
}
}
return sum;
}
inline uint64_t popcount(const uint8_t* data, size_t size) {
uint64_t sum = 0;
size_t offset = 0;
+18 -6
View File
@@ -223,21 +223,25 @@ class TestRaBitQ(unittest.TestCase):
index_rbq.train(ds.get_train())
index_rbq.add(ds.get_database())
for qb in [1, 2, 3, 4, 8]:
params = faiss.RaBitQSearchParameters(qb=qb)
def test(params):
_, I_rbq = index_rbq.search(ds.get_queries(), k, params=params)
# ensure that RaBitQ and PQ are relatively close
loss_rbq = 1 - eval_I(I_rbq)
ratio_threshold = 2 ** (1 / qb)
print(
f"{random_rotate=:1}, {params.qb=}: "
f"{random_rotate=:1}, {params.qb=}, {params.centered=:1}: "
f"{loss_rbq=:5.3f} = loss_pq * {loss_rbq / loss_pq:5.3f}"
f" < {ratio_threshold=:.2f}"
)
np.testing.assert_(loss_rbq < loss_pq * ratio_threshold)
for qb in [1, 2, 3, 4, 8]:
print()
for centered in [False, True]:
test(faiss.RaBitQSearchParameters(qb=qb, centered=centered))
def test_comparison_vs_pq_L2(self):
self.do_comparison_vs_pq_test(faiss.METRIC_L2)
@@ -387,21 +391,29 @@ class TestIVFRaBitQ(unittest.TestCase):
print(f"{random_rotate=:1}, {loss_pq=:5.3f}")
np.testing.assert_(loss_pq < 0.25, f"{loss_pq}")
for qb in [1, 2, 3, 4, 8]:
params = faiss.IVFRaBitQSearchParameters(nprobe=nprobe, qb=qb)
def test(params):
D_rbq, I_rbq = index_rbq.search(xq, k, params=params)
# ensure that RaBitQ and PQ are relatively close
loss_rbq = 1 - eval_I(I_rbq)
ratio_threshold = 2 ** (1 / qb)
print(
f"{random_rotate=:1}, {params.qb=}: "
f"{random_rotate=:1}, {params.qb=}, {params.centered=:1}: "
f"{loss_rbq=:5.3f} = loss_pq * {loss_rbq / loss_pq:5.3f}"
f" < {ratio_threshold=:.2f}"
)
np.testing.assert_(loss_rbq < loss_pq * ratio_threshold)
for qb in [1, 2, 3, 4, 8]:
print()
for centered in [False, True]:
test(
faiss.IVFRaBitQSearchParameters(
nprobe=nprobe, qb=qb, centered=centered
)
)
def test_comparison_vs_pq_L2(self):
self.do_comparison_vs_pq_test(faiss.METRIC_L2)