利用约束优化技术实现自动化最近邻搜索配置|ICLR2023
摘要
《Automating Nearest Neighbor Search Configuration with Constrained Optimization》是一篇关于自动化最近邻搜索配置的研究论文,主要探讨了如何使用约束优化技术来自动配置基于量化的最近邻搜索算法。
一、研究背景
能否解决billion-scale的问题是ANNS算法从理论到工程落地的一个重要衡量标准。在面对billion-scale的问题时,单一的索引结构已经显得“力不从心”,因此向量索引开始展示出了由单一结构向多层结构、融合索引过渡的趋势。可以类比深度学习技术,从逻辑回归到多层感知机再到现在风靡的大模型,数据量剧增势必需要更复杂的数据结构来构建高效索引。如下图,文中给出了一个多层量化索引的例子,在billion-scale的数据集上比较了层数更多更复杂的索引结构(Original)与其两个简化版本(Shallow-Small、Shallow-Large)的检索效果。可以明显看出,结构更复杂的索引在召回率和吞吐量两个角度都表现出了明显的优势。



随着索引结构的复杂化,其参数量也随之成倍增长,不同参数配置带来的效果变化更是成几何倍数增加。目前主流的参数配置手段有:
·网格搜索:遍历所有参数配置可能,虽然能得到最优解,但是计算开销庞大,不切实际;
·人工调参:以工程师丰富的经验来配置参数,这种方法会花费大量人力,且受工程师的经验所限,效果参差;
·黑盒优化:利用优化器来探索参数配置,这种方法同样需要大量计算,且经常只能得到次优解,差强人意。

基于以上两点考量,文章提出了一种基于约束优化的参数配置办法,以更低的计算开销获得更加权衡召回率和吞吐量的参数配置。文章的思路也十分“直截了当”,只需要定义两个函数,对于给定参数配置能够快速预估出近似的召回率和吞吐量表现即可。下面我们对文章中的方法进行详细介绍。
二、预估召回率和吞吐量
下式给出了多层量化索引的召回率计算方法,即,给定参数配置 t构建的 m层量化索引的查询结果 {S}_m(q,t)和真实解 {G}(q)的交集大小比上真实解数量。

衡量整体的多层量化索引的召回率的计算量还是很大,所以文中对上式做了进一步简化。首先,将其拆分为每层量化器的召回率估计的累积,即

将召回率改写为损失的形式(负对数形式),即

因为每后一层量化索引的召回表现严格依赖于上一层的计算,为了削除层级之间的影响,文章对每层量化索引的计算结果进一步近似化,即


此时,一整个多层量化索引的求解遍可以改写前几层量化索引的累积和单独的最后一层索引的结果的乘积形式,即

最终,整个召回率损失的估计函数可以写成下式。实际使用时,可以先采样一部分query用来估计召回率损失。

量化索引的吞吐量其实很好估计,已知每次距离计算的开销是固定的,那么只需要估测出每次query大概的距离计算次数即可估计时延、吞吐量。不像图索引,每个节点的邻居数量不固定,距离计算的次数无法确定,量化索引的距离计算次数是可以计算的,确定的。对于单层的量化索引,距离计算次数就是全部数据量。对于多层的量化索引,设每层搜索都选择前 t_i个区域作为候选集到下一层量化索引继续计算,则可以将多层量化索引的距离计算次数表示为下式,其中 X^(1) 表示第一层量化索引,由于初始时还没确定候选集,所以第一层的所有数据都要和query做一次距离计算。

经实验验证,以上两个估算函数虽然准确率和真实值存在差异,但是都表现出了和真实值一样的变化趋势,可以有效的估计参数变化对索引性能的影响。

三、基于约束优化自动化配置
经过上一节,我们已经获得了快速估算一组参数配置的召回率和吞吐量的方法。根据我们的目标,给出期望的召回率和吞吐量要求,得到一组权衡二者的参数配置,文章给出如下约束优化表达,即,限定召回率上限,最小化召回率损失,

进一步地,文章为了得到一种权衡二者的配置,对上式进行分解,得到

基于上式便可以直接使用优化问题的求解器快速获得一组优质的配置。
下图的实验结果表明,本文的方法与精确的网格搜索结果性能相差甚微,明显优于黑盒优化器的计算结果。


因为在估算召回率时使用的是采样的query,而不是真正全量的query,文章进一步实验来评估方法对于采样率的泛化性。如下图,可以看出,在In-Sample和Out-of-Sample两种情况下,该方法的配置出的量化索引性能差异微乎其微。

总结
数据量激增到亿万级别的如今,向量索引的结构也愈发地趋于复杂。本文以易懂的数学方式合理地计算出权衡召回率和吞吐量的索引参数配置。依赖本文的方法可以直接以性能为指标来装配索引,而无需深入研究索引中每个参数的具体意义、索引的具体实现原理,显著降低了向量索引的使用门槛。因为量化索引的计算特性,召回率和吞吐量易于估量,未来对于结构复杂的图索引、数索引的参数配置方法也值得探索。
参考
Sun, Philip et al. “Automating Nearest Neighbor Search Configuration with Constrained Optimization.” ArXiv abs/2301.01702 (2023): n. pag.
Guo, Ruiqi et al. “Accelerating Large-Scale Inference with Anisotropic Vector Quantization.” International Conference on Machine Learning (2019).

有任何问题可以搜索公众号“向量检索实验室”,关注回复【加入社区】进群和大家一起交流讨论。也可以直接添加负责人进群。
