MS MARCO

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

本页面展示如何在此数据集上训练稀疏编码器模型（更确切地说是Splade模型），以便它能够用于根据查询（关键词、短语或问题）搜索文本段落。

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

有预训练模型可用，您可以直接使用，无需自行训练。欲了解更多信息，请参阅：预训练模型。

本页面描述了一种在MS MARCO数据集上训练Splade模型的策略。

SparseMultipleNegativesRankingLoss

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

为了进一步改善训练，我们使用批内负例。

MultipleNegativesRankingLoss

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

改善训练的一种方法是选择真正好的负例，也称为难负例：负例应该看起来与正例段落非常相似，但它不应与查询相关。

我们通过以下方式找到这些难负例：我们使用现有的检索系统（例如词汇搜索和其他双编码器检索系统），并为每个查询找到最相关的段落。然后，我们使用强大的cross-encoder/ms-marco-MiniLM-L6-v2 交叉编码器对找到的(查询, 段落)对进行评分。我们在MS MARCO 挖掘三元组数据集集合中提供了1.6亿对这样的评分。

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

您可以通过浏览MS MARCO 挖掘三元组数据集集合中的任何数据集并使用triplet-hard子集来找到这些数据。所有数据集总共包含1.757亿个三元组。原始数据可以在这里找到。在我们的示例中，我们只使用了原始的三元组数据集，因为

from datasets import load_dataset

dataset_size = 100_000  # We only use the first 100k samples for training
print("The dataset has not been fully stored as texts on disk yet. We will do this now.")
corpus = load_dataset("sentence-transformers/msmarco", "corpus", split="train")
corpus = dict(zip(corpus["passage_id"], corpus["passage"]))
queries = load_dataset("sentence-transformers/msmarco", "queries", split="train")
queries = dict(zip(queries["query_id"], queries["query"]))
dataset = load_dataset("sentence-transformers/msmarco", "triplets", split="train")
dataset = dataset.select(range(dataset_size))

def id_to_text_map(batch):
    return {
        "query": [queries[qid] for qid in batch["query_id"]],
        "positive": [corpus[pid] for pid in batch["positive_id"]],
        "negative": [corpus[pid] for pid in batch["negative_id"]],
    }

dataset = dataset.map(id_to_text_map, batched=True, remove_columns=["query_id", "positive_id", "negative_id"])
dataset = dataset.train_test_split(test_size=10_000)
train_dataset = dataset["train"]
eval_dataset = dataset["test"]