Learn Where to Store: Query-Aware IVF

Approximate nearest neighbor (ANN) vector search is an essential primitive in modern AI and ML applications, e.g., to support retrieval augmented generation. For very large datasets, one of the most popular ANN methods is the inverted-file (IVF) index, which performs locality-friendly contiguous scans, can run on GPUs and on disk, and can serve as the first filtering step in ANN search pipelines. IVF clusters vectors by $k$-means and stores each vector in the list of its nearest centroid. Achieving high recall typically requires scanning more clusters; an alternative is to place copies of a vector in several lists. We present query-aware IVF, which learns where to store copies from a sample of the query distribution, formulating the allocation as a maximum coverage problem that we solve greedily. On the 24 datasets of the Vector Index Benchmark for Embeddings (VIBE), at recall@100 = 0.9, our query-aware index is faster than the IVF baselines on all 24 datasets: a median 1.63$\times$ faster than the best query-agnostic replication and up to 21.9$\times$ faster, and a median 1.79$\times$ faster than vanilla IVF on the 21 datasets where vanilla IVF reaches that recall, at a median 2.58$\times$ the memory of vanilla IVF.