FlashRAG：可能是最全的、最快搭建RAG的开源框架

发布时间：2024 年 05 月 22 日

RAG

今天这篇论文来自中国人民大学高瓴人工智能学院

FlashRAG: A Modular Toolkit for Efficient Retrieval-Augmented Generation Research

随着大型语言模型（LLMs）的兴起，检索增强生成（RAG）技术因其潜力而备受研究关注。尽管已有多种创新算法和模型被提出以提升 RAG 系统的性能，但由于缺乏统一的标准化框架，加之 RAG 过程的复杂性，研究人员在一致环境中比较和评估这些方法面临挑战。现有的 RAG 工具包，如 LangChain 和 LlamaIndex，虽已存在，但往往过于庞大且不够灵活，难以满足个性化需求。为此，我们开发了 FlashRAG，一个高效且模块化的开源工具包，旨在帮助研究者在统一框架内复现现有 RAG 方法并创新开发新算法。FlashRAG 集成了 12 种先进 RAG 方法，并整合了 32 个基准数据集，提供可定制的模块化框架、丰富的预实现作品、全面数据集、高效预处理脚本及广泛的标准评估指标。所有资源均可在https://github.com/RUC-NLPIR/FlashRAG获取。

https://arxiv.org/abs/2405.13576

1. 为什么要提出FlashRAG？

在大型语言模型（LLMs）的浪潮中，检索增强生成（RAG）技术凭借其利用外部知识库的能力，有效缓解了幻觉问题，成为一项备受瞩目的解决方案。RAG技术的广泛应用前景吸引了众多研究者的目光。近年来，随着众多新算法和模型的涌现，旨在提升RAG系统各方面性能，但要在统一框架下比较和评估这些方法，却变得日益困难。

许多研究工作要么不开源，要么开源代码中设置了固定参数，这使得适应新数据或创新组件变得不易。

此外，研究中使用的数据库和检索语料库常常各不相同，资源分散，导致研究者在预处理上耗费大量时间，而非专注于方法的优化。

RAG系统的复杂性，涉及索引、检索和生成等多个步骤，也使得研究者往往需要自行实现系统的多个部分。尽管市面上已有如LangChain和LlamaIndex等RAG工具包，但它们通常庞大且不易操作，难以满足定制化需求，也未能解决上述问题。因此，迫切需要一个统一的、以研究者为中心的RAG工具包，以简化方法开发和比较研究。

为了应对这一挑战，本篇论文作者推出了FlashRAG，这是一个开源库，旨在帮助研究者轻松复现现有的RAG方法，并开发他们自己的RAG算法。该库允许研究者利用构建好的管道复现现有工作，使用提供的RAG组件构建自己的流程，或利用组织好的数据集和语料库加速他们的RAG工作流程。相较于现有的RAG工具包，FlashRAG更适合研究者的需求。总的来说，FlashRAG库的核心特点和能力可以概括为以下四个方面：

• • 1.广泛且可定制的模块化RAG框架。为了便于RAG过程的扩展，们在两个层面上实现了模块化设计。在组件层面，提供了全面的RAG组件，涵盖裁判、检索器、细化器和生成器四大类别共13个组件。这些组件既可以单独使用，也可以组合成完整的流程。在流程层面，根据RAG发展的现状，实现了8种常见的RAG流程。基于此框架，可以轻松复现现有方法，并在不同设置下运行和评估RAG流程。

• • 2.预先实现的高级RAG算法。FlashRAG提供的现有工作实现是最全面的。已经基于此框架实现了12种高级RAG算法，包括Self-RAG和FLARE等，覆盖了顺序RAG、条件RAG、分支RAG和循环RAG等类别。这些方法已在统一设置下评估，并提供了基准报告。FlashRAG 使研究者能够轻松评估这些方法，并与他们自己的方法进行公平比较，从而提高整体的可复制性。

• • 3.全面的基准数据集。为了提高RAG研究中数据集的一致性和可重用性，编译并预处理了32个常见的RAG基准数据集，统一了格式。其中一些数据集，如asqa和wikiasp，针对RAG场景进行了特别调整，以确保一致性。已将这些数据集托管在Hugging Face平台上，便于访问和使用。

• • 4.高效的RAG辅助脚本。为了缩短RAG实验的设置时间，提供了一套全面的辅助脚本，包括下载和切割维基百科创建语料库、构建检索索引以及提前准备检索结果等。这些步骤对后续流程至关重要，但往往繁琐耗时。我们设计的用户友好脚本直观易用，确保研究者能够轻松完成RAG相关研究的准备阶段。

2. FlashRAG架构

FlashRAG工具包的目的在于简化与RAG相关的研究工作。如上图所示，该工具包由三个层次化的模块构成：环境模块、组件模块和管道模块。

• • 环境模块：The environment module，整个工具包的基石，它负责搭建实验所需的数据集、超参数和评估指标。

• • 组件模块：The component module，建立在环境模块之上，包含多种RAG组件，每个组件都担负着独特的任务，如检索和生成等。

• • 管道模块：The pipeline module，则将不同的组件模块有机结合，以实现完整的RAG流程。

2.1 组件模块

组件模块将RAG流程中涉及的所有要素汇集于一个统一的框架之下。每个组件都具备独立运作的能力，可以单独使用。目前，组件模块主要包括五大组件：裁判器、检索器、重排器、细化器和生成器。

2.1.1 裁判器（Judger）

作为流程的首个环节，负责判断查询是否需要进行检索。由于这一领域的工作和模型相对较少，目前提供了基于SKR方法的裁判器，它利用精心策划的大型语言模型自我知识数据集来判断检索的必要性。

2.1.2 检索器（Retriever）

FlashRAG 对检索器的实现给予了广泛支持。

在稀疏检索方面，整合了Pyserini库以便于应用BM25方法。

对于密集检索，支持包括DPR、E5和BGE在内的多种基于BERT的嵌入模型。

FlashRAG同样支持基于T5架构的模型，例如ANCE。

采用FAISS进行向量数据库计算，确保检索效率，并利用HuggingFace的数据集库来提升语料库加载速度。

为提高检索结果的复用性，并适应非开源检索器，FlashRAG支持使用预先检索的结果，即“检索缓存”。在每次检索时，系统会自动在检索缓存中搜索与当前查询相关的结果，并将其作为返回值。用户可以使用我们的检索器设置自动保存检索缓存为JSONL文件，以备将来使用。对于非开源检索器，用户可以调整检索结果格式，以适配缓存结构进行加载。

2.1.3 重排器（Reranker）

重排器的目标是优化检索器返回的结果顺序，以提升检索的准确性。目前，FlashRAG支持包括bge-reranker和jina-reranker在内的多种使用的交叉编码器模型。在嵌入模型用于重排的情况下（例如，以BM25作为检索器），还支持使用E5等Bi-Encoder模型作为重排器。在实际应用中，重排器通过装饰器模式集成到检索器的检索功能中，实现与任何检索器的无缝配合。用户可以轻松地通过一行代码结合任何检索器和重排器。

2.1.4 细化器（Refiner）

细化器的作用是对输入文本进行细化，供生成器使用，以减少令牌的使用并降低检索文档中的噪声，从而提升最终RAG响应的质量。作为RAG流程中的关键一环，众多研究致力于开发更高效的细化技术。

经过文献综述，实现了四种不同类型的细化器，在处理检索文档时各有所长。

• • 提取型细化器利用嵌入模型从检索文本中提取与查询语义相似度较高的语义单元，如句子或短语。

• • 抽象型细化器则采用seq2seq模型直接对检索文本进行总结，支持如RECOMP这样的专用模型，以及HuggingFace提供的具有相似结构的通用摘要模型。

• • 还支持使用基于困惑度的LLMLingua细化器和选择性上下文细化器。

2.1.5 生成器（Generator）

生成器是RAG流程中的最后一个组件，在FlashRAG的工具包中得到了全面的支持。在生成器模块中，整合了vllm和FastChat这两个领先的大型语言模型加速库，因此支持了众多主流的大型语言模型。还提供了Transformers库的原生接口，以增强整体的鲁棒性。也支持包括Flan-T5在内的各种编码器-解码器模型。对于这些模型支持使用Fusion in Decoder (FiD)技术，进一步优化了处理检索文档时的效率。

2.2 管道模块

依托先前介绍的多样化组件，FlashRAG 成功实现了 RAG 过程算法逻辑与其各个组件具体实现的解耦，大大简化了整个管道的构建工作。该管道能够处理用户提供的数据集，执行相应的 RAG 操作，并且同时输出最终评估结果与中间成果。在构建管道时，用户只需确定整个 RAG 流程所需的组件及其数据流向逻辑。具体来说，每个管道都需要在 init(.) 函数中加载所需的组件，并根据组件接口在 run(.) 函数中实现相应的逻辑。

为了有序执行各类 RAG 任务的操作逻辑，深入调研了 RAG 相关文献。参考 RAG 综述 《Retrieval-augmented generation for large language models: A survey, 2024》的整理，将所有 RAG 流程归纳为四种类型：顺序、分支、条件和循环。目前，已经构建了 8 种不同的管道，覆盖了众多前沿的 RAG 实践。

• • 顺序管道 提供了一条清晰的线性执行路径，即查询 -> 检索器 -> 检索后处理（重排器、细化器）-> 生成器。用户一旦设定好配置，库便会自动加载所需的组件及其逻辑。

• • 分支管道 则针对单个查询并行执行多条路径（通常每个检索到的文档对应一条路径），并汇聚所有路径的结果以形成最终输出。目前支持两种先进的分支方法：REPLUG 管道 和 SuRe 管道 。REPLUG 管道并行处理每个检索文档，汇总所有文档的生成概率以生成最终答案。SuRe 管道则从每个检索文档生成候选答案，然后对所有候选答案进行排序。在实现 SuRe 时，严格遵循原论文的提示和流程，确保结果的精确性和可比性。

• • 条件管道 则利用裁判器根据裁判结果将查询导向不同的执行路径。在当前框架下，需要检索的查询会进入常规顺序流程，而其他查询则跳过检索，直接进行生成。提供了辅助函数，用以根据裁判器的判断拆分和合并输入数据集，确保所有处理工作能够高效批量执行。此外，条件管道还支持与多种管道类型集成，能够根据不同裁判结果执行不同的管道。

• • 循环管道 涉及检索与生成过程之间的复杂交互，通常包含多次检索与生成循环。与前三种管道相比，循环管道提供了更大的灵活性和更优的结果。支持四种广泛认可的方法：迭代 、自我提问、自我 RAG和 。对于这些方法，都支持对检索器和生成器进行灵活调整，以适应不同场景的测试需求。

2.3 数据集与语料库

2.3.1 数据集概览

如上表，FlashRAG精心搜集并预处理了32个标杆数据集，覆盖了 RAG 研究中广泛使用的数据集。大多数数据集的知识库源自维基百科，数据集已被统一转换成 JSONL 格式，主要包括四个关键字段：ID、问题、标准答案和元数据。对于 MMLU和 OpenBookQA这类选择题数据集，还额外提供了“选项”字段。为方便用户，在 HuggingFace 平台上提供了这些经过处理的数据集。

除了数据集，我们还为用户提供了多种数据集筛选工具，以便筛选整个数据集。用户可以随机或顺序地从整个数据集中选取一定数量的样本进行评估，也可以通过数据集的元数据来选取数据集的子集。这些筛选方法都集成在数据集加载功能中，用户可以通过一个标准化的接口来访问。同时，用户也被允许自定义自己的筛选函数。

2.3.2 语料库

检索所用的语料库，也被称作知识库，同样是实验准备中不可或缺的一环。在众多研究中，以下两种类型的语料库被频繁使用：维基百科和 MS MARCO 。

2.4 评估方法

评估指标包括两大类：检索相关指标与生成相关指标。

检索相关指标：提供四种指标来评价检索效果，包括 recall@k、precision@k、F1@k 和平均准确率（MAP）。与传统的独立检索系统评估不同，RAG 流程中的检索文档通常没有黄金标签（例如，相关或不相关的标记）。因此，通过判断检索文档中是否包含正确答案来评估相关性。用户还可以通过继承现有指标并调整计算方法来获取其他类型的指标。

生成相关指标：在评估生成品质方面，支持五种指标，涵盖了 token 级别的 F1 分数、精确度匹配、准确率、BLEU 评分 和 ROUGE-L 评分 。此外，还支持对生成过程中使用的 token 数量进行评估，帮助用户分析整个流程的成本。

为了满足用户自定义评估指标的需求，FlashRAG提供了一个指标模板供用户实现。由于库能自动保存执行过程中的中间结果，用户可以轻松地对中间组件产生的结果进行评估。例如，用户可以对比细化器处理前后的 token 数量变化，或是比较不同轮次检索结果间的精确度差异。

3 效果评估

FlashRAG为研究者提供了一个强大的平台，用以对RAG方法进行基准测试，评估他们独创的RAG方案，以及深入探究RAG领域的各种可能性。为展现FlashRAG的强大功能，我们开展了一系列实验，旨在提供可复现的基准测试和深入的探索性研究。

实验架构： 在核心实验中，采用了最新的LLAMA3-8B-instruct 作为生成组件，E5-base-v2作为检索组件，并以2018年12月的维基百科数据构建检索语料库。将生成模型的最大输入长度限定在4096。对于每条查询，检索并选取了五篇相关文档。对于那些不采用自定义提示的方法，统一使用了默认提示。任何需要特别设置和调整超参数的方法，在表格中以星号标注。所有的实验都在配备8块NVIDIA A100 GPU的系统上执行。

方法论： 对所有支持的RAG方法进行了全面的实验。这些方法根据它们主要优化的RAG组件被分类：

• • AAR 专注于优化检索组件；

• • LongLLMLingua 、RECOMP 和Selective-Context 致力于通过压缩输入提示来提升细化组件的性能；

• • Ret-Robust 和REPLUG 集中于改进生成组件及其解码方法；

• • SKR 用于增强判断查询是否需要检索的裁判组件；-

• • SuRe 、Self-RAG 、FLARE 、Iter-RetGen 和ITRG 则旨在整体优化RAG流程，这包括了多重检索和生成过程的优化。

3.1 主要结果

上表列出了多种方法的主要成效。总体来看，RAG方法相较于直接生成的基线有了显著提升。配备高级检索器和生成器的标准RAG，表现强劲，在六个数据集上均有良好表现。AAR通过训练对抗模型来优化检索器，并在多个数据集上与e5基线相媲美。在细化器方面，三种方法均显示出显著进步，尤其在HotpotQA和2WikiMultihopQA等多跳数据集上表现突出，这可能是因为复杂问题导致文档检索的准确性下降，产生更多噪声，需要细化器进行优化。在生成器优化策略中，Ret-Robust采用Llama2-13B模型搭配lora模块，显著提升了生成器对检索文档的理解，超越了其他无需额外训练的方法。RAG过程的优化效果因数据集而异，在NQ和TriviaQA等较简单数据集上，FLARE和Iter-RetGen与标准RAG不相上下或略逊一筹。然而，在需要复杂推理的HotpotQA等数据集上，改进显著，这表明自适应检索方法更适合解决复杂问题，而在简单任务上，可能成本较高而收益有限。

3.2 检索器的影响

在RAG流程中，检索器扮演着至关重要的角色，它对最终结果有着显著的影响。输入的检索文档数量和品质是决定答案的关键。然而，出于成本等因素的考量，现有研究往往采用固定的检索器和检索文档数量，忽略了这一领域的深入探索。

上面两个图展示了不同检索文档数量的实验结果。第一个图所示，当检索文档数量为3或5时，整体表现达到最佳。过多或过少的检索文档数量都会导致性能显著下降，降幅高达40%。这一趋势在包括密集型和稀疏型检索方法的不同检索器中均保持一致。

此外，当检索文档数量较多时，三种不同品质的检索器结果趋于一致。而在top1结果中，密集型方法（如E5、Bge）与BM25之间存在较大差异，表明检索文档数量越少，检索器的品质对最终结果的影响越大。

第二个图展示了检索文档数量对不同数据集的影响。可以看出，在大多数数据集上，采用top3或top5的检索结果能够带来最佳表现，这暗示了一个在检索文档品质与噪声之间取得良好平衡的策略。

4 局限性

FlashRAG工具箱目前尚有不足之处:

-（1）尽管力求涵盖众多典型的RAG方法，但受时间和成本的限制，未能包含所有现有的RAG研究成果。未来，期望开源社区能够贡献力量 -（2）工具箱目前不支持RAG相关组件的训练功能。虽然在最初设计时考虑了训练功能，但鉴于训练方法的多样性以及众多专门针对检索器和生成器训练的资源库的存在，决定不包含此部分。未来，可能会增加一些辅助脚本，以协助研究人员的训练需求。

• • 论文原文: https://arxiv.org/abs/2405.13576

• • 获取更多最新 Arxiv 论文更新: https://github.com/HuggingAGI/HuggingArxiv!

• • 加入社群，+v: iamxxn886