相似性搜索為何重要？

人工智能和機(jī)器學(xué)習(xí)的興起，催生了大量高維數(shù)據(jù)表示形式，即嵌入（embeddings），它們捕捉數(shù)據(jù)點(diǎn)之間的復(fù)雜關(guān)系，助力強(qiáng)大的分析與理解。然而，在大型數(shù)據(jù)集中查找相似嵌入是一項(xiàng)計(jì)算密集型任務(wù)。相似性搜索在檢索增強(qiáng)生成（Retrieval-Augmented Generation，RAG）領(lǐng)域引發(fā)了變革。RAG 將傳統(tǒng)信息檢索與語言模型相結(jié)合，通過利用相似性搜索查找相關(guān)文檔，使模型能訪問更廣泛的知識(shí)庫，生成更具信息量和上下文豐富的輸出，從而提高生成文本的準(zhǔn)確性和相關(guān)性。

大規(guī)模相似性搜索的挑戰(zhàn)

傳統(tǒng)數(shù)據(jù)庫和搜索引擎難以滿足大規(guī)模相似性搜索的需求。它們依賴結(jié)構(gòu)化查詢和索引方法，無法應(yīng)對(duì)高維數(shù)據(jù)的動(dòng)態(tài)特性。因此，需要專門的技術(shù)來解決這一問題。

Faiss 框架

Faiss 由 Facebook AI Research 開發(fā)，是一個(gè)專為高效相似性搜索設(shè)計(jì)的強(qiáng)大庫。它提供多種索引方法，在性能、準(zhǔn)確性和內(nèi)存使用之間進(jìn)行了不同的權(quán)衡優(yōu)化。Faiss 還支持 GPU 加速，非常適合處理大規(guī)模數(shù)據(jù)集。

基礎(chǔ)準(zhǔn)備

首先安裝和導(dǎo)入必要的依賴項(xiàng)：

pip install faiss-cpu
import time
import faiss
import numpy as np

接著定義一些常量：d? 表示向量維度（128），nb? 表示基礎(chǔ)向量數(shù)量（10000），nq 表示查詢向量數(shù)量（100）。

d = 128
nb = 10000
nq = 100

為保證結(jié)果可復(fù)現(xiàn)，初始化隨機(jī)種子：

np.random.seed(1234)

生成兩組隨機(jī)向量，xb? 代表基礎(chǔ)向量（10000 x 128），xq 代表查詢向量（100 x 128）：

xb = np.random.random((nb, d)).astype('float32')
xq = np.random.random((nq, d)).astype('float32')

這些向量本質(zhì)上就是我們的數(shù)據(jù)點(diǎn)。

分層可導(dǎo)航小世界（Hierarchical Navigable Small World，HNSW）

• 工作原理：HNSW 是一種基于圖的索引方法，向量被組織在小世界圖的層次結(jié)構(gòu)中。圖中的每個(gè)節(jié)點(diǎn)（向量）都與其最近鄰節(jié)點(diǎn)相連。搜索時(shí)，算法在圖中導(dǎo)航，快速收斂到最近的向量。
• 優(yōu)勢(shì)：HNSW 準(zhǔn)確性高且搜索速度快，尤其適用于高維數(shù)據(jù)集。
• 關(guān)鍵參數(shù)：

M（每個(gè)節(jié)點(diǎn)的連接數(shù)）：控制圖中每個(gè)節(jié)點(diǎn)連接的鄰居數(shù)量。數(shù)值越高，準(zhǔn)確性越高，但內(nèi)存消耗也越大。

efConstruction? 和efSearch：分別控制索引構(gòu)建和搜索過程中的探索深度。數(shù)值越高，搜索準(zhǔn)確性越好，但計(jì)算量也更大。

創(chuàng)建兩個(gè) HNSW 索引 —— HNSWFlat? 和 HNSWSQ（標(biāo)量量化）來對(duì)比性能：

# HNSWFlat 是基本的 HNSW 實(shí)現(xiàn)
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 結(jié)合了標(biāo)量量化（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

兩個(gè)索引都使用“扁平”存儲(chǔ)方法，即不壓縮原始向量。

輸出結(jié)果：

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 方法的索引速度明顯較慢，因?yàn)?HNSW 需要構(gòu)建鄰居圖，計(jì)算成本較高。由于標(biāo)量量化（SQ）在索引過程中更緊湊地表示向量，降低了向量維度，所以 HNSWSQ? 比 HNSWFlat? 稍快。HNSWFlat? 和 HNSWSQ 的搜索時(shí)間比基于 IVF 的方法略長(zhǎng)，這是因?yàn)樾枰闅v圖來找到最近鄰居。HNSW 以高精度著稱，尤其在高維空間中，但代價(jià)是索引和搜索時(shí)間較長(zhǎng)。

倒排文件索引（Inverted File Index，IVF）

• 工作原理：IVF 是大規(guī)模數(shù)據(jù)集相似性搜索中另一種常用方法。它將數(shù)據(jù)集劃分為多個(gè)桶（buckets）或“列表”，每次查詢僅搜索其中一部分桶。
• 優(yōu)勢(shì)：與 HNSW 相比，IVF 的主要優(yōu)勢(shì)是內(nèi)存需求較低。
• 關(guān)鍵參數(shù)：

nlist：聚類（或桶）的數(shù)量。數(shù)值越高，精度越高，但索引時(shí)間也會(huì)增加。

nprobe：查詢時(shí)搜索的聚類數(shù)量。增加此參數(shù)可提高召回率，但會(huì)降低搜索速度。

對(duì)比 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 結(jié)合乘積量化（PQ）以提高內(nèi)存效率
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 的關(guān)鍵思想是將數(shù)據(jù)集劃分為聚類（桶），每次查詢僅搜索其中一部分桶。這里設(shè)置 nlist 為 100，即有 100 個(gè)聚類。

輸出結(jié)果：

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? 的索引時(shí)間比 IVFFlat? 長(zhǎng)得多，因?yàn)?nbsp;IVFPQ? 在初始聚類后還涉及額外的量化步驟。它應(yīng)用乘積量化（PQ），需要學(xué)習(xí)一個(gè)碼本，將每個(gè)向量壓縮為多個(gè)量化子向量。IVFPQ 的搜索時(shí)間略長(zhǎng)，這是由于在搜索過程中需要從量化表示中解壓縮和重構(gòu)向量，但差異很小，其內(nèi)存效率的提升通常值得這額外的搜索時(shí)間。

局部敏感哈希（Locality Sensitive Hashing，LSH）

• 工作原理：LSH 將高維向量轉(zhuǎn)換為低維“哈?！敝?。相似向量更有可能具有相同的哈希值，通過關(guān)注包含相關(guān)哈希的桶來實(shí)現(xiàn)高效搜索。
• 優(yōu)勢(shì)：LSH 為相似性搜索提供了一種快速且可擴(kuò)展的方法，尤其適用于大型數(shù)據(jù)集和高維空間。
• 關(guān)鍵參數(shù)：

哈希表數(shù)量：控制準(zhǔn)確性和速度之間的權(quán)衡。哈希表越多，準(zhǔn)確性越高，但搜索時(shí)間也會(huì)增加。

每個(gè)表的哈希函數(shù)數(shù)量：用于生成每個(gè)哈希值的哈希函數(shù)數(shù)量。

使用 IndexLSH（基于哈希的方法，利用隨機(jī)投影為每個(gè)向量創(chuàng)建哈希值）：

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 位哈希。

輸出結(jié)果：

print(f"LSH Indexing Time: {indexing_time_lsh:.4f} seconds")
print(f"LSH Search Time: {search_time_lsh:.4f} seconds")

LSH 的索引速度極快，因?yàn)樗皇腔陔S機(jī)投影將數(shù)據(jù)點(diǎn)哈希到哈希桶中，無需像 IVF 或 HNSW 那樣的訓(xùn)練過程，所以索引幾乎是即時(shí)的。與 IVFFlat? 和 HNSW? 相比，LSH 的搜索相對(duì)較慢，這是因?yàn)?LSH 的隨機(jī)性，可能需要搜索多個(gè)哈希桶才能找到最近鄰居。LSH 通常索引速度快，但與 HNSW 或 IVFPQ 等更復(fù)雜的方法相比，可能會(huì)犧牲準(zhǔn)確性和搜索速度。

本文只是對(duì)大規(guī)模相似性搜索領(lǐng)域的簡(jiǎn)要介紹，僅觸及了基礎(chǔ)知識(shí)，還有更多內(nèi)容有待探索。未來我們將深入研究實(shí)際應(yīng)用，探索更高級(jí)的索引方法，甚至使用 Faiss 構(gòu)建一些有趣的項(xiàng)目。

本文轉(zhuǎn)載自 ??柏企閱文??，作者： 柏企