Quest:一种以查询为中心的长文本数据合成方法

文摘   2024-09-23 23:00   荷兰  

今天给大家带来一篇长文本数据合成方法。

作者: 高超尘
机构: 中国科学院信息工程研究所
Arxiv: https://arxiv.org/abs/2405.19846

大规模语言模型可接受长度愈加得到关注,就像LLaMA模型已从2k(LLaMA)增加到128k(LLaMA3.1),但是目前长文本数据匮乏且分布不均。

正如下图(左)所示,对Pile语料库的分析表明,长文档主要集中在Books3数据集中,并且随着目标上下文长度的增加,分布会进一步倾斜。

本文介绍一种以查询为中心的长文本数据合成方法-Quest,能够有效聚合相关但冗余度低的文档以合成长上下文数据,并且具备高度可扩展性,可构建用于持续训练的大型数据集。

图1 长文本来源分布对比图,对比pile长文本分布(图左),Quest合成数据中的长文分布更均匀(图右)

方法介绍

在此之前,有一些基于检索增强预训练的研究,比如:REALM(基于KNN检索将相似的语句拼接到了一起),ICLM(通过建立了一条相似度路径将数据聚合起来,避免了数据的重复使用)等,但把模型从8k扩展到了32k时反而出现了性能下降。如下表所示,基于检索的方法常常因相似句子的重复性而导致冗余,特别是在大规模语料库中,降低了预测难度和上下文的多样性,从而削弱了长上下文建模的有效性。

表1 不同方法在longbench上的效果,基座模型为pythia

为了聚合相关但不冗余的文档,提出了一种以查询为中心的数据合成方法(Quest),其灵感来源于搜索引擎,当用户查询query时,会返回大量的相关文档。具体来说,会对每一个doc都生成一条query,同时基于相似的query将对应doc 聚集在一起,合成的长文本数据 用来训练模型,扩展模型的窗口长度,整体流程如下图所示:

图2 Quest方法图

Quest方法分为五步:

  • 查询预测:使用开源的doc2query模型[3]为每个文档预测查询。对于超过doc2query模型上下文长度限制的文本,将其分割成多个部分,并为每个部分生成一个查询。因此,对于文档,预测得到的查询列表为
  • 关键词提取:使用高效工具Rake从每个查询中提取关键词。对于具有多个查询的文本,Rake生成多个关键词列表。为了确保提取的关键词质量,采用了两种过滤策略。首先,过滤掉Rake得分低于3.0的关键词。其次,移除频繁出现但无信息价值的关键词。然后,从剩余的关键词中随机选择一个作为该文档的代表关键词。
  • 构建基于关键词的倒排索引:在将每个文档映射到其代表关键词后,可以构建一个基于关键词的倒排索引 。具有相同代表关键词的文档会被一起索引。
  • 索引划分:根据每个倒排索引中包含的文档数量对其进行升序排列,并将排序后的索引分为两组。排名靠前的 关键词倒排索引被分配到短索引集 ,其余部分分配到长索引集
  • 训练过程:从采样的索引中的文档中进行不放回抽样,并将选定的文档在模型的上下文长度 内进行拼接用于训练。对短索引进行过采样,以确保参与训练的短索引和长索引的文档数量分布均衡。

伪代码如下图所示:

实验结果

使用不同方法分别合成了32k,128k的长文训练集,基于pythia的模型训练并在对应的Benchmark上进行了测试。在32k的场景下longbench测试结果如上文表1所示,128k的场景下LongBook QA测试结果如表2所示。可以看出,在32k以及128k的场景下,Quest都取得了更好的性能。有趣的是,在128k的场景下,KNN和ICLM反超了Standard方法,这说明了在更长文本的场景下,对于文本的相似度建模是有效的。

表2 使用pythia模型在不同合成数据集训练之后的longbook qa表现

同样测试了不同方法的短文本保持能力。可以看到,经过Quest训练的模型短文本能力并没有下降。

表3 使用pythia模型在不同合成数据集训练之后的短文本表现

此外,还在llama3 8B上使用Quest 数据进行了训练,longbookQA的指标达到了32.39,达到了开源模型的最好效果。在128k的捞针实验(召回一条文本)上达到了97%的准确率。

表4 使用llama3 8B模型在Quest训练之后longbook qa 的表现
图3 使用llama3 8B模型在Quest训练之后128k捞针的表现

为了进一步探究Quest在超长文本下的表现,将文本长度扩展到了500k。为了便于超长文本的测试,参考了EasyContext以及LWM在召回随机数字串的捞针任务进行了测试,在500k的长度下达到了全绿。在捞针以及longbookqa任务上,达到了同尺寸开源预训练模型(base model)的最好效果。

图4 使用llama3 8B模型在Quest训练之后500k捞针的表现

效果分析

在上文中,提到基于检索的方法常常因相似句子的重复性而导致冗余,为了证明这一点,抽样统计了不同算法合成的每条数据内doc 的相似度,如下图所示,可以发现,随着相似度的升高,在两个Benchmark上的指标都出现了先上升再下降的趋势。

图5 随着聚合文档相似度的增加,长上下文性能呈现先上升再下降趋势,所有结果均在相似度范围内归一化

还将每种方法合成数据内部的文档通过TSNE可视化了出来,如下图所示,标准方法(Standard)随机拼接文档,导致文档分布过于分散,破坏了文档之间的关系,进而导致性能下降,这点在128k的时候尤其显著。在 32k 上下文中,ICLM 由于较短的相似路径长度,导致过度聚类,呈现出类似 KNN 的分布,进而影响了 Longbench 基准上的性能表现。然而,在 128k 上下文中,ICLM 形成了较长的相似路径,使文档分布更加分散,从而提升了 LongbookQA 基准上的表现。

值得注意的是,Quest 在这两种上下文中都保持了均匀分布的文档分布,突显了其在两个基准测试中的卓越表现。这些发现表明,过于分散或过于集中的文档语义分布会损害模型性能,而 Quest 通过对查询相关文档进行聚类,确保了文档的相关性并避免冗余,从而提升了模型性能。

写在最后

本文介绍了Quest,一种通过将相关但冗余度低的文档进行分组和拼接来合成长上下文数据的新方法。该方法缓解了长上下文数据中的数据稀缺性和分布不均问题,从而提升了预训练模型的长上下文建模能力。

PS:看到这里,如果觉得不错,可以来个点赞在看关注。给公众号添加【星标⭐️】不迷路!您的支持是我坚持的最大动力!

欢迎多多关注公众号「NLP工作站」,加入交流群,交个朋友吧,一起学习,一起进步!

NLP工作站
日常分享AIGC前沿知识&落地经验总结,也欢迎关注《ChatGPT原理与实战》、《大型语言模型实战指南》两本大模型相关书籍
 最新文章