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∈tlog(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)