Add refinement of quantized vector scores with fp distance calculations

[jmazanec15]

Description

With quantization techniques that are compressing vectors in memory further and further, because of how much information is lost, recall is going to drop. However, with the current quantization support, we already store the full precision float vector values. With this, we could create a two phased search process that oversamples vectors from the quantized index (i.e. r*k results), and then lazily loads vectors from disk for the r*k results, re-scoring the vectors. This has been discussed directly or indirectly in a couple different places, but I figured itd make sense to create a separate issue for it:

• Can we add configuration on dropping raw vectors from quantized formats after some period of time? #13251
• Expose flat vectors in "user space" #13468
• Should we explore DiskANN for aKNN vector search? #12615

It has been shown to work with binary compression techniques and some PQ in a few different places (https://medium.com/qdrant/hyperfast-vector-search-with-binary-quantization-865052c5c259, https://huggingface.co/blog/embedding-quantization#binary-rescoring). We also have done some experiments in OpenSearch with IVFPQ and re-scoring and its shown pretty strong promise (assuming that the disk IOPS/throughput are strong enough - see opensearch-project/k-NN#1779 (comment)).

That being said, it can provide similar benefits to the DiskANN approach.

One challenge I foresee is the approach requires the quantized ANN index to be resident in memory. However, I am unsure if loading the full precision vectors from disk will cause the page cache to evict the quantized ANN index or if the page cache will naturally adapt to the access pattern. If it is the former, we would probably need some way to pin the quantized ANN index and its vectors in memory.

In the future, it could probably be fine tuned quite a bit with access pattern optimizations and/or some kind of hierarchical quantization and re-scoring.

[benwtrent]

I think the API would be tricky, but I am all for this idea. The ability to "approximately score" and then do second pass to "exact score" is useful for all sorts of levels of quantization.

Whatever the design, it would be most efficient to first gather the nearest vectors from ALL segments with an approximate score, and then do a second pass over all segments with to refine the top k.

Rescoring per segment would be needlessly in-efficient.

[jmazanec15]

I think the API would be tricky, but I am all for this idea

Yes agree, Ill think on this a little bit. Ill start with a PoC and go from there.

Whatever the design, it would be most efficient to first gather the nearest vectors from ALL segments with an approximate score, and then do a second pass over all segments with to refine the top k.

Rescoring per segment would be needlessly in-efficient.

Yes I agree. However, with this, I think we would need to refine not the top k but the top r*k and then reduce to k. Otherwise, I dont think that recall would actually be improved - just ordering might be better.

[benwtrent]

@jmazanec15

I think we would need to refine not the top k but the top r*k and then reduce to k

100%, I wasn't clear. Yes, over all segments, we gather some approximate scoring that is r*k total when collapsed together, and then refine & rescore the total r*k.

[jmazanec15]

Makes sense thanks @benwtrent . Im working on PoC and some experiments. Didnt realize that the full-precision vectors for quantized index are exposed via getFloatVectorValues. That should make things simpler.

I was also looking at this comment for DiskANN from @kevindrosendahl (#12615 (comment)) and noticed that mlock'ing parts of the index in memory became really important in lower mem environments for performance. I dont think this is currently supported in Lucene, but will take a look.

[dungba88]

Hi all, I was looking at this idea for some experimentation ideas (not mean to be intrusive to ongoing effort).

If the full sized vectors are exposed through getFloatVectorValues then it seems we can just call exactSearch over the results of approximateSearch, but this would inefficiently re-rank k * oversample * num_segments instead of just k * oversample (with some improve in recall perhaps).

However I don't know how to get the full-sized vectors without LeafContextReader. Do you have an idea how we can achieve that?

Or maybe we can just reduce oversample to account for the number of segments, e.g: oversample = oversample / leafReaderContexts.size() (which maybe differs)?

[benwtrent]

then it seems we can just call exactSearch over the results of approximateSearch,

Exact search will hopefully just be using the quantized values as well.

However I don't know how to get the full-sized vectors without LeafContextReader. Do you have an idea how we can achieve that?

I don't see why that would be a problem. How I would understand this is you would need to iterate the segments twice regardless.

• Once to get the top k as desired with the oversampling
• Apply the brute-force reranking (by passing the "exactSearch" path as that uses quantized values).

However, the query shouldn't bother attempting to rerank results that are considered "accurate".

So, the collector will likely need to be told "Hey, these are estimations", and then a second pass over the top docs there would be possible via something like if(collector.scoresAreEstimates)

[dungba88]

Hi @benwtrent

Thanks for the reply. It seems I wasn't clear enough in my previous question.

Regarding exactSearch, it was based on this jmazanec15@ comment. getFloatVectorValues() is used in exactSearch(), thus I thought that it would use full-precision vectors, but based on your comment, it seems it also uses quantized vectors? As the re-ranking step needs full-precision vectors, how can I access them?

Didnt realize that the full-precision vectors for quantized index are exposed via getFloatVectorValues

Update: It seems the raw (full precision) vectors are stored in QuantizedVectorValues along with the quantized values. The vectorValue(int ord) will return the raw vector.

[dungba88]

Apply the brute-force reranking (by passing the "exactSearch" path as that uses quantized values).

Maybe I misunderstood something, but I thought the idea is to re-rank with full precision vectors over the oversampled matches?

[dungba88]

I also saw where my misunderstanding comes from now: getFloatVectorValues() returns QuantizedVectorValues, which has 2 methods: vectorValue(int ord) returns the raw vector while scorer(float[] target) uses the quantized vector for scoring.

[jmazanec15]

Hi all, I was looking at this idea for some experimentation ideas (not mean to be intrusive to ongoing effort).

Hey @dungba88 - not intrusive at all! I ended up not having more time to work more on this.

So, the collector will likely need to be told "Hey, these are estimations", and then a second pass over the top docs there would be possible via something like if(collector.scoresAreEstimates)

I like this approach. I think it is a good way to maintain abstraction between knn codec formats and queries, where the query does not need to know any specifics about quantization done.