使用提示进行训练

什么是提示 (Prompts)？

许多现代嵌入模型都遵循 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%。作者总结说，使用指令可能对特定任务微调的质量有重大贡献。

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

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

自 Sentence Transformers v3.3.0 版本起，可以使用 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。

重要提示

如果您使用提示进行训练，强烈建议将提示作为提示名称到提示字符串的映射存储在模型配置中。您可以通过在保存 SentenceTransformer 之前使用 prompts 字典对其进行初始化，在保存已加载的模型之前更新其 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。

重要提示

如果您使用提示进行训练，强烈建议将提示作为提示名称到提示字符串的映射存储在模型配置中。您可以通过在保存 SentenceTransformer 之前使用 prompts 字典对其进行初始化，在保存已加载的模型之前更新其 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]])