Module rabitq 

RaBitQ, the quantiser the searchable form of a vector is written in (10 section 3).

A 768 dimensional embedding is 3072 bytes of f32, and ten million of them is 30 GB. RaBitQ writes the same vector as one bit per dimension, which is 96 bytes and a 32x reduction, and that is the difference between an index that sits in memory and one that does not.

Binary quantisation on its own is old and it is not good enough: keeping only the sign of each coordinate throws away how far the point is from the boundary, and the distances that come back are biased in a way that no amount of rerank hides. RaBitQ’s contribution is the estimator. It stores one extra number per vector, the cosine between the vector and the corner of the cube it was rounded to, and dividing by that turns a biased guess into an unbiased one with an error bound that shrinks as 1/sqrt(D). The ordering that comes out of the codes alone is then good enough that rerank only has to look at a small multiple of k real vectors.

use yo_vector::{Bits, Quantizer};

let q = Quantizer::new(64, Bits::One, 7);
let centroid = vec![0.0f32; 64];
let mut code = vec![0u8; q.code_bytes()];

let v: Vec<f32> = (0..64).map(|i| (i as f32 * 0.37).sin()).collect();
let coded = q.encode(&v, &centroid, &mut code);

// The query is prepared once and then measured against many codes.
let query = q.query(&v, &centroid);
let guess = query.distance(&code, &coded);
// A vector against itself, so the answer should be near zero.
assert!(guess < 0.2, "{guess}");

What is stored

Per vector: the code, and two f32. The first is the length of the residual, which is the vector minus its partition’s centroid, and it is what turns an angle in the unit sphere back into a distance. The second is the correction, which is the estimator, and it is kept as its reciprocal because the estimator only ever divides by it. 10 section 3’s table says one f32 and it is two, because both are needed and neither can be recovered from the other. Eight bytes on top of 96 is still a 30x reduction rather than 32x.

One bit and four

Bits::One rounds each coordinate to a sign and Bits::Four rounds it to one of sixteen levels between the smallest and the largest coordinate the vector has. Four bits is four times the index and roughly a quarter of the error, and which one a collection wants is a question about the embedding family rather than about the engine, which is why both are here and the choice is per collection.

The two share one code path. Reconstructing a code is lo + level * delta either way, with lo and delta fixed by the dimension for one bit and measured per vector for four, so the estimator is written once and the only thing that changes is how many bits a level takes.

A code is bit planes, and that is what makes the scan fast

The obvious layout writes a coordinate’s level in the bits next to it, and then measuring a code against a query is a multiply per dimension. That is 768 multiplies per candidate and it does not fit inside a millisecond search.

So a code is stored transposed. Plane b holds bit b of every coordinate’s level, one bit per coordinate, dim bits rounded up to whole 64 bit words, and the planes run least significant first. A one bit code is one plane and a four bit code is four, and the byte count is the same either way.

The query is quantised to Bits::query_bits and transposed the same way. Then the sum of the code’s levels times the query’s levels is

sum over a, b of 2^(a+b) * popcount(code plane a AND query plane b)

which is four ANDs and four popcounts per word for a one bit code, against 64 float multiplies for the same 64 coordinates. The sums are exact integers, so the arithmetic is also better behaved than the float version it replaces, and nothing is left to round until the end.

The query is quantised finer than the code

Quantising the query costs accuracy, and how much was measured rather than assumed. At one bit the query at four bits is off by about a third of what the code itself is off by, which is lost in the quadrature and does not matter. At four bits the code is ten times more accurate and the same four bit query is off by four times as much as the code, which throws away the entire reason anyone would pay for four bit codes.

So the query width follows the code width: four bits against a one bit code and eight against a four bit one. That puts the query’s error back at about a third of the code’s in both cases, and Query::cosine against Query::cosine_exact is the test that holds it there.

Structs

Coded

What a code needs alongside it to be measured against a query.

Quantizer

The quantiser for one collection: its rotation and its width.

Query

A query, rotated and quantised once and then measured against every code in a partition.

Enums

Bits

How many bits a coordinate is rounded to.