# Efficient K-generalizable Learned Search

arXiv: https://arxiv.org/abs/2603.06159 · HTML: https://mlsanity.com/paper/2603.06159/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 23 of its 50 references that are in the map.
- First posted on arXiv 6 Mar 2026.

## Builds on (most cited first, 23 of 23)

- [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://mlsanity.com/paper/2005.11401/) (2020)
- [Efficient and robust approximate nearest neighbor search using Hierarchical Navigable Small World graphs](https://mlsanity.com/paper/1603.09320/) (2016)
- [The Faiss library](https://mlsanity.com/paper/2401.08281/) (2024)
- [Fast Approximate Nearest Neighbor Search With The Navigating Spreading-out Graph](https://mlsanity.com/paper/1707.00143/) (2017)
- [Embedding-based Product Retrieval in Taobao Search](https://mlsanity.com/paper/2106.09297/) (2021)
- [Results of the NeurIPS'21 Challenge on Billion-Scale Approximate Nearest Neighbor Search](https://mlsanity.com/paper/2205.03763/) (2022)
- [SPFresh: Incremental In-Place Update for Billion-Scale Vector Search](https://mlsanity.com/paper/2410.14452/) (2024)
- [High-Dimensional Approximate Nearest Neighbor Search: with Reliable and Efficient Distance Comparison Operations](https://mlsanity.com/paper/2303.09855/) (2023)
- [Non-autoregressive Generative Models for Reranking Recommendation](https://mlsanity.com/paper/2402.06871/) (2024)
- [Starling: An I/O-Efficient Disk-Resident Graph Index Framework for High-Dimensional Vector Similarity Search on Data Segment](https://mlsanity.com/paper/2401.02116/) (2024)
- [The Design and Implementation of a Real Time Visual Search System on JD E-commerce Platform](https://mlsanity.com/paper/1908.07389/) (2019)
- [Off the Beaten Path: Let's Replace Term-Based Retrieval with k-NN Search](https://mlsanity.com/paper/1610.10001/) (2016)
- [BANG: Billion-Scale Approximate Nearest Neighbor Search using a Single GPU](https://mlsanity.com/paper/2401.11324/) (2024)
- [Effective and General Distance Computation for Approximate Nearest Neighbor Search](https://mlsanity.com/paper/2404.16322/) (2024)
- [boldsymbolSteiner-Hardness: A Query Hardness Measure for Graph-Based ANN Indexes](https://mlsanity.com/paper/2408.13899/) (2024)
- [Quake: Adaptive Indexing for Vector Search](https://mlsanity.com/paper/2506.03437/) (2025)
- [Automating Nearest Neighbor Search Configuration with Constrained Optimization](https://mlsanity.com/paper/2301.01702/) (2023)
- [MS MARCO Web Search: a Large-scale Information-rich Web Dataset with Millions of Real Click Labels](https://mlsanity.com/paper/2405.07526/) (2024)
- [HYRR: Hybrid Infused Reranking for Passage Retrieval](https://mlsanity.com/paper/2212.10528/) (2022)
- [VSAG: An Optimized Search Framework for Graph-based Approximate Nearest Neighbor Search](https://mlsanity.com/paper/2503.17911/) (2025)
- [LIRA: A Learning-based Query-aware Partition Framework for Large-scale ANN Search](https://mlsanity.com/paper/2503.23409/) (2025)
- [Distribution-Aware Exploration for Adaptive HNSW Search](https://mlsanity.com/paper/2512.06636/) (2025)
- [Efficient Data-aware Distance Comparison Operations for High-Dimensional Approximate Nearest Neighbor Search](https://mlsanity.com/paper/2411.17229/) (2024)

## Abstract (arXiv)

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.