使用Sentence Transformers v5训练和微调稀疏嵌入模型

发布于2025年7月1日

在 GitHub 上更新

106

Sentence Transformers是一个Python库，用于使用和训练嵌入模型和重排序模型，适用于广泛的应用，例如检索增强生成、语义搜索、语义文本相似度、释义挖掘等。最近几个主要版本对训练进行了重大改进

v3.0: （改进的）Sentence Transformer（密集嵌入）模型训练
v4.0: （改进的）Cross Encoder（重排序器）模型训练
v5.0: （新增的）稀疏嵌入模型训练

在这篇博文中，我将向您展示如何使用它来微调稀疏编码器/嵌入模型，并解释为什么您可能想要这样做。这得到了sparse-encoder/example-inference-free-splade-distilbert-base-uncased-nq，一个廉价的模型，在混合搜索或检索和重排序场景中表现出色。

微调稀疏嵌入模型涉及几个组件：模型、数据集、损失函数、训练参数、评估器和训练器类。我将详细介绍每个组件，并附上实际示例，说明它们如何用于微调强大的稀疏嵌入模型。

除了训练您自己的模型外，您还可以从Hugging Face Hub上提供的各种预训练稀疏编码器中进行选择。为了帮助您在这个不断增长的领域中导航，我们策划了一个SPLADE模型集合，重点介绍了一些最相关的模型。
我们在文档的预训练模型中列出了最著名的模型及其基准测试结果。

什么是稀疏嵌入模型？

更广义的“嵌入模型”是指将某种输入（通常是文本）转换为向量表示（嵌入），以捕获输入语义的模型。与原始输入不同，您可以对这些嵌入执行数学运算，从而得到可用于各种任务（例如搜索、聚类或分类）的相似性分数。

对于密集嵌入模型（即常见类型），嵌入通常是低维向量（例如384、768或1024维），其中大多数值是非零的。另一方面，稀疏嵌入模型生成高维向量（例如30,000+维），其中大多数值是零。通常，稀疏嵌入中的每个活跃维度（即具有非零值的维度）对应于模型词汇表中的特定标记，从而实现可解释性。

让我们以最先进的稀疏嵌入模型naver/splade-v3为例

from sentence_transformers import SparseEncoder

# Download from the 🤗 Hub
model = SparseEncoder("naver/splade-v3")

# Run inference
sentences = [
    "The weather is lovely today.",
    "It's so sunny outside!",
    "He drove to the stadium.",
]
embeddings = model.encode(sentences)
print(embeddings.shape)
# (3, 30522)

# Get the similarity scores for the embeddings
similarities = model.similarity(embeddings, embeddings)
print(similarities)
# tensor([[   32.4323,     5.8528,     0.0258],
#         [    5.8528,    26.6649,     0.0302],
#         [    0.0258,     0.0302,    24.0839]])

# Let's decode our embeddings to be able to interpret them
decoded = model.decode(embeddings, top_k=10)
for decoded, sentence in zip(decoded, sentences):
    print(f"Sentence: {sentence}")
    print(f"Decoded: {decoded}")
    print()

Sentence: The weather is lovely today.
Decoded: [('weather', 2.754288673400879), ('today', 2.610959529876709), ('lovely', 2.431990623474121), ('currently', 1.5520408153533936), ('beautiful', 1.5046082735061646), ('cool', 1.4664798974990845), ('pretty', 0.8986214995384216), ('yesterday', 0.8603134155273438), ('nice', 0.8322536945343018), ('summer', 0.7702118158340454)]

Sentence: It's so sunny outside!
Decoded: [('outside', 2.6939032077789307), ('sunny', 2.535827398300171), ('so', 2.0600898265838623), ('out', 1.5397940874099731), ('weather', 1.1198079586029053), ('very', 0.9873268604278564), ('cool', 0.9406591057777405), ('it', 0.9026399254798889), ('summer', 0.684999406337738), ('sun', 0.6520509123802185)]

Sentence: He drove to the stadium.
Decoded: [('stadium', 2.7872302532196045), ('drove', 1.8208855390548706), ('driving', 1.6665740013122559), ('drive', 1.5565159320831299), ('he', 1.4721972942352295), ('stadiums', 1.449463129043579), ('to', 1.0441515445709229), ('car', 0.7002660632133484), ('visit', 0.5118278861045837), ('football', 0.502326250076294)]

在此示例中，嵌入是30,522维向量，其中每个维度对应于模型词汇表中的一个标记。`decode`方法返回了嵌入中值最高的10个标记，使我们能够解释哪些标记对嵌入的贡献最大。

我们甚至可以确定嵌入之间的交集或重叠，这对于确定为什么两个文本被认为是相似或不相似非常有用。

# Let's also compute the intersection/overlap of the first two embeddings
intersection_embedding = model.intersection(embeddings[0], embeddings[1])
decoded_intersection = model.decode(intersection_embedding)
print(decoded_intersection)

Decoded: [('weather', 3.0842742919921875), ('cool', 1.379457712173462), ('summer', 0.5275946259498596), ('comfort', 0.3239051103591919), ('sally', 0.22571465373039246), ('julian', 0.14787325263023376), ('nature', 0.08582140505313873), ('beauty', 0.0588383711874485), ('mood', 0.018594780936837196), ('nathan', 0.000752730411477387)]

查询和文档扩展

神经稀疏嵌入模型的一个关键组成部分是**查询/文档扩展**。与BM25等传统词汇方法仅匹配精确标记不同，神经稀疏模型通常会自动用语义相关的术语扩展原始文本。

传统词汇方法（例如BM25）：仅匹配文本中的精确标记。
神经稀疏模型：自动扩展相关术语。

例如，在上面的代码输出中，句子“The weather is lovely today”被扩展为包含“beautiful”、“cool”、“pretty”和“nice”等原始文本中没有的词语。类似地，“It's so sunny outside!”被扩展为包含“weather”、“summer”和“sun”。

这种扩展使得神经稀疏模型即使没有精确的词元匹配也能匹配语义相关的内容或同义词，处理拼写错误，并克服词汇不匹配问题。这就是为什么像SPLADE这样的神经稀疏模型在保持稀疏表示效率优势的同时，通常优于传统的词汇搜索方法。

然而，扩展也有其风险。例如，对“星期二天气如何？”的查询扩展很可能也会扩展到“星期一”、“星期三”等，这可能不是所希望的。

为什么要使用稀疏嵌入模型？

简而言之，神经稀疏嵌入模型在传统词汇方法（如BM25）和密集嵌入模型（如Sentence Transformers）之间占据了一个重要的利基市场。它们具有以下优点：

混合潜力：可与密集模型高效结合，密集模型在词汇匹配重要的搜索中可能表现不佳。
可解释性：您可以准确地看到哪些标记促成了匹配。
性能：在许多检索任务中与密集模型相比具有竞争力或更优。

在本博客中，我将交替使用“稀疏嵌入模型”和“稀疏编码器模型”。

为什么要微调？

大多数（神经）稀疏嵌入模型都采用上述查询/文档扩展，这样您就可以匹配含义几乎相同的文本，即使它们不共享任何单词。简而言之，模型必须识别同义词，以便这些标记可以放置在最终嵌入中。

大多数开箱即用的稀疏嵌入模型都能轻易识别出“supermarket”、“food”和“market”是包含“grocery”的文本的有用扩展，但例如

“The patient complained of severe cephalalgia.”（患者抱怨严重的头痛。）

扩展为

'##lal', 'severe', '##pha', 'ce', '##gia', 'patient', 'complaint', 'patients', 'complained', 'warning', 'suffered', 'had', 'disease', 'complain', 'diagnosis', 'syndrome', 'mild', 'pain', 'hospital', 'injury'

然而，我们希望它扩展为“headache”（头痛），这是“cephalalgia”的常用词。这个例子扩展到许多领域，例如无法识别“Java”是一种编程语言，“Audi”制造汽车，或者“NVIDIA”是一家制造显卡的的公司。

通过微调，模型可以学习专门关注对您重要的领域和/或语言。

训练组件

训练Sentence Transformer模型涉及以下组件

模型：要训练或微调的模型，可以是预训练的稀疏编码器模型或基础模型。
数据集：用于训练和评估的数据。
损失函数：量化模型性能并指导优化过程的函数。
训练参数（可选）：影响训练性能和跟踪/调试的参数。
评估器（可选）：用于在训练前、训练中或训练后评估模型的工具。
训练器：将模型、数据集、损失函数和其他组件整合在一起进行训练。

现在，让我们更详细地探讨这些组件。

模型

稀疏编码器模型由一系列模块、稀疏编码器特定模块或自定义模块组成，具有很大的灵活性。如果您想进一步微调稀疏编码器模型（例如，它有一个modules.json文件），那么您无需担心使用哪些模块

from sentence_transformers import SparseEncoder

model = SparseEncoder("naver/splade-cocondenser-ensembledistil")

但如果您想从其他检查点或从头开始训练，那么这些是最常见的架构

Splade

Splade模型使用MLMTransformer模块，后接SpladePooling模块。前者加载预训练的掩码语言建模Transformer模型（例如BERT、RoBERTa、DistilBERT、ModernBERT等），后者将MLMHead的输出池化以生成词汇表大小的单一稀疏嵌入。

from sentence_transformers import models, SparseEncoder
from sentence_transformers.sparse_encoder.models import MLMTransformer, SpladePooling

# Initialize MLM Transformer (use a fill-mask model)
mlm_transformer = MLMTransformer("google-bert/bert-base-uncased")

# Initialize SpladePooling module
splade_pooling = SpladePooling(pooling_strategy="max")

# Create the Splade model
model = SparseEncoder(modules=[mlm_transformer, splade_pooling])

如果您向SparseEncoder提供一个fill-mask模型架构，此架构是默认的，因此使用快捷方式会更容易

from sentence_transformers import SparseEncoder

model = SparseEncoder("google-bert/bert-base-uncased")
# SparseEncoder(
#   (0): MLMTransformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'BertForMaskedLM'})
#   (1): SpladePooling({'pooling_strategy': 'max', 'activation_function': 'relu', 'word_embedding_dimension': None})
# )

免推理Splade

免推理Splade 使用Router模块，其中查询和文档使用不同的模块。通常，对于这种类型的架构，文档部分是传统的Splade架构（一个MLMTransformer后接一个SpladePooling模块），而查询部分是SparseStaticEmbedding模块，它只返回查询中每个标记的预计算分数。

from sentence_transformers import SparseEncoder
from sentence_transformers.models import Router
from sentence_transformers.sparse_encoder.models import SparseStaticEmbedding, MLMTransformer, SpladePooling

# Initialize MLM Transformer for document encoding
doc_encoder = MLMTransformer("google-bert/bert-base-uncased")

# Create a router model with different paths for queries and documents
router = Router.for_query_document(
    query_modules=[SparseStaticEmbedding(tokenizer=doc_encoder.tokenizer, frozen=False)],
    # Document path: full MLM transformer + pooling
    document_modules=[doc_encoder, SpladePooling("max")],
)

# Create the inference-free model
model = SparseEncoder(modules=[router], similarity_fn_name="dot")
# SparseEncoder(
#   (0): Router(
#     (query_0_SparseStaticEmbedding): SparseStaticEmbedding ({'frozen': False}, dim:30522, tokenizer: BertTokenizerFast)
#     (document_0_MLMTransformer): MLMTransformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'BertForMaskedLM'})
#     (document_1_SpladePooling): SpladePooling({'pooling_strategy': 'max', 'activation_function': 'relu', 'word_embedding_dimension': None})
#   )
# )

这种架构允许使用轻量级的SparseStaticEmbedding方法进行快速查询时处理，该方法可以训练并视为线性权重，而文档则通过完整的MLM transformer和SpladePooling进行处理。

免推理Splade特别适用于查询延迟至关重要的搜索应用，因为它将计算复杂度转移到可以离线完成的文档索引阶段。

训练带有`Router`模块的模型时，您必须在`SparseEncoderTrainingArguments`中使用`router_mapping`参数将训练数据集列映射到正确的路由（“query”或“document”）。例如，如果您的数据集有`["question", "answer"]`列，那么您可以使用以下映射
args = SparseEncoderTrainingArguments(
    ...,
    router_mapping={
        "question": "query",
        "answer": "document",
    }
)
此外，建议对SparseStaticEmbedding模块使用比模型其余部分更高的学习率。为此，您应该在`SparseEncoderTrainingArguments`中使用`learning_rate_mapping`参数将参数模式映射到其学习率。例如，如果您希望SparseStaticEmbedding模块的学习率为`1e-3`，模型其余部分为`2e-5`，您可以这样做
args = SparseEncoderTrainingArguments(
    ...,
    learning_rate=2e-5,
    learning_rate_mapping={
        r"SparseStaticEmbedding\.*": 1e-3,
    }
)

对比稀疏表示 (CSR)

对比稀疏表示（CSR）模型，在Beyond Matryoshka: Revisiting Sparse Coding for Adaptive Representation中引入，在密集Sentence Transformer模型之上应用了SparseAutoEncoder模块，后者通常由Transformer后接Pooling模块组成。您可以像这样从头开始初始化一个：

from sentence_transformers import models, SparseEncoder
from sentence_transformers.sparse_encoder.models import SparseAutoEncoder

# Initialize transformer (can be any dense encoder model)
transformer = models.Transformer("google-bert/bert-base-uncased")

# Initialize pooling
pooling = models.Pooling(transformer.get_word_embedding_dimension(), pooling_mode="mean")

# Initialize SparseAutoEncoder module
sparse_auto_encoder = SparseAutoEncoder(
    input_dim=transformer.get_word_embedding_dimension(),
    hidden_dim=4 * transformer.get_word_embedding_dimension(),
    k=256,  # Number of top values to keep
    k_aux=512,  # Number of top values for auxiliary loss
)
# Create the CSR model
model = SparseEncoder(modules=[transformer, pooling, sparse_auto_encoder])

或者，如果您的基础模型是 1) 一个密集 Sentence Transformer 模型，或者 2) 一个非 MLM Transformer 模型（默认情况下这些模型会作为 Splade 模型加载），那么这个快捷方式将自动为您初始化 CSR 模型

from sentence_transformers import SparseEncoder

model = SparseEncoder("mixedbread-ai/mxbai-embed-large-v1")
# SparseEncoder(
#   (0): Transformer({'max_seq_length': 512, 'do_lower_case': False, 'architecture': 'BertModel'})
#   (1): Pooling({'word_embedding_dimension': 1024, 'pooling_mode_cls_token': True, 'pooling_mode_mean_tokens': False, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False, 'pooling_mode_weightedmean_tokens': False, 'pooling_mode_lasttoken': False, 'include_prompt': True})
#   (2): SparseAutoEncoder({'input_dim': 1024, 'hidden_dim': 4096, 'k': 256, 'k_aux': 512, 'normalize': False, 'dead_threshold': 30})
# )

与（免推理）Splade模型不同，CSR模型生成的稀疏嵌入与基础模型的词汇表大小不同。这意味着您无法像Splade模型那样直接解释嵌入中激活了哪些单词，因为Splade模型中每个维度都对应于词汇表中的特定标记。

此外，CSR模型在那些使用高维表示（例如1024-4096维）的密集编码器模型上最有效。

架构选择指南

如果您不确定使用哪种架构，这里有一个快速指南

您想稀疏化现有的密集嵌入模型吗？如果想，请使用CSR。
您希望查询推理即时完成，但代价是性能略有下降吗？如果是，请使用免推理SPLADE。
否则，请使用SPLADE。

数据集

SparseEncoderTrainer使用datasets.Dataset或datasets.DatasetDict实例进行训练和评估。您可以从Hugging Face数据集中心加载数据，或使用各种格式的本地数据，如CSV、JSON、Parquet、Arrow或SQL。

注意：许多与Sentence Transformers开箱即用的公共数据集都已在Hugging Face Hub上标记为sentence-transformers，因此您可以在https://huggingface.co/datasets?other=sentence-transformers上轻松找到它们。考虑浏览这些数据集，以找到可能对您的任务、领域或语言有用的现成数据集。

Hugging Face Hub上的数据

您可以使用load_dataset函数从Hugging Face Hub中的数据集加载数据。

from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/natural-questions", split="train")

print(train_dataset)
"""
Dataset({
    features: ['query', 'answer'],
    num_rows: 100231
})
"""

一些数据集，例如nthakur/swim-ir-monolingual，具有多个不同数据格式的子集。您需要指定子集名称以及数据集名称，例如dataset = load_dataset("nthakur/swim-ir-monolingual", "de", split="train")。

本地数据（CSV、JSON、Parquet、Arrow、SQL）

您也可以使用load_dataset加载某些文件格式的本地数据

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")
# or
dataset = load_dataset("json", data_files="my_file.json")

需要预处理的本地数据

如果您的本地数据需要预处理，您可以使用datasets.Dataset.from_dict。这允许您使用字典列表初始化数据集。

from datasets import Dataset

queries = []
documents = []
# Open a file, perform preprocessing, filtering, cleaning, etc.
# and append to the lists

dataset = Dataset.from_dict({
    "query": queries,
    "document": documents,
})

字典中的每个键都将成为结果数据集中的一列。

数据集格式

确保数据集格式与您选择的损失函数匹配至关重要。这涉及检查两件事

如果您的损失函数需要一个标签（如损失概述表中所示），则您的数据集必须有一个名为"label"或"score"的列。
除"label"或"score"之外的所有列都被视为输入（如损失概述表中所示）。这些列的数量必须与您所选损失函数的有效输入数量相匹配。列的名称不重要，只有它们的顺序重要。

例如，如果您的损失函数接受(anchor, positive, negative)三元组，那么您的第一个、第二个和第三个数据集列分别对应anchor、positive和negative。这意味着您的第一列和第二列必须包含应该紧密嵌入的文本，而您的第一列和第三列必须包含应该相距遥远的文本。这就是为什么根据您的损失函数，您的数据集列顺序很重要。

考虑一个包含列["text1", "text2", "label"]的数据集，其中"label"列包含浮点相似度分数。此数据集可与SparseCoSENTLoss、SparseAnglELoss和SparseCosineSimilarityLoss一起使用，因为

数据集有一个“label”列，这是这些损失函数所必需的。
数据集有2个非标签列，与这些损失函数所需的输入数量相匹配。

如果您的数据集中的列顺序不正确，请使用Dataset.select_columns重新排序。此外，使用Dataset.remove_columns删除任何多余的列（例如，sample_id、metadata、source、type），否则它们将被视为输入。

损失函数

损失函数衡量模型在给定批量数据上的表现，并指导优化过程。损失函数的选择取决于您可用的数据和目标任务。有关选项的全面列表，请参阅损失概述。

要训练SparseEncoder，您需要一个SpladeLoss或CSRLoss，具体取决于架构。这些是包装损失，它们在主损失函数之上添加稀疏性正则化，主损失函数必须作为参数提供。唯一可以独立使用的损失是SparseMSELoss，因为它执行嵌入级别蒸馏，通过直接复制教师的稀疏嵌入来确保稀疏性。

大多数损失函数可以通过仅用您正在训练的SparseEncoder以及一些可选参数进行初始化，例如

from datasets import load_dataset
from sentence_transformers import SparseEncoder
from sentence_transformers.sparse_encoder.losses import SpladeLoss, SparseMultipleNegativesRankingLoss

# Load a model to train/finetune
model = SparseEncoder("distilbert/distilbert-base-uncased")

# Initialize the SpladeLoss with a SparseMultipleNegativesRankingLoss
# This loss requires pairs of related texts or triplets
loss = SpladeLoss(
    model=model,
    loss=SparseMultipleNegativesRankingLoss(model=model),
    query_regularizer_weight=5e-5,  # Weight for query loss
    document_regularizer_weight=3e-5,
) 

# Load an example training dataset that works with our loss function:
train_dataset = load_dataset("sentence-transformers/natural-questions", split="train")
print(train_dataset)
"""
Dataset({
    features: ['query', 'answer'],
    num_rows: 100231
})
"""

文档

训练参数

SparseEncoderTrainingArguments类允许您指定影响训练性能和跟踪/调试的参数。虽然是可选的，但尝试这些参数有助于提高训练效率并提供对训练过程的洞察。

在Sentence Transformers文档中，我概述了一些最有用的训练参数。我建议您阅读训练概述 > 训练参数。

以下是初始化SparseEncoderTrainingArguments的示例

from sentence_transformers import SparseEncoderTrainingArguments

args = SparseEncoderTrainingArguments(
    # Required parameter:
    output_dir="models/splade-distilbert-base-uncased-nq",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    fp16=True,  # Set to False if your GPU can't handle FP16
    bf16=False,  # Set to True if your GPU supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # Losses using "in-batch negatives" benefit from no duplicates
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    logging_steps=100,
    run_name="splade-distilbert-base-uncased-nq",  # Used in W&B if `wandb` is installed
)

请注意，`eval_strategy`是在`transformers`版本`4.41.0`中引入的。之前的版本应使用`evaluation_strategy`。

评估器

您可以向SparseEncoderTrainer提供eval_dataset以在训练期间获取评估损失，但在训练期间获取更具体的指标也可能很有用。为此，您可以使用评估器在训练前、训练中或训练后使用有用的指标评估模型的性能。您可以同时使用eval_dataset和评估器，或其中一个，或都不使用。它们根据eval_strategy和eval_steps训练参数进行评估。

以下是Sentence Transformers为稀疏编码器模型实现的评估器：

评估器	所需数据
`SparseBinaryClassificationEvaluator`	带有类标签的对。
`SparseEmbeddingSimilarityEvaluator`	带有相似度分数的对。
`SparseInformationRetrievalEvaluator`	查询（qid => 问题）、语料库（cid => 文档）和相关文档（qid => set[cid]）。
`SparseNanoBEIREvaluator`	无需数据。
`SparseMSEEvaluator`	需要用教师模型嵌入的源句子和用学生模型嵌入的目标句子。可以是相同的文本。
`SparseRerankingEvaluator`	列表形式的`{'query': '...', 'positive': [...], 'negative': [...]}`字典。
`SparseTranslationEvaluator`	两种不同语言的句子对。
`SparseTripletEvaluator`	（锚点、正例、负例）对。

此外，应使用SequentialEvaluator将多个评估器组合成一个评估器，该评估器可以传递给SparseEncoderTrainer。

有时，您没有所需的评估数据来自行准备这些评估器，但您仍然希望跟踪模型在某些常见基准上的表现。在这种情况下，您可以将这些评估器与Hugging Face的数据一起使用。

SparseNanoBEIREvaluator

文档

sentence_transformers.sparse_encoder.evaluation.SparseNanoBEIREvaluator

from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator

# Initialize the evaluator. Unlike most other evaluators, this one loads the relevant datasets
# directly from Hugging Face, so there's no mandatory arguments
dev_evaluator = SparseNanoBEIREvaluator()
# You can run evaluation like so:
# results = dev_evaluator(model)

使用STSb的SparseEmbeddingSimilarityEvaluator

文档

from datasets import load_dataset
from sentence_transformers.evaluation import SimilarityFunction
from sentence_transformers.sparse_encoder.evaluation import SparseEmbeddingSimilarityEvaluator

# Load the STSB dataset (https://huggingface.co/datasets/sentence-transformers/stsb)
eval_dataset = load_dataset("sentence-transformers/stsb", split="validation")

# Initialize the evaluator
dev_evaluator = SparseEmbeddingSimilarityEvaluator(
    sentences1=eval_dataset["sentence1"],
    sentences2=eval_dataset["sentence2"],
    scores=eval_dataset["score"],
    main_similarity=SimilarityFunction.COSINE,
    name="sts-dev",
)
# You can run evaluation like so:
# results = dev_evaluator(model)

使用AllNLI的SparseTripletEvaluator

文档

from datasets import load_dataset
from sentence_transformers.evaluation import SimilarityFunction
from sentence_transformers.sparse_encoder.evaluation import SparseTripletEvaluator

# Load triplets from the AllNLI dataset (https://huggingface.co/datasets/sentence-transformers/all-nli)
max_samples = 1000
eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split=f"dev[:{max_samples}]")

# Initialize the evaluator
dev_evaluator = SparseTripletEvaluator(
    anchors=eval_dataset["anchor"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    main_distance_function=SimilarityFunction.DOT,
    name="all-nli-dev",
)
# You can run evaluation like so:
# results = dev_evaluator(model)

如果在训练期间使用较小的eval_steps频繁评估，请考虑使用微小的eval_dataset以最大程度地减少评估开销。如果您担心评估集大小，90-1-9的训练-评估-测试划分可以提供一个平衡，为最终评估保留一个合理大小的测试集。训练结束后，您可以使用trainer.evaluate(test_dataset)评估模型的测试损失，或使用test_evaluator(model)初始化测试评估器以获取详细的测试指标。

如果您在训练后但在保存模型之前进行评估，则自动生成的模型卡仍将包含测试结果。

使用分布式训练时，评估器仅在第一个设备上运行，与训练和评估数据集不同，后者在所有设备上共享。

训练器

SparseEncoderTrainer是所有先前组件的集合。我们只需指定训练器、模型、训练参数（可选）、训练数据集、评估数据集（可选）、损失函数、评估器（可选），然后就可以开始训练了。让我们看看一个将所有这些组件结合在一起的脚本

import logging

from datasets import load_dataset

from sentence_transformers import (
    SparseEncoder,
    SparseEncoderModelCardData,
    SparseEncoderTrainer,
    SparseEncoderTrainingArguments,
)
from sentence_transformers.models import Router
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
from sentence_transformers.sparse_encoder.losses import SparseMultipleNegativesRankingLoss, SpladeLoss
from sentence_transformers.sparse_encoder.models import SparseStaticEmbedding, MLMTransformer, SpladePooling
from sentence_transformers.training_args import BatchSamplers

logging.basicConfig(format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO)

# 1. Load a model to finetune with 2. (Optional) model card data
mlm_transformer = MLMTransformer("distilbert/distilbert-base-uncased", tokenizer_args={"model_max_length": 512})
splade_pooling = SpladePooling(
    pooling_strategy="max", word_embedding_dimension=mlm_transformer.get_sentence_embedding_dimension()
)
router = Router.for_query_document(
    query_modules=[SparseStaticEmbedding(tokenizer=mlm_transformer.tokenizer, frozen=False)],
    document_modules=[mlm_transformer, splade_pooling],
)

model = SparseEncoder(
    modules=[router],
    model_card_data=SparseEncoderModelCardData(
        language="en",
        license="apache-2.0",
        model_name="Inference-free SPLADE distilbert-base-uncased trained on Natural-Questions tuples",
    ),
)

# 3. Load a dataset to finetune on
full_dataset = load_dataset("sentence-transformers/natural-questions", split="train").select(range(100_000))
dataset_dict = full_dataset.train_test_split(test_size=1_000, seed=12)
train_dataset = dataset_dict["train"]
eval_dataset = dataset_dict["test"]
print(train_dataset)
print(train_dataset[0])

# 4. Define a loss function
loss = SpladeLoss(
    model=model,
    loss=SparseMultipleNegativesRankingLoss(model=model),
    query_regularizer_weight=0,
    document_regularizer_weight=3e-3,
)

# 5. (Optional) Specify training arguments
run_name = "inference-free-splade-distilbert-base-uncased-nq"
args = SparseEncoderTrainingArguments(
    # Required parameter:
    output_dir=f"models/{run_name}",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    learning_rate=2e-5,
    learning_rate_mapping={r"SparseStaticEmbedding\.weight": 1e-3},  # Set a higher learning rate for the SparseStaticEmbedding module
    warmup_ratio=0.1,
    fp16=True,  # Set to False if you get an error that your GPU can't run on FP16
    bf16=False,  # Set to True if you have a GPU that supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
    router_mapping={"query": "query", "answer": "document"},  # Map the column names to the routes
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=1000,
    save_strategy="steps",
    save_steps=1000,
    save_total_limit=2,
    logging_steps=200,
    run_name=run_name,  # Will be used in W&B if `wandb` is installed
)

# 6. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = SparseNanoBEIREvaluator(dataset_names=["msmarco", "nfcorpus", "nq"], batch_size=16)

# 7. Create a trainer & train
trainer = SparseEncoderTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=loss,
    evaluator=dev_evaluator,
)
trainer.train()

# 8. Evaluate the model performance again after training
dev_evaluator(model)

# 9. Save the trained model
model.save_pretrained(f"models/{run_name}/final")

# 10. (Optional) Push it to the Hugging Face Hub
model.push_to_hub(run_name)

在这个例子中，我正在从distilbert/distilbert-base-uncased微调，这是一个尚未成为稀疏编码器模型的基础模型。这比微调现有稀疏编码器模型（如naver/splade-cocondenser-ensembledistil）需要更多的训练数据。

运行此脚本后，sparse-encoder/example-inference-free-splade-distilbert-base-uncased-nq模型已为我上传。该模型在NanoMSMARCO上得分0.5241 NDCG@10，在NanoNFCorpus上得分0.3299 NDCG@10，在NanoNQ上得分0.5357 NDCG@10，对于一个基于distilbert的免推理模型，仅在自然语言问答数据集的10万对数据上进行训练，这是一个不错的结果。

该模型在文档的稀疏嵌入中平均使用184个活跃维度，而查询的活跃维度为7.7个（即查询中的平均标记数）。这分别对应于99.39%和99.97%的稀疏度。

所有这些信息都存储在自动生成的模型卡中，包括基础模型、语言、许可证、评估结果、训练和评估数据集信息、超参数、训练日志等等。无需任何努力，您上传的模型将包含您的潜在用户确定模型是否适合他们所需的所有信息。

回调

Sentence Transformers训练器支持各种transformers.TrainerCallback子类，包括

如果安装了wandb，则使用WandbCallback将训练指标记录到W&B。
如果可以访问tensorboard，则使用TensorBoardCallback将训练指标记录到TensorBoard。
如果安装了codecarbon，则使用CodeCarbonCallback跟踪训练期间的碳排放。

只要安装了所需的依赖项，这些功能就会自动使用，您无需进行任何指定。

有关这些回调以及如何创建自己的回调的更多信息，请参阅Transformers回调文档。

多数据集训练

表现出色的模型通常通过同时使用多个数据集进行训练。SparseEncoderTrainer通过允许您使用多个数据集进行训练而无需将它们转换为相同格式来简化此过程。您甚至可以对每个数据集应用不同的损失函数。以下是多数据集训练的步骤

使用datasets.Dataset实例（或datasets.DatasetDict）的字典作为train_dataset和eval_dataset。
（可选）如果想对不同数据集使用不同的损失函数，请使用将数据集名称映射到损失函数的字典。

每个训练/评估批次将仅包含来自一个数据集的样本。从多个数据集中采样批次的顺序由MultiDatasetBatchSamplers枚举决定，该枚举可以通过multi_dataset_batch_sampler传递给SparseEncoderTrainingArguments。有效选项包括

MultiDatasetBatchSamplers.ROUND_ROBIN: 以循环方式从每个数据集中采样，直到其中一个耗尽。此策略可能不会使用每个数据集中的所有样本，但它确保从每个数据集中进行相同的采样。
MultiDatasetBatchSamplers.PROPORTIONAL（默认）：根据每个数据集的大小按比例从每个数据集中采样。此策略确保使用每个数据集中的所有样本，并且更频繁地从更大的数据集中采样。

评估

让我们使用NanoMSMARCO数据集评估我们新训练的免推理SPLADE模型，看看它与密集检索方法相比如何。我们还将探讨结合稀疏和密集向量的混合检索方法，以及重排序以进一步提高搜索质量。

运行我们略作修改的hybrid_search.py脚本后，我们获得了NanoMSMARCO数据集的以下结果，使用了这些模型：

稀疏：sparse-encoder/example-inference-free-splade-distilbert-base-uncased-nq（我们刚刚训练的模型）
密集：sentence-transformers/all-MiniLM-L6-v2
重排序器：cross-encoder/ms-marco-MiniLM-L6-v2

稀疏	密集	重排序器	NDCG@10	MRR@10	MAP
x			52.41	43.06	44.20
	x		55.40	47.96	49.08
x	x		62.22	53.02	53.44
x		x	66.31	59.45	60.36
	x	x	66.28	59.43	60.34
x	x	x	66.28	59.43	60.34

稀疏和密集排名可以使用倒数排名融合（RRF）进行组合，这是一种结合多个排名结果的简单方法。如果应用了重排序器，它将对先前检索步骤的结果进行重排序。

结果表明，对于此数据集，结合密集和稀疏排名表现非常出色，与密集和稀疏基线相比，分别提高了12.3%和18.7%。简而言之，结合稀疏和密集检索方法是提高搜索性能的非常有效的方法。

此外，对任何排名应用重排序器都将性能提高到大约66.3 NDCG@10，这表明无论是稀疏、密集还是混合（密集+稀疏）都在其前100名中找到了相关文档，然后重排序器将其排到前10名。因此，将密集 -> 重排序器管道替换为稀疏 -> 重排序器管道可能会同时提高延迟和成本

稀疏嵌入存储成本更低，例如我们的模型对MS MARCO文档只使用约180个活跃维度，而不是密集模型常用的1024个维度。
一些稀疏编码器支持免推理查询处理，实现近乎即时的第一阶段检索，类似于BM25等词汇解决方案。

训练技巧

稀疏编码器模型在训练时有一些您应该注意的特点：

稀疏编码器模型不应仅通过评估分数进行评估，还应通过嵌入的稀疏性进行评估。毕竟，低稀疏性意味着模型嵌入存储成本高昂且检索速度慢。
更强的稀疏编码器模型几乎完全通过从更强的教师模型（例如CrossEncoder模型）进行蒸馏来训练，而不是直接从文本对或三元组进行训练。例如，参见SPLADE-v3论文，其中使用SparseDistillKLDivLoss和SparseMarginMSELoss进行蒸馏。我们在此博客中没有详细介绍，因为它需要更多的数据准备，但蒸馏设置应该被认真考虑。

向量数据库集成

训练稀疏嵌入模型后，下一个关键步骤是将其有效地部署到生产环境中。向量数据库提供了存储、索引和大规模检索稀疏嵌入所需的基础设施。流行的选项包括Qdrant、OpenSearch、Elasticsearch和Seismic等。

有关上述向量数据库的全面示例，请参阅向量数据库语义搜索文档或下方的Qdrant示例。

Qdrant集成示例

Qdrant为稀疏向量提供了出色的支持，具有高效的存储和快速检索功能。以下是一个全面的实现示例

先决条件：

Qdrant本地运行（或可访问），更多详细信息请参阅Qdrant快速入门。
安装了Python Qdrant客户端
```
pip install qdrant-client
```

此示例演示了如何设置 Qdrant 进行稀疏向量搜索，通过展示如何使用稀疏编码器高效编码和索引文档，使用稀疏向量构建搜索查询，并提供交互式查询界面。请参阅下文

import time

from datasets import load_dataset
from sentence_transformers import SparseEncoder
from sentence_transformers.sparse_encoder.search_engines import semantic_search_qdrant

# 1. Load the natural-questions dataset with 100K answers
dataset = load_dataset("sentence-transformers/natural-questions", split="train")
num_docs = 10_000
corpus = dataset["answer"][:num_docs]

# 2. Come up with some queries
queries = dataset["query"][:2]

# 3. Load the model
sparse_model = SparseEncoder("naver/splade-cocondenser-ensembledistil")

# 4. Encode the corpus
corpus_embeddings = sparse_model.encode_document(
    corpus, convert_to_sparse_tensor=True, batch_size=16, show_progress_bar=True
)

# Initially, we don't have a qdrant index yet
corpus_index = None
while True:
    # 5. Encode the queries using the full precision
    start_time = time.time()
    query_embeddings = sparse_model.encode_query(queries, convert_to_sparse_tensor=True)
    print(f"Encoding time: {time.time() - start_time:.6f} seconds")

    # 6. Perform semantic search using qdrant
    results, search_time, corpus_index = semantic_search_qdrant(
        query_embeddings,
        corpus_index=corpus_index,
        corpus_embeddings=corpus_embeddings if corpus_index is None else None,
        top_k=5,
        output_index=True,
    )

    # 7. Output the results
    print(f"Search time: {search_time:.6f} seconds")
    for query, result in zip(queries, results):
        print(f"Query: {query}")
        for entry in result:
            print(f"(Score: {entry['score']:.4f}) {corpus[entry['corpus_id']]}, corpus_id: {entry['corpus_id']}")
        print("")

    # 8. Prompt for more queries
    queries = [input("Please enter a question: ")]

附加资源

训练示例

以下页面包含带有解释和代码链接的训练示例。我们建议您浏览这些示例，熟悉训练循环

模型蒸馏 - 使模型更小、更快、更轻的示例。
MS MARCO - 在 MS MARCO 信息检索数据集上进行训练的示例训练脚本。
检索器 - 在通用信息检索数据集上进行训练的示例训练脚本。
自然语言推理 - 自然语言推理 (NLI) 数据对于预训练和微调模型以创建有意义的稀疏嵌入非常有帮助。
Quora 重复问题 - Quora 重复问题是一个大型语料库，包含来自 Quora 社区的重复问题。该文件夹包含如何训练模型用于重复问题挖掘和语义搜索的示例。
STS - 训练模型最基本的方法是使用语义文本相似度 (STS) 数据。在这里，我们使用句子对和表示语义相似度的分数。

文档

此外，以下页面可能有助于您了解更多关于 Sentence Transformers 的信息

安装
快速入门
用法
预训练模型
训练概述（这篇博客文章是训练概述文档的精髓）
数据集概述
损失概述
API参考

最后，这里有一些您可能感兴趣的高级页面

更多博客文章

使用 Sentence Transformers v4 训练和微调 Reranker 模型

作者： 2025年3月26日 • 155

使用 Sentence Transformers 训练静态嵌入模型快 400 倍

作者： 2025年1月15日 • 202

社区

pulkitmehtawork

7月4日

很棒的工作。最好的部分是可解释性和速度。 @tomaarsen - 我计划使用以下设置微调一个模型进行文本到代码的检索。请指导一下，这个设置是否适合开始，或者我还可以调整什么以做得更好。目标是在文本到代码方面做得不错，并在 (https://github.com/CoIR-team/coir) 上进行评估。
训练数据集 - claudios/code_search_net .. 过滤 Python 代码 .. 查询是代码的文档字符串，段落是代码 ... 损失 - SparseMultipleNegativesRankingLoss.. 无法想到一个像样的开发评估 .. 我应该使用 SparseTripletEvaluator 吗 .. 此外，只需要查询和正向段落是否可以，因为我相信负向选项将是该批次中的所有其他数据，或者我们必须明确准备数据（我的负向数据）.. 请指导 ..

tomaarsen

文章作者 7月4日

你好！
我认为这是一个不错的设置。我个人会推荐通过 MTEB 使用 COIR (https://github.com/CoIR-team/coir#coconut-mteb-usage)，其余的听起来都很可靠。挖掘困难负例显然是可能的，但我个人建议先从一个简单的设置作为基线开始，然后再尝试使其过于复杂/耗时。

Tom Aarsen

rasyosef

13 天前

这是一个 Python 包，您可以使用它来索引、查询和使用 Sentence-Transformers 中的 SPLADE 模型对文档进行排名。

splade-index: https://github.com/rasyosef/splade-index

SPLADE-Index⚡

SPLADE-Index 是一个超快速的 SPLADE 稀疏检索模型索引，它使用纯 Python 实现，并由 Scipy 稀疏矩阵提供支持。它构建在 BM25s 库之上。

安装

您可以使用 pip 安装 splade-index

pip install splade-index

快速入门

这是一个如何使用 splade-index 的简单示例

from sentence_transformers import SparseEncoder
from splade_index import SPLADE

# Download a SPLADE model from the 🤗 Hub
model = SparseEncoder("rasyosef/splade-tiny")

# Create your corpus here
corpus = [
    "a cat is a feline and likes to purr",
    "a dog is the human's best friend and loves to play",
    "a bird is a beautiful animal that can fly",
    "a fish is a creature that lives in water and swims",
]

# Create the SPLADE retriever and index the corpus
retriever = SPLADE()
retriever.index(model=model, documents=corpus)

# Query the corpus
queries = ["does the fish purr like a cat?"]

# Get top-k results as a tuple of (doc ids, documents, scores). All three are arrays of shape (n_queries, k).
results = retriever.retrieve(queries, k=2)
doc_ids, result_docs, scores = results.doc_ids, results.documents, results.scores

for i in range(doc_ids.shape[1]):
    doc_id, doc, score = doc_ids[0, i], result_docs[0, i], scores[0, i]
    print(f"Rank {i+1} (score: {score:.2f}) (doc_id: {doc_id}): {doc}")

# You can save the index to a directory
retriever.save("animal_index_splade")

# ...and load it when you need it
import splade_index

reloaded_retriever = splade_index.SPLADE.load("animal_index_splade", model=model)

通过拖放到文本输入框、粘贴或点击此处上传图片、音频和视频。

点击或粘贴此处以上传图片

· 注册或登录发表评论

106

使用Sentence Transformers v5训练和微调稀疏嵌入模型

目录

什么是稀疏嵌入模型？

查询和文档扩展

为什么要使用稀疏嵌入模型？

为什么要微调？

训练组件

模型

Splade

免推理Splade

对比稀疏表示 (CSR)

架构选择指南

数据集

Hugging Face Hub上的数据

本地数据（CSV、JSON、Parquet、Arrow、SQL）

需要预处理的本地数据

数据集格式

损失函数

训练参数

评估器

SparseNanoBEIREvaluator

使用STSb的SparseEmbeddingSimilarityEvaluator

使用AllNLI的SparseTripletEvaluator

训练器

回调

多数据集训练

评估

训练技巧

向量数据库集成

Qdrant集成示例

先决条件：

附加资源

训练示例

文档

使用 Sentence Transformers v4 训练和微调 Reranker 模型

使用 Sentence Transformers 训练静态嵌入模型快 400 倍

社区

SPLADE-Index⚡

安装

快速入门