损失函数

sentence_transformers.losses 定义了不同的损失函数，可用于在训练数据上微调嵌入模型。损失函数的选择在微调模型时起着关键作用。它决定了我们的嵌入模型在特定下游任务中的表现。

遗憾的是，没有“一劳永逸”的损失函数。哪种损失函数适合取决于可用的训练数据和目标任务。建议查看损失函数概述，以帮助缩小您的损失函数选择范围。

BatchAllTripletLoss

class sentence_transformers.losses.BatchAllTripletLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, distance_metric=<function BatchHardTripletLossDistanceFunction.eucledian_distance>, margin: float = 5)[source]

BatchAllTripletLoss 接收一批 (句子, 标签) 对，并计算所有可能有效三元组的损失，即锚点和正例必须具有相同的标签，锚点和负例必须具有不同的标签。标签必须是整数，相同的标签表示来自同一类别的句子。您的训练数据集必须每个标签类别至少包含 2 个示例。

参数:

model – SentenceTransformer 模型
distance_metric – 返回两个嵌入之间距离的函数。SiameseDistanceMetric 类包含可用的预定义度量。
margin – 负样本与锚点的距离应至少比正样本远 margin。

参考文献

来源: https://github.com/NegatioN/OnlineMiningTripletLoss/blob/master/online_triplet_loss/losses.py
论文: In Defense of the Triplet Loss for Person Re-Identification, https://hugging-face.cn/papers/1703.07737
博客文章: https://omoindrot.github.io/triplet-loss

要求

每个句子必须用一个类别标记。
您的数据集每个标签类别必须至少包含 2 个示例。

输入

文本	标签
单个句子	类别

建议

使用 BatchSamplers.GROUP_BY_LABEL (docs) 确保每个批次包含每个标签类别 2 个或更多示例。

关系

BatchHardTripletLoss 仅使用最难的正样本和负样本，而不是所有可能的有效三元组。
BatchHardSoftMarginTripletLoss 仅使用最难的正样本和负样本，而不是所有可能的有效三元组。此外，它不需要设置 margin。
BatchSemiHardTripletLoss 仅使用半硬三元组、有效三元组，而不是所有可能的有效三元组。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
# E.g. 0: sports, 1: economy, 2: politics
train_dataset = Dataset.from_dict({
    "sentence": [
        "He played a great game.",
        "The stock is up 20%",
        "They won 2-1.",
        "The last goal was amazing.",
        "They all voted against the bill.",
    ],
    "label": [0, 1, 0, 0, 2],
})
loss = losses.BatchAllTripletLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

BatchHardSoftMarginTripletLoss

class sentence_transformers.losses.BatchHardSoftMarginTripletLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, distance_metric=<function BatchHardTripletLossDistanceFunction.eucledian_distance>)[source]

BatchHardSoftMarginTripletLoss 接收一批 (句子, 标签) 对，并计算所有可能有效三元组的损失，即锚点和正例必须具有相同的标签，锚点和负例必须具有不同的标签。标签必须是整数，相同的标签表示来自同一类别的句子。您的训练数据集每个标签类别必须至少包含 2 个示例。此软边距变体不需要设置 margin。

参数:

model – SentenceTransformer 模型
distance_metric – 返回两个嵌入之间距离的函数。SiameseDistanceMetric 类包含可用的预定义度量。

定义

简单三元组:: 损失为 0 的三元组，因为 distance(anchor, positive) + margin < distance(anchor, negative)。
硬三元组:: 负样本比正样本更接近锚点的三元组，即 distance(anchor, negative) < distance(anchor, positive)。
半硬三元组:: 负样本不比正样本更接近锚点，但仍具有正损失的三元组，即 distance(anchor, positive) < distance(anchor, negative) + margin。

参考文献

来源: https://github.com/NegatioN/OnlineMiningTripletLoss/blob/master/online_triplet_loss/losses.py
论文: In Defense of the Triplet Loss for Person Re-Identification, https://hugging-face.cn/papers/1703.07737
博客文章: https://omoindrot.github.io/triplet-loss

要求

每个句子必须用一个类别标记。
您的数据集每个标签类别必须至少包含 2 个示例。
您的数据集应包含硬正样本和硬负样本。

输入

文本	标签
单个句子	类别

建议

使用 BatchSamplers.GROUP_BY_LABEL (docs) 确保每个批次包含每个标签类别 2 个或更多示例。

关系

BatchHardTripletLoss 使用用户指定的 margin，而此损失不需要设置 margin。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
# E.g. 0: sports, 1: economy, 2: politics
train_dataset = Dataset.from_dict({
    "sentence": [
        "He played a great game.",
        "The stock is up 20%",
        "They won 2-1.",
        "The last goal was amazing.",
        "They all voted against the bill.",
    ],
    "label": [0, 1, 0, 0, 2],
})
loss = losses.BatchHardSoftMarginTripletLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

BatchHardTripletLoss

class sentence_transformers.losses.BatchHardTripletLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, distance_metric=<function BatchHardTripletLossDistanceFunction.eucledian_distance>, margin: float = 5)[source]

BatchHardTripletLoss 接收一批 (句子, 标签) 对，并计算所有可能有效三元组的损失，即锚点和正例必须具有相同的标签，锚点和负例必须具有不同的标签。然后，它寻找最硬的正样本和最硬的负样本。标签必须是整数，相同的标签表示来自同一类别的句子。您的训练数据集每个标签类别必须至少包含 2 个示例。

参数:

model – SentenceTransformer 模型
distance_metric – 返回两个嵌入之间距离的函数。SiameseDistanceMetric 类包含可用的预定义度量。
margin – 负样本与锚点的距离应至少比正样本远 margin。

定义

简单三元组:: 损失为 0 的三元组，因为 distance(anchor, positive) + margin < distance(anchor, negative)。
硬三元组:: 负样本比正样本更接近锚点的三元组，即 distance(anchor, negative) < distance(anchor, positive)。
半硬三元组:: 负样本不比正样本更接近锚点，但仍具有正损失的三元组，即 distance(anchor, positive) < distance(anchor, negative) + margin。

参考文献

来源: https://github.com/NegatioN/OnlineMiningTripletLoss/blob/master/online_triplet_loss/losses.py
论文: In Defense of the Triplet Loss for Person Re-Identification, https://hugging-face.cn/papers/1703.07737
博客文章: https://omoindrot.github.io/triplet-loss

要求

每个句子必须用一个类别标记。
您的数据集每个标签类别必须至少包含 2 个示例。
您的数据集应包含硬正样本和硬负样本。

输入

文本	标签
单个句子	类别

建议

使用 BatchSamplers.GROUP_BY_LABEL (docs) 确保每个批次包含每个标签类别 2 个或更多示例。

关系

BatchAllTripletLoss 使用所有可能的有效三元组，而不是仅使用最硬的正样本和负样本。
BatchSemiHardTripletLoss 仅使用半硬三元组、有效三元组，而不是仅使用最硬的正样本和负样本。
BatchHardSoftMarginTripletLoss 不需要设置 margin，而此损失需要。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
# E.g. 0: sports, 1: economy, 2: politics
train_dataset = Dataset.from_dict({
    "sentence": [
        "He played a great game.",
        "The stock is up 20%",
        "They won 2-1.",
        "The last goal was amazing.",
        "They all voted against the bill.",
    ],
    "label": [0, 1, 0, 0, 2],
})
loss = losses.BatchHardTripletLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

BatchSemiHardTripletLoss

class sentence_transformers.losses.BatchSemiHardTripletLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, distance_metric=<function BatchHardTripletLossDistanceFunction.eucledian_distance>, margin: float = 5)[source]

BatchSemiHardTripletLoss 接收一批 (标签, 句子) 对，并计算所有可能有效三元组的损失，即锚点和正例必须具有相同的标签，锚点和负例必须具有不同的标签。然后，它寻找半硬正样本和负样本。标签必须是整数，相同的标签表示来自同一类别的句子。您的训练数据集每个标签类别必须至少包含 2 个示例。

参数:

model – SentenceTransformer 模型
distance_metric – 返回两个嵌入之间距离的函数。SiameseDistanceMetric 类包含可用的预定义度量。
margin – 负样本与锚点的距离应至少比正样本远 margin。

定义

简单三元组:: 损失为 0 的三元组，因为 distance(anchor, positive) + margin < distance(anchor, negative)。
硬三元组:: 负样本比正样本更接近锚点的三元组，即 distance(anchor, negative) < distance(anchor, positive)。
半硬三元组:: 负样本不比正样本更接近锚点，但仍具有正损失的三元组，即 distance(anchor, positive) < distance(anchor, negative) + margin。

参考文献

来源: https://github.com/NegatioN/OnlineMiningTripletLoss/blob/master/online_triplet_loss/losses.py
论文: In Defense of the Triplet Loss for Person Re-Identification, https://hugging-face.cn/papers/1703.07737
博客文章: https://omoindrot.github.io/triplet-loss

要求

每个句子必须用一个类别标记。
您的数据集每个标签类别必须至少包含 2 个示例。
您的数据集应包含半硬正样本和半硬负样本。

输入

文本	标签
单个句子	类别

建议

使用 BatchSamplers.GROUP_BY_LABEL (docs) 确保每个批次包含每个标签类别 2 个或更多示例。

关系

BatchHardTripletLoss 仅使用最硬的正样本和负样本，而不是仅使用半硬正样本和负样本。
BatchAllTripletLoss 使用所有可能的有效三元组，而不是仅使用半硬正样本和负样本。
BatchHardSoftMarginTripletLoss 仅使用最硬的正样本和负样本，而不是仅使用半硬正样本和负样本。此外，它不需要设置 margin。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
# E.g. 0: sports, 1: economy, 2: politics
train_dataset = Dataset.from_dict({
    "sentence": [
        "He played a great game.",
        "The stock is up 20%",
        "They won 2-1.",
        "The last goal was amazing.",
        "They all voted against the bill.",
    ],
    "label": [0, 1, 0, 0, 2],
})
loss = losses.BatchSemiHardTripletLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

ContrastiveLoss

class sentence_transformers.losses.ContrastiveLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, distance_metric=<function SiameseDistanceMetric.<lambda>>, margin: float = 0.5, size_average: bool = True)[source]

对比损失。期望输入是两个文本和一个标签（0 或 1）。如果标签 == 1，则减小两个嵌入之间的距离。如果标签 == 0，则增大嵌入之间的距离。

参数:

model – SentenceTransformer 模型
distance_metric – 返回两个嵌入之间距离的函数。SiameseDistanceMetric 类包含可用的预定义度量。
margin – 负样本（标签 == 0）的距离应至少为 margin 值。
size_average – 按 mini-batch 的大小进行平均。

参考文献

要求

(锚点, 正例/负例) 对

输入

文本	标签
(锚点, 正例/负例) 对	正例为 1，负例为 0

关系

OnlineContrastiveLoss 类似，但使用硬正样本和硬负样本对。它通常能产生更好的结果。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "sentence1": ["It's nice weather outside today.", "He drove to work."],
    "sentence2": ["It's so sunny.", "She walked to the store."],
    "label": [1, 0],
})
loss = losses.ContrastiveLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

OnlineContrastiveLoss

class sentence_transformers.losses.OnlineContrastiveLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, distance_metric=<function SiameseDistanceMetric.<lambda>>, margin: float = 0.5)[source]

此在线对比损失与 ConstrativeLoss 类似，但它选择硬正样本对（距离较远的正样本）和硬负样本对（距离较近的负样本），并仅对这些对计算损失。此损失通常比 ContrastiveLoss 产生更好的性能。

参数:

model – SentenceTransformer 模型
distance_metric – 返回两个嵌入之间距离的函数。SiameseDistanceMetric 类包含可用的预定义度量。
margin – 负样本（标签 == 0）的距离应至少为 margin 值。

参考文献

训练示例 > Quora 重复问题

要求

(锚点, 正例/负例) 对
数据应包括硬正样本和硬负样本。

输入

文本	标签
(锚点, 正例/负例) 对	正例为 1，负例为 0

关系

ContrastiveLoss 类似，但不使用硬正样本和硬负样本对。OnlineContrastiveLoss 通常能产生更好的结果。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "sentence1": ["It's nice weather outside today.", "He drove to work."],
    "sentence2": ["It's so sunny.", "She walked to the store."],
    "label": [1, 0],
})
loss = losses.OnlineContrastiveLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

ContrastiveTensionLoss

class sentence_transformers.losses.ContrastiveTensionLoss(model: SentenceTransformer)[source]

此损失仅接受单个句子，不带任何标签。正样本对和负样本对通过随机采样自动创建，其中正样本对包含两个相同的句子，负样本对包含两个不同的句子。会创建一个独立于编码器模型的副本，用于编码每个对的第一个句子。原始编码器模型编码第二个句子。嵌入通过生成的标签（如果为正则为 1，如果为负则为 0）使用二元交叉熵目标进行比较和评分。

请注意，您必须将 ContrastiveTensionDataLoader 用于此损失。ContrastiveTensionDataLoader 的 pos_neg_ratio 可用于确定每个正样本对的负样本对数量。

通常，建议使用 ContrastiveTensionLossInBatchNegatives 而不是此损失，因为它提供了更强的训练信号。

参数:: model – SentenceTransformer 模型

参考文献

语义重调与对比张力: https://openreview.net/pdf?id=Ov_sMNau-PF
无监督学习 > CT

输入

文本	标签
单个句子	无

关系

ContrastiveTensionLossInBatchNegatives 使用批内负采样，这比此损失提供了更强的训练信号。

示例

from sentence_transformers import SentenceTransformer, losses
from sentence_transformers.losses import ContrastiveTensionDataLoader

model = SentenceTransformer('all-MiniLM-L6-v2')
train_examples = [
    'This is the 1st sentence',
    'This is the 2nd sentence',
    'This is the 3rd sentence',
    'This is the 4th sentence',
    'This is the 5th sentence',
    'This is the 6th sentence',
    'This is the 7th sentence',
    'This is the 8th sentence',
    'This is the 9th sentence',
    'This is the final sentence',
]

train_dataloader = ContrastiveTensionDataLoader(train_examples, batch_size=3, pos_neg_ratio=3)
train_loss = losses.ContrastiveTensionLoss(model=model)

model.fit(
    [(train_dataloader, train_loss)],
    epochs=10,
)

初始化内部 Module 状态，由 nn.Module 和 ScriptModule 共享。

ContrastiveTensionLossInBatchNegatives

class sentence_transformers.losses.ContrastiveTensionLossInBatchNegatives(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, scale: float = 20.0, similarity_fct=<function cos_sim>)[source]

此损失仅接受单个句子，不带任何标签。正样本对和负样本对通过随机采样自动创建，其中正样本对包含两个相同的句子，负样本对包含两个不同的句子。会创建一个独立于编码器模型的副本，用于编码每个对的第一个句子。原始编码器模型编码第二个句子。与 ContrastiveTensionLoss 不同，此损失使用批次负采样策略，即负样本对从批次中采样。使用批次内负采样比原始 ContrastiveTensionLoss 提供了更强的训练信号。性能通常随批次大小的增加而提高。

请注意，您不应将 ContrastiveTensionDataLoader 用于此损失，而应使用带有 InputExample 实例的普通 DataLoader。每个 InputExample 实例的两个文本应相同。

参数:

model – SentenceTransformer 模型
scale – 相似度函数的输出乘以比例值。
similarity_fct – 句子嵌入之间的相似度函数。默认为 cos_sim。也可以设置为点积（然后将 scale 设置为 1）。

参考文献

语义重调与对比张力: https://openreview.net/pdf?id=Ov_sMNau-PF
无监督学习 > CT (批内负样本)

关系

ContrastiveTensionLoss 不在批次内选择负样本对，导致训练信号弱于此损失。

输入

文本	标签
(锚点, 锚点) 对	无

示例

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

model = SentenceTransformer('all-MiniLM-L6-v2')
train_examples = [
    InputExample(texts=['This is a positive pair', 'Where the distance will be minimized'], label=1),
    InputExample(texts=['This is a negative pair', 'Their distance will be increased'], label=0),
]
train_examples = [
    InputExample(texts=['This is the 1st sentence', 'This is the 1st sentence']),
    InputExample(texts=['This is the 2nd sentence', 'This is the 2nd sentence']),
    InputExample(texts=['This is the 3rd sentence', 'This is the 3rd sentence']),
    InputExample(texts=['This is the 4th sentence', 'This is the 4th sentence']),
    InputExample(texts=['This is the 5th sentence', 'This is the 5th sentence']),
]

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=32)
train_loss = losses.ContrastiveTensionLossInBatchNegatives(model=model)

model.fit(
    [(train_dataloader, train_loss)],
    epochs=10,
)

CoSENTLoss

class sentence_transformers.losses.CoSENTLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, scale: float = 20.0, similarity_fct=<function pairwise_cos_sim>)[source]

此类实现了 CoSENT (Consistent SENTence embedding) 损失。它期望每个 InputExamples 包含一对文本和一个浮点值标签，表示该对的预期相似度分数。

它计算以下损失函数

loss = logsum(1+exp(s(k,l)-s(i,j))+exp...)，其中 (i,j) 和 (k,l) 是批次中任意输入对，使得 (i,j) 的预期相似度大于 (k,l)。求和是对批次中所有符合此条件的输入对的所有可能对进行的。

经验证，此损失函数比 CosineSimilarityLoss 产生更强大的训练信号，从而实现更快的收敛和最终模型的卓越性能。因此，CoSENTLoss 可用作任何训练脚本中 CosineSimilarityLoss 的直接替代品。

参数:

model – SentenceTransformer模型
similarity_fct – 计算嵌入之间成对相似度的函数。默认为 util.pairwise_cos_sim。
scale – 相似度函数的输出乘以比例值。表示逆温度。

参考文献

更多详情，请参见：https://penghao-bdsc.github.io/papers/CoSENT_TASLP2024.pdf

要求

句子对以及在相似度函数范围内的相应相似度分数。默认为 [-1,1]。

输入

文本	标签
(句子_A, 句子_B) 对	浮点相似度分数

关系

AnglELoss 是 CoSENTLoss，但使用 pairwise_angle_sim 作为度量，而不是 pairwise_cos_sim。
CosineSimilarityLoss 似乎比 CoSENTLoss 产生更弱的训练信号。在我们的实验中，推荐使用 CoSENTLoss。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "sentence1": ["It's nice weather outside today.", "He drove to work."],
    "sentence2": ["It's so sunny.", "She walked to the store."],
    "score": [1.0, 0.3],
})
loss = losses.CoSENTLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

AnglELoss

class sentence_transformers.losses.AnglELoss(model: SentenceTransformer, scale: float = 20.0)[source]

此类实现了 AnglE (Angle Optimized) 损失。这是对 CoSENTLoss 的修改，旨在解决以下问题：余弦函数的梯度在波形接近其顶部或底部时趋近于 0。这可能会阻碍优化过程，因此 AnglE 提出通过优化复数空间中的角度差来减轻这种影响。

它期望每个 InputExamples 包含一对文本和一个浮点值标签，表示该对的预期相似度分数。

它计算以下损失函数

loss = logsum(1+exp(s(i,j)-s(k,l))+exp...)，其中 (i,j) 和 (k,l) 是批次中任意输入对，使得 (i,j) 的预期相似度大于 (k,l)。求和是对批次中所有符合此条件的输入对的所有可能对进行的。这与 CoSENTLoss 相同，只是相似度函数不同。

参数:

model – SentenceTransformer模型
scale – 相似度函数的输出乘以比例值。表示逆温度。

参考文献

更多详情，请参见：https://aclanthology.org/2024.acl-long.101/

要求

句子对以及在相似度函数范围内的相应相似度分数。默认为 [-1,1]。

输入

文本	标签
(句子_A, 句子_B) 对	浮点相似度分数

关系

CoSENTLoss 是 AnglELoss，但使用 pairwise_cos_sim 作为度量，而不是 pairwise_angle_sim。
CosineSimilarityLoss 似乎比 CoSENTLoss 或 AnglELoss 产生更弱的训练信号。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "sentence1": ["It's nice weather outside today.", "He drove to work."],
    "sentence2": ["It's so sunny.", "She walked to the store."],
    "score": [1.0, 0.3],
})
loss = losses.AnglELoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

CosineSimilarityLoss

对于每个句子对，我们将句子 A 和句子 B 通过我们的网络，产生嵌入 u 和 v。使用余弦相似度计算这些嵌入的相似度，并将结果与黄金相似度分数进行比较。

这使得我们的网络可以针对识别句子相似度进行微调。

class sentence_transformers.losses.CosineSimilarityLoss(model: SentenceTransformer, loss_fct: Module = MSELoss(), cos_score_transformation: Module = Identity())[source]

CosineSimilarityLoss 期望 InputExamples 包含两个文本和一个浮点标签。它计算向量 u = model(sentence_A) 和 v = model(sentence_B)，并测量两者之间的余弦相似度。默认情况下，它最小化以下损失：||input_label - cos_score_transformation(cosine_sim(u,v))||_2。

参数:

model – SentenceTransformer 模型
loss_fct – 应该使用哪个 PyTorch 损失函数来比较 cosine_similarity(u, v) 和 input_label？默认使用 MSE：||input_label - cosine_sim(u, v)||_2
cos_score_transformation – cos_score_transformation 函数应用于余弦相似度之上。默认情况下，使用恒等函数（即不改变）。

参考文献

训练示例 > 语义文本相似度

要求

句子对及相应的相似度分数，范围 [0, 1]

输入

文本	标签
(句子_A, 句子_B) 对	浮点相似度分数

关系

CoSENTLoss 似乎比 CosineSimilarityLoss 产生更强的训练信号。在我们的实验中，推荐使用 CoSENTLoss。
AnglELoss 是 CoSENTLoss，但使用 pairwise_angle_sim 作为度量，而不是 pairwise_cos_sim。它也比 CosineSimilarityLoss 产生更强的训练信号。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "sentence1": ["It's nice weather outside today.", "He drove to work."],
    "sentence2": ["It's so sunny.", "She walked to the store."],
    "score": [1.0, 0.3],
})
loss = losses.CosineSimilarityLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

DenoisingAutoEncoderLoss

class sentence_transformers.losses.DenoisingAutoEncoderLoss(model: SentenceTransformer, decoder_name_or_path: str | None = None, tie_encoder_decoder: bool = True)[source]

此损失期望输入为受损句子对和相应的原始句子。在训练期间，解码器从编码的句子嵌入中重建原始句子。这里的参数“decoder_name_or_path”指示用于解码器的预训练模型（由 Hugging Face 支持）。由于包含解码过程，因此这里的解码器应该有一个名为 XXXLMHead 的类（在 Hugging Face 的 Transformers 的上下文中）。“tie_encoder_decoder”标志指示是否绑定编码器和解码器的可训练参数，这被证明有利于模型性能，同时限制所需的内存量。只有当编码器和解码器来自相同的架构时，“tie_encoder_decoder”标志才能起作用。

数据生成过程（即“损坏”过程）已在 DenoisingAutoEncoderDataset 中实现，允许您仅提供常规句子。

参数:

model (SentenceTransformer) – SentenceTransformer 模型。
decoder_name_or_path (str, optional) – 用于初始化解码器的模型名称或路径（与 Hugging Face 的 Transformers 兼容）。默认为 None。
tie_encoder_decoder (bool) – 是否绑定编码器和解码器的可训练参数。默认为 True。

参考文献

TSDAE 论文: https://hugging-face.cn/papers/2104.06979
无监督学习 > TSDAE

要求

解码器应该有一个名为 XXXLMHead 的类（在 Hugging Face 的 Transformers 的上下文中）
应该使用大型语料库

输入

文本	标签
(受损句子, 原始句子) 对	无
通过 `DenoisingAutoEncoderDataset` 输入的句子	无

示例

from sentence_transformers import SentenceTransformer, losses
from sentence_transformers.datasets import DenoisingAutoEncoderDataset
from torch.utils.data import DataLoader

model_name = "bert-base-cased"
model = SentenceTransformer(model_name)
train_sentences = [
    "First training sentence", "Second training sentence", "Third training sentence", "Fourth training sentence",
]
batch_size = 2
train_dataset = DenoisingAutoEncoderDataset(train_sentences)
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
train_loss = losses.DenoisingAutoEncoderLoss(
    model, decoder_name_or_path=model_name, tie_encoder_decoder=True
)
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=10,
)

GISTEmbedLoss

class sentence_transformers.losses.GISTEmbedLoss(model: SentenceTransformer, guide: SentenceTransformer, temperature: float = 0.01, margin_strategy: Literal['absolute', 'relative'] = 'absolute', margin: float = 0.0, contrast_anchors: bool = True, contrast_positives: bool = True, gather_across_devices: bool = False)[source]

此损失用于使用 GISTEmbed 算法训练 SentenceTransformer 模型。它将一个模型和一个指导模型作为输入，并使用指导模型来指导批内负样本选择。余弦相似度用于计算损失，温度参数用于缩放余弦相似度。

您可以应用不同的假阴性过滤策略来丢弃与正样本过于相似的硬负样本。支持两种策略：

“absolute”：丢弃相似度分数大于或等于 positive_score - margin 的负样本。

“relative”：丢弃相似度分数大于或等于 positive_score * (1 - margin) 的负样本。

参数:

model – 基于 transformers 模型的 SentenceTransformer 模型。
guide – 用于指导批内负样本选择的 SentenceTransformer 模型。
temperature – 用于缩放余弦相似度的温度参数。MultipleNegativesRankingLoss 中 scale 参数的倒数。
margin_strategy – 用于假阴性过滤的策略。{"absolute", "relative"} 之一。
margin – 过滤负样本的 margin 值。默认为 0.0，与“absolute”策略一起使用时，仅移除比正样本更相似于查询的负样本。
contrast_anchors – 如果为 True，则在损失计算中包含锚点-锚点对，从而使锚点的嵌入彼此推开。默认为 True，遵循原始 GISTEmbed 论文。
contrast_positives – 如果为 True，则在损失计算中包含正样本-正样本对，从而使正样本的嵌入彼此推开。默认为 True，遵循原始 GISTEmbed 论文，但在某些检索任务中设置为 False 可能会产生更好的结果。
gather_across_devices – 如果为 True，则在计算损失之前聚合所有设备上的嵌入。在多 GPU 训练时推荐使用，因为它允许更大的批次大小，但由于通信开销可能会减慢训练速度，并可能导致内存不足错误。

参考文献

更多详情，请参见：https://hugging-face.cn/papers/2402.16829

要求

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

输入

文本	标签
(锚点, 正例, 负例) 三元组	无
(锚点, 正例) 对	无

建议

使用 BatchSamplers.NO_DUPLICATES (docs) 确保批内负样本中没有与锚点或正样本重复的。

关系

MultipleNegativesRankingLoss 与此损失类似，但它不使用指导模型来指导批内负样本选择。GISTEmbedLoss 以一些训练开销为代价产生更强的训练信号。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
guide = SentenceTransformer("all-MiniLM-L6-v2")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.GISTEmbedLoss(model, guide)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

CachedGISTEmbedLoss

class sentence_transformers.losses.CachedGISTEmbedLoss(model: SentenceTransformer, guide: SentenceTransformer, temperature: float = 0.01, mini_batch_size: int = 32, show_progress_bar: bool = False, margin_strategy: Literal['absolute', 'relative'] = 'absolute', margin: float = 0.0, contrast_anchors: bool = True, contrast_positives: bool = True, gather_across_devices: bool = False)[source]

此损失是 GISTEmbedLoss 和 CachedMultipleNegativesRankingLoss 的组合。通常，MultipleNegativesRankingLoss 需要更大的批次大小才能获得更好的性能。GISTEmbedLoss 由于使用指导模型进行批内负样本选择，因此比 MultipleNegativesRankingLoss 产生更强的训练信号。同时，CachedMultipleNegativesRankingLoss 允许通过将计算分为嵌入和损失计算两个阶段来扩展批次大小，这两个阶段都可以通过 mini-batches 进行缩放（https://hugging-face.cn/papers/2101.06983）。

通过结合 GISTEmbedLoss 的引导选择和 CachedMultipleNegativesRankingLoss 的梯度缓存，可以减少内存使用，同时保持与 GISTEmbedLoss 相当的性能水平。

您可以应用不同的假阴性过滤策略来丢弃与正样本过于相似的硬负样本。支持两种策略：

“absolute”：丢弃相似度分数大于或等于 positive_score - margin 的负样本。

“relative”：丢弃相似度分数大于或等于 positive_score * (1 - margin) 的负样本。

参数:

model – SentenceTransformer 模型
guide – 用于指导批内负样本选择的 SentenceTransformer 模型。
temperature – 温度参数，用于缩放余弦相似度。
mini_batch_size – 前向传播的 mini-batch 大小，表示训练和评估期间实际使用的内存量。mini-batch 大小越大，训练的内存效率越高，但训练速度会越慢。建议将其设置为 GPU 内存允许的最大值。默认值为 32。
show_progress_bar – 如果为 True，则在训练期间显示 mini-batch 的进度条。默认值为 False。
margin_strategy – 用于假阴性过滤的策略。{"absolute", "relative"} 之一。
margin – 过滤负样本的 margin 值。默认为 0.0，与“absolute”策略一起使用时，仅移除比正样本更相似于查询的负样本。
contrast_anchors – 如果为 True，则在损失计算中包含锚点-锚点对，从而使锚点的嵌入彼此推开。默认为 True，遵循原始 GISTEmbed 论文。
contrast_positives – 如果为 True，则在损失计算中包含正样本-正样本对，从而使正样本的嵌入彼此推开。默认为 True，遵循原始 GISTEmbed 论文，但在某些检索任务中设置为 False 可能会产生更好的结果。
gather_across_devices – 如果为 True，则在计算损失之前聚合所有设备上的嵌入。在多 GPU 训练时推荐使用，因为它允许更大的批次大小，但由于通信开销可能会减慢训练速度，并可能导致内存不足错误。

参考文献

智能回复的高效自然语言回复建议，第 4.4 节: https://hugging-face.cn/papers/1705.00652
内存受限设置下扩展深度对比学习批次大小: https://hugging-face.cn/papers/2101.06983
GISTEmbed: 用于文本嵌入微调的样本内负样本引导选择 https://hugging-face.cn/papers/2402.16829

要求

(锚点, 正样本) 对或 (锚点, 正样本, 负样本对)
应与大批次大小一起使用以获得卓越性能，但训练时间比 MultipleNegativesRankingLoss 慢。

输入

文本	标签
(锚点, 正例) 对	无
(锚点, 正例, 负例) 三元组	无
(anchor, positive, negative_1, ..., negative_n)	无

建议

使用 BatchSamplers.NO_DUPLICATES (docs) 确保批内负样本中没有与锚点或正样本重复的。

关系

等效于 GISTEmbedLoss，但具有缓存功能，允许更高的批次大小。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
guide = SentenceTransformer("all-MiniLM-L6-v2")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.CachedGISTEmbedLoss(
    model,
    guide,
    mini_batch_size=64,
    margin_strategy="absolute",   # or "relative" (e.g., margin=0.05 for max. 95% of positive similarity)
    margin=0.1
)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

MSELoss

class sentence_transformers.losses.MSELoss(model: SentenceTransformer)[source]

计算计算出的句子嵌入与目标句子嵌入之间的 MSE 损失。此损失用于将句子嵌入扩展到新语言，如我们发表的《使用知识蒸馏使单语句子嵌入多语言化》中所述。

有关示例，请参阅蒸馏文档中关于将语言模型扩展到新语言的部分。

参数:: model – SentenceTransformer模型

参考文献

使用知识蒸馏使单语句子嵌入多语言化: https://hugging-face.cn/papers/2004.09813
训练 > 模型蒸馏
训练 > 多语言模型

要求

通常在知识蒸馏设置中，使用经过微调的教师模型 M

输入

文本	标签
句子	模型句子嵌入
句子_1, 句子_2, …, 句子_N	模型句子嵌入

关系

MarginMSELoss 等同于此损失，但通过负样本对添加了 margin。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

student_model = SentenceTransformer("microsoft/mpnet-base")
teacher_model = SentenceTransformer("all-mpnet-base-v2")
train_dataset = Dataset.from_dict({
    "english": ["The first sentence",  "The second sentence", "The third sentence",  "The fourth sentence"],
    "french": ["La première phrase",  "La deuxième phrase", "La troisième phrase",  "La quatrième phrase"],
})

def compute_labels(batch):
    return {
        "label": teacher_model.encode(batch["english"])
    }

train_dataset = train_dataset.map(compute_labels, batched=True)
loss = losses.MSELoss(student_model)

trainer = SentenceTransformerTrainer(
    model=student_model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

MarginMSELoss

class sentence_transformers.losses.MarginMSELoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, similarity_fct=<function pairwise_dot_score>)[source]

计算 |sim(Query, Pos) - sim(Query, Neg)| 与 |gold_sim(Query, Pos) - gold_sim(Query, Neg)| 之间的 MSE 损失。默认情况下，sim() 是点积。gold_sim 通常是来自教师模型的相似度分数。

与 MultipleNegativesRankingLoss 相比，这两个段落不必严格为正样本和负样本，两者都可以与给定查询相关或不相关。这可能是 MarginMSELoss 相对于 MultipleNegativesRankingLoss 的一个优势，但请注意 MarginMSELoss 训练速度慢得多。使用 MultipleNegativesRankingLoss，在批次大小为 64 时，我们将一个查询与 128 个段落进行比较。使用 MarginMSELoss，我们仅将一个查询与两个段落进行比较。MarginMSELoss 也可以使用多个负样本，但训练速度会更慢。

参数:

model – SentenceTransformer模型
similarity_fct – 使用哪种相似度函数。

参考文献

有关更多详情，请参阅 https://hugging-face.cn/papers/2010.02666。
训练示例 > MS MARCO
无监督学习 > 领域自适应

要求

(查询, 段落一, 段落二) 三元组或 (查询, 正样本, 负样本_1, …, 负样本_n)
通常在知识蒸馏设置中，使用经过微调的教师模型 M

输入

文本	标签
(查询, 段落_一, 段落_二) 三元组	M(查询, 段落一) - M(查询, 段落二)
(查询, 段落_一, 段落_二) 三元组	[M(查询, 段落一), M(查询, 段落二)]
(查询, 正例, 负例_1, ..., 负例_n)	[M(查询, 正样本) - M(查询, 负样本_i) for i in 1..n]
(查询, 正例, 负例_1, ..., 负例_n)	[M(查询, 正样本), M(查询, 负样本_1), …, M(查询, 负样本_n)]

关系

MSELoss 与此损失类似，但没有通过负样本对添加 margin。

示例

带有金标签，例如，如果您有句子硬分数。想象一下，您希望模型将具有相似“质量”的句子嵌入在一起。如果“文本 1”的质量为 5 分（满分 5 分），“文本 2”的质量为 1 分（满分 5 分），“文本 3”的质量为 3 分（满分 5 分），那么一对句子的相似度可以定义为质量分数的差异。因此，“文本 1”和“文本 2”之间的相似度为 4，“文本 1”和“文本 3”之间的相似度为 2。如果我们将此作为“教师模型”，则标签变为 similarity(“文本 1”, “文本 2”) - similarity(“文本 1”, “文本 3”) = 4 - 2 = 2。

正值表示第一个段落比第二个段落与查询更相似，而负值表示相反。

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "text1": ["It's nice weather outside today.", "He drove to work."],
    "text2": ["It's so sunny.", "He took the car to work."],
    "text3": ["It's very sunny.", "She walked to the store."],
    "label": [0.1, 0.8],
})
loss = losses.MarginMSELoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

我们也可以使用教师模型来计算相似度分数。在这种情况下，我们可以使用教师模型来计算相似度分数并将其用作银标签。这通常用于知识蒸馏。

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

student_model = SentenceTransformer("microsoft/mpnet-base")
teacher_model = SentenceTransformer("all-mpnet-base-v2")
train_dataset = Dataset.from_dict({
    "query": ["It's nice weather outside today.", "He drove to work."],
    "passage1": ["It's so sunny.", "He took the car to work."],
    "passage2": ["It's very sunny.", "She walked to the store."],
})

def compute_labels(batch):
    emb_queries = teacher_model.encode(batch["query"])
    emb_passages1 = teacher_model.encode(batch["passage1"])
    emb_passages2 = teacher_model.encode(batch["passage2"])
    return {
        "label": teacher_model.similarity_pairwise(emb_queries, emb_passages1) - teacher_model.similarity_pairwise(emb_queries, emb_passages2)
    }

train_dataset = train_dataset.map(compute_labels, batched=True)
loss = losses.MarginMSELoss(student_model)

trainer = SentenceTransformerTrainer(
    model=student_model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

我们也可以在知识蒸馏过程中使用多个负样本。

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset
import torch

student_model = SentenceTransformer("microsoft/mpnet-base")
teacher_model = SentenceTransformer("all-mpnet-base-v2")

train_dataset = Dataset.from_dict(
    {
        "query": ["It's nice weather outside today.", "He drove to work."],
        "passage1": ["It's so sunny.", "He took the car to work."],
        "passage2": ["It's very cold.", "She walked to the store."],
        "passage3": ["Its rainy", "She took the bus"],
    }
)


def compute_labels(batch):
    emb_queries = teacher_model.encode(batch["query"])
    emb_passages1 = teacher_model.encode(batch["passage1"])
    emb_passages2 = teacher_model.encode(batch["passage2"])
    emb_passages3 = teacher_model.encode(batch["passage3"])
    return {
        "label": torch.stack(
            [
                teacher_model.similarity_pairwise(emb_queries, emb_passages1)
                - teacher_model.similarity_pairwise(emb_queries, emb_passages2),
                teacher_model.similarity_pairwise(emb_queries, emb_passages1)
                - teacher_model.similarity_pairwise(emb_queries, emb_passages3),
            ],
            dim=1,
        )
    }


train_dataset = train_dataset.map(compute_labels, batched=True)
loss = losses.MarginMSELoss(student_model)

trainer = SentenceTransformerTrainer(model=student_model, train_dataset=train_dataset, loss=loss)
trainer.train()

MatryoshkaLoss

class sentence_transformers.losses.MatryoshkaLoss(model: SentenceTransformer, loss: Module, matryoshka_dims: Sequence[int], matryoshka_weights: Sequence[float] | Sequence[int] | None = None, n_dims_per_step: int = -1)[source]

MatryoshkaLoss 可被视为一种损失“修饰符”，它允许您在各种不同的嵌入维度上使用其他损失函数。当您想要训练一个模型时，用户可以选择降低嵌入维度以提高其嵌入比较速度和成本，这非常有用。

此损失还与 Cached... 损失兼容，后者是批内负样本损失，允许更大的批次大小。更大的批次大小允许更多的负样本，并且通常会产生更强的模型。

参数:

model – SentenceTransformer 模型
loss – 要使用的损失函数，例如 MultipleNegativesRankingLoss、CoSENTLoss 等。
matryoshka_dims – 用于损失函数的一系列嵌入维度，例如 [768, 512, 256, 128, 64]。
matryoshka_weights – 用于损失函数的一系列权重，例如 [1, 1, 1, 1, 1]。如果为 None，则所有维度的权重将设置为 1。
n_dims_per_step – 每个步骤要使用的维度数量。如果为 -1，则使用所有维度。如果 > 0，则每个步骤使用 n_dims_per_step 个随机采样的维度。默认值为 -1。

参考文献

该概念在此论文中引入：https://hugging-face.cn/papers/2205.13147
套娃嵌入 (Matryoshka Embeddings)

输入

文本	标签
任何	任何

关系

Matryoshka2dLoss 将此损失与 AdaptiveLayerLoss 结合使用，从而允许
层减少以实现更快的推理。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.MultipleNegativesRankingLoss(model)
loss = losses.MatryoshkaLoss(model, loss, [768, 512, 256, 128, 64])

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

Matryoshka2dLoss

class sentence_transformers.losses.Matryoshka2dLoss(model: SentenceTransformer, loss: Module, matryoshka_dims: list[int], matryoshka_weights: list[float | int] | None = None, n_layers_per_step: int = 1, n_dims_per_step: int = 1, last_layer_weight: float = 1.0, prior_layers_weight: float = 1.0, kl_div_weight: float = 1.0, kl_temperature: float = 0.3)[source]

Matryoshka2dLoss 可被视为一种损失“修饰符”，它结合了 AdaptiveLayerLoss 和 MatryoshkaLoss。这允许您训练一个嵌入模型，该模型 1) 允许用户指定要使用的模型层数，以及 2) 允许用户指定要使用的输出维度。

前者对于希望用户可以选择降低使用的层数以提高其推理速度和内存使用的情况很有用，后者对于希望用户可以选择降低输出维度以提高其下游任务（例如检索）效率或降低存储成本的情况很有用。

请注意，此损失默认使用 n_layers_per_step=1 和 n_dims_per_step=1，遵循原始 2DMSE 实现。

参数:

model – SentenceTransformer 模型
loss – 要使用的损失函数，例如 MultipleNegativesRankingLoss、CoSENTLoss 等。
matryoshka_dims – 用于损失函数的一系列嵌入维度，例如 [768, 512, 256, 128, 64]。
matryoshka_weights – 用于损失函数的一系列权重，例如 [1, 1, 1, 1, 1]。如果为 None，则所有维度的权重将设置为 1。
n_layers_per_step – 每个步骤要使用的层数。如果为 -1，则使用所有层。如果 > 0，则每个步骤使用 n_layers_per_step 个随机采样的层，与始终使用的最后一层分开。2DMSE 论文使用 n_layers_per_step=1。默认值为 -1。
n_dims_per_step – 每个步骤要使用的维度数量。如果为 -1，则使用所有维度。如果 > 0，则每个步骤使用 n_dims_per_step 个随机采样的维度。默认值为 -1。
last_layer_weight – 最后一层损失的权重。增加此值可更侧重于使用所有层时的性能。默认值为 1.0。
prior_layers_weight – 前面层损失的权重。增加此值可更侧重于使用较少层时的性能。默认值为 1.0。
kl_div_weight – 用于 KL 散度损失的权重，该损失用于使前面层与最后一层匹配。增加此值可更侧重于使用较少层时的性能。默认值为 1.0。
kl_temperature – 用于 KL 散度损失的温度。如果为 0，则不使用 KL 散度损失。默认值为 1.0。

参考文献

参见 2D Matryoshka Sentence Embeddings (2DMSE) 论文: https://hugging-face.cn/papers/2402.14776
套娃嵌入 (Matryoshka Embeddings)
自适应层 (Adaptive Layers)

要求

基础损失不能是 CachedMultipleNegativesRankingLoss、CachedMultipleNegativesSymmetricRankingLoss 或 CachedGISTEmbedLoss。

输入

文本	标签
任何	任何

关系

MatryoshkaLoss 用于此损失，它负责维度缩减。
AdaptiveLayerLoss 用于此损失，它负责层缩减。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.MultipleNegativesRankingLoss(model)
loss = losses.Matryoshka2dLoss(model, loss, [768, 512, 256, 128, 64])

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

AdaptiveLayerLoss

class sentence_transformers.losses.AdaptiveLayerLoss(model: SentenceTransformer, loss: nn.Module, n_layers_per_step: int = 1, last_layer_weight: float = 1.0, prior_layers_weight: float = 1.0, kl_div_weight: float = 1.0, kl_temperature: float = 0.3)[source]

AdaptiveLayerLoss 可被视为一种损失“修饰符”，它允许您在 Sentence Transformer 模型的非最终层上使用其他损失函数。当您想要训练一个模型时，用户可以选择降低使用的层数以提高其推理速度和内存使用，这非常有用。

参数:

model – SentenceTransformer 模型
loss – 要使用的损失函数，例如 MultipleNegativesRankingLoss、CoSENTLoss 等。
n_layers_per_step – 每个步骤要使用的层数。如果为 -1，则使用所有层。如果 > 0，则每个步骤使用 n_layers_per_step 个随机采样的层，与始终使用的最后一层分开。2DMSE 论文使用 n_layers_per_step=1。默认值为 1。
last_layer_weight – 最后一层损失的权重。增加此值可更侧重于使用所有层时的性能。默认值为 1.0。
prior_layers_weight – 前面层损失的权重。增加此值可更侧重于使用较少层时的性能。默认值为 1.0。
kl_div_weight – 用于 KL 散度损失的权重，该损失用于使前面层与最后一层匹配。增加此值可更侧重于使用较少层时的性能。默认值为 1.0。
kl_temperature – 用于 KL 散度损失的温度。如果为 0，则不使用 KL 散度损失。默认值为 1.0。

参考文献

该概念受 2DMSE 论文启发: https://hugging-face.cn/papers/2402.14776
自适应层 (Adaptive Layers)

要求

基础损失不能是 CachedMultipleNegativesRankingLoss、CachedMultipleNegativesSymmetricRankingLoss 或 CachedGISTEmbedLoss。

输入

文本	标签
任何	任何

关系

Matryoshka2dLoss 将此损失与 MatryoshkaLoss 结合使用，从而允许
输出维度缩减以实现更快的下游任务（例如检索）。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.MultipleNegativesRankingLoss(model=model)
loss = losses.AdaptiveLayerLoss(model, loss)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

MegaBatchMarginLoss

class sentence_transformers.losses.MegaBatchMarginLoss(model: SentenceTransformer, positive_margin: float = 0.8, negative_margin: float = 0.3, use_mini_batched_version: bool = True, mini_batch_size: int = 50)[source]

给定一个大型批次（例如 500 个或更多示例）的 (锚点_i, 正样本_i) 对，为批次中的每对找到最硬的负样本，即找到 j != i 使得 cos_sim(锚点_i, 正样本_j) 最大。然后从中创建一个三元组 (锚点_i, 正样本_i, 正样本_j)，其中正样本_j 作为此三元组的负样本。

然后像使用三元组损失一样进行训练。

参数:

model – SentenceTransformer模型
positive_margin – 正向间隔，cos(anchor, positive) 应该 > positive_margin
negative_margin – 负向间隔，cos(anchor, negative) 应该 < negative_margin
use_mini_batched_version – 由于大批量需要大量内存，我们可以使用小批量版本。我们将大批量分解为包含更少示例的小批量。
mini_batch_size – 小批量的大小。应该是数据加载器中批量大小的除数。

参考文献

此损失函数受到 ParaNMT 论文的启发：https://www.aclweb.org/anthology/P18-1042/

要求

(锚点, 正例) 对
大批量（500 个或更多示例）

输入

文本	标签
(锚点, 正例) 对	无

建议

使用 BatchSamplers.NO_DUPLICATES (docs) 确保批内负样本中没有与锚点或正样本重复的。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainingArguments, SentenceTransformerTrainer, losses
from datasets import Dataset

train_batch_size = 250
train_mini_batch_size = 32

model = SentenceTransformer('all-MiniLM-L6-v2')
train_dataset = Dataset.from_dict({
    "anchor": [f"This is sentence number {i}" for i in range(500)],
    "positive": [f"This is sentence number {i}" for i in range(1, 501)],
})
loss = losses.MegaBatchMarginLoss(model=model, mini_batch_size=train_mini_batch_size)

args = SentenceTransformerTrainingArguments(
    output_dir="output",
    per_device_train_batch_size=train_batch_size,
)
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

MultipleNegativesRankingLoss (多负例排序损失)

如果您只有正向对，例如，只有相似文本对（如释义对、重复问题对、(查询, 响应) 对或 (源语言, 目标语言) 对），那么 MultipleNegativesRankingLoss 是一个很棒的损失函数。

class sentence_transformers.losses.MultipleNegativesRankingLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, scale: float = 20.0, similarity_fct=<function cos_sim>, gather_across_devices: bool = False)[source]

给定一个 (anchor, positive) 对或 (anchor, positive, negative) 三元组列表，此损失函数优化以下目标

给定一个锚点（例如一个问题），在批次中的每一个正例和负例（例如所有答案）中，为对应的正例（即答案）赋予最高的相似度。

如果您提供可选的负例，它们都将用作额外的选项，模型必须从中选择正确的正例。在合理范围内，这种“选择”越困难，模型就会变得越强大。因此，更高的批量大小会导致更多的批量内负例，从而提高性能（达到一定程度）。

此损失函数非常适合训练检索设置的嵌入，其中您有正向对（例如（查询，答案）），因为它将在每个批量中随机采样 n-1 个负向文档。

此损失函数也称为 InfoNCE 损失、SimCSE 损失、带批内负例的交叉熵损失，或简称为批内负例损失。

参数:

model – SentenceTransformer 模型
scale – 相似度函数的输出乘以比例值。在某些文献中，缩放参数被称为温度，它是比例的倒数。简而言之：scale = 1 / temperature，因此 scale=20.0 等价于 temperature=0.05。
similarity_fct – 句子嵌入之间的相似度函数。默认为 cos_sim。也可以设置为点积（然后将 scale 设置为 1）。
gather_across_devices – 如果为 True，则在计算损失之前聚合所有设备上的嵌入。在多 GPU 训练时推荐使用，因为它允许更大的批次大小，但由于通信开销可能会减慢训练速度，并可能导致内存不足错误。

参考文献

智能回复的高效自然语言回复建议，第 4.4 节: https://hugging-face.cn/papers/1705.00652
训练示例 > 自然语言推理
训练示例 > 释义数据
训练示例 > Quora 重复问题
训练示例 > MS MARCO
无监督学习 > SimCSE
无监督学习 > GenQ

要求

(anchor, positive) 对或 (anchor, positive, negative) 三元组

输入

文本	标签
(锚点, 正例) 对	无
(锚点, 正例, 负例) 三元组	无
(anchor, positive, negative_1, ..., negative_n)	无

建议

使用 BatchSamplers.NO_DUPLICATES (docs) 确保批内负样本中没有与锚点或正样本重复的。

关系

CachedMultipleNegativesRankingLoss 等价于此损失，但它使用缓存，允许在不增加额外内存使用的情况下实现更高的批量大小（从而获得更好的性能）。然而，它稍微慢一些。
MultipleNegativesSymmetricRankingLoss 等价于此损失，但额外包含一个损失项。
GISTEmbedLoss 等价于此损失，但使用引导模型来引导批量内负样本选择。GISTEmbedLoss 以一些训练开销为代价产生更强的训练信号。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.MultipleNegativesRankingLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

CachedMultipleNegativesRankingLoss

class sentence_transformers.losses.CachedMultipleNegativesRankingLoss(model: SentenceTransformer, scale: float = 20.0, similarity_fct: callable[[Tensor, Tensor], Tensor] = <function cos_sim>, mini_batch_size: int = 32, gather_across_devices: bool = False, show_progress_bar: bool = False)[source]

GradCache (https://hugging-face.cn/papers/2101.06983) 优化版 MultipleNegativesRankingLoss (https://hugging-face.cn/papers/1705.00652)。

对比学习（此处为我们的 MNRL 损失）与批量内负例在处理大批量时通常因（GPU）内存限制而难以工作。即使使用梯度缩放等批量缩放方法也无法奏效。这是因为批量内负例使得同一批量内的数据点不独立，因此无法将批量分解为小批量。GradCache 是一种智能解决此问题的方法。它通过将计算分为嵌入和损失计算两个阶段来实现目标，这两个阶段都可以通过小批量进行缩放。结果是，恒定大小的内存（例如，适用于批量大小 = 32 的内存）现在可以处理更大的批量（例如 65536）。

详细信息

它首先执行一个没有梯度/计算图的快速嵌入步骤以获取所有嵌入；

计算损失，反向传播到嵌入并缓存相对于嵌入的梯度；

执行第二个带有梯度/计算图的嵌入步骤，并将缓存的梯度连接到反向链中。

注意：所有步骤都以小批量完成。在 GradCache 的原始实现中，(2) 未以小批量完成，并且在批量大小较大时需要大量内存。一个缺点是速度。根据论文，梯度缓存会牺牲大约 20% 的计算时间。

参数:

model – SentenceTransformer 模型
scale – 相似度函数的输出乘以比例值。在某些文献中，缩放参数被称为温度，它是比例的倒数。简而言之：scale = 1 / temperature，因此 scale=20.0 等价于 temperature=0.05。
similarity_fct – 句子嵌入之间的相似度函数。默认为 cos_sim。也可以设置为点积（然后将 scale 设置为 1）。
mini_batch_size – 前向传播的 mini-batch 大小，表示训练和评估期间实际使用的内存量。mini-batch 大小越大，训练的内存效率越高，但训练速度会越慢。建议将其设置为 GPU 内存允许的最大值。默认值为 32。
gather_across_devices – 如果为 True，则在计算损失之前聚合所有设备上的嵌入。在多 GPU 训练时推荐使用，因为它允许更大的批次大小，但由于通信开销可能会减慢训练速度，并可能导致内存不足错误。
show_progress_bar – 如果为 True，则在训练期间显示 mini-batch 的进度条。默认值为 False。

参考文献

智能回复的高效自然语言回复建议，第 4.4 节: https://hugging-face.cn/papers/1705.00652
内存受限设置下扩展深度对比学习批次大小: https://hugging-face.cn/papers/2101.06983

要求

(锚点, 正样本) 对或 (锚点, 正样本, 负样本对)
应与大的 per_device_train_batch_size 和小的 mini_batch_size 一起使用以获得卓越性能，但训练时间比 MultipleNegativesRankingLoss 慢。

输入

文本	标签
(锚点, 正例) 对	无
(锚点, 正例, 负例) 三元组	无
(anchor, positive, negative_1, ..., negative_n)	无

建议

使用 BatchSamplers.NO_DUPLICATES (docs) 确保批内负样本中没有与锚点或正样本重复的。

关系

等价于 MultipleNegativesRankingLoss，但增加了缓存机制，可在不增加额外内存使用的情况下实现更高的批量大小（从而获得更好的性能）。此损失的训练速度也比 MultipleNegativesRankingLoss 慢大约 2 到 2.4 倍。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.CachedMultipleNegativesRankingLoss(model, mini_batch_size=64)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

MultipleNegativesSymmetricRankingLoss

class sentence_transformers.losses.MultipleNegativesSymmetricRankingLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, scale: float = 20.0, similarity_fct=<function cos_sim>, gather_across_devices: bool = False)[source]

给定一对 (anchor, positive) 列表，此损失函数将以下两个损失相加

前向损失：给定一个锚点，从批量中所有正样本中找到相似度最高的样本。这等价于 MultipleNegativesRankingLoss。
后向损失：给定一个正样本，从批量中所有锚点中找到相似度最高的样本。

例如，对于问答对，MultipleNegativesRankingLoss 只计算在给定问题的情况下找到答案的损失，但 MultipleNegativesSymmetricRankingLoss 额外计算在给定答案的情况下找到问题的损失。

注意：如果传递三元组，负向条目将被忽略。仅为正向条目搜索锚点。

参数:

model – SentenceTransformer 模型
scale – 相似度函数的输出乘以比例值。在某些文献中，缩放参数被称为温度，它是比例的倒数。简而言之：scale = 1 / temperature，因此 scale=20.0 等价于 temperature=0.05。
similarity_fct – 句子嵌入之间的相似度函数。默认为 cos_sim。也可以设置为点积（然后将 scale 设置为 1）。
gather_across_devices – 如果为 True，则在计算损失之前聚合所有设备上的嵌入。在多 GPU 训练时推荐使用，因为它允许更大的批次大小，但由于通信开销可能会减慢训练速度，并可能导致内存不足错误。

要求

(锚点, 正例) 对

输入

文本	标签
(锚点, 正例) 对	无

建议

使用 BatchSamplers.NO_DUPLICATES (docs) 确保批内负样本中没有与锚点或正样本重复的。

关系

与 MultipleNegativesRankingLoss 类似，但额外包含一个损失项。
CachedMultipleNegativesSymmetricRankingLoss 等价于此损失，但它使用缓存，允许在不增加额外内存使用的情况下实现更高的批量大小（从而获得更好的性能）。然而，它稍微慢一些。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.MultipleNegativesSymmetricRankingLoss(model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

CachedMultipleNegativesSymmetricRankingLoss

class sentence_transformers.losses.CachedMultipleNegativesSymmetricRankingLoss(model: SentenceTransformer, scale: float = 20.0, similarity_fct: callable[[Tensor, Tensor], Tensor] = <function cos_sim>, mini_batch_size: int = 32, gather_across_devices: bool = False, show_progress_bar: bool = False)[source]

GradCache (https://hugging-face.cn/papers/2101.06983) 优化的 MultipleNegativesSymmetricRankingLoss (MNSRL) 版本。

给定一个 (anchor, positive) 对列表，MNSRL 累加以下两个损失

前向损失：给定一个锚点，从批次中所有正样本中找到相似度最高的样本。
后向损失：给定一个正样本，从批量中所有锚点中找到相似度最高的样本。

例如，对于问答对，前向损失找到给定问题的答案，后向损失找到给定答案的问题。这种损失在对称任务中很常见，例如语义文本相似性。

缓存修改允许在不增加内存使用的情况下实现大批量（提供更好的训练信号），从而使用常规硬件即可达到最佳训练信号。

注意：如果传递三元组，负向条目将被忽略。仅为正向条目搜索锚点。

参数:

model – SentenceTransformer 模型
scale – 相似度函数的输出乘以比例值。在某些文献中，缩放参数被称为温度，它是比例的倒数。简而言之：scale = 1 / temperature，因此 scale=20.0 等价于 temperature=0.05。
similarity_fct – 句子嵌入之间的相似度函数。默认为 cos_sim。也可以设置为点积（然后将 scale 设置为 1）。
mini_batch_size – 前向传播的 mini-batch 大小，表示训练和评估期间实际使用的内存量。mini-batch 大小越大，训练的内存效率越高，但训练速度会越慢。建议将其设置为 GPU 内存允许的最大值。默认值为 32。
gather_across_devices – 如果为 True，则在计算损失之前聚合所有设备上的嵌入。在多 GPU 训练时推荐使用，因为它允许更大的批次大小，但由于通信开销可能会减慢训练速度，并可能导致内存不足错误。
show_progress_bar – 如果为 True，则在训练期间显示 mini-batch 的进度条。默认值为 False。

要求

(锚点, 正例) 对
应与大批量一起使用以获得卓越性能，但训练时间比非缓存版本慢

输入

文本	标签
(锚点, 正例) 对	无

建议

使用 BatchSamplers.NO_DUPLICATES (docs) 确保批内负样本中没有与锚点或正样本重复的。

关系

类似于 MultipleNegativesRankingLoss，但额外增加了一个对称损失项和缓存机制。
灵感来源于 CachedMultipleNegativesRankingLoss，并根据对称损失计算进行了调整。

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
})
loss = losses.CachedMultipleNegativesSymmetricRankingLoss(model, mini_batch_size=32)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

参考文献

智能回复的高效自然语言回复建议，第 4.4 节: https://hugging-face.cn/papers/1705.00652
内存受限设置下扩展深度对比学习批次大小: https://hugging-face.cn/papers/2101.06983

SoftmaxLoss

class sentence_transformers.losses.SoftmaxLoss(model: SentenceTransformer, sentence_embedding_dimension: int, num_labels: int, concatenation_sent_rep: bool = True, concatenation_sent_difference: bool = True, concatenation_sent_multiplication: bool = False, loss_fct: Callable = CrossEntropyLoss())[source]

此损失函数用于我们的 SBERT 出版物 (https://hugging-face.cn/papers/1908.10084) 中，用于在 NLI 数据上训练 SentenceTransformer 模型。它在两个 Transformer 网络的输出之上添加了一个 softmax 分类器。

根据 https://hugging-face.cn/papers/2004.09813，MultipleNegativesRankingLoss 是一种替代损失函数，通常能产生更好的结果。

参数:

model (SentenceTransformer) – SentenceTransformer 模型。
sentence_embedding_dimension (int) – 句子嵌入的维度。
num_labels (int) – 不同标签的数量。
concatenation_sent_rep (bool) – 是否将向量 u,v 拼接用于 softmax 分类器。默认为 True。
concatenation_sent_difference (bool) – 是否将 abs(u-v) 添加到 softmax 分类器。默认为 True。
concatenation_sent_multiplication (bool) – 是否将 u*v 添加到 softmax 分类器。默认为 False。
loss_fct (Callable) – 自定义 pytorch 损失函数。如果未设置，则使用 nn.CrossEntropyLoss()。默认为 nn.CrossEntropyLoss()。

参考文献

Sentence-BERT：使用 Siamese BERT-Networks 的句子嵌入：https://hugging-face.cn/papers/1908.10084
训练示例 > 自然语言推理

要求

带有类标签的句子对

输入

文本	标签
(句子_A, 句子_B) 对	类别

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "sentence1": [
        "A person on a horse jumps over a broken down airplane.",
        "A person on a horse jumps over a broken down airplane.",
        "A person on a horse jumps over a broken down airplane.",
        "Children smiling and waving at camera",
    ],
    "sentence2": [
        "A person is training his horse for a competition.",
        "A person is at a diner, ordering an omelette.",
        "A person is outdoors, on a horse.",
        "There are children present.",
    ],
    "label": [1, 2, 0, 0],
})
loss = losses.SoftmaxLoss(model, model.get_sentence_embedding_dimension(), num_labels=3)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

TripletLoss

class sentence_transformers.losses.TripletLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, distance_metric=<function TripletDistanceMetric.<lambda>>, triplet_margin: float = 5)[source]

此类实现了三元组损失。给定一个 (锚点, 正样本, 负样本) 三元组，损失函数会最小化锚点与正样本之间的距离，同时最大化锚点与负样本之间的距离。它计算以下损失函数

loss = max(||anchor - positive|| - ||anchor - negative|| + margin, 0).

裕度是一个重要的超参数，需要相应调整。

参数:

model – SentenceTransformer模型
distance_metric – 计算两个嵌入之间距离的函数。TripletDistanceMetric 类包含可使用的常见距离度量。
triplet_margin – 负例与锚点的距离应至少比正例与锚点的距离远这么多。

参考文献

欲了解更多详情，请参阅：https://en.wikipedia.org/wiki/Triplet_loss

要求

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

输入

文本	标签
(锚点, 正例, 负例) 三元组	无

示例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset

model = SentenceTransformer("microsoft/mpnet-base")
train_dataset = Dataset.from_dict({
    "anchor": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to the office."],
    "negative": ["It's quite rainy, sadly.", "She walked to the store."],
})
loss = losses.TripletLoss(model=model)

trainer = SentenceTransformerTrainer(
    model=model,
    train_dataset=train_dataset,
    loss=loss,
)
trainer.train()

DistillKLDivLoss

class sentence_transformers.losses.DistillKLDivLoss(model: ~sentence_transformers.SentenceTransformer.SentenceTransformer, similarity_fct=<function pairwise_dot_score>, temperature: float = 1.0)[source]

计算学生模型和教师模型相似度得分导出的概率分布之间的 KL 散度损失。默认情况下，相似度使用点积计算。此损失旨在用于知识蒸馏，其中较小的学生模型从更强大的教师模型学习。

该损失计算教师相似度得分的 softmax 概率和学生模型的 log-softmax 概率，然后计算这些分布之间的 KL 散度。

参数:

model – SentenceTransformer 模型（学生模型）
similarity_fct – 学生模型使用的相似度函数
temperature – 用于软化概率分布的温度参数（温度越高 = 分布越平滑）。温度为 1.0 不会缩放分数。注意：在 v5.0.1 版本中，默认温度从 2.0 更改为 1.0。

参考文献

有关更多详细信息，请参阅 https://hugging-face.cn/papers/2010.11386

要求

(query, positive, negative_1, …, negative_n) 示例
包含教师模型在查询-正例和查询-负例对之间分数的标签

输入

文本	标签
(查询, 正例, 负例)	[Teacher(query, positive), Teacher(query, negative)]
(查询, 正例, 负例_1, ..., 负例_n)	[Teacher(query, positive), Teacher(query, negative_i)…]

关系

类似于 MarginMSELoss 但使用 KL 散度而非 MSE
更适合保留排名很重要的蒸馏任务

示例

使用教师模型计算相似度分数进行蒸馏

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset
import torch

student_model = SentenceTransformer("microsoft/mpnet-base")
teacher_model = SentenceTransformer("all-mpnet-base-v2")
train_dataset = Dataset.from_dict({
    "query": ["It's nice weather outside today.", "He drove to work."],
    "positive": ["It's so sunny.", "He took the car to work."],
    "negative": ["It's very cold.", "She walked to the store."],
})

def compute_labels(batch):
    emb_queries = teacher_model.encode(batch["query"])
    emb_positives = teacher_model.encode(batch["positive"])
    emb_negatives = teacher_model.encode(batch["negative"])

    pos_scores = teacher_model.similarity_pairwise(emb_queries, emb_positives)
    neg_scores = teacher_model.similarity_pairwise(emb_queries, emb_negatives)

    # Stack the scores for positive and negative pairs
    return {
        "label": torch.stack([pos_scores, neg_scores], dim=1)
    }

train_dataset = train_dataset.map(compute_labels, batched=True)
loss = losses.DistillKLDivLoss(student_model)

trainer = SentenceTransformerTrainer(model=student_model, train_dataset=train_dataset, loss=loss)
trainer.train()

有多个负例

from sentence_transformers import SentenceTransformer, SentenceTransformerTrainer, losses
from datasets import Dataset
import torch

student_model = SentenceTransformer("microsoft/mpnet-base")
teacher_model = SentenceTransformer("all-mpnet-base-v2")

train_dataset = Dataset.from_dict(
    {
        "query": ["It's nice weather outside today.", "He drove to work."],
        "positive": ["It's so sunny.", "He took the car to work."],
        "negative1": ["It's very cold.", "She walked to the store."],
        "negative2": ["Its rainy", "She took the bus"],
    }
)


def compute_labels(batch):
    emb_queries = teacher_model.encode(batch["query"])
    emb_positives = teacher_model.encode(batch["positive"])
    emb_negatives1 = teacher_model.encode(batch["negative1"])
    emb_negatives2 = teacher_model.encode(batch["negative2"])

    pos_scores = teacher_model.similarity_pairwise(emb_queries, emb_positives)
    neg_scores1 = teacher_model.similarity_pairwise(emb_queries, emb_negatives1)
    neg_scores2 = teacher_model.similarity_pairwise(emb_queries, emb_negatives2)

    # Stack the scores for positive and multiple negative pairs
    return {
        "label": torch.stack([pos_scores, neg_scores1, neg_scores2], dim=1)
    }

train_dataset = train_dataset.map(compute_labels, batched=True)
loss = losses.DistillKLDivLoss(student_model)

trainer = SentenceTransformerTrainer(model=student_model, train_dataset=train_dataset, loss=loss)
trainer.train()