稠密检索的规模艺术：模型、数据与性能的精准匹配

• 论文：https://arxiv.org/pdf/2403.18684
• 代码：GitHub - jingtaozhan/DRScale
• 机构：清华大学
• 领域：稠密检索、Scaling Laws
• 发表：SIGIR2024最佳论文

Abstract

Scaling Laws已经在广泛的任务中被观察到，特别是在语言生成中。先前的研究发现，大型语言模型的性能和模型、数据集的大小存在可预测的模式关系。这有助于我们有效地设计训练策略，特别是在大规模训练越来越需要丰富的计算资源的时候。然而，在密集检索中，这种Scaling Laws尚未得到充分的探索。在本研究中，我们研究了Scaling如何影响密集检索模型的性能。我们使用不同参数量来实现密集检索模型，并使用不同数量的标注数据来训练它们。本文使用对比熵作为评价指标，与离散排序指标相比，它是连续的，因此可以准确地反映模型的性能。结果表明，密集检索模型的性能与模型大小和标注数据量大小具有精确的幂律Scaling关系。此外，我们还证明了Scaling Laws有助于优化训练过程，如解决预算限制下的资源分配问题。我们认为，这些发现有助于理解密集检索模型的Scaling效应，并为未来的研究提供了有意义的指导。

对比熵（Contrastive Entropy）是信息论中的一个概念，用于衡量两个不同概率分布之间的相似性或差异性。与传统的熵（Entropy）不同，熵是一个单一分布的度量，用来衡量随机变量的不确定性。对比熵则关注两个分布之间的关系。对比熵的计算通常涉及将一个分布作为参考分布，然后计算另一个分布相对于参考分布的差异（这里是对比查询和文档的分布差异）。

Methods

问题形式

对于给定的一个语料库，密集检索模型的目标是对于一个具体的查询确定最相关的段落，密集检索模型通过将查询和候选段落映射到一个共享的密集表征空间来实现这一目标。获得表征后，再使用一个打分函数，比如内积或余弦相似度来计算相关度分数。f代表映射函数，θ代表模型参数，则相似度分数可以用以下公式计算：

本文中，仅考虑查询和段落共享编码器的情况。训练数据采用最流行的query-positive passage对的标注方式，获得了n个数据对，其中每个查询可能对应多个匹配段落。

模型

采用transformer架构，得到向量的函数如下，其中W和b为映射头的参数：

为了避免不同预训练任务对实验的影响，对于英文检索，本文选用了Google开源的24个BERT系列模型，从最小的0.5M参数（BERT-Tiny）到最大的82M参数（BERT-Base）。对于中文检索，本文选用了文心系列模型，对于每个模型，我们都统一映射到768维的表征空间从而保持公平对比。

训练数据

对于英文，使用MS MARCO，包含从英文网站上收集的8.8M个段落以及0.5M条查询，每条查询都标注了一个正例段落。同时MS MARCO提供了7000条查询-段落对用于评估。对于中文，本文使用了T2Ranking，包含从真实搜索引擎中收集的300k条查询以及超过2M个段落。

训练设置

本文采用最经典的随机批内负采样策略进行对比学习，损失如下：

本文训练固定的10000步，批大小设置为256。

训练评估

由于一般的协议MAP@K，NDCG@K是离散的，在模型变化时不一定能准确反映出全部的变化（比如排序从12到11对于R@10没有区别），因此本文提出使用对比熵作为评价指标，从而可以得到连续的结果。

在信息检索领域，通常使用的一些评价指标，如 MAP@K（平均准确率）和 NDCG@K（归一化折损累计增益），是基于离散的排名位置的指标。这些指标在评估模型性能时，有时可能无法捕捉到细微的变化。例如，如果一个文档的排名从第12位提高到第11位，对于 R@10（前10个结果的召回率）来说，这种变化是无意义的，因为这两个排名都在前10名之外，所以 R@10 的值不会发生变化。

为了克服这种离散指标的局限性，文中提出使用对比熵作为一种评价指标。对比熵可以提供一种连续的度量方式，能够更敏感地捕捉到模型变化带来的细微影响。这意味着，使用对比熵作为评价指标，即使是排名位置之间的小变化，也可以在整体的评估中被反映出来，从而提供更精细的性能评估。

对于测试集中的每个查询-段落对，我们随机选取了256个负段落，并定义对比熵如下（越小越好）：

我们将对比熵与几个离散指标进行了对比，如图，可以看到对比熵与这几个评价指标呈线性相关。同时我们还可以发现，在对比熵0.25左右时，三种离散排序指标都会断层式地增大，语言模型中也发现了类似现象，即生成能力与某个具体的损失值紧密相关。

实验

Scaling For Model Size

使用全部的训练集，报告训练过程中测试集效果最好的结果，避免欠拟合或过拟合。

模型大小与效果的关系图如下：

结论表明，模型的效果随着模型规模的增加而提升，且模型大小的缩放规律（Scaling Law）可以用以下公式表示：即模型效果与参数量倒数的某个幂指数成正比。其中，δ 是一个不可约的损失项，也就是说，即使模型规模无限增大，损失也只能缩小到这个值，而不会降至零。这一现象是合理的，因为在人类标注过程中，由于可能存在错误标注，导致标注存在一定的局限性。

• L(N)：模型的损失函数，表示在参数量为 N 时的损失值。
• A：比例因子，它影响模型损失随参数量 N 增长时的变化幅度。
• N：模型的参数量。
• α：幂指数，描述模型性能提升的速率。
• δN：不可约损失项，表示当模型的参数量无限增大时的最小损失。

Scaling For Data Size

接下来本文使用固定的BERT-Base模型来探究数据大小的影响。

结论仍然是数据量越大，效果越好。

这两条规律说明：

1. 可以先在一些小尺度模型上获得相应的效果，拟合模型Scaling曲线，从而直接计算得到更大的模型的效果，而不需要训练。

2. 可以先在一些少量的标注数据上获得相应的效果，拟合数据Scaling曲线，从而直接计算得到更多标注数据时的效果，而不需要训练。

Annotation Quality

对于不同的标注质量，Scaling Laws是否仍然成立。

本文使用三种方式对MSMARCO的段落进行了查询生成从而构造不同标注类型数据集：

1. Inverse Cloze Task (ICT)：直接使用段落中的某个句子作为查询，一般质量较低。

2. Supervised Generation Models：使用docT5query为每个段落生成查询，质量一般高于ICT。

3. Large Language Models (LLMs)：提示ChatGLM3为每个段落生成相关的查询，本文认为这种方式会生成比ICR和T5更高质量的数据集。

ChatGLM3的提示词为：

请根据给定的段落生成5个相关查询，用于搜索目的。
1. 每个查询都应该与段落相关。
2. 每个查询应该包含大约10到20个单词。
3. 请生成多样化的查询。
4. 以JSON格式输出，包含以下键："query1"、"query2"、"query3"、"query4"、"query5"。
5. 请用英文回复。不要使用中文。

对于ICT和ChatGLM3，本文为原训练集中每个人类标注的正例段落都生成了一个查询。对于docT5query，由于他是使用人类标注的正例段落训练得来，因此本文随机从语料库中随机选取了0.5M个段落用于生成查询。评估仍然使用人类标注的验证集。

根据图5，可以看到，四种数据集基本都满足幂律尺度，人类标注效果最好，ICT数据效果最差。其中docT5query效果好于ChatGLM3，原因是docT5query在MSMARCO训练集上微调过。（如果使用更大的闭源模型比如GPT-4 GLM-4，那么效果应该比较好）

Mixed Scaling Law

模型和数据联合起来与效果的关系如下：

使用一个包含模型大小N和数据大小D的公式来拟合：

预算分配中的应用

本文做了一个有趣的预算分配应用。首先根据常见的数据标注、云服务和实际应用中的文档数量，推算出单位大小的数据和模型对应的数据标注、训练和推理的开销（美元），基于此来估计总开销。总开销公式如下：

这里，𝑍data、𝑍train、𝑍infer 分别代表标注、训练、推理对应的成本因子

其中每一项的单位开销如下：

由于随着模型的增大，训练开销会越来越大，留给数据的部分就会越来越少，但数据标注本身是昂贵的，增大模型到一定程度后就几乎无法再标注数据了。因此给定预算下，将模型大小尽可能设置合理会带来最优选择。如果不考虑推理开销，对于5000美元来说，训练一个1B左右模型比较合适；对于20000美元成本来说，训练一个13B的模型比较合适。

Limitation

• 本文使用的对比熵无法完全和排序指标一致，未来可以研究更好的连续指标

• 本文仅考虑随机负采样，没有考虑更精确的训练技术比如难负采样，蒸馏和对比预训练等，他们也可能对Scaling Laws带来不同影响

• 本文仅考虑密集检索的经典架构，将查询与文档映射到固定长度的向量，没有考虑其他架构，比如映射到不同长度、多个向量、稀疏向量等。

• 本文仅考虑域内数据集，尽管本文也做了少量域外实验，但由于数据量少效果不稳定，本文并未报告域外数据集相关结果

• 本文探究的模型大小与数据集大小未来可以进一步扩大

结论

本文系统地研究了密集检索的Scaling Laws，在中文和英语的数据集上进行了实验，以评估模型大小、数据大小和标注方法对检索性能的影响。

通过利用对比熵作为度量，我们观察到相对于不同标注方法和数据集的模型效果与模型大小/数据大小之间的幂律关系。本文还表明，Scaling Laws有助于优化训练过程。例如，正如我们的实验所示的那样，Scaling Laws对于预算分配管理很重要。此外，Scaling Laws允许评估不同标注方法的有效性。本文的实验表明，使用LLM来生成相关性标注仍然有很大的改进空间。

本文相信Scaling Laws提供了一个系统的方法来评估和改进排序模型的训练过程。虽然本研究为这一领域的未来探索奠定了基础，但还需要进一步的研究来扩展我们对不同的领域、尺度、架构和评估协议的Scaling Laws的理解。