训练概览

为什么要进行微调?

微调Sentence Transformer模型通常能显著提升模型在您的用例上的性能,因为每个任务都需要不同的相似性概念。例如,给定新闻文章

  • “苹果发布了新款iPad”

  • “英伟达正准备推出下一代GPU”

那么在以下用例中,我们可能对相似性有不同的理解:

  • 用于分类新闻文章(如经济、体育、科技、政治等)的模型,应为这些文本生成相似的嵌入

  • 用于语义文本相似性的模型,应为这些文本生成不相似的嵌入,因为它们含义不同。

  • 用于语义搜索的模型则不需要对两个文档之间的相似性概念,因为它只需要比较查询和文档。

另请参阅训练示例,了解您可以采纳的常见实际应用的众多训练脚本。

训练组件

训练Sentence Transformer模型涉及4到6个组件

模型

Sentence Transformer模型由一系列模块自定义模块组成,这提供了很大的灵活性。如果您想进一步微调一个SentenceTransformer模型(例如,它有一个modules.json文件),那么您无需担心使用了哪些模块

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

但如果您想从另一个检查点训练,或者从头开始训练,那么这些是最常用的架构

大多数Sentence Transformer模型使用TransformerPooling模块。前者加载一个预训练的Transformer模型(例如BERTRoBERTaDistilBERTModernBERT等),后者则对Transformer的输出进行池化,为每个输入句子生成一个单一的向量表示。

from sentence_transformers import models, SentenceTransformer

transformer = models.Transformer("google-bert/bert-base-uncased")
pooling = models.Pooling(transformer.get_word_embedding_dimension(), pooling_mode="mean")

model = SentenceTransformer(modules=[transformer, pooling])

这是Sentence Transformers中的默认选项,所以使用快捷方式会更简单

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("google-bert/bert-base-uncased")

提示

最强大的基础模型通常是“编码器模型”,即训练用于为输入生成有意义的token嵌入的模型。您可以在这里找到强劲的候选模型

考虑寻找针对您的语言和/或领域设计的预训练模型。例如,FacebookAI/xlm-roberta-base 在土耳其语上的表现会优于 google-bert/bert-base-uncased

静态嵌入模型(博客文章)使用StaticEmbedding模块,它们是编码器模型,不使用缓慢的Transformer或注意力机制。对于这些模型,计算嵌入非常简单:给定输入token,返回预计算的token嵌入。这些模型速度快好几个数量级,但无法捕捉复杂的语义,因为token嵌入是与上下文分开计算的。

from sentence_transformers import models, SentenceTransformer
from tokenizers import Tokenizer

# Load any Tokenizer from Hugging Face
tokenizer = Tokenizer.from_pretrained("google-bert/bert-base-uncased")
# The `embedding_dim` is the dimensionality (size) of the token embeddings
static_embedding = StaticEmbedding(tokenizer, embedding_dim=512)

model = SentenceTransformer(modules=[static_embedding])

数据集

SentenceTransformerTrainer使用datasets.Dataset(一个数据集)或datasets.DatasetDict实例(多个数据集,另请参阅多数据集训练)进行训练和评估。

如果您想从Hugging Face Datasets加载数据,那么您应该使用datasets.load_dataset()

from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/all-nli", "pair-class", split="train")
eval_dataset = load_dataset("sentence-transformers/all-nli", "pair-class", split="dev")

print(train_dataset)
"""
Dataset({
    features: ['premise', 'hypothesis', 'label'],
    num_rows: 942069
})
"""

一些数据集(包括sentence-transformers/all-nli)需要您提供一个“子集”以及数据集名称。sentence-transformers/all-nli有4个子集,每个子集都有不同的数据格式:pairpair-classpair-scoretriplet

注意

许多Hugging Face数据集通过sentence-transformers标签与Sentence Transformers无缝集成,您可以轻松地通过浏览https://hugging-face.cn/datasets?other=sentence-transformers找到它们。我们强烈建议您浏览这些数据集,以找到对您的任务有用的训练数据集。

如果您有常用文件格式的本地数据,那么您可以使用datasets.load_dataset()轻松加载这些数据

from datasets import load_dataset

dataset = load_dataset("csv", data_files="my_file.csv")

from datasets import load_dataset

dataset = load_dataset("json", data_files="my_file.json")

如果您有需要额外预处理的本地数据,我的建议是使用datasets.Dataset.from_dict()和一个列表字典来初始化数据集,如下所示

from datasets import Dataset

anchors = []
positives = []
# Open a file, do preprocessing, filtering, cleaning, etc.
# and append to the lists

dataset = Dataset.from_dict({
    "anchor": anchors,
    "positive": positives,
})

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

数据集格式

重要的是,您的数据集格式必须与您的损失函数匹配(或者您选择的损失函数必须与您的数据集格式匹配)。验证数据集格式是否与损失函数兼容需要两个步骤

  1. 如果您的损失函数根据损失概述表需要标签,那么您的数据集必须包含名为“label”、“labels”、“score”或“scores”的列。此列将自动作为标签。

  2. 根据损失概述表,所有未命名为“label”、“labels”、“score”或“scores”的列都被视为输入。剩余列的数量必须与您选择的损失函数的有效输入数量匹配。这些列的名称不重要,只有顺序重要

例如,给定一个包含["text1", "text2", "label"]列的数据集,其中“label”列包含0到1之间的浮点相似度分数,我们可以将其与CoSENTLossAnglELossCosineSimilarityLoss一起使用,因为它

  1. 包含这些损失函数所需的“label”列。

  2. 有2个非标签列,正好是这些损失函数所需的数量。

如果您的列顺序不正确,请务必使用Dataset.select_columns重新排序您的数据集列。例如,如果您的数据集有["good_answer", "bad_answer", "question"]作为列,那么该数据集理论上可以与需要(锚点、正例、负例)三元组的损失函数一起使用,但good_answer列将被视为锚点,bad_answer视为正例,question视为负例。

此外,如果您的数据集包含无关列(例如 sample_id、metadata、source、type),您应该使用Dataset.remove_columns将其删除,否则它们将被用作输入。您也可以使用Dataset.select_columns仅保留所需的列。

损失函数

损失函数量化了模型在给定一批数据上的表现,从而允许优化器更新模型权重以产生更优(即更低)的损失值。这是训练过程的核心。

遗憾的是,没有一个损失函数能适用于所有用例。相反,使用哪个损失函数在很大程度上取决于您可用的数据和目标任务。请参阅数据集格式,了解哪些数据集对哪些损失函数有效。此外,损失概述将是您了解选项的最佳助手。

大多数损失函数只需使用您正在训练的SentenceTransformer以及一些可选参数进行初始化,例如

from datasets import load_dataset
from sentence_transformers import SentenceTransformer
from sentence_transformers.losses import CoSENTLoss

# Load a model to train/finetune
model = SentenceTransformer("xlm-roberta-base")

# Initialize the CoSENTLoss
# This loss requires pairs of text and a float similarity score as a label
loss = CoSENTLoss(model)

# Load an example training dataset that works with our loss function:
train_dataset = load_dataset("sentence-transformers/all-nli", "pair-score", split="train")
"""
Dataset({
    features: ['sentence1', 'sentence2', 'label'],
    num_rows: 942069
})
"""

训练参数

SentenceTransformerTrainingArguments类可用于指定影响训练性能的参数以及定义跟踪/调试参数。尽管它是可选的,但强烈建议尝试各种有用的参数。



以下是SentenceTransformerTrainingArguments如何初始化的示例

args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir="models/mpnet-base-all-nli-triplet",
    # 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 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,  # losses that use "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="mpnet-base-all-nli-triplet",  # Will be used in W&B if `wandb` is installed
)

评估器

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

以下是Sentence Transformers中已实现的评估器

评估器

所需数据

BinaryClassificationEvaluator

带有类别标签的对。

EmbeddingSimilarityEvaluator

带相似度分数的句子对。

InformationRetrievalEvaluator

查询 (qid => 问题), 语料库 (cid => 文档), 以及相关文档 (qid => 集合[cid])。

NanoBEIREvaluator

无需数据。

MSEEvaluator

用教师模型嵌入源句子,用学生模型嵌入目标句子。可以是相同的文本。

ParaphraseMiningEvaluator

ID 到句子的映射 & 带有重复句子 ID 的对。

RerankingEvaluator

字典列表,格式为 {'query': '...', 'positive': [...], 'negative': [...]}

TranslationEvaluator

两种不同语言的句子对。

TripletEvaluator

(锚点,正例,负例)三元组。

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

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

from datasets import load_dataset
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator, SimilarityFunction

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

# Initialize the evaluator
dev_evaluator = EmbeddingSimilarityEvaluator(
    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)
from datasets import load_dataset
from sentence_transformers.evaluation import TripletEvaluator, SimilarityFunction

# Load triplets from the AllNLI dataset (https://hugging-face.cn/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 = TripletEvaluator(
    anchors=eval_dataset["anchor"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    main_distance_function=SimilarityFunction.COSINE,
    name="all-nli-dev",
)
# You can run evaluation like so:
# results = dev_evaluator(model)
from sentence_transformers.evaluation import NanoBEIREvaluator

# 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 = NanoBEIREvaluator()
# You can run evaluation like so:
# results = dev_evaluator(model)

提示

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

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

警告

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

训练器

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

from datasets import load_dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    SentenceTransformerModelCardData,
)
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers.evaluation import TripletEvaluator

# 1. Load a model to finetune with 2. (Optional) model card data
model = SentenceTransformer(
    "microsoft/mpnet-base",
    model_card_data=SentenceTransformerModelCardData(
        language="en",
        license="apache-2.0",
        model_name="MPNet base trained on AllNLI triplets",
    )
)

# 3. Load a dataset to finetune on
dataset = load_dataset("sentence-transformers/all-nli", "triplet")
train_dataset = dataset["train"].select(range(100_000))
eval_dataset = dataset["dev"]
test_dataset = dataset["test"]

# 4. Define a loss function
loss = MultipleNegativesRankingLoss(model)

# 5. (Optional) Specify training arguments
args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir="models/mpnet-base-all-nli-triplet",
    # 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 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
    # 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="mpnet-base-all-nli-triplet",  # Will be used in W&B if `wandb` is installed
)

# 6. (Optional) Create an evaluator & evaluate the base model
dev_evaluator = TripletEvaluator(
    anchors=eval_dataset["anchor"],
    positives=eval_dataset["positive"],
    negatives=eval_dataset["negative"],
    name="all-nli-dev",
)
dev_evaluator(model)

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

# (Optional) Evaluate the trained model on the test set
test_evaluator = TripletEvaluator(
    anchors=test_dataset["anchor"],
    positives=test_dataset["positive"],
    negatives=test_dataset["negative"],
    name="all-nli-test",
)
test_evaluator(model)

# 8. Save the trained model
model.save_pretrained("models/mpnet-base-all-nli-triplet/final")

# 9. (Optional) Push it to the Hugging Face Hub
model.push_to_hub("mpnet-base-all-nli-triplet")

回调

此Sentence Transformers训练器集成了对各种transformers.TrainerCallback子类的支持,例如

  • WandbCallback,如果安装了wandb,则自动将训练指标记录到W&B

  • TensorBoardCallback,如果tensorboard可访问,则将训练指标记录到TensorBoard。

  • CodeCarbonCallback,如果安装了codecarbon,则在训练期间跟踪模型的碳排放。

    • 注意:这些碳排放量将被包含在您自动生成的模型卡片中。

有关集成回调以及如何编写您自己的回调的更多信息,请参阅 Transformers 回调文档。

多数据集训练

性能最佳的模型是同时使用多个数据集训练出来的。通常,这相当棘手,因为每个数据集的格式都不同。然而,sentence_transformers.trainer.SentenceTransformerTrainer可以训练多个数据集,而无需将每个数据集转换为相同的格式。它甚至可以对每个数据集应用不同的损失函数。使用多个数据集进行训练的步骤是

  • 使用Dataset实例的字典(或DatasetDict)作为train_dataset(可选地,也作为eval_dataset)。

  • (可选)使用一个损失函数字典,将数据集名称映射到损失。仅当您希望为不同的数据集使用不同的损失函数时才需要。

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

  • MultiDatasetBatchSamplers.ROUND_ROBIN: 从每个数据集中轮流采样,直到其中一个耗尽。使用此策略,可能不会使用每个数据集中的所有样本,但每个数据集都被同等地采样。

  • MultiDatasetBatchSamplers.PROPORTIONAL(默认):按其大小比例从每个数据集采样。使用此策略,每个数据集中的所有样本都会被使用,并且更大的数据集会更频繁地被采样。

这种多任务训练已被证明非常有效,例如Huang et al. 采用了MultipleNegativesRankingLossCoSENTLoss以及MultipleNegativesRankingLoss的一种变体(不带批内负例,只带难负例)来在中国数据集上达到最先进的性能。他们甚至应用了MatryoshkaLoss,使模型能够生成Matryoshka Embeddings

在多个数据集上训练看起来像这样

from datasets import load_dataset
from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer
from sentence_transformers.losses import CoSENTLoss, MultipleNegativesRankingLoss, SoftmaxLoss

# 1. Load a model to finetune
model = SentenceTransformer("bert-base-uncased")

# 2. Load several Datasets to train with
# (anchor, positive)
all_nli_pair_train = load_dataset("sentence-transformers/all-nli", "pair", split="train[:10000]")
# (premise, hypothesis) + label
all_nli_pair_class_train = load_dataset("sentence-transformers/all-nli", "pair-class", split="train[:10000]")
# (sentence1, sentence2) + score
all_nli_pair_score_train = load_dataset("sentence-transformers/all-nli", "pair-score", split="train[:10000]")
# (anchor, positive, negative)
all_nli_triplet_train = load_dataset("sentence-transformers/all-nli", "triplet", split="train[:10000]")
# (sentence1, sentence2) + score
stsb_pair_score_train = load_dataset("sentence-transformers/stsb", split="train[:10000]")
# (anchor, positive)
quora_pair_train = load_dataset("sentence-transformers/quora-duplicates", "pair", split="train[:10000]")
# (query, answer)
natural_questions_train = load_dataset("sentence-transformers/natural-questions", split="train[:10000]")

# We can combine all datasets into a dictionary with dataset names to datasets
train_dataset = {
    "all-nli-pair": all_nli_pair_train,
    "all-nli-pair-class": all_nli_pair_class_train,
    "all-nli-pair-score": all_nli_pair_score_train,
    "all-nli-triplet": all_nli_triplet_train,
    "stsb": stsb_pair_score_train,
    "quora": quora_pair_train,
    "natural-questions": natural_questions_train,
}

# 3. Load several Datasets to evaluate with
# (anchor, positive, negative)
all_nli_triplet_dev = load_dataset("sentence-transformers/all-nli", "triplet", split="dev")
# (sentence1, sentence2, score)
stsb_pair_score_dev = load_dataset("sentence-transformers/stsb", split="validation")
# (anchor, positive)
quora_pair_dev = load_dataset("sentence-transformers/quora-duplicates", "pair", split="train[10000:11000]")
# (query, answer)
natural_questions_dev = load_dataset("sentence-transformers/natural-questions", split="train[10000:11000]")

# We can use a dictionary for the evaluation dataset too, but we don't have to. We could also just use
# no evaluation dataset, or one dataset.
eval_dataset = {
    "all-nli-triplet": all_nli_triplet_dev,
    "stsb": stsb_pair_score_dev,
    "quora": quora_pair_dev,
    "natural-questions": natural_questions_dev,
}

# 4. Load several loss functions to train with
# (anchor, positive), (anchor, positive, negative)
mnrl_loss = MultipleNegativesRankingLoss(model)
# (sentence_A, sentence_B) + class
softmax_loss = SoftmaxLoss(model, model.get_sentence_embedding_dimension(), 3)
# (sentence_A, sentence_B) + score
cosent_loss = CoSENTLoss(model)

# Create a mapping with dataset names to loss functions, so the trainer knows which loss to apply where.
# Note that you can also just use one loss if all of your training/evaluation datasets use the same loss
losses = {
    "all-nli-pair": mnrl_loss,
    "all-nli-pair-class": softmax_loss,
    "all-nli-pair-score": cosent_loss,
    "all-nli-triplet": mnrl_loss,
    "stsb": cosent_loss,
    "quora": mnrl_loss,
    "natural-questions": mnrl_loss,
}

# 5. Define a simple trainer, although it's recommended to use one with args & evaluators
trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss=losses,
)
trainer.train()

# 6. save the trained model and optionally push it to the Hugging Face Hub
model.save_pretrained("bert-base-all-nli-stsb-quora-nq")
model.push_to_hub("bert-base-all-nli-stsb-quora-nq")

已弃用的训练

在Sentence Transformers v3.0版本发布之前,模型将使用SentenceTransformer.fit()方法和InputExampleDataLoader进行训练,其大致如下所示

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader

# Define the model. Either from scratch of by loading a pre-trained model
model = SentenceTransformer("distilbert/distilbert-base-uncased")

# Define your train examples. You need more than just two examples...
train_examples = [
    InputExample(texts=["My first sentence", "My second sentence"], label=0.8),
    InputExample(texts=["Another pair", "Unrelated sentence"], label=0.3),
]

# Define your train dataset, the dataloader and the train loss
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
train_loss = losses.CosineSimilarityLoss(model)

# Tune the model
model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=1, warmup_steps=100)

自v3.0版本发布以来,仍然可以使用SentenceTransformer.fit(),但它会在幕后初始化一个SentenceTransformerTrainer。建议直接使用Trainer,因为通过SentenceTransformerTrainingArguments您将拥有更多控制权,但依赖SentenceTransformer.fit()的现有训练脚本应该仍然有效。

如果更新的SentenceTransformer.fit()存在问题,您也可以通过调用SentenceTransformer.old_fit()来获得完全相同的旧行为,但此方法计划在未来完全弃用。

最佳基础嵌入模型

您的文本嵌入模型的质量取决于您选择的Transformer模型。遗憾的是,我们无法从GLUE或SuperGLUE基准测试中更好的性能推断该模型也将产生更好的表示。

为了测试Transformer模型的适用性,我使用training_nli_v2.py脚本,并在560k(锚点,正例,负例)三元组上训练1个epoch,批大小为64。然后我在来自各个领域的14个不同的文本相似性任务(聚类,语义搜索,重复检测等)上进行评估。

下表列出了不同模型在该基准测试上的性能

模型 性能(14个句子相似性任务)
microsoft/mpnet-base 60.99
nghuyong/ernie-2.0-en 60.73
microsoft/deberta-base 60.21
roberta-base 59.63
t5-base 59.21
bert-base-uncased 59.17
distilbert-base-uncased 59.03
nreimers/TinyBERT_L-6_H-768_v2 58.27
google/t5-v1_1-base 57.63
nreimers/MiniLMv2-L6-H768-distilled-from-BERT-Large 57.31
albert-base-v2 57.14
microsoft/MiniLM-L12-H384-uncased 56.79
microsoft/deberta-v3-base 54.46

与CrossEncoder训练的比较

训练SentenceTransformer模型与训练CrossEncoder模型非常相似,但存在一些关键差异

  • 对于CrossEncoder训练,您可以在一列中使用(大小可变的)文本列表。在SentenceTransformer训练中,您不能在训练/评估数据集的列中使用输入(例如文本)列表。简而言之:不支持可变数量的负样本训练。

有关训练CrossEncoder模型的更多详细信息,请参阅Cross Encoder > 训练概述文档。