PSP團隊 投稿量子位 | 公眾號 QbitAI
只需修改兩行程式碼,RAG向量檢索效率暴漲30%!
不僅適用於文搜文”、“圖搜圖”、“文搜圖”、“推薦系統召回”多種任務;而且具備良好擴充套件性,適合十億、百億級別大規模應用。
浙江大學高雲君、柯翔宇團隊聯手向量檢索領域大佬傅聰,開源新方法PSP(Proximity graph with Spherical Pathway),突破RAG兩大難題。

簡單來說,主流向量檢索方法都是基於歐幾里得距離設計,主要看“誰離你最近”;但有時AI其實更需要比較“語義相關性”,也就是最大內積、看誰最相似。
以往的內積檢索辦法,不能像歐式距離檢索方法那樣滿足數學上的三角關係,所以很多老方法失效。
PSP發現,只要進行微小改動,老圖結構也能找到最大內積最優解。
而且PSP還設定了提前停止策略,能判斷檢索是否應該提前結束,避免浪費算力,讓搜尋更快。
AI產品背後的技術核心
向量檢索,是支撐起明星AI產品的核心技術元件。它不僅大大拓寬傳統語義檢索(關鍵詞檢索)的邊界,和大模型的配合更是渾然天成。
如何發揮這項技術的真正潛力,讓向量模型和向量資料庫的組合真正跑出效果,關鍵在於——選對“度量空間”。
儘管基於圖的向量檢索演算法,如HNSW、NSG等,因其優秀的檢索速度備受青睞,但往往被忽視的是,它們都是面向歐式空間設計的向量檢索演算法。
“度量錯配”在很多場景下是毀滅性的,很多適合用“最大內積”檢索的向量資料,搭配歐式向量演算法,往往會出現“檢索結果和query語義無關”的問題。
回看最大內積檢索領域,其實還沒有出現類似HNSW、NSG這樣現象級的檢索演算法。之前的很多工作往往只在某些資料集上面表現良好,但換了資料集,效果就會劇烈退化。
破局關鍵:僅需修改2行程式碼,實現全域性最大內積解
研究團隊透過理論探索發現,在最大內積檢索領域的研究涇渭分明地分成兩種正規化:
一是把最大內積轉換為最小歐式距離,進而可以用HNSW、NSG來解決。但這種轉化往往會伴隨著資訊損失或者拓撲空間的非線性轉換,而這些問題會對搜尋效果帶來不同程度的負面影響。
二是不進行空間轉化,直接在內積空間進行檢索。這樣做的好處是避免了資訊損失或空間扭曲,但相對應的痛點是,缺少有效手段對無效檢索空間進行裁剪,進而難以達到更好的檢索速度。
為什麼在內積空間直接做檢索這麼難呢?
最核心原因在於內積空間並不是一個嚴格意義上的“度量空間”。從數學上來說,一個空間可以稱之為“度量空間”,需要滿足諸多條件,典型地,我們最常接觸的歐式空間就是一個度量空間。而作為一個“殘缺空間”,內積空間缺少的最重要的屬性就是“三角不等式”。
根據NSG論文的理論部分,HNSW、NSG、SSG等state-of-the-art的向量檢索演算法之所以能如此高效,就是因為他們都利用了三角不等式對索引結構(圖結構)進行了高效的裁剪。
而以內積作為距離度量,構建的三角形,不滿足我們耳熟能詳的那句口訣“三角形中任意兩邊之和大於第三邊,而任意兩邊之差小於第三邊”。正是這一屬性的缺失,阻礙了最大內積檢索演算法進一步發展。
PSP研究團隊對這一問題進行了深入研究,從理論上證明了一件事情:對任意搜尋請求,即Query點q,它在一個為歐式距離設計的圖索引結構上,可以透過簡單的貪心演算法找到全域性最優的最大內積解。
基於圖的向量檢索演算法都利用貪心演算法進行檢索:當我們從隨機點開始在圖上游走時,NSG這類演算法會從路徑上的點的鄰居中,尋找一個距離目標“最近”的鄰居進行跳轉,這樣從鄰居的鄰居逐步跳轉到全域性最優解。
而這種貪心演算法曾經隱含的理論要求的是,如果構建圖用的是歐式距離表達“遠和近”,那麼貪婪遊走也需要用歐式距離來定義遠和近。
而PSP團隊的研究成果意義在於,如果構建圖用的是歐式距離,在貪婪遊走的時候可以用內積來定義遠和近,最終到達的終點就是全域性最優的最大內積解!
因此,研究團隊可以透過僅修改檢索(貪婪遊走)演算法中的兩行程式碼,就實現將一個現有的歐式演算法向最大內積的適配:

△實操中改變候選點佇列的“最大堆”、“最小堆”設定,以及距離度量
最佳化:合理引導搜尋行為規避冗餘計算
PSP研究團隊發現,最大內積檢索的過程中,會存在大量冗餘計算,而這些冗餘是可以透過合理引導搜尋行為來規避的。
最大內積中的搜尋行為與歐式空間中的搜尋行為有極大差異,如下圖所示:

左圖中,綠色方框(query)的最近歐式近鄰是紅色三角,但它的最大內積近鄰是橙色方塊。因此,在搜尋query的最近歐式鄰居的時候,遊走行為會很快在三角形附近停止,但搜尋他的最大內積鄰居會繼續走到“外圍”橙色方塊附近。
從更宏觀的角度看,研究團隊發現,最大內積檢索的解空間往往在資料集“外圍”(不同於歐式距離最近鄰,可以存在於資料空間的任意位置)。因此,最大內積的搜尋行為往往服從一種“由內而外,再外圍擴充套件”的模式(如上圖右圖)。
針對這種特性,PSP會設計針對性的策略,讓圖上搜索的起始點就儘量分佈在距離“答案”更近的區域。
同時,冗餘不僅僅發生在搜尋過程的前段,也非常多地集中在搜尋過程的後段。

如上圖,PSP研究團隊發現,在圖索引上搜索到精確解的“最少步數”因Query而異,呈現明顯的長尾分佈(圖a),而他們也透過大量實驗挖掘出四類“特徵”幫助我們判斷搜尋應該在什麼時候停下來(圖b)。這四類特徵可以在搜尋過程中以非常低的成本被計算和記錄,實現自適應的“早停”策略。
具體來說,可以在資料庫中隨機取樣一部分點作為query,透過對它們進行搜尋來收集最優停止步數前後的資料構成可分類的樣本,再用這個樣本去訓練一顆決策樹,就可以輔助搜尋過程判斷停止條件:

如上圖,研究團隊透過對決策樹剪枝,可以讓整棵決策樹保留較小的高度。選擇決策樹作為分類器,可以有效擬合少量樣本,並直接翻譯為if-else語句嵌入搜尋程式碼中,實現高效的“停止判斷”。
效能實測:穩定、高效、可擴充套件性
研究團隊為了充分測試PSP演算法的效果,在8個大規模、高維度的資料集上進行了充分測試。從維度看,DBpedia100K和DBpedia1M分別高達1536和3072維,用OpenAI text-embedding-3-large模型抽取;從數量看,最大的資料集Commerce100M包含1億資料庫點。

比較向量檢索演算法,往往關注相同召回率下的檢索速度,即Query-Per-Second(QPS)。從上圖中可看出,PSP相對於現有state-of-the-art的方法有著穩定、明顯的提升。在MNIST資料上,甚至超過第二名4倍之多。
值得注意的是,baseline的方法裡,往往有一些會在圖中“缺席”。這是因為它們效能遠差於其它方法,而很難和其它方法畫到同一張圖中。比如ip-HNSW在MNIST資料集中缺席;ScaNN在Laion10M和Commerce100M上缺席等等,這突出了PSP的表現穩定性。
另外,所使用的資料集包含了“文搜文”“圖搜圖”“文搜圖”“推薦系統召回”等諸多資料模態,體現出PSP強大的泛化性。
除了比較檢索效能,另外一個考察向量檢索演算法的應用價值的重要維度是scalability。好的檢索需要遠低於線性增長的時間複雜度(time complexity)。

上圖可以看出,PSP在Top-1近鄰上表現出log(N)速率增長的時間複雜度。而在Top-K檢索上表現出接近log(N)的複雜度。這體現出PSP優秀的可擴充套件性,即在十億乃至百億級別的資料上進行高效檢索的潛力。
論文連結: https://arxiv.org/pdf/2503.06882Github連結:https://github.com/ZJU-DAILY/PSP
一鍵三連「點贊」「轉發」「小心心」
歡迎在評論區留下你的想法!
— 完 —

🌟 點亮星標 🌟