又来一个RAG：RankRAG，英伟达RAG新思路
发布时间：2024 年 07 月 02 日
RAG
RankRAG: Unifying Context Ranking with Retrieval-Augmented Generation in LLMs
我们提出了一种名为 RankRAG 的创新指令微调框架，该框架使单个 LLM 能够同时进行上下文排序和答案生成，显著提升了 RAG 的性能。通过在训练中融入少量排序数据，RankRAG 不仅超越了专门优化的排序模型，还在生成任务上表现卓越，击败了包括 GPT-4 在内的多个顶尖模型。特别是在知识密集型和生物医学领域的基准测试中，RankRAG 展现了其强大的泛化能力，无需特定领域的微调即可与 GPT-4 媲美。
1. 背景
检索增强生成无疑是当前最热门的LLM研究领域了，但是传统RAG存在一些局限性：
- • 检索器容量的限制。考虑到处理效率，现有的RAG一般采用稀疏检索（比如BM25）或中等规模的嵌入模型（比如Bert）作为检索器。
- • 只选择前K个文档。尽管最新的大语言模型扩大了上下文长度的限制，能够接受更多的上下文作为输入，但是实际性能会随着K的增加而迅速达到饱和。比如在长问答任务中，最佳的分块上下文数量大约是10。虽然更大的K可以提高召回率，但是同时也引入了更多无关的内容，干扰大语言模型生成准确回答。
因此，作者设计了一个新的RAG：RankRAG，利用单一大语言模型来实现高召回率的上下文提取和高质量内容生成。通过对单一大语言模型进行指令调优，使其可以同时进行上下文排序和答案生成，进一步提升LLM在RAG检索和生成阶段排除不相关上下文的能力。
2. RankRAG
RankRAG整体包括两个阶段：指令调优阶段、排名与生成综合指令调优阶段。
2.1 第一阶段：指令调优（SFT）阶段
指令调优（或监督式微调）可以显著提升大语言模型指令遵循的能力，从而提高下游任务的零样本性能。所以，作者先对大语言模型进行指令调优。
总共准备了128K SFT样本，确保样本中没有数据泄露，没有与评估任务重的数据重叠，并采用多轮对话格式。
第一阶段的指令调优数据集包括：公共对话数据集、长篇问答数据集、Chain of Thought数据集以及合成数据集。
2.2 排名与生成的综合指令调优阶段
第一阶段的SFT目的是为大模型奠定指令遵循的基础能力，但是在RAG任务上，这些模型仍然还有提升空间。为了提升这些模型在RAG任务上的性能，作者进行了第二阶段的指令调优。第二阶段的指令调优包括五个部分：
2.2.1 第一阶段的SFT数据
这部分数据的纳入是为了保持LLM遵循指令的能力。
2.2.2 富含上下文的QA数据
在数据集中加入多个富含上下文的QA任务数据来加强LLM利用上下文回答问题的能力。采用混合训练数据集：
- • 标准的QA和阅读理解数据集
- • 对话QA数据集
2.2.3 检索增强问答数据
引入两个新数据集，不仅包含标准答案的上下文，还包括通过BM25算法检索出的最相关上下文。
增强LLM在生成回答时，对无关上下文的抗干扰能力也非常重要，所以作者特别考虑了SQuAD和WebQuestion两个问答任务数据集。对于每个问题，作者将标准答案的上下文和BM25检索出的上下文进行融合，确保每个问题都能至少关联5个上下文。其中部分检索出的上下文可能是不包含答案的，这些被称为负样本。
2.2.4 上下文排名数据
为了使模型具有LLM排名能力，作者采用了大家普遍认可的MS MARCO段落排名数据集。将标准查询-段落对（q，d+）定义为正样本，而通过BM25算法挖掘出的硬负样本段落（q，d-）定义为负样本。LLM需要根据给定的查询-段落对生成“真”或“假”的回答，任务指令是“针对问题{question}，请判断该段落是否与问题相关。”
虽然MS MARCO覆盖了众多主题，但问题均为单轮简短句子。然而，针对对话问答的排名数据极为稀缺。
为解决这一局限，将对话问答对重新设计，以生成伪相关对。鉴于每段对话仅关联一个文档，将每个文档切割成150字的段落（d1，d2，...，dn）。计算每个段落di与标准答案间的4-gram召回率，将召回率超过0.5的段落判定为与对话相关，低于0.1的则判定为不相关。每个样本包含一对问题-上下文，用于排名数据集。
合成后的数据加上开始的MS MARCO数据，一共约50K数据用于本阶段的指令微调。
2.2.5 检索增强型排名数据
为了使模型能够针对一个问题评估多个上下文的相关性，采用了SQuAD和WebQuestions两个问答数据集，通过BM25算法，将标准答案上下文与检索出的最相关上下文相结合，确保每个问题关联五个上下文。凡包含答案的上下文均被视为相关，训练LLM识别与问题相关的所有上下文。
2.3 RankRAG 推理流程：检索-重排-生成三部曲
融入额外的重排环节，RankRAG 的推理流程转变为三步走策略：检索-重排-生成，具体步骤如下：
(1) 检索器 ℛ 率先在文本库中筛选出 top-k 个相关上下文。
(2) RankRAG 模型进而基于上表中的提示，评估问题与这些检索到的上下文之间的相关性得分，以此作为生成正确答案的概率，随后对上下文进行重排，精挑细选出 top-k（k'远小于k）个最为贴切的上下文，作为下一步生成答案的依据。
(3) 精选出的 top-k' 个上下文与问题串联，再次输入 RankRAG 模型，以此生成最终的答复。
3. 效果对比
3.1 主要结果
上图展示了RankRAG与基线在九个数据集上的比较结果。所有结果均在零样本评估条件下得出（除非特别标注），未附加额外示例。
RankRAG在性能上超越了现有的RAG方法： 以8B模型规模来看，RankRAG持续领先于ChatQA-1.5 8B，后者是最近开源且在众多RAG基准测试中表现卓越的模型。即使面对参数量多得多的基线模型，RankRAG 8B依然展现出强劲的竞争力。举例来说，它显著超越了参数量为其五倍的InstructRetro、参数量为其八倍的RA-DIT 65B，甚至在NQ和TriviaQA任务上超越了参数量为其八倍的Llama3-instruct 70B。参数量更多的RankRAG 70B不仅战胜了强劲的ChatQA-1.5 70B模型，还大幅超越了以InstructGPT为基础LLM的先前RAG基线。
RankRAG在更具挑战性的数据集上取得了更显著的进步： RankRAG在处理更具挑战性的QA数据集时，性能提升更为显著。比如，在长尾QA（PopQA）和多跳QA（2WikimQA）任务中，相较于ChatQA-1.5，实现了超过10%的性能提升。这些结果表明，在挑战性的开放问答数据集中，当检索器返回的顶部文档与答案的相关性不高时，上下文排名能显著增强性能。今天这篇论文专注于提升QA任务的单次检索效果。如何有效地将多轮RAG流程与RankRAG结合，是未来研究值得探索的有趣方向。
5.3 组件效果分析
通过在九个广泛领域的数据集上以Llama3 8B作为基础模型，展示了对RankRAG进行的组件切除实验。总体来看，发现所提出的各个组件均对提升最终性能起到了积极作用。
若去除上下文排名功能，将在所有任务中造成性能下降，这证实了其在筛选与目标问题最密切相关上下文中的关键作用。
同时，为指令微调特别设计的检索增强问答（RQA）和检索增强排名（RAR）通过辅助模型明确识别相关上下文，在大多数任务中带来了性能提升。
性能随不同检索器的变化。如上图，对比了RankRAG和ChatQA-1.5在三个典型任务中搭配不同密集检索器的表现，特别考虑了DPR与Contriever-MS MARCO这两种检索器变体。可以看到，即便初次检索结果不尽人意，RankRAG在平均性能上依然超越了ChatQA-1.5，幅度超过10%。综上，RankRAG对于检索器的选择展现出了良好的适应性和鲁棒性。
5.4 领域特定RAG基准测试的实验验证
为验证RankRAG对专业领域的适应能力，在生物医学领域的最新RAG基准测试Mirage上进行了实验。采用MedCPT（Jin等人，2023年）作为检索器ℛ，并以MedCorp2作为语料库𝒟。
上图展示了RankRAG与基线的实验结果，即便未针对生物医学领域进行微调，RankRAG在医学问答任务上依然表现卓越。特别是RankRAG 8B在性能上超越了医疗领域领先的开源LLM Meditron 70B达6.3%。
此外，RankRAG 70B的性能更是达到了GPT-4的98%以上。充分证明了RankRAG具备快速适应新领域并无需额外训练的能力。
5.5 细究排名模块
RankRAG在数据使用上极为高效：
传统RAG流程的上下文排名方法通常需要一个独立的重排模型。作者评估了四款在完整MS MARCO段落排名数据集上经过微调的模型（BERT（Glass等人，2022年）/T5（Nogueira等人，2020年）/Llama3（Ma等人，2023年）），一个强大的现成重排模型BGE-ranker，以及两款OpenAI的GPT系列模型。对于GPT系列，用‘True’标记的概率来代理相关性得分。
如上表，RankRAG在多数情况下，即使在排名数据量仅为十分之一，也能实现比专用排名模型更好的召回率。此外，RankRAG在大多数任务上仍能超越经过超过100万个排名对训练的BGE-ranker。这一优势可能源于我们模型训练的适应性，排名数据与一般RAG微调数据高度相似。直接使用ChatQA-1.5进行段落排名会降低性能，这表明在指令微调中纳入排名数据的必要性。
作者验证了上下文排名数据量与最终性能之间的关联。仅用5k排名数据（约MS MARCO数据集的1%），RankRAG便能取得非常出色的成果，而将排名数据量增加至50k，能带来额外的显著增益。这验证了RankRAG的数据高效性——仅需少量排名数据即可达到高效性能，并在多样的任务中维持其适应性。
RankRAG的性能与时间效率权衡：
扩大模型规模时需注意的一个细节是，延迟开销的增加——模型需要对每个样本进行排名，这会耗费更多时间。
为了探究时间效率与性能之间的联系，调整了重排名过程中使用的样本数k，上图展示了k值与最终准确性之间的关系。结果表明，即便k值设为20，RankRAG的性能依然超过了未进行重排名的基线模型。当k值从20增加至100，三个任务的完全匹配准确率提升了5.9%至9.1%，而所需时间仅增加了0.9至6.0倍，远低于人们可能预期的20至100倍的增长。
