分享一篇SIGIR2024的最佳论文,探究了密集检索的Scaling Laws,尽管结论是数据和模型都越大越好,不过这篇文章的行文和严谨实验设置还是值得一看的,对于评估协议、不同的数据标注方法、预算分配也提出了一些较深刻的见解。
论文题目:Scaling Laws For Dense Retrieval
来源:SIGIR2024/清华
方向:文本检索
开源地址:https://github.com/jingtaozhan/DRScale
1.摘要
Scaling Laws已经在广泛的任务中被观察到,特别是在语言生成中。先前的研究发现,大型语言模型的性能和模型、数据集的大小存在可预测的模式关系。这有助于我们有效地设计训练策略,特别是在大规模训练越来越需要丰富的计算资源的时候。然而,在密集检索中,这种Scaling Laws尚未得到充分的探索。在本研究中,我们研究了Scaling如何影响密集检索模型的性能。我们使用不同参数量来实现密集检索模型,并使用不同数量的标注数据来训练它们。
本文使用对比熵作为评价指标,与离散排序指标相比,它是连续的,因此可以准确地反映模型的性能。结果表明,密集检索模型的性能与模型大小和标注数据量大小具有精确的幂律Scaling关系。此外,我们还证明了Scaling Laws有助于优化训练过程,如解决预算限制下的资源分配问题。我们认为,这些发现有助于理解密集检索模型的Scaling效应,并为未来的研究提供了有意义的指导。
3.方法
3.1 问题形式化
对于给定的一个语料库,密集检索模型的目标是对于一个具体的查询确定最相关的段落,密集检索模型通过将查询和候选段落映射到一个共享的密集表征空间来实现这一目标。获得表征后,再使用一个打分函数,比如内积或余弦相似度来计算相关度分数。f代表映射函数,θ代表模型参数,则相似度分数可以用以下公式计算:
本文中,仅考虑查询和段落共享编码器的情况。训练数据采用最流行的query-positive passage对的标注方式,获得了n个数据对,其中每个查询可能对应多个匹配段落。
3.2 模型架构
本文采用了最流行的transformer架构,得到向量的函数如下,其中W和b为映射头的参数
为了避免不同预训练任务对实验的影响,对于英文检索,本文选用了Google开源的24个BERT系列模型,从最小的0.5M参数(BERT-Tiny)到最大的82M参数(BERT-Base)。对于中文检索,本文选用了文心系列模型,对于每个模型,我们都统一映射到768维的表征空间从而保持公平对比。
3.3 训练数据
对于英文,本文使用了MS MARCO,包含从英文网站上收集的8.8M个段落以及0.5M条查询,每条查询都标注了一个正例段落。同时MS MARCO提供了7000条查询-段落对用于评估。对于中文,本文使用了T2Ranking,包含从真实搜索引擎中收集的300k条查询以及超过2M个段落。
3.4 训练设置
本文采用最经典的随机批内负采样策略进行对比学习,损失如下:
本文训练固定的10000步,批大小设置为256。
注:由于截至2024年7月22日作者还未开源代码且原文的负采样策略介绍有点不详细,这里的介绍可能不一定和作者的实现完全一致。
3.5 评估协议
由于一般的协议MAP@K,NDCG@K是离散排序评估协议,在模型变化时不一定能准确反映出全部的变化(比如排序从12到11对于R@10没有区别),因此本文提出使用对比熵作为评价指标,从而可以得到连续的结果。
对于测试集中的每个查询-段落对,我们随机选取了256个负段落,并定义对比熵如下(越小越好):
我们将对比熵与几个离散指标进行了对比,如图2,可以看到对比熵与这几个评价指标呈线性相关。同时我们还可以发现,在对比熵0.25左右时,三种离散排序指标都会断层式地增大,语言模型中也发现了类似现象,即生成能力与某个具体的损失值紧密相关。
4.实验
这一节主要回答下列三个问题:
- 模型大小如何影响密集检索性能
- 训练数据大小如何影响密集检索性能
- 不同类型的标注数据是否会给密集检索模型带来不同的Scaling效应
4.1 模型大小Scaling
使用全部的训练集,报告训练过程中测试集效果最好的结果,避免欠拟合或过拟合。
模型大小与效果的关系图如下:
结论是模型越大效果越好,模型大小的Scaling Law可用如下公式表示,即和参数量的倒数的幂指数成正比。其中δ是一个不可约的损失项,即一个足够大的模型(趋近无穷)也只能将损失缩小到这个值而不是0。由于人类标注的局限性(可能存在很多假负例),这一项是合理的。
本文使用最小二乘法来拟合这个曲线,发现模型的Scaling Law确实满足幂律尺度。这一发现为未来的研究提供了新视角,即我们可以先在一些小尺度模型上获得相应的效果,拟合模型Scaling曲线,从而直接计算得到更大的模型的效果,而不需要训练。
4.2 数据大小Scaling
接下来本文使用固定的BERT-Base模型来探究数据大小的影响。
数据大小与效果的关系图如下:
结论仍然是数据量越大,效果越好。同样也用一个相同公式进行拟合:
本文使用最小二乘法来拟合这个曲线,发现数据的Scaling Law确实满足幂律尺度。这一发现也为未来的研究提供了新视角,即我们可以先在一些少量的标注数据上获得相应的效果,拟合数据Scaling曲线,从而直接计算得到更多标注数据时的效果,而不需要训练。
4.3 标注质量
接下来本文需要探究对于不同的标注质量,Scaling Laws是否仍然成立。本文使用三种方式对MSMARCO的段落进行了查询生成从而构造不同标注类型数据集:
- Inverse Cloze Task (ICT):直接使用段落中的某个句子作为查询,一般质量较低
- Supervised Generation Models:使用docT5query为每个段落生成查询,质量一般高于ICT
- Large Language Models (LLMs):提示ChatGLM3为每个段落生成相关的查询,本文认为这种方式会生成比ICR和T5更高质量的数据集
ChatGLM3的提示词如下:
对于ICT和ChatGLM3,本文为原训练集中每个人类标注的正例段落都生成了一个查询。对于docT5query,由于他是使用人类标注的正例段落训练得来,因此本文随机从语料库中随机选取了0.5M个段落用于生成查询。评估仍然使用人类标注的验证集。
根据图5,可以看到,四种数据集基本都满足幂律尺度,人类标注效果最好,ICT数据效果最差。其中docT5query效果好于ChatGLM3,原因是docT5query在MSMARCO训练集上微调过。(如果使用更大的闭源模型比如GPT-4 GLM-4,那么效果应该比较好)
4.4 模型数据联合定律
模型和数据联合起来与效果的关系如下:
使用一个包含模型大小N和数据大小D的公式来拟合:
5.预算分配中的应用
本文做了一个有趣的预算分配应用。首先根据常见的数据标注、云服务和实际应用中的文档数量,推算出单位大小的数据和模型对应的数据标注、训练和推理的开销(美元),基于此来估计总开销。总开销公式如下:
其中每一项的单位开销如下:
由于随着模型的增大,训练开销会越来越大,留给数据的部分就会越来越少,但数据标注本身是昂贵的,增大模型到一定程度后就几乎无法再标注数据了。因此给定预算下,将模型大小尽可能设置合理会带来最优选择。如果不考虑推理开销,对于5000美元来说,训练一个1B左右模型比较合适;对于20000美元成本来说,训练一个13B的模型比较合适。
考虑推理开销,由于需要将30trillion级别的文档编码,可以训练的模型大小就比较小了。(笔者认为开销这里有待商榷,假设使用自己的卡,那么训练和推理成本会带来相应的减少,且推理主要关注的应当是对于查询的embedding,以及还要考虑到具体的向量数据库算法等)
6.不足和未来工作
- 本文使用的对比熵无法完全和排序指标一致,未来可以研究更好的连续指标
- 本文仅考虑随机负采样,没有考虑更精确的训练技术比如难负采样,蒸馏和对比预训练等,他们也可能对Scaling Laws带来不同影响
- 本文仅考虑密集检索的经典架构,将查询与文档映射到固定长度的向量,没有考虑其他架构,比如映射到不同长度、多个向量、稀疏向量等。
- 本文仅考虑域内数据集,尽管本文也做了少量域外实验,但由于数据量少效果不稳定,本文并未报告域外数据集相关结果
- 本文探究的模型大小与数据集大小未来可以进一步扩大
结论
本文系统地研究了密集检索的Scaling Laws,在中文和英语的数据集上进行了实验,以评估模型大小、数据大小和标注方法对检索性能的影响。
通过利用对比熵作为度量,我们观察到相对于不同标注方法和数据集的模型效果与模型大小/数据大小之间的幂律关系。本文还表明,Scaling Laws有助于优化训练过程。例如,正如我们的实验所示的那样,Scaling Laws对于预算分配管理很重要。此外,Scaling Laws允许评估不同标注方法的有效性。本文的实验表明,使用LLM来生成相关性标注仍然有很大的改进空间。
本文相信Scaling Laws提供了一个系统的方法来评估和改进排序模型的训练过程。虽然本研究为这一领域的未来探索奠定了基础,但还需要进一步的研究来扩展我们对不同的领域、尺度、架构和评估协议的Scaling Laws的理解。
大家好,我是NLP研究者BrownSearch,如果你觉得本文对你有帮助的话,不妨点赞或收藏支持我的创作,您的正反馈是我持续更新的动力!如果想了解更多LLM/检索的知识,记得关注我!