動態超級塊剪枝:加速稀疏檢索的革命性技術 精華
突破性能瓶頸:動態超級塊剪枝如何重塑信息檢索效率
在當今數據爆炸的時代,高效的信息檢索系統對于各類應用至關重要,從搜索引擎到基于檢索增強的大語言模型(RAG)。隨著學習型稀疏表示模型的興起,如何在保持高檢索質量的同時提升檢索速度成為研究熱點。本文深入探討一項革命性技術——動態超級塊剪枝(Superblock Pruning,簡稱SP),這一創新方法在保持高相關性的前提下,顯著提升了稀疏檢索的效率。
稀疏檢索的挑戰與機遇
稀疏檢索模型如BM25和學習型稀疏表示(如SPLADE、E-SPLADE等)在僅使用CPU的服務器環境中廣受歡迎,主要得益于它們能夠充分利用高效的倒排索引實現。傳統的稀疏檢索速度優化通常采用動態秩安全索引剪枝技術,該技術能夠準確跳過那些得分較低、不可能出現在最終top-k結果中的文檔。
近年來,基于塊的檢索方法成為研究熱點,這類方法將文檔分配到塊(或稱為簇)中,并利用塊級信息改進索引遍歷順序,同時剪枝低分文檔組。然而,這些方法在處理大規模數據集時仍然面臨效率挑戰,尤其是當需要保持高相關性時。
超級塊剪枝:創新的兩級剪枝策略
超級塊剪枝(SP)技術在現有基于塊的剪枝方法基礎上進行了創新性擴展。SP方法將一系列連續的文檔塊均勻聚合成超級塊,然后以自上而下的方式進行在線索引遍歷。這種設計為每個超級塊分配固定數量的文檔塊,簡化了向量化和緩存優化過程,同時提供了具有概率安全保證的兩級剪枝機制。
SP的核心創新在于其兩級剪枝策略:首先計算所有超級塊的邊界信息并進行剪枝,然后再計算塊的邊界并進行剪枝。具體來說,SP執行以下動態剪枝步驟:
- 對于超級塊X,計算該超級塊內文檔的最大和平均排名分數邊界
- 當超級塊的最大和平均超級塊邊界滿足特定條件時,該超級塊被剪枝
- 對于文檔塊B,如果其邊界和滿足特定條件,則剪枝該塊
- 對于所有未被剪枝的塊,按照其邊界和值的降序對相應的文檔塊進行排序和評分
這種方法允許SP更有效地跳過文檔塊,以排名安全或概率排名安全的方式加速檢索。剪枝一個超級塊不僅避免了計算子塊的最大分數,還避免了對其子塊內文檔的評分,從而大幅提升檢索效率。
理論保證與實現優化
SP具有與ASC類似的排名安全μ-競爭性質。可以證明,SP的平均top-k'排名分數與任何排名安全檢索算法R在μ因子內相同。作為額外保障,如果我們假設文檔的排名分數在每個超級塊內獨立同分布,SP還提供概率安全性。
在實現層面,SP采用了多項優化策略:
CPU緩存使用優化
SP使用SIMD指令計算相關公式。當順序計算所有查詢項的這些公式而不進行塊跳過時,現代編譯器可以輕松向量化其實現,現代CPU可以有效地預取數據。然而,由于超級塊剪枝導致的不規則和非連續數據訪問,編譯器難以優化塊級邊界計算。因此,SP需要顯式控制CPU緩存在計算塊級邊界時的重用模式。
SP采用了超級塊優先的邊界和計算方式,即對每個未剪枝的超級塊,先對該超級塊內的所有塊進行完整評分,然后再處理下一個未剪枝的超級塊。這種方法允許在內部循環中重用累積寄存器以獲得更好的L1緩存性能。實驗表明,這種方式比傳統的項優先方法最高可提速1.89倍。
實驗評估與性能對比
研究團隊在MS MARCO段落排名數據集上進行了全面評估,該數據集包含880萬個英文段落。評估采用標準指標:平均倒數排名(MRR@10)和位置1000(k=1000)或10(k=10)的召回率。所有實驗在配備Intel i7-1260P、64GB RAM和AVX2指令的Linux系統上使用單線程運行。
SP與三種最先進的基于塊的檢索算法進行了比較:BMP、Seismic和ASC。實驗結果表明,在高相關性預算要求下,SP在SPLADE和E-SPLADE上的表現顯著優于這些基線方法。
SPLADE模型上的性能比較
在SPLADE上的排名安全搜索中,SP比BMP在k=10時快32%,在k=1000時快25%。與ASC相比,SP在k=10和k=1000時都快約3.3倍。對于99%的召回預算,SP比BMP快至多2.9倍,比Seismic快3.3倍,比ASC快9.1倍。
上圖展示了SP和BMP在塊大小b從128減小到8時的總延遲(上圖)和成本細分(下圖)。當b變小時,BMP能夠獲得更緊密的邊界和,但塊過濾開銷增加。SP在評估小塊的同時減少了塊和超級塊過濾的開銷。
超級塊剪枝的有效性
實驗數據顯示,即使在安全搜索(μ=1)情況下,SP也能剪枝24%的超級塊(k=10)。隨著μ減小,超級塊級別的剪枝量顯著增加,而被剪枝的塊數量大致相同。這是因為塊對其內部文檔形成了緊密邊界;SP能夠避開不太可能包含相關文檔的塊組,從而減少計算塊邊界和的開銷。
在100%概率安全性(η=1)下,即使在μ=0.4(k=1000)時,對Dev集、DL 19和DL 20的相關性指標影響也可忽略不計,盡管當μ=0.4時召回率開始下降。相比之下,即使在BMP中使用低估計閾值也會導致相關性大幅下降。
E-SPLADE模型上的性能
在E-SPLADE上,SP在不同高相關性召回預算下的表現也優于其他基線,比Seismic快至多16倍,比BMP快1.4倍。
超級塊剪枝的優勢與局限
與現有方法相比,SP具有以下顯著優勢:
- 更高效的塊跳過:與BMP相比,SP利用其超級塊結構快速跳過大量塊,同時提供額外的η保障以確保概率安全性。
- 更好的緩存利用:與Anytime Ranking、ASC和Seismic相比,SP能夠處理更多的塊數量,并通過緩存優化的超級塊剪枝克服額外開銷,自然導致更緊密的邊界估計。
- 高相關性保證:在保持高相關性的同時顯著提升檢索速度,特別適合對檢索質量要求較高的應用場景。
然而,SP也存在一些局限性:
- 額外空間成本:與BMP相比,SP需要額外空間來維護每個超級塊的最大和平均項權重。在MS MARCO評估中,當c=64、b=8時,額外空間約為2GB;當b=16時,額外空間約為1GB。
- 靜態索引剪枝的缺失:與Seismic不同,SP沒有利用靜態索引剪枝,這可能在某些情況下限制其性能。
應用場景與未來展望
SP技術特別適合需要高相關性的應用場景。對于此類應用,建議將η設置接近1.0,并將μ從0.4變化到1。檢索是大規模搜索系統和基于檢索增強的大語言模型(如RAG)的關鍵組件,在低成本CPU上實現高相關性的快速檢索可以產生積極影響。
未來研究方向包括:
- 探索靜態索引剪枝、自定義摘要和文檔鄰近圖等技術與SP的結合
- 研究SP與索引壓縮方案的結合
- 開發針對輸入復雜性的動態縮放策略
- 整合自適應推理深度控制以在推理期間平衡效率和安全性能
結論
動態超級塊剪枝(SP)是一種創新的動態剪枝方案,除標準塊級別外,還在超級塊級別進行剪枝,并設計為利用CPU緩存局部性。實驗評估表明,在SPLADE上99%或更高的召回預算下,SP比Seismic快2.3倍至3.8倍,比ASC快3.2倍至9.4倍,比BMP快至多2.9倍。對于安全搜索,SP比BMP快至多1.3倍。
隨著信息檢索需求的不斷增長和大語言模型對高效檢索系統的依賴加深,SP技術有望在未來發揮更加重要的作用,特別是在需要在保持高相關性的同時提高檢索效率的場景中。
參考資料
論文:???https://arxiv.org/abs/2504.17045??
GitHub:???https://github.com/thefxperson/hierarchical_pruning??
本文轉載自???頓數AI????,作者:可可
