Module muvera 

MUVERA: a set of token vectors as one vector, so late interaction retrieval runs on the index that is already here (10 section 6).

Late interaction is where retrieval quality went. A ColBERT style model does not give a document one embedding, it gives every token one, and it scores a query against a document with Chamfer similarity: for each query token, the best match anywhere in the document, averaged over the query. That is a much better score than one vector against one vector, and it is much more expensive, because there is no single vector to put in an index and the score is a loop over two sets rather than a dot product.

The usual answer is a second index over every token of every document, so a collection of a million documents at a hundred tokens each becomes a hundred million vector index, plus a gather and a scoring pass on top. That is a whole second system with its own memory, its own tuning and its own failure modes.

MUVERA (NeurIPS 2024, arXiv:2405.19504) does away with it. It maps a set of token vectors to one fixed length vector, a Fixed Dimensional Encoding, such that the dot product of a query’s encoding with a document’s approximates the Chamfer similarity between the two sets, with a proven bound on the error. So multi vector retrieval costs an encode at write time, the index that is already here, and a different rerank function. No second index.

How it works

The trick is to cut the vector space into buckets, with random hyperplanes, and then compare a query’s tokens only against the document tokens that landed in the same bucket.

A document’s encoding holds, for each bucket, the average of the document tokens that fell in it. A query’s encoding holds, for each bucket, the sum of the query tokens that fell in it. Take the dot product of the two and the bucket contributes, for each query token in it, that token against the average of the document tokens near it. If the hyperplanes did their job then the document token that maximises the dot product is in the same bucket, and the bucket average is close enough to it, so the sum over buckets is close to the sum of maxima that Chamfer wanted.

Two details make the difference between that being an argument and it being true.

An empty bucket in a document is filled with the document tokens whose own bucket is nearest in Hamming distance. Without it, a query token in a bucket the document did not reach contributes nothing, when the truth is that it still has a best match somewhere in the document, and short documents lose badly. With it, a one token document has that token in every bucket, which is exactly right, because that token is the best match for every query token.

The whole construction is repeated with independent hyperplanes and the results are laid end to end. One repetition is a coin toss on whether a query token and its true best match landed together. Several repetitions average that away, and this is the knob that actually buys accuracy.

What the repetitions buy

Three hundred documents of twenty four tokens each, forty queries, where a query is six of one document’s own tokens with noise on them and the answer is the document it came from. How often the encoding alone, with no rerank, puts that document first:

repetitions      1      2      4      8     16     32
ranked first  .825   .900   .950   .950  1.000  1.000

That is the shape to expect. Nothing else moves it nearly as much: buckets and block width change how long the encoding is far more than how good it is, which is why the default puts eight repetitions on sixteen buckets rather than the other way round.

What is different here

The paper ends with an optional random projection of the whole encoding down to a smaller dimension, to make it cheap to store and search. There is no point doing that here. crate::Quantizer already turns a vector into one bit a dimension with an error bound that a random projection does not have, so projecting first would throw away accuracy to save space that RaBitQ was going to save anyway, and better. The per bucket projection that shrinks each block from the token dimension down to Shape::dproj is still here, because that one is what keeps the encoding from being buckets times token dimension long.

The encodings come out unit length. The index ranks by squared distance, and for vectors of equal length that ordering is exactly the dot product ordering the approximation is stated in, so normalising is what makes the two agree. What it costs is the length of a document’s raw encoding, which mostly says how many tokens crowded into each bucket rather than anything about whether the document is a good answer, and chamfer on the candidates puts back any ordering that lost.

use yo_vector::muvera::{Encoder, Shape, chamfer};

let dim = 16;
let enc = Encoder::new(dim, Shape::default(), 7);

// Two tokens for the document, one for the query, laid out end to end.
let doc: Vec<f32> = (0..2 * dim).map(|i| if i % dim == i / dim { 1.0 } else { 0.0 }).collect();
let query: Vec<f32> = (0..dim).map(|i| f32::from(u8::from(i == 1))).collect();

// The query token is the second document token, so Chamfer is 1.
assert!((chamfer(&query, &doc, dim) - 1.0).abs() < 1e-6);

// And both sides encode to one vector of the same fixed length.
assert_eq!(enc.document(&doc).len(), enc.fde_dim());
assert_eq!(enc.query(&query).len(), enc.fde_dim());

Structs

Encoder

Turns a set of token vectors into one vector.

Shape

How big an encoding is and how much accuracy it buys.

Functions

chamfer

The Chamfer similarity of a query’s tokens against a document’s: for each query token, the best it does against any document token, averaged over the query.