Cluster with Auctions for Vector Search

2026-07-15

AI Summary

This work addresses the suboptimality in large-scale approximate nearest neighbor search arising from the tight coupling between database partitioning and query probing functions, particularly under distributional shifts. To resolve this, the authors propose CwA, a novel approach that decouples and jointly optimizes these two components for the first time. CwA alternately optimizes a neural probing function and a large-scale balanced clustering assignment derived via a parallel auction algorithm, directly tailoring the retrieval structure to the query distribution. It further enhances partition granularity through a Cartesian product clustering architecture. Experiments demonstrate that under out-of-distribution settings, CwA achieves up to a 4.7× throughput improvement at equal recall compared to existing methods; notably, even with a linear probing function in in-distribution scenarios, it surpasses state-of-the-art deep learning approaches.

Abstract

Large-scale approximate nearest neighbor search commonly relies on partitions for indexing: database vectors are partitioned into clusters, and for each query a probing function selects the clusters to be scanned. The query probing function and the database partition are rarely treated as separate entities: most techniques assign queries with the same assignment function as the database vectors, which is suboptimal especially when database and query distributions differ. This paper introduces CwA (Cluster with Auctions), which addresses this limitation by jointly learning a balanced database partition and a neural probing function. CwA optimizes search performance directly for the query distribution. It minimizes its objective by alternating two steps: (i) gradient descent on the neural network of the probing function, and (ii) a large-scale combinatorial optimization of the cluster assignment for the database vectors. We solve the latter with a parallelizable auction algorithm that balances the partition by design. To further scale CwA, we extend the method to a Cartesian product of clusters that increases the partition's granularity. When database and query distributions differ, CwA achieves up to 4.7$\times$ throughput over the state-of-the-art at equal recall. In the in-distribution (ID) setting, even a simple linear probing function trained with CwA outperforms competing deep neural methods.

Problem

• approximate nearest neighbor search
• vector clustering
• query distribution mismatch
• database partitioning
• search efficiency

Innovation

• vector search
• approximate nearest neighbor
• auction algorithm
• neural probing function
• balanced clustering