MS MARCO

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

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

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

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

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

SparseMultipleNegativesRankingLoss

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

为了进一步改进训练，我们使用批内负样本

MultipleNegativesRankingLoss

我们将所有 查询、正样本段落 和 负样本段落 嵌入到向量空间中。匹配的 (query_i, positive_passage_i) 应该接近，而一个 query 与批次中所有其他三元组的所有其他（正/负）段落之间应该有很大的距离。对于批次大小为 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"]