Efficient Vector Search in the Wild: One Model for Multi-K Queries
This work provides a more efficient and accurate solution for serving vector search queries with varying K values, which is a common problem for systems that rely on vector search.
The paper addresses the challenge of efficiently serving multi-K vector queries in learned top-K search, where existing models struggle with accuracy for larger K and performance for smaller K. They introduce OMEGA, a method that uses a K=1 trained base model with trajectory-based features and a dynamic refinement procedure to achieve 6-33% lower average latency compared to state-of-the-art methods, while maintaining the same recall.
Learned top-K search is a promising approach for serving vector queries with both high accuracy and performance. However, current models trained for a specific K value fail to generalize to real-world multi-K queries: they suffer from accuracy degradation (for larger Ks) and performance loss (for smaller Ks). Training the model to generalize on different Ks requires orders of magnitude more preprocessing time and is not suitable for serving vector queries in the wild. We present OMEGA, a K-generalizable learned top-K search method that simultaneously achieves high accuracy, high performance, and low preprocessing cost for multi-K vector queries. The key idea is that a base model properly trained on K=1 with our trajectory-based features can be used to accurately predict larger Ks with a dynamic refinement procedure and smaller Ks with minimal performance loss. To make our refinements efficient, we further leverage the statistical properties of top-K searches to reduce excessive model invocations. Extensive evaluations on multiple public and production datasets show that, under the same preprocessing budgets, OMEGA achieves 6-33% lower average latency compared to state-of-the-art learned search methods, while all systems achieve the same recall target. With only 16-30% of the preprocessing time, OMEGA attains 1.01-1.28x of the optimal average latency of these baselines.