使用提示进行训练

什么是提示？

许多现代嵌入模型都遵循 INSTRUCTOR 论文，使用“指令”或“提示”进行训练。这些提示是字符串，前置于每个要嵌入的文本，允许模型区分不同类型的文本。

例如，mixedbread-ai/mxbai-embed-large-v1 模型训练时对所有查询都使用了提示 Represent this sentence for searching relevant passages: 。此提示以提示名称 "query" 存储在模型配置中，因此用户可以在 model.encode 中指定 prompt_name。

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
query_embedding = model.encode("What are Pandas?", prompt_name="query")
# or
# query_embedding = model.encode("What are Pandas?", prompt="Represent this sentence for searching relevant passages: ")
document_embeddings = model.encode([
    "Pandas is a software library written for the Python programming language for data manipulation and analysis.",
    "Pandas are a species of bear native to South Central China. They are also known as the giant panda or simply panda.",
    "Koala bears are not actually bears, they are marsupials native to Australia.",
])
similarity = model.similarity(query_embedding, document_embeddings)
print(similarity)
# => tensor([[0.7594, 0.7560, 0.4674]])

有关使用提示进行推理的更多信息，请参阅提示模板。

我们为什么要使用提示进行训练？

INSTRUCTOR 论文表明，在每个文本之前添加提示或指令可以将模型性能平均提高约 6%，尤其是在分类、聚类和语义文本相似度方面取得了显著的收益。请注意，检索方面的性能提升明显较低，小型模型和大型模型分别为 0.4% 和 2.7%。

最近，BGE 论文也发现了类似的结果，表明如果查询前缀为 Represent this sentence for searching relevant passages: ，检索性能会提高约 1.4%。作者总结道，使用指令可能对任务特定微调的质量有重大贡献。

本质上，只要在训练和推理过程中都使用指令或提示，它们就能提高性能。

我们如何使用提示进行训练？

自 v3.3.0 Sentence Transformers 版本发布以来，可以通过 SentenceTransformerTrainingArguments 类中的 prompts 参数来微调带提示的嵌入模型。此参数有 4 种接受的格式：

1. str：一个用于所有数据集中所有列的单一提示。例如：

   args = SentenceTransformerTrainingArguments(
       ...,
       prompts="text: ",
       ...,
   )
   

2. Dict[str, str]：一个将列名映射到提示的字典，应用于所有数据集。例如：

   args = SentenceTransformerTrainingArguments(
       ...,
       prompts={
           "query": "query: ",
           "answer": "document: ",
       },
       ...,
   )
   

3. Dict[str, str]：一个将数据集名称映射到提示的字典。这仅应在您的训练/评估/测试数据集是 DatasetDict 或 Dataset 字典时使用。例如：

   args = SentenceTransformerTrainingArguments(
       ...,
       prompts={
           "stsb": "Represent this text for semantic similarity search: ",
           "nq": "Represent this text for retrieval: ",
       },
       ...,
   )
   

4. Dict[str, Dict[str, str]]：一个将数据集名称映射到字典的字典，该字典又将列名映射到提示。这仅应在您的训练/评估/测试数据集是 DatasetDict 或 Dataset 字典时使用。例如：

   args = SentenceTransformerTrainingArguments(
       ...,
       prompts={
           "stsb": {
               "sentence1": "sts: ",
               "sentence2": "sts: ",
           },
           "nq": {
               "query": "query: ",
               "document": "document: ",
           },
       },
       ...,
   )
   

此外，一些研究论文（INSTRUCTOR、NV-Embed）将提示从平均池化步骤中排除，使其仅在 Transformer 块中使用。在 Sentence Transformers 中，这可以通过 Pooling 模块中的 include_prompt 参数/属性或通过 SentenceTransformer.set_pooling_include_prompt() 方法进行配置。根据我个人经验，在池化中包含提示的模型往往表现更好。

训练脚本

请参阅以下脚本，作为实际使用提示进行训练的示例：

• training_nq_prompts.py：此脚本使用 CachedMultipleNegativesRankingLoss 损失，对来自 natural-questions 数据集的 10 万个查询-答案对进行 mpnet-base 微调。模型在训练期间使用 NanoBEIREvaluator 进行评估。

此脚本有两个变量，分别影响 1) 是否使用提示，以及 2) 提示是否包含在池化中。我已在各种不同设置下微调了 mpnet-base 和 bert-base-uncased，导致 NDCG@10 相对改进了 0.66% 和 0.90%，且没有额外成本。

mpnet-base 实验

在各种设置下运行脚本产生了以下检查点：

• tomaarsen/mpnet-base-nq

• tomaarsen/mpnet-base-nq-prompts

注意

mpnet-base 在使用提示且将这些提示排除在池化之外进行训练时似乎有点不稳定：损失在某个时刻会出现峰值，这种现象在例如 bert-base-uncased 中没有观察到。

对于这两个模型，使用提示训练的模型在整个训练过程中始终优于基线模型。

此外，经过提示训练的模型在自动生成的模型卡中包含了训练数据集的这些提示详情：tomaarsen/mpnet-base-nq-prompts#natural-questions。

重要

如果您使用提示进行训练，强烈建议将提示作为提示名称到提示字符串的映射存储在模型配置中。您可以通过以下方式实现：在保存模型之前使用 prompts 字典初始化 SentenceTransformer；在保存加载的模型之前更新 model.prompts；和/或更新保存模型的 config_sentence_transformers.json 文件。

在模型配置中添加提示后，提示训练模型的最终用法变为：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("tomaarsen/mpnet-base-nq-prompts")
query_embedding = model.encode("What are Pandas?", prompt_name="query")
document_embeddings = model.encode([
    "Pandas is a software library written for the Python programming language for data manipulation and analysis.",
    "Pandas are a species of bear native to South Central China. They are also known as the giant panda or simply panda.",
    "Koala bears are not actually bears, they are marsupials native to Australia.",
    ],
    prompt_name="document",
)
similarity = model.similarity(query_embedding, document_embeddings)
print(similarity)
# => tensor([[0.4725, 0.7339, 0.4369]])

bert-base-uncased 实验

在各种设置下运行脚本产生了以下检查点：

• tomaarsen/bert-base-nq

• tomaarsen/bert-base-nq-prompts

• tomaarsen/bert-base-nq-prompts-exclude-pooling-prompts

对于这三个模型，除了第一次评估外，使用提示训练的模型在整个训练过程中始终优于基线模型。在平均池化中排除提示的模型始终明显差于其他两个模型。

此外，经过提示训练的模型在自动生成的模型卡中包含了训练数据集的这些提示详情：tomaarsen/bert-base-nq-prompts#natural-questions。

重要

如果您使用提示进行训练，强烈建议将提示作为提示名称到提示字符串的映射存储在模型配置中。您可以通过以下方式实现：在保存模型之前使用 prompts 字典初始化 SentenceTransformer；在保存加载的模型之前更新 model.prompts；和/或更新保存模型的 config_sentence_transformers.json 文件。

在模型配置中添加提示后，提示训练模型的最终用法变为：

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("tomaarsen/bert-base-nq-prompts")
query_embedding = model.encode("What are Pandas?", prompt_name="query")
document_embeddings = model.encode([
    "Pandas is a software library written for the Python programming language for data manipulation and analysis.",
    "Pandas are a species of bear native to South Central China. They are also known as the giant panda or simply panda.",
    "Koala bears are not actually bears, they are marsupials native to Australia.",
    ],
    prompt_name="document",
)
similarity = model.similarity(query_embedding, document_embeddings)
print(similarity)
# => tensor([[0.3955, 0.8226, 0.5706]])