Efficient Vector Search in the Wild: One Model for Multi-K Queries

TL;DR

OMEGA is a novel vector search method that generalizes across different K values, achieving high accuracy and efficiency in multi-K queries with minimal preprocessing, outperforming existing methods.

Contribution

The paper introduces OMEGA, a K-generalizable learned top-K search model that uses trajectory-based features and statistical properties to efficiently handle multiple K values.

Findings

OMEGA reduces average latency by 6-33% compared to state-of-the-art methods.

It achieves similar recall with only 16-30% of the preprocessing time.

OMEGA maintains high accuracy across various K values in diverse datasets.

Abstract

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.