摘要

《Automating Nearest Neighbor Search Configuration with Constrained Optimization》是一篇關(guān)于自動化最近鄰搜索配置的研究論文，主要探討了如何使用約束優(yōu)化技術(shù)來自動配置基于量化的最近鄰搜索算法。

一、研究背景

????能否解決billion-scale的問題是ANNS算法從理論到工程落地的一個重要衡量標準。在面對billion-scale的問題時，單一的索引結(jié)構(gòu)已經(jīng)顯得“力不從心”，因此向量索引開始展示出了由單一結(jié)構(gòu)向多層結(jié)構(gòu)、融合索引過渡的趨勢?？梢灶惐壬疃葘W習技術(shù)，從邏輯回歸到多層感知機再到現(xiàn)在風靡的大模型，數(shù)據(jù)量劇增勢必需要更復雜的數(shù)據(jù)結(jié)構(gòu)來構(gòu)建高效索引。如下圖，文中給出了一個多層量化索引的例子，在billion-scale的數(shù)據(jù)集上比較了層數(shù)更多更復雜的索引結(jié)構(gòu)（Original）與其兩個簡化版本（Shallow-Small、Shallow-Large）的檢索效果?？梢悦黠@看出，結(jié)構(gòu)更復雜的索引在召回率和吞吐量兩個角度都表現(xiàn)出了明顯的優(yōu)勢。

隨著索引結(jié)構(gòu)的復雜化，其參數(shù)量也隨之成倍增長，不同參數(shù)配置帶來的效果變化更是成幾何倍數(shù)增加。目前主流的參數(shù)配置手段有：

·網(wǎng)格搜索：遍歷所有參數(shù)配置可能，雖然能得到最優(yōu)解，但是計算開銷龐大，不切實際；

·人工調(diào)參：以工程師豐富的經(jīng)驗來配置參數(shù)，這種方法會花費大量人力，且受工程師的經(jīng)驗所限，效果參差；

·黑盒優(yōu)化：利用優(yōu)化器來探索參數(shù)配置，這種方法同樣需要大量計算，且經(jīng)常只能得到次優(yōu)解，差強人意。

基于以上兩點考量，文章提出了一種基于約束優(yōu)化的參數(shù)配置辦法，以更低的計算開銷獲得更加權(quán)衡召回率和吞吐量的參數(shù)配置。文章的思路也十分“直截了當”，只需要定義兩個函數(shù)，對于給定參數(shù)配置能夠快速預估出近似的召回率和吞吐量表現(xiàn)即可。下面我們對文章中的方法進行詳細介紹。

二、預估召回率和吞吐量

????下式給出了多層量化索引的召回率計算方法，即，給定參數(shù)配置?t構(gòu)建的?m層量化索引的查詢結(jié)果?{S}_m(q,t)和真實解?{G}(q)的交集大小比上真實解數(shù)量。

????衡量整體的多層量化索引的召回率的計算量還是很大，所以文中對上式做了進一步簡化。首先，將其拆分為每層量化器的召回率估計的累積，即

將召回率改寫為損失的形式（負對數(shù)形式），即

因為每后一層量化索引的召回表現(xiàn)嚴格依賴于上一層的計算，為了削除層級之間的影響，文章對每層量化索引的計算結(jié)果進一步近似化，即

此時，一整個多層量化索引的求解遍可以改寫前幾層量化索引的累積和單獨的最后一層索引的結(jié)果的乘積形式，即

最終，整個召回率損失的估計函數(shù)可以寫成下式。實際使用時，可以先采樣一部分query用來估計召回率損失。

????量化索引的吞吐量其實很好估計，已知每次距離計算的開銷是固定的，那么只需要估測出每次query大概的距離計算次數(shù)即可估計時延、吞吐量。不像圖索引，每個節(jié)點的鄰居數(shù)量不固定，距離計算的次數(shù)無法確定，量化索引的距離計算次數(shù)是可以計算的，確定的。對于單層的量化索引，距離計算次數(shù)就是全部數(shù)據(jù)量。對于多層的量化索引，設(shè)每層搜索都選擇前?t_i個區(qū)域作為候選集到下一層量化索引繼續(xù)計算，則可以將多層量化索引的距離計算次數(shù)表示為下式，其中?X^(1)?表示第一層量化索引，由于初始時還沒確定候選集，所以第一層的所有數(shù)據(jù)都要和query做一次距離計算。

????經(jīng)實驗驗證，以上兩個估算函數(shù)雖然準確率和真實值存在差異，但是都表現(xiàn)出了和真實值一樣的變化趨勢，可以有效的估計參數(shù)變化對索引性能的影響。

三、基于約束優(yōu)化自動化配置

????經(jīng)過上一節(jié)，我們已經(jīng)獲得了快速估算一組參數(shù)配置的召回率和吞吐量的方法。根據(jù)我們的目標，給出期望的召回率和吞吐量要求，得到一組權(quán)衡二者的參數(shù)配置，文章給出如下約束優(yōu)化表達，即，限定召回率上限，最小化召回率損失，

進一步地，文章為了得到一種權(quán)衡二者的配置，對上式進行分解，得到

基于上式便可以直接使用優(yōu)化問題的求解器快速獲得一組優(yōu)質(zhì)的配置。

????下圖的實驗結(jié)果表明，本文的方法與精確的網(wǎng)格搜索結(jié)果性能相差甚微，明顯優(yōu)于黑盒優(yōu)化器的計算結(jié)果。

????因為在估算召回率時使用的是采樣的query，而不是真正全量的query，文章進一步實驗來評估方法對于采樣率的泛化性。如下圖，可以看出，在In-Sample和Out-of-Sample兩種情況下，該方法的配置出的量化索引性能差異微乎其微。

總結(jié)

????數(shù)據(jù)量激增到億萬級別的如今，向量索引的結(jié)構(gòu)也愈發(fā)地趨于復雜。本文以易懂的數(shù)學方式合理地計算出權(quán)衡召回率和吞吐量的索引參數(shù)配置。依賴本文的方法可以直接以性能為指標來裝配索引，而無需深入研究索引中每個參數(shù)的具體意義、索引的具體實現(xiàn)原理，顯著降低了向量索引的使用門檻。因為量化索引的計算特性，召回率和吞吐量易于估量，未來對于結(jié)構(gòu)復雜的圖索引、數(shù)索引的參數(shù)配置方法也值得探索。

參考

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).

有任何問題可以搜索公眾號“向量檢索實驗室”，關(guān)注回復【加入社區(qū)】進群和大家一起交流討論。也可以直接添加負責人進群。