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.

翻译：基于学习的Top-K搜索是一种兼具高精度与高性能的向量查询服务方法。然而，当前针对特定K值训练的模型难以泛化至真实场景中的多K查询：它们在大K值下会出现精度下降，在小K值下则面临性能损失。训练能够泛化不同K值的模型需要数量级增长的预处理时间，因此不适用于实际环境中的向量查询服务。本文提出OMEGA，一种具备K值泛化能力的Top-K学习搜索方法，能够同时为多K向量查询实现高精度、高性能与低预处理成本。其核心思想在于：基于我们提出的轨迹特征在K=1条件下训练的基础模型，可通过动态优化程序准确预测更大K值，并在预测更小K值时仅产生微小的性能损失。为实现高效优化，我们进一步利用Top-K搜索的统计特性来减少过度的模型调用。在多个公开数据集与生产数据集上的广泛实验表明：在相同预处理资源约束下，相较于最先进的基于学习的搜索方法，OMEGA实现了平均延迟降低6-33%，同时所有系统均达到相同的召回率目标。仅需基线方法16-30%的预处理时间，OMEGA即可达到这些基线最优平均延迟的1.01-1.28倍。