MS MARCO¶
MS MARCO Passage Ranking 是一个大型数据集,用于训练信息检索模型。它包含约50万条来自Bing搜索引擎的真实搜索查询,以及回答这些查询的相关文本段落。
本页展示如何训练 Sentence Transformer 模型在这个数据集上,以便它可以用于根据查询(关键词、短语或问题)搜索文本段落。
如果你对如何使用这些模型感兴趣,请参阅 应用 - 检索与重排。
有预训练模型可用,你可以直接使用,无需训练自己的模型。更多信息请参阅:预训练模型 > MSMARCO Passage 模型。
双编码器¶
为了从大型集合中检索合适的文档,我们必须使用 Sentence Transformer(又名双编码器)模型。文档被独立编码为固定大小的嵌入。查询被嵌入到相同的向量空间中。然后可以使用余弦相似度或点积来找到相关文档。
本页描述了两种策略来在 MS MARCO 数据集上训练双编码器:
MultipleNegativesRankingLoss¶
当我们使用 MultipleNegativesRankingLoss
时,我们提供三元组:(query, positive_passage, negative_passage)
,其中 positive_passage
是与查询相关的段落,而 negative_passage
是与查询无关的段落。我们计算语料库中所有查询、正样本段落和负样本段落的嵌入,然后优化以下目标:(query, positive_passage)
对必须在向量空间中接近,而 (query, negative_passage)
应在向量空间中远离。
为了进一步改进训练,我们使用 批次内负样本:
我们将所有 queries
、positive_passages
和 negative_passages
嵌入到向量空间中。匹配的 (query_i, positive_passage_i)
应该接近,而一个 query
与批次中其他所有三元组的(正/负)段落应该有较大的距离。对于批次大小为64,我们将一个查询与64+64=128个段落进行比较,其中只有一个段落应该接近,其余127个段落应在向量空间中远离。
改进训练的一种方法是选择真正好的负样本,也称为 硬负样本:负样本应该看起来非常类似于正样本段落,但它不应该与查询相关。
我们以以下方式找到这些硬负样本:我们使用现有的检索系统(例如词法搜索和其他双编码器检索系统),并为每个查询找到最相关的段落。然后我们使用强大的 cross-encoder/ms-marco-MiniLM-L-6-v2 交叉编码器 对找到的 (query, passage)
对进行评分。我们在 MS MARCO 挖掘三元组数据集集合 中提供了1.6亿对这样的对的评分。
对于 MultipleNegativesRankingLoss
,我们必须确保在三元组 (query, positive_passage, negative_passage)
中,negative_passage
确实与查询无关。遗憾的是,MS MARCO 数据集**高度冗余**,尽管平均每个查询只有一个段落被标记为相关,但实际上它包含许多人类会认为是相关的段落。我们必须确保这些段落**不会被作为负样本传递**:我们通过确保相关段落和挖掘的硬负样本之间在交叉编码器评分中达到一定阈值来实现这一点。默认情况下,我们设置阈值为3:如果 (query, positive_passage)
从交叉编码器得到9分,那么我们将只考虑交叉编码器评分低于6的负样本。此阈值确保我们实际上在三元组中使用了负样本。
你可以通过访问 MS MARCO 挖掘三元组数据集集合 中的任何数据集并使用 triplet-hard
子集来找到这些数据。所有数据集中,这指的是1.757亿个三元组。原始数据可以在这里找到 here。使用以下命令加载部分数据:
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': '什么是文科?', 'positive': '文科。1. 学院提供的旨在提供一般知识的学术课程,包括艺术、人文、自然科学和社会科学,与专业或技术科目相对。', 'negative': "与其为学生准备特定职业,文科课程更注重文化素养,并培养沟通和分析技能。它们通常涵盖各种学科,从人文科学到社会科学。1 文科的课程级别:副学士学位、学士学位、硕士学位。"}
MarginMSE¶
训练代码: train_bi-encoder_margin-mse.py
MarginMSELoss
基于 Hofstätter 等人 的论文。与使用 MultipleNegativesRankingLoss
进行训练时类似,我们可以使用三元组: (查询, 段落1, 段落2)
。然而,与 MultipleNegativesRankingLoss
不同,段落1 和 段落2 不必严格为正/负,两者都可以与给定查询相关或不相关。
然后我们计算 Cross-Encoder 分数,分别为 (查询, 段落1)
和 (查询, 段落2)
。我们在 msmarco-hard-negatives 数据集 中提供了 1.6 亿对这样的分数。然后我们计算距离: CE_distance = CEScore(查询, 段落1) - CEScore(查询, 段落2)
。
对于我们的 Sentence Transformer(例如双编码器)训练,我们将 查询
、段落1
和 段落2
编码为嵌入,然后测量 (查询, 段落1)
和 (查询, 段落2)
之间的点积。同样,我们测量距离: BE_distance = DotScore(查询, 段落1) - DotScore(查询, 段落2)
然后我们希望确保双编码器预测的距离接近交叉编码器预测的距离,即我们优化 CE_distance
和 BE_distance
之间的均方误差 (MSE)。
与 MultipleNegativesRankingLoss
相比,MarginMSELoss
的一个 优势 是我们 不需要 正
和 负
段落。如前所述,MS MARCO 是冗余的,许多段落包含相同或相似的内容。使用 MarginMSELoss
,我们可以在两个相关段落上进行训练而不会出现问题:在这种情况下,CE_distance
将更小,我们期望我们的双编码器也将这两个段落在向量空间中放置得更近。
而 MarginMSELoss
的一个 劣势 是训练时间较慢:我们需要更多的 epoch 才能获得良好的结果。在 MultipleNegativesRankingLoss
中,批量大小为 64 时,我们将一个查询与 128 个段落进行比较。使用 MarginMSELoss
,我们仅将一个查询与两个段落进行比较。