SPLADE 在稀疏向量搜索中的原理與應用詳解

今天看到Sentence Transformers v5.0 集成了許多稀疏嵌入模型。為了搞清楚什么稀疏嵌入模型以及應用，查到了SPLADE，比較巧合的是在paper reading分享的時候看到有同學分享了一片ACL 2025的工作也是基于SPLADE去做的。下面結合一些資料分享關于SPLADE 在稀疏向量搜索中的原理以及應用。

主要內容來自：https://www.pinecone.io/learn/splade/

在現代向量搜索出現之前，我們主要采用“傳統”的詞袋模型（ Bag of Words, BOW）方法。基于這些方法，我們將待檢索的文檔（例如谷歌的網頁）轉化為一個詞集（即詞袋），進而生成一個表示詞頻的稀疏向量。TF-IDF 和 BM25 是其中典型的算法。

稀疏向量因其高效性、可解釋性和精確的詞語匹配特性，在信息檢索領域曾廣受歡迎。然而，它們遠非完美。

稀疏向量搜索的工作方式與人類的自然表達存在脫節。我們在搜索信息時，通常很難預測目標文檔中會包含哪些確切的詞語。

稠密嵌入模型（Dense Embedding Models）在這方面提供了一定的幫助。利用稠密模型，我們可以基于“語義含義”進行搜索，而非僅僅依賴于詞語匹配。然而，這些模型仍然有不足之處。

稠密嵌入模型需要大量數據進行微調（fine-tune），否則其性能可能不如稀疏方法。這對于難以獲取數據且領域特定術語很重要的小眾領域來說，是一個棘手的問題。

過去曾涌現出各種臨時性的補救方案來應對這些挑戰，包括復雜且（仍然不完美）的兩階段檢索系統，以及查詢和文檔擴展或改寫方法（我們將在后文探討）。然而，這些方案都未能提供真正魯棒、持久的解決方案。

幸運的是，該領域已取得顯著進展，有望結合兩者的優勢。通過混合搜索 (Hybrid Search)技術，稀疏和稠密檢索得以融合；而可學習的稀疏嵌入 (Learnable Sparse Embeddings)則有助于克服傳統稀疏檢索的不足。

本文將深入探討可學習稀疏嵌入領域的最新進展——SPLADE (Sparse Lexical and Expansion model)模型 [1]。

稀疏向量與稠密向量

在信息檢索中，向量嵌入（Vector Embeddings）將文檔和查詢表示為數值向量格式。這種格式使得我們能夠在向量數據庫中通過計算相似度來檢索相似的向量。

稀疏向量和稠密向量是向量表示的兩種不同形式，各有優缺點。

稀疏向量包含很多零值，非零值比例非常小。

TF-IDF 或 BM25 等稀疏向量具有高維度，但包含的非零值非常少（因此得名“稀疏”）。稀疏向量已有數十年的研究歷史，由此產生了緊湊的數據結構和許多專門針對這類向量設計的高效檢索算法。

稠密向量維度較低，但包含豐富信息，大多數或全部維度都包含非零值。這些向量通常由神經網絡模型（如 Transformer）構建，因此能夠表示更抽象的信息，例如文本的語義含義。

總的來說，這兩種方法的優缺點可以概括如下表：

稀疏檢索

優點	缺點
通常檢索速度更快	性能無法相比基線顯著提升
具有良好的基線性能	性能無法相比基線顯著提升
不需要模型微調	存在詞匯不匹配問題
詞匯精確匹配

密集檢索

優點	缺點
通過微調可以超越稀疏檢索性能	需要訓練數據，在低資源場景下較困難
可以搜索類似人類的抽象概念	泛化能力不強，特別是對于細分術語
支持多模態（文本、圖像、音頻等）和跨模態搜索（如文本到圖像）	比稀疏檢索需要更多計算和內存資源
	無法精確匹配
	不易解釋

理想情況下，我們希望結合兩者的優勢，但這很難實現。

兩階段檢索

一種常見的處理方法是實現兩階段檢索和排序系統（Two-stage Retrieval and Ranking）。在這種方案中，系統使用兩個不同的階段來檢索和排序與給定查詢相關的文檔。

在第一階段，系統使用稀疏檢索方法召回大量候選文檔。然后，這些文檔被傳遞到第二階段，使用稠密模型根據它們與查詢的相關性重新排序結果。

兩階段檢索系統包含一個稀疏檢索器和一個稠密重排序器。

這種方法有一些優點：（1）對完整的文檔集應用稀疏模型進行召回更高效；（2）對召回后較小的文檔集使用較慢的稠密模型進行重排序可以更準確。通過這種方式，我們可以向用戶返回更相關的結果。另一個優點是，重排序階段與檢索系統是分離的，這對于多用途的檢索系統很有幫助。

然而，它并非完美。兩階段的檢索和重排序可能比使用近似搜索算法的單階段系統要慢。擁有兩個階段會帶來更高的工程復雜性。最后，系統的性能依賴于第一階段檢索器能否返回相關的結果；如果第一階段未能召回有用的內容，第二階段的重排序也無濟于事。

改進單階段系統

由于兩階段檢索存在固有的不足，大量研究致力于改進單階段檢索系統。

單階段檢索系統。注意，檢索器可以是稀疏、稠密，甚至兩者兼具。

這方面的研究成果之一就是更魯棒、更可學習的稀疏嵌入模型——其中性能最優的模型之一就是 SPLADE。

SPLADE（SparseLexical and Expansion model）模型的理念是：一個預訓練語言模型（如 BERT）可以識別詞語/子詞（在本文中稱為“詞片段”或“詞項”）之間的聯系，并利用這些知識來增強我們的稀疏向量嵌入。

這通過兩種方式實現：它允許我們為不同詞項賦予相關性權重（例如，“the”這樣的詞項權重較低，而“orangutan”等不常用詞權重較高）；同時，它支持詞項擴展（Term Expansion）：包含除原始文本中出現詞項之外的、相關但不同的備選詞項。

詞項擴展允許我們識別相關但不同的詞項，并在稀疏向量檢索步驟中使用它們。

SPLADE 最顯著的優勢在于它能夠學習詞項擴展，而非僅僅執行詞項擴展。傳統方法需要基于規則進行詞項擴展，這既耗時又本質受限。而 SPLADE可以利用最優秀的語言模型來學習詞項擴展，甚至可以根據句子的上下文對其進行微調。

盡管查詢和文檔包含許多相關的詞項，但由于它們不是“精確匹配”，因此未能被識別。

詞項擴展對于緩解詞匯不匹配問題至關重要——這是查詢和相關文檔之間典型缺乏詞項重疊的現象。

通過對查詢進行詞項擴展，我們將獲得更大的重疊度，因為現在能夠識別相似的詞語。

由于語言的復雜性以及描述同一事物的多種方式，相關文檔與查詢之間可能存在很少甚至沒有詞項重疊，這是可以預期的。詞項擴展正是為了解決這一問題。

SPLADE 嵌入構建過程

SPLADE 構建稀疏嵌入的過程是易于理解的。我們首先使用一個帶有**掩碼語言模型（ Masked-Language Modeling, MLM）**頭的 Transformer 模型，例如 BERT。

MLM 是許多 Transformer 常用的預訓練方法。我們可以直接使用一個現成的預訓練 BERT 模型。### BERT 模型介紹

如前所述，我們將使用帶有 MLM 頭的 BERT 模型。如果您熟悉 BERT 和 MLM，那很好；如果不熟悉，下面我們將進行分解。

BERT 是一種流行的 Transformer 模型。與所有 Transformer 一樣，其核心功能是生成信息豐富的詞元嵌入（Token Embeddings）。這具體意味著什么呢？

我們以文本 "Orangutans are native to the rainforests of Indonesia and Malaysia"（猩猩原產于印度尼西亞和馬來西亞的雨林）為例。首先，我們會將這段文本**詞元化（tokenize）**為 BERT 特定的子詞詞元：

text = ("Orangutans are native tothe rainforests of ""Indonesia and Malaysia"
)# create the tokens that will be input into the model
tokens = tokenizer(text, returntensors="pt")
tokens

{'inputids': tensor([[  101,  2030,  5654, 13210,  3619,  2024,  3128,2000,  1996, 18951,\2015,  1997,  6239,  1998,  6027,   102]]), 'tokentypeids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,0, 0, 0, 0, 0]]), 'attentionmask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1,1, 1, 1, 1, 1, 1, 1]])}

# we transform the inputids to human-readable tokens
tokenizer.convertidstotokens(tokens["inputids"][0])

['[CLS]',\'or',\'##ang',\'##uta',\'##ns',\'are',\'native',\'to',\'the',\'rainforest',\'##s',\'of',\'indonesia',\'and',\'malaysia',\'[SEP]']

詞元 ID 被映射到嵌入矩陣中學習到的詞元嵌入。

這些詞元會與一個**“嵌入矩陣（EmbeddingMatrix）”**相匹配，嵌入矩陣是 BERT 模型的第一層。在這個嵌入矩陣中，我們可以找到學習到的“向量嵌入（Vector Embeddings）”，它們是這些詞語/子詞詞元的“數值表示”。

嵌入矩陣中的向量在有意義的向量空間中分別代表一個詞元。

接下來，原始文本的詞元表示會通過多個**“編碼器塊（Encoder Blocks）”**。這些編碼器塊基于文本的其余上下文，將越來越多的上下文信息編碼到每個向量嵌入中。

在此之后，我們就得到了 Transformer 的**“輸出（Output）”**：信息豐富的向量嵌入。每個嵌入都代表了先前的詞元，但融入了從原始句子中提取的其他詞元向量嵌入中獲取的信息。

通過多個注意力編碼器塊處理初始詞元嵌入，可以編碼更多上下文信息，從而生成信息豐富的嵌入。

這個過程是 BERT 和其他所有 Transformer 模型的核心。然而，Transformer 的強大之處在于，這些信息豐富的向量可以用于眾多不同的任務。通常，我們會給 Transformer 添加一個任務特定的“頭（Head）”，將這些向量轉換為其他形式，例如預測結果或稀疏向量。

掩碼語言模型頭 (MLM Head)

MLM 頭是 BERT 模型常用的眾多“頭”之一。與大多數頭不同，MLM 頭用于 BERT 的初始預訓練階段。

其工作原理是：輸入一個句子，例如 "Orangutans are native to the rainforests of Indonesia and Malaysia"。我們將文本進行詞元化，然后用特殊的 [MASK] 詞元隨機替換部分詞元。

可以使用 [MASK] 詞元屏蔽任何詞語或子詞詞元。

這個經過掩碼處理的詞元序列作為輸入傳遞給 BERT。在另一端，我們將原始句子提供給 MLM 頭。然后，對 BERT 和 MLM 頭進行優化，使其能夠預測被 [MASK] 詞元替換的原始詞語/子詞詞元。

MLM 頭從每個輸出logits 生成一個概率分布。這些概率表示對詞匯表中每個詞元代表 [MASK] 的預測。

為實現這一功能，MLM 頭為每個詞元位置輸出 30522 個值。這30522 個值代表 BERT 的詞匯表大小，并構成一個在詞匯表上的概率分布。激活度最高的值對應的詞元，即為該詞元位置的詞元預測結果。

MLM與稀疏向量

這 30522 個概率分布 wijw{ij}wij? 指示了詞匯表 VVV 中哪些詞語/詞元 jjj 最為重要。MLM 頭為模型輸入的每個詞元 iii 輸出這些分布。

MLM 頭為每個詞元（無論是否被掩碼）提供一個概率分布。這些分布被聚合起來，得到重要性估計。

SPLADE 將所有這些概率分布聚合成一個單一的分布，稱為重要性估計（Importance Estimation） wjw\jwj?。這個重要性估計就是 SPLADE 生成的稀疏向量。我們可以將所有這些概率分布組合成一個單一的分布，它告訴我們詞匯表中的每個詞元與輸入句子的相關性。

其計算公式如下：

$w_j = \sum_{i \in t} \log(1 + \text{ReLU}(w_{ij}))$

其中：

$\in t$ ：表示輸入詞元集合 $t$ 中的每一個詞元 $i$ 。

$w_{ij}$ ：表示對于每個詞元 $i$ ，模型預測的詞匯表 $V$ 中所有詞元 $j$ 的權重值。

這使得我們能夠識別輸入句子中不存在但相關的詞元。例如，如果我們掩碼了詞語 rainforest（雨林），模型可能會對 jungle（叢林）、land（土地）和 forest（森林）等詞返回較高的預測概率 $w_j$ 。這些詞語及其相關的概率隨后會在 SPLADE 構建的稀疏向量中得到體現。

這種**“學習到的”查詢/文檔擴展能力，即包含其他相關詞項的能力，是 SPLADE 相較于傳統稀疏方法的一個關鍵優勢**。它基于學習到的詞項關系和上下文，幫助我們最大程度地緩解詞匯不匹配問題。

查詢中的詞項擴展可以大大增加查詢與相關文檔之間的重疊度，從而幫助我們緩解詞匯不匹配問題。

由于許多 Transformer 模型都使用 MLM 進行預訓練，因此有大量模型在預訓練階段學習了 MLM 頭的權重，這些權重可以用于后續的 SPLADE 微調。

SPLADE 的不足之處

SPLADE 是緩解稀疏向量方法常見詞匯不匹配問題的一種優秀方法。然而，我們還需要考慮它的一些局限性。

相較于其他稀疏方法，使用SPLADE 進行檢索速度相對較慢。這主要有三個原因：

SPLADE 查詢和文檔向量中的非零值數量通常多于傳統稀疏向量，而現有的稀疏檢索系統并未針對這一點進行優化。
非零值的分布偏離了傳統稀疏檢索系統所預期的分布，這也會導致速度變慢。
大多數稀疏檢索系統不原生支持 SPLADE 向量，這意味著我們必須執行多步預處理和后處理，例如權重離散化等。

幸運的是，所有這些問題都有解決方案。針對原因 (1)，SPLADE 的作者在模型的后續版本 (SPLADEv2) 中解決了這個問題，該版本最小化了查詢向量中的非零值數量 [2]。

減少查詢向量中的非零值數量是通過兩個步驟實現的。首先，通過對原始池化策略進行最大池化（Max Pooling）修改，提高了 SPLADE 文檔編碼的性能：

wj=maxi∈tlog(1+ReLU(wij))w\j = max{i \in t}log(1 + ReLU(w{ij}))wj?=maxi∈t?log(1+ReLU(wij?))

其次，將詞項擴展僅限于文檔編碼。得益于改進后的文檔編碼性能，即使去掉了查詢擴展，性能依然優于原始的 SPLADE 模型。

原因 (2) 和 (3) 則可以通過使用 Pinecone 向量數據庫解決。(2) 的解決方案在于 Pinecone 的檢索引擎從頭設計時就不依賴數據分布。Pinecone 支持實數值的稀疏向量，這意味著 SPLADE 向量天然就能得到支持。

SPLADE 實現示例

實現 SPLADE 有兩種選擇：直接使用 Hugging Face 的 Transformer 和 PyTorch，或者使用封裝程度更高的官方 SPLADE 庫。我們將演示這兩種方法，先從 Hugging Face 和 PyTorch 的實現開始，以便理解其內部工作原理。

使用Hugging Face 和 PyTorch

首先，安裝所有必需的庫：

!pip install -U transformers torch

然后，初始化 BERT 的分詞器（Tokenizer）和帶有掩碼語言模型（MLM）頭的 BERT 模型。我們加載 naver/splade-cocondenser-ensembledistil 中經過微調的 SPLADE 模型權重。

from transformers import AutoModelForMaskedLM, AutoTokenizermodelid = 'naver/splade-cocondenser-ensembledistil'tokenizer = AutoTokenizer.frompretrained(modelid)
model = AutoModelForMaskedLM.frompretrained(modelid)

接下來，我們可以創建一個輸入文檔文本 text，對其進行詞元化，并通過 model 處理，以生成 MLM 頭的輸出 logits。

tokens = tokenizer(text, returntensors='pt')
output = model(**tokens)
output

MaskedLMOutput(loss=None, logits=tensor([[[ -6.9833,  -8.2131,  -8.1693,  ...,  -8.1552,  -7.8168,  -5.8152],\[-13.6888, -11.7828, -12.5595,  ..., -12.4415, -11.5789, -12.0632],\[ -8.7075,  -8.7019,  -9.0092,  ...,  -9.1933,  -8.4834,  -6.8165],\...,\[ -5.1051,  -7.7245,  -7.0402,  ...,-7.5713,  -6.9855,  -5.0462],\[-23.5020, -18.8779, -17.7931,  ..., -18.2811, -17.2806, -19.4826],\[-21.6329,-17.7142, -16.6525,  ..., -17.1870, -16.1865, -17.9581]]],gradfn=<ViewBackward0>), hiddenstates=None, attentions=None)

output.logits.shape

torch.Size([1, 91, 30522])

我們得到了 91 個概率分布，每個分布的維度是 30522。要將其轉換為 SPLADE 稀疏向量，我們執行以下操作：

importtorchvec = torch.max(torch.log(1 + torch.relu(output.logits)) * tokens.attentionmask.unsqueeze(-1),
dim=1)[0].squeeze()vec.shape

torch.Size([30522])

vec

tensor([0., 0., 0.,  ..., 0.,0., 0.], gradfn=<SqueezeBackward0>)

由于我們的向量是稀疏的，我們可以將其轉換為更緊湊的字典格式，只保留非零值的位置和權重。

cols = vec.nonzero().squeeze().cpu().tolist()
print(f"非零值數量: {len(cols)}")# extract the non-zero values
weights = vec[cols].cpu().tolist()
# use to create a dictionary of token ID to weight
sparsedict = dict(zip(cols, weights))
sparsedict

非零值數量: 174```

{1000: 0.6246446967124939,
1039: 0.45678916573524475,
1052: 0.3088974058628082,
1997: 0.15812619030475616,
1999: 0.07194626331329346,
2003: 0.6496524810791016,
2024: 0.9411943554878235,
…,
29215: 0.3594200909137726,
29278: 2.276832342147827}


這是稀疏向量的最終格式，但還不太易于解釋。我們可以將詞元 ID 鍵轉換為人類可讀的純文本詞元。操作如下：```python
# extract the ID position to text token mappings
idx2token = {idx: token for token, idx in tokenizer.getvocab().items()
}

# map tokenIDs to human-readable tokens
sparsedicttokens = {idx2token[idx]: round(weight, 2) for idx, weight in zip(cols, weights)
}
# sort so we cansee most relevant tokens first
sparsedicttokens = {k: v for k, v in sorted(sparsedicttokens.items(),key=lambda item: item[1],reverse=True)
}
sparsedicttokens

{'pc': 3.02,'lace': 2.95,'programmed': 2.36,'##for': 2.28,'madagascar': 2.26,'death': 1.96,'##d': 1.95,'lattice':1.81,...,'carter': 0.0,'reg': 0.0}

現在我們可以看到稀疏向量中得分最高的詞元，包括一些重要的領域特定詞項，如 programmed（編程的）、cell（細胞）、lattice（晶格）、regulated（被調節的）等等。

使用 Naver Labs SPLADE 庫

另一個更高級的替代方案是直接使用 SPLADE 官方庫。我們可以通過 pip 安裝它：pip install git+https://github.com/naver/splade.git。然后使用以下代碼初始化相同的模型和構建向量的步驟：


sparsemodelid = 'naver/splade-cocondenser-ensembledistil'sparsemodel = Splade(sparsemodelid, agg='max')
sparsemodel.eval()

我們仍然需要使用 Hugging Face 分詞器對輸入文本進行詞元化以獲取 tokens，然后使用以下代碼創建稀疏向量：

    sparseemb = sparsemodel(dkwargs=tokens)['drep'].squeeze()
sparseemb.shape

torch.Size([30522])

這些嵌入可以像之前使用 Hugging Face 和 PyTorch 方法那樣，被處理成一個更小的稀疏向量字典。最終得到的數據是相同的。

向量比較示例

讓我們看看如何實際比較我們的稀疏向量。我們將定義三個短文本。

texts = [\"Programmed cell death (PCD) is the regulated death of cells within an organism",\"How is thescheduled death of cells within a living thing regulated?",\"Photosynthesis is the process of storing light energy as chemical energy in cells"\
]

像之前一樣，我們使用 tokenizer 對所有文本進行編碼，使用 model 生成輸出 logits，并將詞元級別的向量轉換為單一的稀疏向量。

tokens = tokenizer(texts, returntensors='pt',padding=True, truncation=True
)output = model(**tokens)
# aggregate the token-level vecs and transform to sparse
vecs = torch.max(torch.log(1 + torch.relu(output.logits)) * tokens.attentionmask.unsqueeze(-1), dim=1
)[0].squeeze().detach().cpu().numpy()
vecs.shape

(3, 30522)

現在我們得到了三個30522 維的稀疏向量。為了比較它們，我們可以使用余弦相似度（Cosine Similarity）或點積（Dot Product）計算。使用余弦相似度，我們執行以下操作：

importnumpy as npsim = np.zeros((vecs.shape[0], vecs.shape[0]))for i, vec in enumerate(vecs):sim[i,:] = np.dot(vec, vecs.T) / (np.linalg.norm(vec) * np.linalg.norm(vecs, axis=1))

sim

array([[1.        , 0.54609376, 0.20535842],\[0.54609376, 0.99999988, 0.20411882],\[0.2053584 , 0.20411879, 1.]])

最終得到以下相似度矩陣：

使用上面計算出的相似度值生成的相似度熱力圖。句子 1 和句子 2 共享最高的相似度（對角線除外，它們代表每個句子與自身的比較）。

可以看到，兩個內容相似的句子（句 1 和句 2）的相似度得分自然高于與第三個不相關句子（句 3）的相似度得分。

—以上便是對 SPLADE 學習型稀疏嵌入的介紹。通過 SPLADE，我們可以使用高效的稀疏向量嵌入來表示文本，有助于解決詞匯不匹配問題，同時兼顧精確詞語匹配的能力。

我們也探討了 SPLADE 在傳統檢索系統中的局限性，以及 SPLADEv2 和 Pinecone 這類數據分布無關檢索系統如何克服這些問題。

該領域仍有許多工作可以開展。更多的研究和近期成果表明，結合稠密和稀疏表示的混合搜索 (Hybrid Search)索引能帶來更多優勢。通過這些以及其他眾多進展，我們可以看到向量搜索正變得越來越精確和易用。

參考文獻

[1] T. Formal, B. Piwowarski, S. Clinchant, SPLADE: Sparse Lexical and Expansion Model for First Stage Ranking (2021), SIGIR 21

[2] T. Formal, C. Lassance, B. Piwowarski, S. Clinchant, SPLADE v2: Sparse Lexical and Expansion Model for Information Retrieval (2021)