facebookresearch/faiss
1
0
Fork 0
mirror of https://github.com/facebookresearch/faiss.git synced 2026-10-12 10:59:57 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add staged RaBitQ storage to HNSW (#5526)

Browse Source 
Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5526

Reviewed By: junjieqi, mnorris11

Differential Revision: D116650696

Pulled By: alibeklfc

fbshipit-source-id: 68ec701da84bd72a71e83ae5fadc041b2890eefe
This commit is contained in:
lyang24
2026-09-08 08:53:40 -07:00
committed by meta-codesync[bot]
parent 2ed4c106e9
commit 7f4cedda9c
16 changed files with 770 additions and 88 deletions
Download Patch File Download Diff File Expand all files Collapse all files
faiss/IndexHNSW.cpp
+58 -2
View File
@@ -7,6 +7,8 @@
#include <faiss/IndexHNSW.h>
#include <faiss/IndexRaBitQ.h>
#include <omp.h>
#include <atomic>
#include <cinttypes>
@@ -577,6 +579,7 @@ void hnsw_search(
}
}
size_t n1 = 0, n2 = 0, ndis = 0, nhops = 0;
size_t n_rabitq_1bit = 0, n_rabitq_refine = 0;
idx_t check_period = InterruptCallback::get_period_hint(
hnsw.max_level * index->d * efSearch);
@@ -602,7 +605,9 @@ void hnsw_search(
omp_capture_exception(ex, [&] { interrupt = true; });
}
#pragma omp for reduction(+ : n1, n2, ndis, nhops) schedule(guided)
#pragma omp for reduction( \
+ : n1, n2, ndis, nhops, n_rabitq_1bit, n_rabitq_refine) \
schedule(guided)
for (idx_t i = i0; i < i1; i++) {
if (interrupt.load(std::memory_order_relaxed)) {
continue;
@@ -610,6 +615,10 @@ void hnsw_search(
try {
res->begin(i);
dis->set_query(x + i * index->d);
auto* rq = dynamic_cast<RaBitQDistanceComputer*>(dis.get());
if (rq) {
rq->stats.reset();
}
HNSWStats stats =
hnsw.search(*dis, index, *res, *vt, params);
@@ -617,6 +626,10 @@ void hnsw_search(
n2 += stats.n2;
ndis += stats.ndis;
nhops += stats.nhops;
if (rq) {
n_rabitq_1bit += rq->stats.n_1bit;
n_rabitq_refine += rq->stats.n_refine;
}
res->end();
vt->advance();
} catch (...) {
@@ -629,6 +642,7 @@ void hnsw_search(
}
hnsw_stats.combine({n1, n2, ndis, nhops});
rabitq_stats.add({n_rabitq_1bit, n_rabitq_refine});
}
} // anonymous namespace
@@ -799,6 +813,7 @@ void IndexHNSW::search_level_0(
{
std::unique_ptr<DistanceComputer> qdis;
HNSWStats search_stats;
RaBitQStats rq_search_stats;
VisitedTable* vt = nullptr;
std::unique_ptr<typename RH::SingleResultHandler> res;
try {
@@ -818,6 +833,11 @@ void IndexHNSW::search_level_0(
try {
res->begin(i);
qdis->set_query(x + i * d);
auto* rq =
dynamic_cast<RaBitQDistanceComputer*>(qdis.get());
if (rq) {
rq->stats.reset();
}
hnsw.search_level_0(
*qdis.get(),
@@ -829,6 +849,9 @@ void IndexHNSW::search_level_0(
search_stats,
*vt,
params);
if (rq) {
rq_search_stats.add(rq->stats);
}
res->end();
vt->advance();
} catch (...) {
@@ -838,6 +861,7 @@ void IndexHNSW::search_level_0(
#pragma omp critical
{
hnsw_stats.combine(search_stats);
rabitq_stats.add(rq_search_stats);
}
}
omp_rethrow_if_exception(ex);
@@ -1061,7 +1085,7 @@ IndexHNSWFlatPanorama::IndexHNSWFlatPanorama(
// Enable Panorama search mode.
// This is not ideal, but is still more simple than making a subclass of
// HNSW and overriding the search logic.
hnsw.is_panorama = true;
hnsw.search_method = HNSW::SM_PANORAMA;
}
void IndexHNSWFlatPanorama::add(idx_t n, const float* x) {
@@ -1128,6 +1152,38 @@ IndexHNSWSQ::IndexHNSWSQ(
IndexHNSWSQ::IndexHNSWSQ() = default;
/**************************************************************
* IndexHNSWRaBitQ implementation
**************************************************************/
IndexHNSWRaBitQ::IndexHNSWRaBitQ() = default;
namespace {
IndexRaBitQ* make_hnsw_rabitq_storage(
int d,
uint8_t nb_bits,
MetricType metric) {
FAISS_THROW_IF_NOT_MSG(
metric == METRIC_L2, "IndexHNSWRaBitQ supports only the L2 metric");
return new IndexRaBitQ(d, metric, nb_bits);
}
} // namespace
IndexHNSWRaBitQ::IndexHNSWRaBitQ(
int d,
int M,
uint8_t nb_bits,
MetricType metric)
: IndexHNSW(make_hnsw_rabitq_storage(d, nb_bits, metric), M) {
own_fields = true;
is_trained = storage->is_trained;
// 1-bit codes store plain SignBitFactors with no f_error, so there is no
// bound to prune with and the staged path does not apply.
hnsw.search_method = nb_bits >= 2 ? HNSW::SM_RABITQ : HNSW::SM_DEFAULT;
}
/**************************************************************
* IndexHNSW2Level implementation
**************************************************************/
faiss/IndexHNSW.h
+30
View File
@@ -216,6 +216,36 @@ struct IndexHNSWSQ : IndexHNSW {
MetricType metric = METRIC_L2);
};
/** HNSW index whose storage is RaBitQ-compressed.
*
* The graph is built directly from the compressed codes. Neighbor-diversity
* pruning uses a fast 1-bit symmetric L2 estimator, including when the storage
* contains multi-bit codes.
*
* With nb_bits >= 2 the codes carry a per-vector error factor, so search uses
* the staged search method: a 1-bit estimate for every neighbor and the full
* multi-bit distance only for candidates the error bound cannot rule out.
* nb_bits = 1 has no error factor and uses ordinary HNSW search.
*/
struct IndexHNSWRaBitQ : IndexHNSW {
IndexHNSWRaBitQ();
IndexHNSWRaBitQ(
int d,
int M,
uint8_t nb_bits = 1,
MetricType metric = METRIC_L2);
IndexHNSWRaBitQ& operator=(const IndexHNSWRaBitQ&) = delete;
private:
// clone_index() replaces the shallow-copied storage with a deep copy
// before returning it. Keep ordinary C++ copies from sharing ownership.
IndexHNSWRaBitQ(const IndexHNSWRaBitQ&) = default;
#ifndef SWIG
friend IndexHNSW* clone_IndexHNSW(const IndexHNSW* index);
#endif
};
/** 2-level code structure with fast random access
*/
struct IndexHNSW2Level : IndexHNSW {
faiss/clone_index.cpp
+3
View File
@@ -38,6 +38,7 @@
#include <faiss/IndexPQ.h>
#include <faiss/IndexPQFastScan.h>
#include <faiss/IndexPreTransform.h>
#include <faiss/IndexRaBitQ.h>
#include <faiss/IndexRefine.h>
#include <faiss/IndexRowwiseMinMax.h>
#include <faiss/IndexScalarQuantizer.h>
@@ -139,6 +140,7 @@ IndexIDMap* clone_IndexIDMap(const IndexIDMap* im) {
IndexHNSW* clone_IndexHNSW(const IndexHNSW* ihnsw) {
TRYCLONE(IndexHNSW2Level, ihnsw)
TRYCLONE(IndexHNSWFlatPanorama, ihnsw)
TRYCLONE(IndexHNSWRaBitQ, ihnsw)
TRYCLONE(IndexHNSWFlat, ihnsw)
TRYCLONE(IndexHNSWPQ, ihnsw)
TRYCLONE(IndexHNSWSQ, ihnsw)
@@ -294,6 +296,7 @@ Index* Cloner::clone_Index(const Index* index) {
TRYCLONE(IndexEDEN, index)
TRYCLONE(IndexScalarQuantizer, index)
TRYCLONE(IndexRaBitQ, index)
TRYCLONE(MultiIndexQuantizer, index)
if (const IndexIVF* ivf = dynamic_cast<const IndexIVF*>(index)) {
faiss/impl/HNSW.cpp
+226 -71
View File
@@ -17,6 +17,7 @@
#include <faiss/impl/DistanceComputer.h>
#include <faiss/impl/IDSelector.h>
#include <faiss/impl/RaBitQuantizer.h>
#include <faiss/impl/ResultHandler.h>
#include <faiss/impl/VisitedTable.h>
#include <faiss/impl/hnsw/LockVector.h>
@@ -1116,13 +1117,82 @@ inline void extract_search_params(
}
}
/** Templated body of `search_from_candidates` — instantiated once per
* VisitedTable subclass × comparator.
struct DefaultCandidateDistanceEvaluator {
DistanceComputer& qdis;
template <typename GetThreshold, typename AddResult>
size_t evaluate(
const storage_idx_t* ids,
int count,
GetThreshold&& /* get_threshold */,
AddResult&& add_result) {
if (count == 4) {
float distances[4];
qdis.distances_batch_4(
ids[0],
ids[1],
ids[2],
ids[3],
distances[0],
distances[1],
distances[2],
distances[3]);
for (int i = 0; i < 4; ++i) {
add_result(ids[i], distances[i]);
}
} else {
for (int i = 0; i < count; ++i) {
add_result(ids[i], qdis(ids[i]));
}
}
return count;
}
};
/** RaBitQ's staged evaluator plugs into the same graph traversal as the
* default distance computer. Each candidate first receives a 1-bit estimate;
* the full code is evaluated only when its error bound can still beat the
* current result threshold.
*/
template <typename VTType, class C>
struct RaBitQCandidateDistanceEvaluator {
RaBitQDistanceComputer& rq;
const bool is_similarity;
template <typename GetThreshold, typename AddResult>
size_t evaluate(
const storage_idx_t* ids,
int count,
GetThreshold&& get_threshold,
AddResult&& add_result) {
size_t ndis = 0;
for (int i = 0; i < count; ++i) {
const uint8_t* code =
rq.codes + static_cast<size_t>(ids[i]) * rq.code_size;
const float estimate = rq.distance_to_code_1bit(code);
rq.stats.n_1bit++;
ndis++;
float distance = estimate;
if (rq.should_refine(
code, estimate, get_threshold(), is_similarity)) {
distance = rq.distance_to_code_full(code);
rq.stats.n_refine++;
ndis++;
}
add_result(ids[i], distance);
}
return ndis;
}
};
/** Templated body shared by the default and staged RaBitQ searches. The
* evaluator changes only how a group of unvisited neighbors receives its
* distance; traversal, result handling and stopping conditions stay identical.
*/
template <typename VTType, class C, typename DistanceEvaluator>
int search_from_candidates_fixVT(
const HNSW& hnsw,
DistanceComputer& qdis,
DistanceEvaluator& evaluator,
ResultHandler& res,
MinimaxHeapT<HC_for<C>>& candidates,
VTType& vt,
@@ -1131,7 +1201,7 @@ int search_from_candidates_fixVT(
int nres_in,
const SearchParameters* params) {
int nres = nres_in;
int ndis = 0;
size_t ndis = 0;
bool do_dis_check;
int efSearch;
@@ -1189,7 +1259,7 @@ int search_from_candidates_fixVT(
}
int counter = 0;
size_t saved_j[4];
storage_idx_t saved_j[4];
threshold = res.threshold;
@@ -1212,32 +1282,18 @@ int search_from_candidates_fixVT(
counter += vt.set(v1) ? 1 : 0;
if (counter == 4) {
float dis[4];
qdis.distances_batch_4(
saved_j[0],
saved_j[1],
saved_j[2],
saved_j[3],
dis[0],
dis[1],
dis[2],
dis[3]);
for (size_t id4 = 0; id4 < 4; id4++) {
add_to_heap(saved_j[id4], dis[id4]);
}
ndis += 4;
ndis += evaluator.evaluate(
saved_j,
counter,
[&] { return threshold; },
add_to_heap);
counter = 0;
}
}
for (int icnt = 0; icnt < counter; icnt++) {
float dis = qdis(saved_j[icnt]);
add_to_heap(saved_j[icnt], dis);
ndis += 1;
if (counter > 0) {
ndis += evaluator.evaluate(
saved_j, counter, [&] { return threshold; }, add_to_heap);
}
nstep++;
@@ -1258,12 +1314,12 @@ int search_from_candidates_fixVT(
return nres;
}
/// Dispatches the VisitedTable concrete type for a given C, then calls
/// the templated `search_from_candidates_fixVT<VTType, C>`.
template <class C>
int search_from_candidates_dispatch(
/// Dispatches the VisitedTable concrete type while keeping the distance
/// evaluator statically selected outside the graph traversal loop.
template <class C, typename DistanceEvaluator>
int search_from_candidates_evaluator_dispatch(
const HNSW& hnsw,
DistanceComputer& qdis,
DistanceEvaluator& evaluator,
ResultHandler& res,
MinimaxHeapT<HC_for<C>>& candidates,
VisitedTable& vt,
@@ -1274,7 +1330,7 @@ int search_from_candidates_dispatch(
auto call = [&]<typename VTType>(VTType& vt_concrete) -> int {
return search_from_candidates_fixVT<VTType, C>(
hnsw,
qdis,
evaluator,
res,
candidates,
vt_concrete,
@@ -1290,6 +1346,59 @@ int search_from_candidates_dispatch(
return call(vts);
}
template <class C>
int search_from_candidates_dispatch(
const HNSW& hnsw,
DistanceComputer& qdis,
ResultHandler& res,
MinimaxHeapT<HC_for<C>>& candidates,
VisitedTable& vt,
HNSWStats& stats,
int level,
int nres_in,
const SearchParameters* params) {
DefaultCandidateDistanceEvaluator evaluator{qdis};
return search_from_candidates_evaluator_dispatch<C>(
hnsw,
evaluator,
res,
candidates,
vt,
stats,
level,
nres_in,
params);
}
template <class C>
int search_from_candidates_rabitq_dispatch(
const HNSW& hnsw,
DistanceComputer& qdis,
ResultHandler& res,
MinimaxHeapT<HC_for<C>>& candidates,
VisitedTable& vt,
HNSWStats& stats,
int level,
int nres_in,
const SearchParameters* params) {
auto* rq = dynamic_cast<RaBitQDistanceComputer*>(&qdis);
FAISS_THROW_IF_NOT_MSG(
rq, "staged RaBitQ search requires an IndexRaBitQ storage");
FAISS_THROW_IF_NOT_MSG(
rq->nb_bits >= 2, "staged RaBitQ search requires nb_bits >= 2");
RaBitQCandidateDistanceEvaluator evaluator{*rq, hnsw.is_similarity};
return search_from_candidates_evaluator_dispatch<C>(
hnsw,
evaluator,
res,
candidates,
vt,
stats,
level,
nres_in,
params);
}
} // namespace
/** Do a BFS on the candidates list. Public dispatcher: only handles the
@@ -1716,37 +1825,49 @@ HNSWStats search_impl(
}
int ef = std::max(cur_efSearch, k);
if (bounded_queue) { // this is the most common branch, for now we only
// support Panorama search in this branch
if (bounded_queue) {
MinimaxHeapT<HC_for<C>> candidates(ef);
candidates.push(nearest, d_nearest);
if (!hnsw.is_panorama) {
search_from_candidates_dispatch<C>(
hnsw, qdis, res, candidates, vt, stats, 0, 0, params);
} else {
// Panorama is L2-specific and is only valid for C_distance.
// The public dispatch ensures we never reach this code path
// with C != C_distance, but assert in debug builds.
if constexpr (std::is_same_v<C, HNSW::C_distance>) {
hnsw_detail::search_from_candidates_panorama(
hnsw,
index,
qdis,
res,
candidates,
vt,
stats,
0,
0,
params);
} else {
FAISS_THROW_MSG(
"Panorama search is not supported with is_similarity=true");
}
switch (hnsw.search_method) {
case HNSW::SM_DEFAULT:
search_from_candidates_dispatch<C>(
hnsw, qdis, res, candidates, vt, stats, 0, 0, params);
break;
case HNSW::SM_RABITQ:
search_from_candidates_rabitq_dispatch<C>(
hnsw, qdis, res, candidates, vt, stats, 0, 0, params);
break;
case HNSW::SM_PANORAMA:
if constexpr (std::is_same_v<C, HNSW::C_distance>) {
hnsw_detail::search_from_candidates_panorama(
hnsw,
index,
qdis,
res,
candidates,
vt,
stats,
0,
0,
params);
} else {
FAISS_THROW_MSG(
"Panorama search does not support similarity metrics");
}
break;
default:
FAISS_THROW_MSG("invalid HNSW search method");
}
} else {
FAISS_THROW_IF_NOT_MSG(
hnsw.search_method != HNSW::SM_RABITQ,
"staged RaBitQ search requires bounded_queue=true");
FAISS_THROW_IF_NOT_MSG(
hnsw.search_method == HNSW::SM_DEFAULT ||
hnsw.search_method == HNSW::SM_PANORAMA,
"invalid HNSW search method");
auto call = [&]<typename VTType>(VTType& vt_concrete) {
return search_from_candidate_unbounded_fixVT<VTType, C>(
hnsw,
@@ -1804,6 +1925,50 @@ void search_level_0_impl(
int k = extract_k_from_ResultHandler<C>(res);
auto search_candidates = [&](MinimaxHeapT<HC_for<C>>& candidates,
int nres) {
switch (hnsw.search_method) {
case HNSW::SM_DEFAULT:
return search_from_candidates_dispatch<C>(
hnsw,
qdis,
res,
candidates,
vt,
search_stats,
0,
nres,
params);
case HNSW::SM_RABITQ:
return search_from_candidates_rabitq_dispatch<C>(
hnsw,
qdis,
res,
candidates,
vt,
search_stats,
0,
nres,
params);
case HNSW::SM_PANORAMA:
// Preserve the pre-enum behavior: search_level_0 did not use
// Panorama's progressive traversal and fell back to the
// ordinary HNSW candidate search.
return search_from_candidates_dispatch<C>(
hnsw,
qdis,
res,
candidates,
vt,
search_stats,
0,
nres,
params);
default:
FAISS_THROW_MSG("invalid HNSW search method");
}
};
if (search_type == 1) {
int nres = 0;
@@ -1823,16 +1988,7 @@ void search_level_0_impl(
candidates.push(cj, nearest_d[j]);
nres = search_from_candidates_dispatch<C>(
hnsw,
qdis,
res,
candidates,
vt,
search_stats,
0,
nres,
params);
nres = search_candidates(candidates, nres);
nres = std::min(nres, candidates_size);
}
} else if (search_type == 2) {
@@ -1849,8 +2005,7 @@ void search_level_0_impl(
candidates.push(cj, nearest_d[j]);
}
search_from_candidates_dispatch<C>(
hnsw, qdis, res, candidates, vt, search_stats, 0, 0, params);
search_candidates(candidates, 0);
}
}
faiss/impl/HNSW.h
+8 -2
View File
@@ -154,8 +154,14 @@ struct HNSW {
/// use bounded queue during exploration
bool search_bounded_queue = true;
/// use Panorama progressive pruning in search
bool is_panorama = false;
/// Specialized level-0 search implementations. This state is derived from
/// the owning IndexHNSW subtype after construction or deserialization.
enum Search_method_t : uint8_t {
SM_DEFAULT,
SM_PANORAMA,
SM_RABITQ,
};
Search_method_t search_method = SM_DEFAULT;
/// distance comparison semantics: when true, distances are treated as
/// similarity scores (larger is better). Default false matches the
faiss/impl/RaBitQuantizer.cpp
+56
View File
@@ -17,6 +17,7 @@
#include <faiss/utils/distances.h>
#include <faiss/utils/rabitq_simd.h>
#include <algorithm>
#include <cmath>
#include <cstring>
#include <memory>
@@ -24,6 +25,8 @@
namespace faiss {
RaBitQStats rabitq_stats;
// Import shared utilities from RaBitQUtils
using rabitq_utils::ExtraBitsFactors;
using rabitq_utils::QueryFactorsData;
@@ -224,6 +227,51 @@ void RaBitQuantizer::decode_core(
}
}
template <SIMDLevel SL>
float symmetric_dis_1bit(const RaBitQDistanceComputer& dc, idx_t i, idx_t j) {
FAISS_THROW_IF_NOT_MSG(
dc.metric_type == MetricType::METRIC_L2,
"RaBitQ symmetric distance supports only L2");
FAISS_ASSERT(i >= 0 && j >= 0);
FAISS_ASSERT(dc.codes != nullptr);
const size_t sign_bytes = (dc.d + 7) / 8;
const uint8_t* code_i = dc.codes + static_cast<size_t>(i) * dc.code_size;
const uint8_t* code_j = dc.codes + static_cast<size_t>(j) * dc.code_size;
const auto* factors_i =
reinterpret_cast<const SignBitFactors*>(code_i + sign_bytes);
const auto* factors_j =
reinterpret_cast<const SignBitFactors*>(code_j + sign_bytes);
const uint64_t xor_popcount =
rabitq::bitwise_xor_dot_product<SL>(code_i, code_j, sign_bytes, 1);
const float sign_dot =
static_cast<float>(dc.d) - 2.0f * static_cast<float>(xor_popcount);
// The L2-optimal reconstruction of residual r is alpha * sign(r), where
// alpha = ||r||_1 / d. The stored factors give
// alpha_i * alpha_j = ||r_i||^2 * ||r_j||^2 /
// (d * dp_multiplier_i * dp_multiplier_j).
float cross_term = 0.0f;
if (factors_i->dp_multiplier != 0.0f && factors_j->dp_multiplier != 0.0f) {
// Dividing each norm first avoids overflowing the product of two
// squared norms even when the final distance is representable.
const float scaled_norm_i =
factors_i->or_minus_c_l2sqr / factors_i->dp_multiplier;
const float scaled_norm_j =
factors_j->or_minus_c_l2sqr / factors_j->dp_multiplier;
cross_term = (scaled_norm_i * (sign_dot / static_cast<float>(dc.d))) *
scaled_norm_j;
}
const float distance = factors_i->or_minus_c_l2sqr +
factors_j->or_minus_c_l2sqr - 2.0f * cross_term;
return std::max(0.0f, distance);
}
float RaBitQDistanceComputer::symmetric_dis(idx_t i, idx_t j) {
return symmetric_dis_1bit<SIMDLevel::NONE>(*this, i, j);
}
namespace {
// Distance computers templatized on SIMDLevel to avoid per-call dynamic
@@ -240,6 +288,10 @@ struct RaBitQDistanceComputerNotQ final : RaBitQDistanceComputer {
RaBitQDistanceComputerNotQ() = default;
float symmetric_dis(idx_t i, idx_t j) final {
return symmetric_dis_1bit<SL>(*this, i, j);
}
// Compute distance using only 1-bit codes (fast)
float distance_to_code_1bit_impl(
const uint8_t* binary_data,
@@ -455,6 +507,10 @@ struct RaBitQDistanceComputerQ final : RaBitQDistanceComputer {
RaBitQDistanceComputerQ() = default;
float symmetric_dis(idx_t i, idx_t j) final {
return symmetric_dis_1bit<SL>(*this, i, j);
}
// Compute distance using only 1-bit codes (fast)
float distance_to_code_1bit_impl(
const uint8_t* binary_data,
faiss/impl/RaBitQuantizer.h
+47 -4
View File
@@ -13,6 +13,8 @@
#include <faiss/MetricType.h>
#include <faiss/impl/DistanceComputer.h>
#include <faiss/impl/Quantizer.h>
#include <faiss/impl/RaBitQUtils.h>
#include <faiss/impl/platform_macros.h>
namespace faiss {
@@ -20,6 +22,31 @@ struct IDSelector;
template <typename T, typename TI>
struct ResultHandlerUnordered;
/** Statistics for staged RaBitQ distance evaluation.
*
* These counters belong to the quantizer rather than to a particular ANN
* graph. They are currently populated by HNSW's staged search path.
*/
struct RaBitQStats {
uint64_t n_1bit = 0;
uint64_t n_refine = 0;
void reset() {
n_1bit = n_refine = 0;
}
void add(const RaBitQStats& other) {
n_1bit += other.n_1bit;
n_refine += other.n_refine;
}
double refine_ratio() const {
return n_1bit ? double(n_refine) / double(n_1bit) : 0.0;
}
};
FAISS_API extern RaBitQStats rabitq_stats;
// the 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".
@@ -122,10 +149,10 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer {
// Used with f_error to compute error bounds for two-stage filtering
float g_error = 0.0f;
float symmetric_dis(idx_t /*i*/, idx_t /*j*/) override {
// Not used for RaBitQ
FAISS_THROW_MSG("Not implemented");
}
/// Per-distance-computer counters, aggregated by the owning index.
RaBitQStats stats;
float symmetric_dis(idx_t i, idx_t j) override;
// Compute 1-bit distance estimate (fast)
virtual float distance_to_code_1bit(const uint8_t* code) = 0;
@@ -133,6 +160,22 @@ struct RaBitQDistanceComputer : FlatCodesDistanceComputer {
// Compute full multi-bit distance (accurate)
virtual float distance_to_code_full(const uint8_t* code) = 0;
/// Apply the RaBitQ error bound without exposing the packed code layout to
/// ANN consumers. Precondition: nb_bits >= 2.
#ifndef SWIG
FAISS_ALWAYS_INLINE bool should_refine(
const uint8_t* code,
float estimate,
float threshold,
bool is_similarity) const {
const auto* factors =
reinterpret_cast<const rabitq_utils::SignBitFactorsWithError*>(
code + (d + 7) / 8);
return rabitq_utils::should_refine_candidate(
estimate, factors->f_error, g_error, threshold, is_similarity);
}
#endif
virtual void set_centroid(const float* centroid_in) {
centroid = centroid_in;
}
faiss/impl/ResultHandler.h
+9 -3
View File
@@ -123,14 +123,20 @@ struct SingleQueryBlockResultHandler : BlockResultHandler<C, use_sel> {
using ResultHandlerT<C>::threshold;
explicit SingleResultHandler(SingleQueryBlockResultHandler& hr)
: the_handler(hr.the_handler) {}
: the_handler(hr.the_handler) {
threshold = the_handler.threshold;
}
/// begin results for query # i
void begin(const size_t /* qid */) {}
void begin(const size_t /* qid */) {
threshold = the_handler.threshold;
}
/// add one result for query i
bool add_result(T dis, TI idx) final {
return the_handler.add_result(dis, idx);
bool updated = the_handler.add_result(dis, idx);
threshold = the_handler.threshold;
return updated;
}
/// series of results for query i is done
faiss/impl/index_read.cpp
+58 -4
View File
@@ -2507,7 +2507,7 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
} else if (
h == fourcc("IHNf") || h == fourcc("IHNp") || h == fourcc("IHNs") ||
h == fourcc("IHN2") || h == fourcc("IHNc") || h == fourcc("IHc2") ||
h == fourcc("IHfP") || h == fourcc("IH00")) {
h == fourcc("IHfP") || h == fourcc("IHNr") || h == fourcc("IH00")) {
std::unique_ptr<IndexHNSW> idxhnsw;
if (h == fourcc("IH00")) {
idxhnsw = std::make_unique<IndexHNSW>();
@@ -2525,6 +2525,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
idxhnsw = std::make_unique<IndexHNSWCagra>();
} else if (h == fourcc("IHc2")) {
idxhnsw = std::make_unique<IndexHNSWCagra>();
} else if (h == fourcc("IHNr")) {
idxhnsw = std::make_unique<IndexHNSWRaBitQ>();
}
read_index_header(*idxhnsw, f);
if (h == fourcc("IHfP")) {
@@ -2559,7 +2561,8 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
"HNSW levels size %zu != index ntotal %" PRId64,
idxhnsw->hnsw.levels.size(),
idxhnsw->ntotal);
idxhnsw->hnsw.is_panorama = (h == fourcc("IHfP"));
idxhnsw->hnsw.search_method =
h == fourcc("IHfP") ? HNSW::SM_PANORAMA : HNSW::SM_DEFAULT;
// `HNSW::is_similarity` is intentionally not serialized, so we
// re-derive it here from the persisted metric type. Without this,
// a saved IP/similarity index would come back configured as a
@@ -2581,6 +2584,57 @@ std::unique_ptr<Index> read_index_up(IOReader* f, int io_flags) {
idxhnsw->storage->d,
idxhnsw->d);
}
if (h == fourcc("IHNr")) {
auto* idx_rabitq = dynamic_cast<IndexHNSWRaBitQ*>(idxhnsw.get());
FAISS_THROW_IF_NOT_MSG(
idx_rabitq, "IHNr must deserialize to an IndexHNSWRaBitQ");
FAISS_THROW_IF_NOT_MSG(
idxhnsw->metric_type == METRIC_L2,
"IndexHNSWRaBitQ supports only the L2 metric");
bool staged;
READ1_BOOL(staged);
idxhnsw->hnsw.search_method =
staged ? HNSW::SM_RABITQ : HNSW::SM_DEFAULT;
if (idxhnsw->storage) {
auto* rq = dynamic_cast<IndexRaBitQ*>(idxhnsw->storage);
FAISS_THROW_IF_NOT_MSG(
rq, "IndexHNSWRaBitQ storage must be an IndexRaBitQ");
FAISS_THROW_IF_NOT_MSG(
rq->metric_type == idxhnsw->metric_type &&
rq->rabitq.metric_type == idxhnsw->metric_type,
"IndexHNSWRaBitQ storage metric mismatch");
FAISS_THROW_IF_NOT_MSG(
rq->is_trained == idxhnsw->is_trained,
"IndexHNSWRaBitQ storage training state mismatch");
FAISS_THROW_IF_NOT_FMT(
rq->rabitq.nb_bits >= 1 && rq->rabitq.nb_bits <= 9,
"invalid RaBitQ nb_bits=%zu",
rq->rabitq.nb_bits);
const size_t expected_code_size =
rq->rabitq.compute_code_size(rq->d, rq->rabitq.nb_bits);
validate_code_size_match(
rq->rabitq.code_size,
expected_code_size,
"IndexHNSWRaBitQ quantizer");
validate_code_size_match(
rq->code_size,
expected_code_size,
"IndexHNSWRaBitQ storage");
FAISS_THROW_IF_NOT(
rq->codes.size() ==
mul_no_overflow(
static_cast<size_t>(rq->ntotal),
rq->code_size,
"IndexHNSWRaBitQ codes"));
FAISS_THROW_IF_NOT_MSG(
!rq->is_trained ||
rq->center.size() == static_cast<size_t>(rq->d),
"IndexHNSWRaBitQ center size mismatch");
FAISS_THROW_IF_NOT_MSG(
staged == (rq->rabitq.nb_bits >= 2),
"IndexHNSWRaBitQ staged-search metadata mismatch");
}
}
if (h == fourcc("IHN2")) {
FAISS_THROW_IF_NOT_MSG(
idxhnsw->storage,
@@ -3402,7 +3456,7 @@ std::unique_ptr<IndexBinary> read_index_binary_up(IOReader* f, int io_flags) {
auto idxhnsw = std::make_unique<IndexBinaryHNSW>();
read_index_binary_header(*idxhnsw, f);
read_HNSW(idxhnsw->hnsw, f);
idxhnsw->hnsw.is_panorama = false;
idxhnsw->hnsw.search_method = HNSW::SM_DEFAULT;
FAISS_THROW_IF_NOT_FMT(
idxhnsw->hnsw.levels.size() == (size_t)idxhnsw->ntotal,
"IndexBinaryHNSW HNSW levels size %zu != ntotal %" PRId64,
@@ -3426,7 +3480,7 @@ std::unique_ptr<IndexBinary> read_index_binary_up(IOReader* f, int io_flags) {
READ1_BOOL(idxhnsw->base_level_only);
READ1(idxhnsw->num_base_level_search_entrypoints);
read_HNSW(idxhnsw->hnsw, f);
idxhnsw->hnsw.is_panorama = false;
idxhnsw->hnsw.search_method = HNSW::SM_DEFAULT;
FAISS_THROW_IF_NOT_FMT(
idxhnsw->hnsw.levels.size() == (size_t)idxhnsw->ntotal,
"IndexBinaryHNSWCagra HNSW levels size %zu != ntotal %" PRId64,
faiss/impl/index_write.cpp
+20
View File
@@ -873,12 +873,22 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
: dynamic_cast<const IndexHNSWSQ*>(idx) ? fourcc("IHNs")
: dynamic_cast<const IndexHNSW2Level*>(idx) ? fourcc("IHN2")
: dynamic_cast<const IndexHNSWCagra*>(idx) ? fourcc("IHc2")
: dynamic_cast<const IndexHNSWRaBitQ*>(idx) ? fourcc("IHNr")
: typeid(*idx) == typeid(IndexHNSW) ? fourcc("IH00")
: 0;
FAISS_THROW_IF_NOT_FMT(
h != 0,
"don't know how to serialize this IndexHNSW subtype: %s",
typeid(*idx).name());
const IndexRaBitQ* storage_rabitq = nullptr;
if (h == fourcc("IHNr")) {
storage_rabitq = dynamic_cast<const IndexRaBitQ*>(idxhnsw->storage);
FAISS_THROW_IF_NOT_MSG(
storage_rabitq ||
((io_flags & IO_FLAG_SKIP_STORAGE) &&
idxhnsw->storage == nullptr),
"IndexHNSWRaBitQ requires IndexRaBitQ storage");
}
WRITE1(h);
write_index_header(idxhnsw, f);
if (h == fourcc("IHfP")) {
@@ -901,6 +911,16 @@ void write_index(const Index* idx, IOWriter* f, int io_flags) {
} else {
write_index(idxhnsw->storage, f);
}
if (h == fourcc("IHNr")) {
// The staged flag is graph-traversal state, so it has to live here:
// with IO_FLAG_SKIP_STORAGE there is no storage to derive it from.
// Storage-owned settings are not duplicated in this payload;
// IndexRaBitQ serializes whatever it owns.
const bool staged = storage_rabitq
? storage_rabitq->rabitq.nb_bits >= 2
: idxhnsw->hnsw.search_method == HNSW::SM_RABITQ;
WRITE1(staged);
}
} else if (const IndexNSG* idxnsg = dynamic_cast<const IndexNSG*>(idx)) {
uint32_t h = dynamic_cast<const IndexNSGFlat*>(idx) ? fourcc("INSf")
: dynamic_cast<const IndexNSGPQ*>(idx) ? fourcc("INSp")
faiss/index_factory.cpp
+7
View File
@@ -575,6 +575,13 @@ IndexHNSW* parse_IndexHNSW(
if (match(sq_pattern)) {
return new IndexHNSWSQ(d, sq_types[sm[1].str()], hnsw_M, mt);
}
// Keep the bare RaBitQ token consistent with Flat and IVF: it means 1 bit.
// Use an explicit width such as RaBitQ4 to enable staged refinement.
if (match("RaBitQ([1-9])?")) {
// the capture is the bare digit, so no substr offset here
int nb_bits = mres_to_int(sm[1], 1);
return new IndexHNSWRaBitQ(d, hnsw_M, nb_bits, mt);
}
if (match("([0-9]+)\\+PQ([0-9]+)?")) {
int ncent = mres_to_int(sm[1]);
int pq_m = mres_to_int(sm[2]);
faiss/python/__init__.pyi
+20
View File
@@ -1939,12 +1939,23 @@ class HNSWStats:
n3: int
ndis: int
nreorder: int
class RaBitQStats:
n_1bit: int
n_refine: int
def reset(self) -> None: ...
def add(self, other: RaBitQStats) -> None: ...
def refine_ratio(self) -> float: ...
class HNSW:
SM_DEFAULT: int
SM_PANORAMA: int
SM_RABITQ: int
max_level: int
entry_point: int
efConstruction: int
efSearch: int
search_method: int
hnsw_stats: HNSWStats
assign_probas: Float32Vector
cum_nneighbor_per_level: Int32Vector
@@ -2008,6 +2019,15 @@ class IndexHNSWSQ(IndexHNSW):
metric: MetricType = METRIC_L2,
) -> None: ...
class IndexHNSWRaBitQ(IndexHNSW):
def __init__(
self,
d: int,
M: int,
nb_bits: int = 1,
metric: MetricType = METRIC_L2,
) -> None: ...
class IndexHNSW2Level(IndexHNSW):
q1: Any # Level1Quantizer
pq: ProductQuantizer # from Index2Layer
faiss/python/swigfaiss.swig
+1
View File
@@ -914,6 +914,7 @@ typedef int cudaDataType_t;
DOWNCAST ( IndexHNSWPQ )
DOWNCAST ( IndexHNSWSQ )
DOWNCAST ( IndexHNSWCagra )
DOWNCAST ( IndexHNSWRaBitQ )
DOWNCAST ( IndexHNSW )
DOWNCAST ( IndexHNSW2Level )
DOWNCAST ( IndexNNDescentFlat )
tests/test_custom_result_handler.cpp
+20
View File
@@ -80,8 +80,28 @@ struct CollectAllResultHandler : ResultHandler {
}
};
struct ThresholdUpdatingResultHandler : ResultHandler {
bool add_result(float distance, idx_t /* i */) override {
threshold = distance;
return true;
}
};
namespace {
TEST(TestSingleQueryBlockResultHandler, SynchronizesWrappedThreshold) {
ThresholdUpdatingResultHandler handler;
handler.threshold = 10.0f;
SingleQueryBlockResultHandler<CMax<float, idx_t>, false> block(handler);
SingleQueryBlockResultHandler<CMax<float, idx_t>, false>::
SingleResultHandler single(block);
single.begin(0);
EXPECT_FLOAT_EQ(single.threshold, 10.0f);
single.add_result(4.0f, 0);
EXPECT_FLOAT_EQ(single.threshold, 4.0f);
}
// dimension of the vectors to index
constexpr int d = 64;
tests/test_hnsw_rabitq.py
+205
View File
@@ -0,0 +1,205 @@
# 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.
import unittest
import faiss
import numpy as np
class TestHNSWRaBitQ(unittest.TestCase):
def make_data(self, d=32, nt=400, nb=600, nq=20):
rs = np.random.RandomState(123)
xt = rs.randn(nt, d).astype("float32")
xb = rs.randn(nb, d).astype("float32")
xq = rs.randn(nq, d).astype("float32")
return xt, xb, xq
def make_index(self, nb_bits=3):
xt, xb, xq = self.make_data()
index = faiss.IndexHNSWRaBitQ(
xt.shape[1], 8, nb_bits, faiss.METRIC_L2
)
index.hnsw.efConstruction = 40
index.hnsw.efSearch = 64
index.train(xt)
index.add(xb)
return index, xb, xq
@staticmethod
def recall_at_k(actual, expected):
matches = sum(
len(set(actual[i]).intersection(expected[i]))
for i in range(actual.shape[0])
)
return matches / actual.size
def test_staged_search_quality(self):
index, xb, xq = self.make_index(nb_bits=3)
faiss.cvar.rabitq_stats.reset()
D, I = index.search(xq, 10)
stats = faiss.cvar.rabitq_stats
self.assertGreater(stats.n_1bit, 0)
self.assertGreater(stats.n_refine, 0)
self.assertLess(stats.n_refine, stats.n_1bit)
exact = faiss.IndexFlatL2(xb.shape[1])
exact.add(xb)
_, Iexact = exact.search(xq, 10)
self.assertGreaterEqual(self.recall_at_k(I, Iexact), 0.8)
self.assertTrue(np.all(I >= 0))
self.assertTrue(np.all(np.isfinite(D)))
def test_symmetric_distance_uses_sign_plane(self):
d = 35 # Exercise padding bits in the final code byte.
rs = np.random.RandomState(456)
xt = np.full((80, d), 2.0, dtype="float32")
residuals = rs.randn(3, d).astype("float32") * 1e10
residuals[2] = 0
xb = (residuals + 2.0).astype("float32")
residuals = xb - 2.0
norms = np.sum(residuals * residuals, axis=1)
alphas = np.sum(np.abs(residuals), axis=1) / d
signs = np.where(residuals > 0, 1.0, -1.0)
pairs = ((0, 1), (0, 2))
expected = [
norms[i]
+ norms[j]
- 2 * (alphas[i] * alphas[j] * np.dot(signs[i], signs[j]))
for i, j in pairs
]
distances = []
for nb_bits in (1, 4):
storage = faiss.IndexRaBitQ(d, faiss.METRIC_L2, nb_bits)
storage.train(xt)
storage.add(xb)
dc = storage.get_FlatCodesDistanceComputer()
distances.append([dc.symmetric_dis(i, j) for i, j in pairs])
for row in distances:
np.testing.assert_allclose(row, expected, rtol=2e-5, atol=2e-5)
np.testing.assert_array_equal(distances[0], distances[1])
def test_clone_and_io_remain_mutable(self):
index, xb, xq = self.make_index(nb_bits=3)
storage = faiss.downcast_index(index.storage)
Dref, Iref = index.search(xq, 10)
cloned = faiss.clone_index(index)
self.assertIsInstance(cloned, faiss.IndexHNSWRaBitQ)
Dclone, Iclone = cloned.search(xq, 10)
np.testing.assert_array_equal(Iclone, Iref)
np.testing.assert_array_equal(Dclone, Dref)
cloned.add(xb[:1])
self.assertEqual(cloned.ntotal, len(xb) + 1)
self.assertEqual(index.ntotal, len(xb))
self.assertEqual(storage.ntotal, len(xb))
loaded = faiss.deserialize_index(faiss.serialize_index(index))
self.assertIsInstance(loaded, faiss.IndexHNSWRaBitQ)
loaded_storage = faiss.downcast_index(loaded.storage)
self.assertEqual(loaded_storage.rabitq.nb_bits, 3)
Dloaded, Iloaded = loaded.search(xq, 10)
np.testing.assert_array_equal(Iloaded, Iref)
np.testing.assert_array_equal(Dloaded, Dref)
loaded.add(xb[:1])
self.assertEqual(loaded.ntotal, len(xb) + 1)
def test_io_and_reset(self):
index, xb, xq = self.make_index(nb_bits=4)
cloned = faiss.clone_index(index)
cloned.add(xb[:1])
self.assertEqual(cloned.ntotal, len(xb) + 1)
loaded = faiss.deserialize_index(faiss.serialize_index(index))
loaded.add(xb[:1])
self.assertEqual(loaded.ntotal, len(xb) + 1)
loaded.reset()
loaded.add(xb)
self.assertEqual(loaded.ntotal, len(xb))
D, I = loaded.search(xq, 10)
self.assertTrue(np.all(I >= 0))
self.assertTrue(np.all(np.isfinite(D)))
def test_one_bit_fallback(self):
index, _, xq = self.make_index(nb_bits=1)
self.assertEqual(index.hnsw.search_method, 0)
faiss.cvar.rabitq_stats.reset()
D, I = index.search(xq, 10)
self.assertTrue(np.all(I >= 0))
self.assertTrue(np.all(np.isfinite(D)))
self.assertEqual(faiss.cvar.rabitq_stats.n_1bit, 0)
self.assertEqual(faiss.cvar.rabitq_stats.n_refine, 0)
loaded = faiss.deserialize_index(faiss.serialize_index(index))
loaded_storage = faiss.downcast_index(loaded.storage)
self.assertEqual(loaded_storage.rabitq.nb_bits, 1)
self.assertEqual(loaded.hnsw.search_method, 0)
Dloaded, Iloaded = loaded.search(xq, 10)
np.testing.assert_array_equal(Dloaded, D)
np.testing.assert_array_equal(Iloaded, I)
def test_staged_search_requires_bounded_queue(self):
index, _, xq = self.make_index(nb_bits=3)
index.hnsw.search_bounded_queue = False
with self.assertRaises(RuntimeError):
index.search(xq, 10)
def test_permute_then_add(self):
index, xb, xq = self.make_index(nb_bits=3)
perm = np.arange(len(xb) - 1, -1, -1, dtype=np.int64)
index.permute_entries(perm)
index.add(xb[:2])
self.assertEqual(index.ntotal, len(xb) + 2)
D, I = index.search(xq, 10)
self.assertTrue(np.all(I >= 0))
self.assertTrue(np.all(np.isfinite(D)))
def test_factory(self):
default_index = faiss.index_factory(32, "HNSW8,RaBitQ")
self.assertIsInstance(default_index, faiss.IndexHNSWRaBitQ)
self.assertFalse(default_index.is_trained)
default_storage = faiss.downcast_index(default_index.storage)
self.assertIsInstance(default_storage, faiss.IndexRaBitQ)
self.assertFalse(default_storage.is_trained)
self.assertEqual(default_storage.rabitq.nb_bits, 1)
self.assertEqual(default_index.hnsw.search_method, 0)
with self.assertRaises(RuntimeError):
default_index.add(np.zeros((1, 32), dtype="float32"))
four_bit = faiss.index_factory(32, "HNSW8,RaBitQ4")
four_bit_storage = faiss.downcast_index(four_bit.storage)
self.assertEqual(four_bit_storage.rabitq.nb_bits, 4)
self.assertNotEqual(four_bit.hnsw.search_method, 0)
def test_skip_storage_can_be_reserialized(self):
index, _, _ = self.make_index(nb_bits=3)
data = faiss.serialize_index(index, faiss.IO_FLAG_SKIP_STORAGE)
metadata_only = faiss.deserialize_index(data)
self.assertIsNone(metadata_only.storage)
self.assertNotEqual(metadata_only.hnsw.search_method, 0)
data_again = faiss.serialize_index(
metadata_only, faiss.IO_FLAG_SKIP_STORAGE
)
metadata_again = faiss.deserialize_index(data_again)
self.assertIsNone(metadata_again.storage)
self.assertEqual(
metadata_again.hnsw.search_method,
metadata_only.hnsw.search_method,
)
def test_unsupported_metric_throws(self):
for metric in (faiss.METRIC_INNER_PRODUCT, faiss.METRIC_L1):
with self.assertRaises(RuntimeError):
faiss.IndexHNSWRaBitQ(32, 8, 1, metric)
if __name__ == "__main__":
unittest.main()
tests/test_io.py
+2 -2
View File
@@ -775,8 +775,8 @@ class Test_IO_HNSW(unittest.TestCase):
index = faiss.IndexHNSWFlat(d, 16)
index.train(self.xt)
index.add(self.xb)
# Note: RaBitQ lacks symmetric_distance, so it can be used for
# `search()`, but not `add()`.
# Build the graph with Flat storage, then verify that swapping in an
# equivalent RaBitQ storage preserves useful search results.
I_flat = index.storage.assign(self.xq, 10)
I_hnsw_flat = index.assign(self.xq, 10)
index.storage = faiss.IndexRaBitQ(d)