大規模相似性搜索:原理、技術與 Faiss 實踐
相似性搜索為何重要?
人工智能和機器學習的興起,催生了大量高維數據表示形式,即嵌入(embeddings),它們捕捉數據點之間的復雜關系,助力強大的分析與理解。然而,在大型數據集中查找相似嵌入是一項計算密集型任務。相似性搜索在檢索增強生成(Retrieval-Augmented Generation,RAG)領域引發了變革。RAG 將傳統信息檢索與語言模型相結合,通過利用相似性搜索查找相關文檔,使模型能訪問更廣泛的知識庫,生成更具信息量和上下文豐富的輸出,從而提高生成文本的準確性和相關性。
大規模相似性搜索的挑戰
傳統數據庫和搜索引擎難以滿足大規模相似性搜索的需求。它們依賴結構化查詢和索引方法,無法應對高維數據的動態特性。因此,需要專門的技術來解決這一問題。
Faiss 框架
Faiss 由 Facebook AI Research 開發,是一個專為高效相似性搜索設計的強大庫。它提供多種索引方法,在性能、準確性和內存使用之間進行了不同的權衡優化。Faiss 還支持 GPU 加速,非常適合處理大規模數據集。
基礎準備
首先安裝和導入必要的依賴項:
pip install faiss-cpu
import time
import faiss
import numpy as np
接著定義一些常量:d? 表示向量維度(128),nb? 表示基礎向量數量(10000),nq 表示查詢向量數量(100)。
d = 128
nb = 10000
nq = 100
為保證結果可復現,初始化隨機種子:
np.random.seed(1234)
生成兩組隨機向量,xb? 代表基礎向量(10000 x 128),xq 代表查詢向量(100 x 128):
xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')
這些向量本質上就是我們的數據點。
分層可導航小世界(Hierarchical Navigable Small World,HNSW)
- 工作原理:HNSW 是一種基于圖的索引方法,向量被組織在小世界圖的層次結構中。圖中的每個節點(向量)都與其最近鄰節點相連。搜索時,算法在圖中導航,快速收斂到最近的向量。
- 優勢:HNSW 準確性高且搜索速度快,尤其適用于高維數據集。
- 關鍵參數:
M(每個節點的連接數):控制圖中每個節點連接的鄰居數量。數值越高,準確性越高,但內存消耗也越大。
efConstruction? 和efSearch:分別控制索引構建和搜索過程中的探索深度。數值越高,搜索準確性越好,但計算量也更大。
創建兩個 HNSW 索引 —— HNSWFlat? 和 HNSWSQ(標量量化)來對比性能:
# HNSWFlat 是基本的 HNSW 實現
index_hnswflat = faiss.IndexHNSWFlat(d, 32)
start = time.time()
index_hnswflat.add(xb)
indexing_time_hnswflat = time.time() - start
start = time.time()
D, I = index_hnswflat.search(xq, 5)
search_time_hnswflat = time.time() - start
# HNSWSQ 結合了標量量化(SQ)以加快索引速度
quantizer_sq = faiss.IndexScalarQuantizer(d, faiss.ScalarQuantizer.QT_8bit)
index_hnswsq = faiss.IndexHNSWFlat(d, 32)
start = time.time()
index_hnswsq.add(xb)
indexing_time_hnswsq = time.time() - start
start = time.time()
D, I = index_hnswsq.search(xq, 5)
search_time_hnswsq = time.time() - start
兩個索引都使用“扁平”存儲方法,即不壓縮原始向量。
輸出結果:
print(f"HNSWFlat Indexing Time: {indexing_time_hnswflat:.4f} seconds")
print(f"HNSWFlat Search Time: {search_time_hnswflat:.4f} seconds")
print(f"HNSWSQ Indexing Time: {indexing_time_hnswsq:.4f} seconds")
print(f"HNSWSQ Search Time: {search_time_hnswsq:.4f} seconds")
與基于 IVF 的方法相比,HNSW 方法的索引速度明顯較慢,因為 HNSW 需要構建鄰居圖,計算成本較高。由于標量量化(SQ)在索引過程中更緊湊地表示向量,降低了向量維度,所以 HNSWSQ? 比 HNSWFlat? 稍快。HNSWFlat? 和 HNSWSQ 的搜索時間比基于 IVF 的方法略長,這是因為需要遍歷圖來找到最近鄰居。HNSW 以高精度著稱,尤其在高維空間中,但代價是索引和搜索時間較長。
倒排文件索引(Inverted File Index,IVF)
- 工作原理:IVF 是大規模數據集相似性搜索中另一種常用方法。它將數據集劃分為多個桶(buckets)或“列表”,每次查詢僅搜索其中一部分桶。
- 優勢:與 HNSW 相比,IVF 的主要優勢是內存需求較低。
- 關鍵參數:
nlist:聚類(或桶)的數量。數值越高,精度越高,但索引時間也會增加。
nprobe:查詢時搜索的聚類數量。增加此參數可提高召回率,但會降低搜索速度。
對比 IndexIVFFlat?(無乘積量化)和 IndexIVFPQ(有乘積量化):
# IVFFlat 使用無量化的扁平索引
nlist = 100
quantizer = faiss.IndexFlatL2(d)
index_ivfflat = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)
start = time.time()
index_ivfflat.train(xb)
index_ivfflat.add(xb)
indexing_time_ivfflat = time.time() - start
index_ivfflat.nprobe = 10
start = time.time()
D, I = index_ivfflat.search(xq, 5)
search_time_ivfflat = time.time() - start
# IVFPQ 結合乘積量化(PQ)以提高內存效率
m = 8
nbits = 8
index_ivfpq = faiss.IndexIVFPQ(quantizer, d, nlist, m, nbits)
start = time.time()
index_ivfpq.train(xb)
index_ivfpq.add(xb)
indexing_time_ivfpq = time.time() - start
index_ivfpq.nprobe = 10
start = time.time()
D, I = index_ivfpq.search(xq, 5)
search_time_ivfpq = time.time() - start
IVF 的關鍵思想是將數據集劃分為聚類(桶),每次查詢僅搜索其中一部分桶。這里設置 nlist 為 100,即有 100 個聚類。
輸出結果:
print(f"IVFFlat Indexing Time: {indexing_time_ivfflat:.4f} seconds")
print(f"IVFFlat Search Time: {search_time_ivfflat:.4f} seconds")
print(f"IVFPQ Indexing Time: {indexing_time_ivfpq:.4f} seconds")
print(f"IVFPQ Search Time: {search_time_ivfpq:.4f} seconds")
IVFPQ? 的索引時間比 IVFFlat? 長得多,因為 IVFPQ? 在初始聚類后還涉及額外的量化步驟。它應用乘積量化(PQ),需要學習一個碼本,將每個向量壓縮為多個量化子向量。IVFPQ 的搜索時間略長,這是由于在搜索過程中需要從量化表示中解壓縮和重構向量,但差異很小,其內存效率的提升通常值得這額外的搜索時間。
局部敏感哈希(Locality Sensitive Hashing,LSH)
- 工作原理:LSH 將高維向量轉換為低維“哈希”值。相似向量更有可能具有相同的哈希值,通過關注包含相關哈希的桶來實現高效搜索。
- 優勢:LSH 為相似性搜索提供了一種快速且可擴展的方法,尤其適用于大型數據集和高維空間。
- 關鍵參數:
哈希表數量:控制準確性和速度之間的權衡。哈希表越多,準確性越高,但搜索時間也會增加。
每個表的哈希函數數量:用于生成每個哈希值的哈希函數數量。
使用 IndexLSH(基于哈希的方法,利用隨機投影為每個向量創建哈希值):
nbits = 16
index_lsh = faiss.IndexLSH(d, nbits)
start = time.time()
index_lsh.add(xb)
indexing_time_lsh = time.time() - start
start = time.time()
D, I = index_lsh.search(xq, 5)
search_time_lsh = time.time() - start
相似向量更有可能具有相同的哈希值,我們使用 16 位哈希。
輸出結果:
print(f"LSH Indexing Time: {indexing_time_lsh:.4f} seconds")
print(f"LSH Search Time: {search_time_lsh:.4f} seconds")
LSH 的索引速度極快,因為它只是基于隨機投影將數據點哈希到哈希桶中,無需像 IVF 或 HNSW 那樣的訓練過程,所以索引幾乎是即時的。與 IVFFlat? 和 HNSW? 相比,LSH 的搜索相對較慢,這是因為 LSH 的隨機性,可能需要搜索多個哈希桶才能找到最近鄰居。LSH 通常索引速度快,但與 HNSW 或 IVFPQ 等更復雜的方法相比,可能會犧牲準確性和搜索速度。
本文只是對大規模相似性搜索領域的簡要介紹,僅觸及了基礎知識,還有更多內容有待探索。未來我們將深入研究實際應用,探索更高級的索引方法,甚至使用 Faiss 構建一些有趣的項目。
本文轉載自 ??柏企閱文??,作者: 柏企
