Efficient K-generalizable Learned Search

Learned top-K search improves the accuracy-latency trade-off of graph-based vector search, but existing methods target a fixed K. Serving real-world workloads with various K values thus incurs preprocessing cost proportional to the number of distinct Ks. We present OMEGA, which supports efficient arbitrary K learned search with the preprocessing cost of one top-1 model. It achieves so by reducing learned top-K search to repeated masked learned top-1 search with refinement, using carefully crafted trajectory features that generalize across Ks. OMEGA further employs rank-wise confidence allocation to tame error accumulation during refinements and applies a statistical forecast method to reduce model invocation overhead without compromising recall. Across nine dataset-scale configurations, OMEGA meets the 0.95 recall target and reduces mean latency by up to 65% (average: 23%; median: 20%) compared with learned search methods (including LAET, DARTH, MultiK-DARTH, ANNiE, and Ada-ef) at their lowest-preprocessing configurations. On production traces, it reduces total CPU time by up to 28% versus traditional HNSW search.