Sentence Transformers 发布稀疏嵌入模型训练与微调教程

Sentence Transformers 是一个 Python 库，用于使用和训练稠密嵌入、重排序器（交叉编码器）和稀疏嵌入模型，适用于检索增强生成、语义搜索、语义文本相似度、复述挖掘等广泛的应用场景。在这篇博客文章中，我将向你展示如何使用它来微调稀疏编码器/嵌入模型，并解释为什么你可能想要这样做。这将产生 sparse-encoder/example-inference-free-splade-distilbert-base-uncased-nq，这是一个成本低廉的模型，在混合搜索或检索与重排序场景中表现尤为出色。

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

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

目录

• What are Sparse Embedding models?
  • 查询和文档扩展
  • 为什么要使用稀疏嵌入模型？
• 为什么要微调？
• 训练组件
• Model
  • Splade
  • Inference-free Splade
  • 对比稀疏表示（CSR）
  • 架构选择指南
• Dataset
  • Hugging Face Hub 上的数据
  • 本地数据（CSV、JSON、Parquet、Arrow、SQL）
  • 需要预处理的本地数据
  • 数据集格式
• 损失函数
• 训练参数
• Evaluator
  • SparseNanoBEIREvaluator
  • 使用 STSb 的 SparseEmbeddingSimilarityEvaluator
  • 使用 AllNLI 的 SparseTripletEvaluator
• Trainer
  • 回调
  • 多数据集训练
• 评估
• 训练技巧
• 向量数据库集成
• Additional Resources
  • 训练示例
  • 文档

什么是稀疏嵌入模型？

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

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

让我们以 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 维的向量，其中每一维对应模型词表中的一个 token。decode 方法返回了嵌入中数值最高的前 10 个 token，让我们得以解读哪些 token 对嵌入贡献最大。

我们甚至可以确定嵌入之间的交集或重叠部分，这对于判断两段文本为何被认为相似或不相似非常有用：

# 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 这类仅匹配精确 token 的传统词汇方法不同，神经稀疏模型通常会使用语义相关的词项自动扩展原始文本：

• 传统的词汇方法（例如 BM25）：仅匹配文本中的精确 token
• 神经稀疏模型：使用相关词项自动扩展

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

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

然而，扩展也有其风险。例如，对“What is the weather on Tuesday?”进行查询扩展时，很可能也会扩展到“monday”、“wednesday”等，而这可能并非我们所期望的。

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

简而言之，神经稀疏嵌入模型处于 BM25 等传统词汇方法与 Sentence Transformers 等稠密嵌入模型之间的一个有价值的生态位。它们具有以下优势：

• 混合潜力：能与稠密模型非常有效地结合，而稠密模型在词汇匹配很重要的搜索中可能表现不佳
• 可解释性：你可以确切看到哪些 token 对匹配有贡献
• 性能：在许多检索任务中与稠密模型相当或更优

在本篇博文中，我会交替使用“稀疏嵌入模型”和“稀疏编码器模型”这两个说法。

为什么要微调？

大多数（神经）稀疏嵌入模型都采用前述的查询/文档扩展，这样即使两段文本没有任何共同的词，你也能匹配含义几乎相同的文本。简而言之，模型必须识别同义词，以便将这些 token 放入最终的嵌入中。

大多数开箱即用的稀疏嵌入模型都能轻松识别出“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 模型涉及以下组件：

1. Model：要训练或微调的模型，可以是预训练的 Sparse Encoder 模型或基础模型。
2. Dataset：用于训练和评估的数据。
3. Loss Function：量化模型性能并指导优化过程的函数。
4. Training Arguments（可选）：影响训练性能以及跟踪/调试的参数。
5. Evaluator（可选）：用于在训练前、训练中或训练后评估模型的工具。
6. Trainer：将模型、数据集、损失函数和其他组件整合在一起进行训练。

现在，让我们更详细地深入探讨每个组件。

Model

Sparse Encoder 模型由一系列 Modules、Sparse Encoder 特定 Modules 或 Custom Modules 组成，具有很大的灵活性。如果你想进一步微调 Sparse Encoder 模型（例如它有一个 modules.json 文件），那么你不必担心使用了哪些模块：

from sentence_transformers import SparseEncoder

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

但如果你想从另一个检查点或从头开始训练，那么这些是你可以使用的最常见的架构：

Splade

Splade 模型使用 MLMTransformer 后跟 SpladePooling 模块。前者加载预训练的 Masked Language Modeling 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})
# )

Inference-free Splade

Inference-free Splade 使用一个 Router 模块，为查询和文档使用不同的模块。通常对于这种类型的架构，文档部分是传统的 Splade 架构（一个 MLMTransformer 后跟一个 SpladePooling 模块），查询部分是一个 SparseStaticEmbedding 模块，它只为查询中的每个 token 返回预先计算的分数。

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 进行处理。

Inference-free 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,
    }
)

Contrastive Sparse Representation (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 模型中每个维度对应词表中的一个特定 token。

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

架构选择指南

如果你不确定该使用哪种架构，这里有一份快速指南：

• 你想对现有的稠密嵌入模型进行稀疏化吗？如果是，请使用 CSR。
• 你希望查询推理是瞬时的，代价是略微的性能损失吗？如果是，请使用 无需推理的 SPLADE。
• 否则，请使用 SPLADE。

数据集

SparseEncoderTrainer 使用 datasets.Dataset 或 datasets.DatasetDict 实例进行训练和评估。你可以从 Hugging Face Datasets Hub 加载数据，或使用 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,
})

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

数据集格式

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

1. 如果你的损失函数需要一个 Label（如 损失函数概览 表中所示），你的数据集必须有一个名为 "label" 或 "score" 的列。
2. 除 "label" 或 "score" 之外的所有列都被视为 输入（如 Loss Overview 表中所示）。这些列的数量必须与你所选损失函数的有效输入数量相匹配。列的名称无关紧要，只有它们的顺序重要。

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

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

1. 该数据集有一个 "label" 列，这是这些损失函数所要求的。
2. 该数据集有 2 个非标签列，与这些损失函数所需的输入数量相匹配。

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

损失函数

损失函数衡量模型在给定批次数据上的表现，并指导优化过程。损失函数的选择取决于你可用的数据和目标任务。请参阅 Loss Overview 以获取完整的选项列表。

要训练 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
})
"""

文档

• SpladeLoss
• CSRLoss
• Losses API 参考

训练参数

SparseEncoderTrainingArguments 类允许你指定影响训练性能以及跟踪/调试的参数。虽然这些参数是可选的，但尝试使用它们有助于提高训练效率，并深入了解训练过程。

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

以下是如何初始化 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 Training Arguments 进行评估。

以下是 Sentence Transformers 为 Sparse Encoder 模型提供的已实现评估器：

此外，应使用 SequentialEvaluator 将多个评估器组合成一个 Evaluator，以便传递给 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)

SparseEmbeddingSimilarityEvaluator 与 STSb

文档

• sentence-transformers/stsb
• sentence_transformers.sparse_encoder.evaluation.SparseEmbeddingSimilarityEvaluator
• sentence_transformers.SimilarityFunction

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)

SparseTripletEvaluator 与 AllNLI

文档

• sentence-transformers/all-nli
• sentence_transformers.sparse_encoder.evaluation.SparseTripletEvaluator
• sentence_transformers.SimilarityFunction

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) 初始化一个测试评估器以获取详细的测试指标。

如果你在训练后、保存模型前进行评估，自动生成的模型卡片仍会包含测试结果。

使用 分布式训练 时，评估器仅在第一个设备上运行，而训练和评估数据集则在所有设备间共享。

Trainer

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 进行微调，这是一个尚未成为 Sparse Encoder 模型的基础模型。这需要比微调现有的 Sparse Encoder 模型（如 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，对于仅使用 Natural Questions 数据集中 10 万对数据训练的基于 distilbert 的无推理模型来说，这是一个不错的结果。

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

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

回调函数

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

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

只要安装了所需的依赖项，这些回调就会自动使用，无需你指定任何内容。

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

多数据集训练

表现最佳的模型通常同时使用多个数据集进行训练。SparseEncoderTrainer 简化了这一过程，允许你使用多个数据集进行训练，而无需将它们转换为相同的格式。你甚至可以为每个数据集应用不同的损失函数。以下是多数据集训练的步骤：

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

每个训练/评估批次将仅包含来自其中一个数据集的样本。从多个数据集中采样批次的顺序由 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

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

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

此外，对任意一种排名应用重排器都能将性能提升至约 66.3 NDCG@10，这表明无论是 Sparse、Dense 还是 Hybrid（Dense + Sparse）都能在其前 100 个结果中找到相关文档，随后重排器再将其排到前 10。因此，将 Dense -> Reranker 流水线替换为 Sparse -> Reranker 流水线可能会同时改善延迟和成本：

• 稀疏嵌入的存储成本更低，例如我们的模型对 MS MARCO 文档仅使用约 180 个活跃维度，而稠密模型通常使用 1024 个维度。
• 一些 Sparse Encoder 支持无需推理的查询处理，可实现近乎即时的第一阶段检索，类似于 BM25 等词汇匹配方案。

训练技巧

Sparse Encoder 模型有一些特性，你在训练时应当注意：

1. Sparse Encoder 模型不应仅通过评估分数来评价，还应考虑嵌入的稀疏度。毕竟，稀疏度低意味着模型嵌入的存储成本高、检索速度慢。
2. 更强的 Sparse Encoder 模型几乎都是通过从更强的教师模型（例如 CrossEncoder 模型）进行蒸馏来训练的，而不是直接从文本对或三元组训练。例如参见 SPLADE-v3 论文，它使用 SparseDistillKLDivLoss 和 SparseMarginMSELoss 进行蒸馏。我们在本博客中不会详细讨论这一点，因为它需要更多的数据准备工作，但蒸馏方案应当被认真考虑。

向量数据库集成

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

关于上述向量数据库的完整示例，请参阅 使用向量数据库进行语义搜索的文档，或参见下方的 Qdrant 示例。

Qdrant 集成示例

Qdrant 对稀疏向量提供了出色的支持，具备高效的存储和快速检索能力。以下是一个完整的实现示例：

前置条件：

• Qdrant 在本地运行（或可访问），更多详情请参阅 Qdrant 快速入门。
• Python Qdrant Client installed:

  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 模型类型，或补充稀疏检索的技术：

• 使用 Sentence Transformers 训练和微调嵌入模型：稠密双编码器嵌入模型的对应指南。
• 使用 Sentence Transformers 训练和微调重排序模型：训练 Cross Encoder 模型，通常用作稀疏检索之上的第二阶段。
• 使用 Sentence Transformers 的多模态嵌入与重排序模型：在推理时使用文本、图像、音频和视频模型。
• 使用 Sentence Transformers 训练和微调多模态嵌入与重排序模型：多模态稠密模型的对应指南。
• 🪆 Matryoshka 嵌入模型简介：用于混合搜索设置的变尺寸稠密嵌入。
• 使用 Sentence Transformers 训练速度提升 400 倍的静态嵌入模型：对 CPU 友好的稠密检索器，可在混合流水线中补充稀疏检索器。
• 用于显著更快、更便宜检索的二进制和标量嵌入量化：为混合检索的稠密侧压缩稠密嵌入。