论文链接:https://arxiv.org/pdf/2407.02485v1
摘要
本篇文章提出了一个新的指令微调框架RankRAG,该框架通过指令微调单个大型语言模型,使其同时具备上下文排序和答案生成的能力,从而提升检索增强生成任务的效率。文章在多个数据集进行了广泛测试,结果表明在多种任务上RankRAG相较于其他基线模型,展现出了更为卓越的性能表现。RankRAG框架
RankRAG框架主要分为两个阶段:训练(Training)阶段与推理(Inference)阶段。其中训练阶段又可以分为:监督微调与统一指令调优排序和生成两个阶段,下面将详细介绍每个阶段。1.训练阶段
1.1 监督微调(SFT)
作者使用高质量的指令遵循数据集对语言模型进行微调,提升了大语言模型(LLMs)的指令遵循能力,改善其在各种下游任务中的零样本表现。训练数据包括众包对话数据集、长篇问答数据集以及由LLM生成的指令数据集。1.2 统一指令调优排序和生成
该阶段的重点在于通过统一指令调优,增强大语言模型(LLM)在排序和生成任务上的能力,特别是对于检索增强生成(RAG)任务的表现。该阶段包括五个关键部分的数据集和任务设计:- 利用多个具有丰富上下文的问答任务来增强大语言模型(LLM)使用上下文进行生成的能力。训练数据集包括:
- 为了提升模型对无关上下文的鲁棒性。作者采用了SQuAD和WebQuestions这两个数据集,这些数据集不仅包含答案上下文,还包括使用BM25检索到的最佳上下文。对于每个问题,作者将答案上下文与使用BM25检索到的最佳上下文结合,确保每个问题都能关联五个上下文。其中一些检索到的上下文可能不包含答案,这些被称为负样本。
- 为提升LLM的上下文排序能力,作者使用了MS MARCO段落排序数据集。该数据集将问题-相关段落(q,c+)视为正例,而使用BM25检索到的困难负例段落(q,c-)视为负例。LLM需要根据问题判断段落是否与其相关,生成“True”或“False”。此外,由于针对对话问答的排名数据极为稀缺,作者利用对话问答对生成伪相关对,由于每个对话只与一个文档相关联,作者将每个文档切分为150字的片段。然后计算每个片段与真实答案之间的4-gram召回率,召回率高于0.5的片段被视为与对话相关,低于0.1的则被视为不相关。每个样本包含一个问题-上下文对用于该排序数据集。最终作者将两种数据结合,用于训练模型的上下文排序能力。
- 为了提升模型同时评估多个上下文与问题相关性的能力,作者将SQuAD和WebQuestions中的答案上下文与BM25检索到的最佳上下文结合,确保每个问题能关联到五个上下文。其中包含答案的上下文被认为是相关的,训练LLM识别所有与问题相关的上下文。
作者将上述五部分的数据和任务成统一为标准化的问答格式 (x, c, y),其中 x 是问题,c 是上下文,y 是目标答案。表1展示了如何将不同任务转换为统一格式。这带来了以下好处:- 通过增加相对较少的排序数据,增强了LLM的排序能力。
- 通过将这些任务标准化为统一格式,它们可以相互增强。
2.推理阶段
由于RankRAG包含了额外的重新排序步骤,因此每个问题的推理管道被修改为检索-重新排序-生成管道,描述如下:- RankRAG模型计算问题与检索到的N个上下文之间的相关性分数,作为使用表1中的提示符生成答案为True的概率,然后对上下文重新排序,只保留前k个(k≪N)上下文,然后将其用作生成步骤的输入。
- 将top-k上下文与问题连接并反馈到RankRAG模型中,以生成最终答案。