IVF_PQ re-ranking

[jparismorgan]

What

Adds support for re-ranking to IVF_PQ. This is controlled by k_factor. We first query for k_factor * k vector IDs and indices, then fetch the original (not PQ-encoded) vectors and re-rank those.

Testing

• Adds several new tests

Benchmarks

Here are the results with local-benchmarks SIFT 1 million. Note that this is with local benchmarks so doesn't include network latency. But for the everything is local case we get a nice accuracy improvement with a fairly small speed penality:

[NikolaosPapailiou]

Great work on implementing this!

Can we also get the benchmark results for SIFT-1M or SIFT-10M and S3 storage? This will give us a good idea of whether we need more code optimizations on the reranking part of the code.

[NikolaosPapailiou]

Should we also load all feature_vectors_ for the infinite memory case during the constructor?

[jparismorgan]

Yep so that will happen (we hit that code in the second call to rerank_query() in rerank().

But this was just to cover the case where you create an empty index and then try to query with re-ranking. But I realized we actually don't need to throw here, so updated by moving the check for group_ below and only reading from the URI if we have the group. Also added a test for this case.

[NikolaosPapailiou]

Sorry but I still don't understand where we load the full feature_vectors_ array. Shouldn't this be inside read_index_infinite?

[jparismorgan]

Oh, oops, I thought you meant loading feature_vectors_ during add(). You are correct, we do not load the full feature_vectors_ anywhere except during add().

In the constructor we don't know if we're opening for infinite or finite queries, that is per-query. We could perhaps load the full feature_vectors_ during the call to query_infinite() (which calls read_index_infinite()), but then we will perhaps be loading way more vectors than we actually need to, and than we would if we just loaded them in rerank(). The benefit though would be we don't need to do any load during the second call to query_infinite(), whereas today we do.

[NikolaosPapailiou]

There are multiple options here, so maybe add some query/constructor parameters to make sure that we can leverage all of them and understand what happens in each case. For example, for small arrays and infinite memory, it would make sense to load the full feature_vectors_ array once.

[jparismorgan]

Talked offline - will do this along with moving memory_budget into the index constructor in a follow-up.

[NikolaosPapailiou]

Please add a TODO or shortcut for this

[jparismorgan]

Started on it here! #518

[jparismorgan]

Can we also get the benchmark results for SIFT-1M or SIFT-10M and S3 storage? This will give us a good idea of whether we need more code optimizations on the reranking part of the code.

Yep, will do, just opened up #514 which will make this comparison easy. Would prefer to not block this PR on that though, if it's okay with you? The k_factor = 1.f case is unaffected by this, so seems fine to get it merged even if it does turn out to be slow for remote URIs.

[NikolaosPapailiou]

Can we also get the benchmark results for SIFT-1M or SIFT-10M and S3 storage? This will give us a good idea of whether we need more code optimizations on the reranking part of the code.

Yep, will do, just opened up #514 which will make this comparison easy. Would prefer to not block this PR on that though, if it's okay with you? The k_factor = 1.f case is unaffected by this, so seems fine to get it merged even if it does turn out to be slow for remote URIs.

As this is new functionality mostly relevant for larger datasets stored in object stores, I would prefer to see the performance of it to understand the current state. And since #514 is ready I think we should try to produce it.

[jparismorgan]

Here are the results of local-benchmarks.py on SIFT 1M using k=100. It is run from my local laptop with an index_uri that is on AWS.

query_infinite_ram()

query_finite_ram() with memory_budget=1000000 for both IVF_FLAT and IVF_PQ:

Observations:

• IVF_PQ re-ranking seems to be a good improvement versus no re-ranking. With it, IVF_FLAT is still better, but IVF_PQ is now reasonably close.
• IVF_PQ infinite queries seem to do better than finite queries.

[NikolaosPapailiou]

I think we need to perform more benchmarks/optimizations to make sure that IVF_PQ is materially better than IVF_FLAT for large objectstore-based datasets and OOC query execution. However, I am ok to merge this as a first version of reranking and follow up in a more specific PR.