MS MARCO

MS MARCO Passage Ranking 是一个大型数据集,用于训练信息检索模型。它包含约 50 万个来自 Bing 搜索引擎的真实搜索查询,以及回答这些查询的相关文本段落。

本页展示了如何在此数据集上训练 Sentence Transformer 模型,以便它能用于根据查询(关键词、短语或问题)搜索文本段落。

如果您对如何使用这些模型感兴趣,请参阅应用 - 检索与重排序

预训练模型可用,您可以直接使用它们,无需训练自己的模型。欲了解更多信息,请参阅:预训练模型 > MSMARCO 段落模型

Bi-Encoder

为了从大量文档集合中检索合适的文档,我们必须使用 Sentence Transformer (也称为 bi-encoder) 模型。文档被独立编码为固定大小的嵌入。查询被嵌入到相同的向量空间中。然后可以通过使用余弦相似度或点积来找到相关文档。

BiEncoder

本页描述了在 MS MARCO 数据集上训练 bi-encoder 的两种策略

MultipleNegativesRankingLoss

训练代码:train_bi-encoder_mnrl.py

当我们使用 MultipleNegativesRankingLoss 时,我们提供三元组:(query, positive_passage, negative_passage),其中 positive_passage 是与查询相关的段落,negative_passage 是与查询不相关的段落。我们计算语料库中所有查询、正例段落和负例段落的嵌入,然后优化以下目标:(query, positive_passage) 对在向量空间中必须接近,而 (query, negative_passage) 应在向量空间中相距较远。

为了进一步改进训练,我们使用批内负样本(in-batch negatives)

MultipleNegativesRankingLoss

我们将所有 queriespositive_passagesnegative_passages 嵌入到向量空间中。匹配的 (query_i, positive_passage_i) 应该接近,而一个 query 与批次中所有其他三元组的所有其他(正/负)段落之间应该有很大的距离。对于批次大小为 64 的情况,我们将一个查询与 64+64=128 个段落进行比较,其中只有一个段落应该接近,而其他 127 个段落应该在向量空间中相距较远。

一种改进训练的方法是选择真正好的负样本,也称为困难负样本(hard negative):负样本应该看起来与正样本非常相似,但它不应与查询相关。

我们通过以下方式找到这些困难负样本:我们使用现有的检索系统(例如词汇搜索和其他双编码器检索系统),并为每个查询找到最相关的段落。然后,我们使用强大的 cross-encoder/ms-marco-MiniLM-L6-v2 Cross-Encoder 来对找到的 (query, passage) 对进行评分。我们在MS MARCO Mined Triplet 数据集集合中提供了 1.6 亿对这样的评分。

对于 MultipleNegativesRankingLoss,我们必须确保在三元组 (query, positive_passage, negative_passage) 中,negative_passage 确实与查询不相关。MS MARCO 数据集不幸地高度冗余,尽管平均而言只有一个段落被标记为与查询相关,但它实际上包含许多人类会认为相关的段落。我们必须确保这些段落不作为负样本传递:我们通过确保相关段落和挖掘的困难负样本之间存在交叉编码器分数阈值来实现这一点。默认情况下,我们将阈值设置为 3:如果 (query, positive_passage) 从交叉编码器获得 9 分,那么我们只会考虑交叉编码器得分低于 6 分的负样本。这个阈值确保我们实际在三元组中使用负样本。

您可以通过访问MS MARCO Mined Triplet 数据集集合中的任何数据集并使用 triplet-hard 子集来找到此数据。所有数据集加起来,这指的是 1.757 亿个三元组。原始数据可以在这里找到。使用以下方式加载一些:

from datasets import load_dataset

train_dataset = load_dataset("sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1", "triplet-hard", split="train")
# Dataset({
#     features: ['query', 'positive', 'negative'],
#     num_rows: 11662655
# })
print(train_dataset[0])
# {'query': 'what are the liberal arts?', 'positive': 'liberal arts. 1. the academic course of instruction at a college intended to provide general knowledge and comprising the arts, humanities, natural sciences, and social sciences, as opposed to professional or technical subjects.', 'negative': "Rather than preparing students for a specific career, liberal arts programs focus on cultural literacy and hone communication and analytical skills. They often cover various disciplines, ranging from the humanities to social sciences. 1  Program Levels in Liberal Arts: Associate degree, Bachelor's degree, Master's degree."}

MarginMSE

训练代码:train_bi-encoder_margin-mse.py

MarginMSELoss 基于 Hofstätter 等人的论文。与使用 MultipleNegativesRankingLoss 训练时类似,我们可以使用三元组:(query, passage1, passage2)。然而,与 MultipleNegativesRankingLoss 不同,passage1passage2 不必严格地是正例/负例,两者都可以与给定查询相关或不相关。

然后我们计算 Cross-Encoder(query, passage1)(query, passage2) 的分数。我们在 msmarco-hard-negatives 数据集中提供了 1.6 亿对这样的分数。然后我们计算距离:CE_distance = CEScore(query, passage1) - CEScore(query, passage2)

对于我们的 Sentence Transformer(例如 bi-encoder)训练,我们将 querypassage1passage2 编码为嵌入,然后测量 (query, passage1)(query, passage2) 之间的点积。同样,我们测量距离:BE_distance = DotScore(query, passage1) - DotScore(query, passage2)

然后我们希望确保双编码器预测的距离接近交叉编码器预测的距离,即,我们优化 CE_distanceBE_distance 之间的均方误差 (MSE)。

MarginMSELoss 相较于 MultipleNegativesRankingLoss 的一个优势是我们不需要一个 positivenegative 段落。如前所述,MS MARCO 是冗余的,许多段落包含相同或相似的内容。使用 MarginMSELoss,我们可以毫无问题地在两个相关段落上进行训练:在这种情况下,CE_distance 会更小,我们期望我们的双编码器也将这两个段落在向量空间中放置得更近。

MarginMSELoss 的一个缺点是训练时间更慢:我们需要更多的 epoch 才能获得好的结果。在 MultipleNegativesRankingLoss 中,当批量大小为 64 时,我们将一个查询与 128 个段落进行比较。而使用 MarginMSELoss,我们只将一个查询与两个段落进行比较。