Efficient K-generalizable Learned Search
TL;DR
This paper shows that learned search can support arbitrary K with the preprocessing cost of a single top-1 model, and reduces total serving and preprocessing computation by up to 28%.
Abstract
Learned top-K search improves the accuracy-latency trade-off of graph-based vector search, but existing methods are designed for a fixed K: serving production workloads with varying K values requires preprocessing cost proportional to the number of distinct Ks served - prohibitive in practice. This paper shows that learned search can support arbitrary K with the preprocessing cost of a single top-1 model. The key idea is to reduce top-K learned search to repeated masked top-1 refinement, which works because the distance-reduction trajectory for discovering the next top-1 vector is largely invariant to the number of results already found. We therefore train the model on trajectory features that remain effective under masking. To make repeated refinement robust and efficient, OMEGA counters error accumulation across iterations with rank-wise confidence allocation, and skips unnecessary model invocations with a statistical forecast of recall from partial results. Across nine dataset-scale configurations, OMEGA meets the 0.95 recall target with one K-independent model. Under the lowest-preprocessing configuration of each learned baseline,it reduces mean latency by 7-36% versus DARTH, 3-25% versus MultiK-DARTH, and 8-21% versus LAET on BIGANN, BIGANN-1B, DEEP, and three production workloads. On GIST, Text2Image, and MS MARCO, its latency remains within 9% of DARTH and MultiK-DARTH. On production traces, OMEGA further reduces total serving and preprocessing computation by up to 28%.
