CORAG 的主要贡献总结
CORAG(Cost-Constrained Retrieval Optimization for Retrieval-Augmented Generation)是一种创新的检索增强生成(RAG)系统,旨在解决现有 RAG 方法中的关键挑战。以下是 CORAG 的主要贡献:
- 全面考虑块间关联:
- 创新点:CORAG 是首个在 RAG 任务中引入块组合顺序优化的框架。与传统方法不同,CORAG 不再独立处理每个文本块或仅在聚类级别考虑块,而是采用 蒙特卡洛树搜索(MCTS) 来顺序搜索最优的块组合顺序。这种方法能够充分捕捉块之间的复杂关联,避免冗余信息,并确保所选块组合能够全面回答用户查询。
- 优势:通过这种方式,CORAG 能够生成更准确、更相关的响应,提升生成质量。
- 解决块效用非单调性问题:
- 创新点:CORAG 将 预算约束 整合到块组合的优化过程中,而不是将预算耗尽视为终止条件。这种方法考虑了块效用的 非单调性,即添加更多块并不总是能提高最终结果的质量。
- 优势:通过在优化过程中动态调整预算使用,CORAG 能够避免过度包含无关或冗余的块,从而提高生成响应的准确性和相关性。
- 适应不同查询领域:
- 创新点:CORAG 引入了一个 配置代理,该代理使用 对比学习 来动态调整 MCTS 配置,以适应不同的查询领域。该代理推荐针对每个查询的最优 重排器模型 和 MCTS 配置。
- 优势:这种方法使 CORAG 能够灵活地处理各种类型的查询,从简单的关键字查询到复杂的推理问题,确保在不同的应用场景中都能提供高质量的响应。
- 高效且可扩展的搜索策略:
- 创新点:CORAG 采用 MCTS 进行节点搜索,并通过 并行扩展技术 加速搜索过程。这种方法将指数级的搜索空间缩减为线性,并有效地平衡了 探索 和 利用。
- 优势:CORAG 在保持高效检索的同时,能够处理大规模数据集,并在计算成本和检索质量之间取得平衡。
- 显著的性能提升:
- 实验验证:实验结果表明,CORAG 在多个基准数据集上均优于现有的基线方法,ROUGE 分数提高了约 30%。此外,CORAG 在效率上也表现出色,能够在严格的成本限制内提供高质量的响应。
CORAG 工作流示例
为了帮助读者快速理解 CORAG 的工作原理,以下是一个完整的工作流示例,展示了 CORAG 如何处理用户查询并生成最终响应。每一步都包含输入和输出,以体现工作流的前后衔接。
步骤 1:用户查询输入
- 输入:用户向 CORAG 系统提交一个自然语言查询,例如:
"请解释光合作用的过程,并列出影响其效率的因素。"
- 输出:查询被传递给 查询嵌入模块 进行处理。
步骤 2:查询嵌入生成
- 输入:用户查询文本。
- 处理:使用预训练的 嵌入模型(例如 BGE-M3)将查询文本转换为向量表示。
- 输出:查询嵌入向量(例如,1024 维向量)。
Query Embedding: [0.123, -0.456, 0.789, ..., -0.012]
步骤 3:配置代理预测
- 输入:查询嵌入向量
[0.123, -0.456, 0.789, ..., -0.012]
。 - 处理:
- 特征提取:嵌入向量被输入到 配置代理 的 编码网络 中进行特征提取。
- 重排器预测:编码网络输出预测的最优 重排器模型,例如
bge-reranker-large
。 - MCTS 配置预测:同时,编码网络预测 MCTS 的最优 配置参数,例如迭代次数、成本系数和探索系数。
- 输出:
- 最优重排器模型:
bge-reranker-large
- MCTS 配置参数:
- 迭代次数:15
- 成本系数:0.2
- 探索系数:2.5
Optimal Reranker: bge-reranker-large MCTS Configuration: - Iterations: 15 - Cost Coefficient: 0.2 - Exploration Coefficient: 2.5
- 最优重排器模型:
步骤 4:检索潜在文本块
- 输入:
- 查询嵌入向量
[0.123, -0.456, 0.789, ..., -0.012]
。 - 向量数据库(例如,包含预分割的文本块及其嵌入向量)。
- 查询嵌入向量
- 处理:使用查询嵌入向量进行 相似性搜索,从向量数据库中检索最相关的潜在文本块。以下是检索到的前五个文本块示例:文本块 1:
光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为葡萄糖和氧气的过程。
文本块 2:
光合作用主要发生在植物叶片的叶绿体中,分为光反应和碳反应两个阶段。
文本块 3:
影响光合作用效率的因素包括光照强度、二氧化碳浓度、温度和水分供应等。
文本块 4:
光照强度对光合作用的影响呈正相关,但过强的光照会导致光抑制现象,降低光合作用效率。
文本块 5:
二氧化碳是光合作用的原料之一,其浓度直接影响光合作用的速率。
- 输出:潜在文本块列表及其嵌入向量。
Retrieved Chunks: - Chunk 1: "光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为葡萄糖和氧气的过程。" - Chunk 2: "光合作用主要发生在植物叶片的叶绿体中,分为光反应和碳反应两个阶段。" - Chunk 3: "影响光合作用效率的因素包括光照强度、二氧化碳浓度、温度和水分供应等。" - Chunk 4: "光照强度对光合作用的影响呈正相关,但过强的光照会导致光抑制现象,降低光合作用效率。" - Chunk 5: "二氧化碳是光合作用的原料之一,其浓度直接影响光合作用的速率。"
步骤 5:MCTS 树搜索
- 输入:
- 潜在文本块列表:
- Chunk 1: ...
- Chunk 2: ...
- Chunk 3: ...
- Chunk 4: ...
- Chunk 5: ...
- MCTS 配置参数:
- 迭代次数:15
- 成本系数:0.2
- 探索系数:2.5
- 重排器模型:
bge-reranker-large
- 潜在文本块列表:
- 处理:
- 初始化根节点:根节点代表一个空状态,没有任何文本块被选中。
- 迭代扩展:
- 选择:使用 UCB 算法 选择当前效用最高的节点。例如,第一次迭代选择 Chunk 3,因为其与查询的相关性最高。
- 扩展:生成所有可能的子节点(新的文本块组合)。例如,Chunk 3 的子节点可能是 Chunk 1、Chunk 2、Chunk 4 和 Chunk 5。
- 评估:使用 重排器模型
bge-reranker-large
并行评估所有新组合的效用。例如,评估 Chunk 3 + Chunk 1、Chunk 3 + Chunk 2 等组合的效用。 - 更新:更新节点效用和访问计数,并向上传播更新。
- 重复迭代:重复选择、扩展、评估和更新步骤,直到达到最大迭代次数(15 次)。
- 终止条件:达到最大迭代次数。
- 输出:最优的文本块组合顺序。
Optimal Chunk Combination: - Chunk 3: "影响光合作用效率的因素包括光照强度、二氧化碳浓度、温度和水分供应等。" - Chunk 1: "光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为葡萄糖和氧气的过程。" - Chunk 2: "光合作用主要发生在植物叶片的叶绿体中,分为光反应和碳反应两个阶段。"
注意:在 MCTS 过程中,CORAG 会根据效用动态调整组合顺序,最终选择最优的组合顺序。
步骤 6:生成最终响应
- 输入:
- 最优的文本块组合顺序:
- Chunk 3: "影响光合作用效率的因素包括光照强度、二氧化碳浓度、温度和水分供应等。"
- Chunk 1: "光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为葡萄糖和氧气的过程。"
- Chunk 2: "光合作用主要发生在植物叶片的叶绿体中,分为光反应和碳反应两个阶段。"
- 用户查询:“请解释光合作用的过程,并列出影响其效率的因素。”
- 最优的文本块组合顺序:
- 处理:
- 构建提示:将最优的文本块组合顺序与用户查询结合,构建 LLM 的输入提示。例如:
用户查询:请解释光合作用的过程,并列出影响其效率的因素。 相关信息: 1. 影响光合作用效率的因素包括光照强度、二氧化碳浓度、温度和水分供应等。 2. 光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为葡萄糖和氧气的过程。 3. 光合作用主要发生在植物叶片的叶绿体中,分为光反应和碳反应两个阶段。
- 生成响应:使用 LLM(例如 Llama3)根据提示生成最终响应。
- 构建提示:将最优的文本块组合顺序与用户查询结合,构建 LLM 的输入提示。例如:
- 输出:最终的自然语言响应。
Final Response: "光合作用是植物、藻类和某些细菌利用光能将二氧化碳和水转化为葡萄糖和氧气的过程。它主要发生在植物叶片的叶绿体中,分为光反应和碳反应两个阶段。影响光合作用效率的因素包括光照强度、二氧化碳浓度、温度和水分供应等。"
步骤 7:响应输出
- 输入:最终的自然语言响应。
- 处理:将响应呈现给用户。
- 输出:用户界面显示响应结果。
总结
通过上述工作流示例,CORAG 的工作原理可以概括为以下步骤:
- 查询嵌入生成:将用户查询转换为向量表示。
- 配置代理预测:预测最优的重排器模型和 MCTS 配置参数。
- 检索潜在文本块:从向量数据库中检索最相关的文本块。
- MCTS 树搜索:使用 MCTS 搜索最优的文本块组合顺序。
- 生成最终响应:使用 LLM 根据最优文本块组合生成最终响应。
- 响应输出:将响应呈现给用户。
这些步骤展示了 CORAG 如何有效地结合检索、增强和生成三个阶段,以提供高质量的自然语言响应。通过详细的数据示例,读者可以更清晰地理解 CORAG 的数据处理过程及其工作原理。
原文:https://arxiv.org/pdf/2411.00744
标题:《CORAG:一种用于检索增强生成的带成本约束的检索优化系统》
作者:Ziting Wang, Haitao Yuan, Wei Dong (南洋理工大学), Gao Cong (南洋理工大学), Feifei Li (阿里巴巴集团)
摘要
大型语言模型(LLMs)在生成任务中展现了卓越的能力,但往往难以获取最新的信息,这可能导致幻觉的产生。检索增强生成(RAG)通过整合外部数据库中的知识来解决这一问题,从而实现更准确和相关的响应。由于LLMs的上下文窗口限制,直接将整个外部数据库上下文输入模型是不切实际的。相反,只有最相关的信息,即“块”(chunks),被选择性地检索。然而,当前的RAG研究面临三个关键挑战。首先,现有的解决方案通常独立选择每个块,忽略了它们之间潜在的关联。其次,在实践中,块的效用是“非单调的”,这意味着添加更多块可能会降低整体效用。传统方法强调最大化包含的块数,这可能会无意中损害性能。第三,每种类型的用户查询都有其独特的特性,需要量身定制的处理方式,而当前的解决方案并未充分考虑这一点。
为了克服这些挑战,我们提出了一种带成本约束的检索优化系统CORAG,用于检索增强生成。我们采用基于蒙特卡洛树搜索(MCTS)的策略框架,依次找到最优的块组合,从而全面考虑块之间的关联。此外,我们不是将预算耗尽视为终止条件,而是将预算约束整合到块组合的优化过程中,有效解决了块效用的非单调性问题。此外,通过设计配置代理,我们的系统预测每种查询类型的最佳配置,提高了适应性和效率。实验结果表明,与基线模型相比,我们的框架在性能上有高达30%的提升,突显了该框架的有效性、可扩展性以及适用于长上下文应用的适用性。
PVLDB 参考文献格式:
Ziting Wang, Haitao Yuan, Wei Dong, Gao Cong, 和 Feifei Li. CORAG: A Cost-Constrained Retrieval Optimization System for Retrieval-Augmented Generation. PVLDB, 14(1): XXX-XXX, 2020.
doi:XX.XX/XXX.XX
1 引言
尽管LLMs在生成任务中展现了卓越的能力,但它们在获取最新信息方面往往遇到困难,这可能导致幻觉的产生[10, 38]。为了应对这些挑战,RAG已成为一种关键的解决方案。通过将外部数据源整合到LLM中,RAG可以提供更准确、相关和最新的信息。如今,RAG在LLMs的背景下得到了广泛的研究,特别是在需要更新外部知识的任务中,如问答任务[2, 22, 29]、医疗信息检索[1, 32]和时间序列分析[12, 26, 40]。外部数据源通常非常庞大,直接将它们输入LLM是不切实际的。为了解决这个问题,数据通常被分割成不重叠的块并存储在向量数据库中,然后用户查询最有用的块以构建LLM的提示。因此,设计高效且准确的搜索结构和算法以找到最相关的块已成为一个突出的研究课题,并在数据库[39, 48]和机器学习社区[2, 35, 43]中得到了广泛研究。
然而,现有的方法存在三个关键挑战。
挑战1:块之间的关联。 目前,有两种主要方法用于识别最相关的块。第一种方法将问题表述为近似k近邻(AKNN)任务[41, 45],其中每个块被分配一个分数,然后选择按分数排名的近似前k个块。第二种方法对块进行聚类,响应查询时返回最相关聚类中的所有块[22, 29]。然而,这两种方法都忽略了块之间的潜在关联:第一种方法完全忽略了关联,而第二种方法仅通过将每个聚类中的所有块视为同等相关来表面地考虑它们。因此,当多个块传达相似或冗余的信息时,这些方法在选定的块中引入了大量冗余。
例如,如图1所示,当查询埃菲尔铁塔的高度和历史时,如果每个块都被独立对待,贪婪方法会选择块χ3和χ1,因为它们具有前两个最高分数。然而,这两个块仅提供历史信息,这不足以完全回答查询。为了更好地回答查询,有必要包括包含建造者姓名的块,例如χ4。另一方面,聚类方法将返回所有χ1、χ2、χ3和χ4,导致冗余。最优解决方案将选择χ3和χ4,因为它们提供了所需的信息而没有冗余。此外,研究[11, 19, 42]表明,块的顺序会影响LLM的性能,这是现有方法也忽略的一个因素。以埃菲尔铁塔为例,当选择块χ3和χ4时,将χ4放在第一位与相反的顺序相比会产生更高的分数。但是,确定最优的块组合顺序是一项具有挑战性的任务,因为两者都需要搜索空间随可用块数的增加而呈指数级增长。在本文中,我们进一步证明了这个问题是NP难的(见第2.1节)。
挑战2:效用的非单调性。 当前的解决方案基于包含更多块总是会产生更好最终结果的假设。具体来说,在基于AKNN的方法中,每次都确定性地选择恰好k个块。在基于聚类的方法中,设置了聚类与查询之间的距离阈值,并返回此阈值内的所有聚类。两者都尽可能多地返回块。然而,在实践中,块的效用不是单调的。更具体地说,过多的块会通过添加边缘相关的内容来稀释关键信息,产生降低清晰度的噪声。此外,块之间的冲突或细微差异可能会使模型感到困惑,降低响应的质量。例如,如图1所示,当选择χ3和χ4时,添加块χ1会降低效用,这突显了实践中效用分数通常是非单调的。
挑战3:查询的多样性。 用户查询有不同类型,每种类型都需要根据其独特特性采用不同的排名策略[47]。在当前的RAG系统中,块的效用分数通常由分配的重排名器模型确定。到目前为止,存在各种重排名器模型,但我们观察到它们在不同查询类型上的性能差异很大,没有一个单一固定的重排名器模型能够在所有查询变化中始终优于其他模型(更多详情见第6.3.4节的实验)。当前的方法[20, 46]通常依赖于静态重排名器模型进行块排名,缺乏适应不同查询上下文的灵活性。
问题陈述: 是否存在一个RAG系统,它充分考虑了块之间的关联和效用的非单调性,同时能够适应所有类型的查询?
1.1 我们的贡献
在本文中,我们通过提出一种新颖的基于MCTS的策略树框架来优化RAG系统中的块检索,从而肯定地回答了这个问题。总的来说,我们的贡献可以概括如下:
• 我们提出了第一个考虑RAG任务中块组合顺序的RAG框架。我们不是独立考虑每个块或以聚类级别考虑,而是使用MCTS来帮助依次搜索最优的块组合顺序。高层次的想法如下:首先,我们初始化根节点。然后,在迭代过程中,我们通过选择最高效用的节点并计算其扩展节点的效用来扩展树。每次扩展后,我们更新整个策略树中的效用。在此过程中,每次迭代的决策取决于已选择的块,使我们能够充分考虑块之间的关联。此外,MCTS将指数搜索空间缩减为线性,我们应用并行扩展技术进一步增强计算效率。通过这样的设计,我们解决了挑战1。
• 与之前将预算耗尽视为终止条件的RAG框架不同,我们提出了一种新颖的公式,其中预算约束被整合到优化块组合的过程中,以充分考虑块效用的非单调性,从而解决了挑战2。此外,通过优先考虑高相关性、低成本的块并考虑标记长度,我们进一步降低了计算成本。
• 我们提出了一种基于对比学习的代理,它根据查询动态调整MCTS配置,使重排名器模型和配置适应特定的查询领域。这种方法为动态、领域特定的查询提供了灵活性和稳健性,适应了挑战3。
• 此外,我们进行了全面的实验,将我们的框架与几种最先进的方法进行了比较。结果验证了我们方法的有效性、效率和可扩展性,并且与基线相比,性能提高了30%。
2 预备知识
在本节中,我们首先在第2.1节中介绍了一些关键概念的定义,例如块和块组合顺序。接下来,我们给出了块顺序优化问题的NP难证明。最后,我们在第2.3节中讨论了相关工作。
2.1 关键概念
RAG & Chunks(块)。 RAG是一种通过从外部语料库中检索相关上下文来提高生成模型性能的有效方法。在这种方法中,语料库首先被分割成更小、更易管理的单元,称为块,这些块存储在向量数据库中。因此,我们可以给出块的正式定义如下:
定义2.1(块)。 让C代表文档的语料库,块χ定义为从C中提取的连续文本块。形式上,块χ由一系列标记(t1, t2, ..., tn)组成,其中每个ti都是C中的标记,大小n由用户设置。
在RAG系统中,每个块都使用嵌入模型嵌入到向量表示中,该模型捕获块的语义含义,并能够检索上下文相似的内容。当接收到新查询时,向量数据库执行相似性搜索以识别与查询在语义上最相关的块。然后将这些检索到的块传递给生成器(例如,大型语言模型)以基于检索到的内容生成最终响应。具体来说,块包含的标记越多,生成器产生的成本就越高。因此,我们定义块的代价为cost(χ)=|χ|,等于块中标记的数量。
块组合顺序。 在RAG系统中,向量数据库的检索结果可能包括多个块。然而,由于生成模型的输入限制,使用所有这些块是不切实际的。因此,有必要选择最优的块子集,称为块组合,以适应给定的成本预算。此外,组合内块的顺序对生成模型的性能有显著影响。目标是识别具有最佳顺序的块组合顺序,正式定义如下:
定义2.2(最优块组合顺序选择)。 让{χ1, χ2, ..., χk}是一组潜在块,ℛ是成本预算,Φ=⟨χφ1, ..., χφm⟩代表一个潜在块组合顺序,其中每个χφi是一个块,索引φi表示其在Φ中的位置。让U(Φ)是由重排名器模型分配的效用分数,它可能是任意的或复合的。我们的目标是找到在将它们输入LLMs以生成最终响应的成本约束下,最大化效用分数的块组合顺序,即搜索
2.2 NP难的证明
为了证明块组合顺序选择是NP难的,我们将最大加权超团问题(MWHP)归约到它。由于MWHP是NP难的,我们展示了任何MWHP实例都可以在多项式时间内转换为块组合优化实例。
2.2.1 MWHP的问题定义
给定一个超图ℋ=(V, E, w1, w2),其中V是顶点集,E是超边集,每个超边包含V的一个子集。w1:v→ℝ和w2:e→ℝ是权重函数,分别给每个顶点和超边分配权重。给定一个顶点子集V'⊆V,我们说一个超边e属于V',即e∈V',如果V'覆盖了e的所有顶点。目标是找到k个顶点,最大化这些顶点和它们覆盖的超边的权重之和:
2.2.2 归约过程
现在我们从给定的MWHP实例构造一个相应的块组合优化问题实例。对于每个节点v∈V,我们创建一个相应的块Xv。我们定义其成本cost(Xv)≡1。然后,一个块组合顺序Φ对应于V的一个顶点子集,表示为V(Φ)⊆V。我们定义其效用为
最后,我们设置B=k,我们的目的是
用Φ表示(4)的解,那么,显然V(Φ)是(2)的解,归约可以在O(|V|·|E|)时间内完成。
请注意,这个归约的前提是,在我们的块组合优化问题中,我们允许重排名器是任意的,这意味着效用分数也可以任意分配。如果对重排名器做出某些假设,可以显著降低找到最优块组合顺序的复杂性。例如,如果重排名器不考虑关联,只是线性地求和各个块的效用分数,则每个块可以独立评估。然而,在本文中,我们处理的是最一般的情况,不对重排名器模型做任何假设。
2.3 相关工作
2.3.1 检索增强生成
RAG[14, 20]被广泛用于处理知识密集型NLP任务。在典型的RAG流程中,基于密集嵌入的检索器从外部数据库中搜索相关信息,然后在生成过程中由LLM使用。为了改进这一流程,一些研究[5, 18, 22, 35]侧重于调整检索器以更好地适应LLM的生成需求,开发多步骤检索方法,并过滤掉不相关的信息。尽管有许多先进的检索器[8, 9, 15, 16, 27, 34],但将检索器和LLM一起优化在一个端到端的过程中[25, 31]更有前景。例如,研究[30]侧重于同时或分阶段共同训练检索器和LLM。然而,这需要代理损失来进行优化,并使训练流程复杂化,特别是当嵌入数据库需要频繁重新索引时,这将带来高计算成本。因此,像[5]这样的方法将复杂的多步骤查询分解为更小的子意图,以提高响应的全面性,而无需频繁重新索引。然而,这些方法通常忽略了块组合顺序的关键作用,这可以显著影响LLM的整体响应质量。据我们所知,本文是第一个在RAG任务中考虑块组合顺序的方法。
2.3.2 RAG的重排名
重排名方法对于提高RAG流程中的检索性能至关重要[43, 44, 51]
传统的重排名方法[33, 50]通常依赖于中等规模的语言模型,如BERT或T5,对检索到的上下文进行排名。然而,这些模型通常难以捕捉查询和上下文之间的语义关系,特别是在零样本或少量样本设置中。因此,最近的研究[43]强调了指令调整的LLM在改进上下文重排名方面的潜力,即使在存在噪声或不相关信息的情况下也是如此。尽管有这些进展,LLM在RAG系统中的重排名能力仍未得到充分利用。特别是,研究表明,块排列会影响LLM的性能[19],强调了需要在RAG任务中考虑块组合顺序。然而,现有的模型并不适合需要特定序列或组合的块才能实现最佳检索的情况,而不是孤立的块。因此,未来的研究需要更好地利用LLM来更有效地排列块以响应RAG框架内的查询。
2.3.3 大型语言模型强化学习
最近,强化学习(RL)越来越多地被用于各种数据管理和RAG任务。RL技术可以使大型语言模型通过利用外部知识来源(如搜索引擎)[13, 23]来提高其生成能力。特别是,可以整合人类反馈[4, 36, 37]来帮助模型通过RL框架产生更准确和上下文相关的响应。此外,一些查询优化方法[17, 21, 49]进一步改进了检索过程,允许模型性能通知查询调整,并最终提高下游任务的结果。在这项工作中,我们应用了一种轻量级的RL技术MCTS来优化RAG系统中块组合顺序搜索过程。我们还引入了一个配置代理来指导MCTS搜索过程。据我们所知,这是第一个解决这一特定问题的方法。
3 系统概述
如前所述,现有的RAG框架面临三个关键挑战:如何充分考虑块之间的关联以及块组合顺序效用的非单调性,并适应不同的查询领域。这些挑战导致输出相关性降低。为了解决这些问题,我们引入了CORAG,这是一个旨在检索最优块组合的系统,同时考虑查询领域和用户预算。作为我们系统最重要的组成部分,我们介绍了最优块组合搜索模型。该模型采用基于MCTS的策略树进行块组合顺序的顺序搜索,允许我们充分考虑块之间的关联(挑战1)以及块组合顺序效用的非单调性(挑战2)。此外,我们提出了一个配置推理模块,它推荐针对各种查询领域的最佳MCTS配置和重排名器,从而解决挑战3。下面,我们对这些两个模块进行简要描述。
最优块组合搜索: 一种考虑块关联的直接方法包括从向量数据库中检索潜在块(如步骤1所示),然后详尽地探索所有可能的块组合。然而,这种方法会导致显著的延迟和计算成本。为了缓解这种情况,我们构建了一个策略树(如步骤2所示),将最优块组合搜索重新构建为树内节点搜索问题。具体来说,策略树的根节点代表一个初始的空状态,每个子节点对应于特定组合的块。例如,如果根节点有子节点代表块χ1,其一个子节点可能代表组合χ1+χ2,而另一个可能代表χ1+χ3
我们设计了一种基于MCTS的搜索算法来解决这个问题。与传统的MCTS不同,我们的方法在每次迭代中扩展最高效用的节点,同时评估所有可能的子节点。此外,在策略树搜索过程中,我们还考虑了成本和预算约束。节点效用是通过平衡探索与成本控制来计算的,优化了效率和准确性。
配置推理: 配置调整的一个简单解决方案是枚举每一种可能的配置或重排名器,并并行计算结果,然后选择最佳配置。然而,这将导致RAG系统的不切实际的成本。为了优化策略树搜索过程的配置(即迭代次数、成本系数和探索系数),我们引入了一个配置代理,它根据查询领域动态生成配置。为了确保模型的有效性,我们采用了一种对比学习方法,使用正负标签对:正标签对应于来自同一最佳重排名器的查询嵌入,而负标签则来自不同的最佳重排名器。联合损失函数用于同时优化回归(用于参数调整)和对比学习(以增强标签区分)。
总结。 我们框架的流程如图2所示。我们首先为输入查询生成嵌入,然后使用它从向量数据库中检索潜在块。这些查询嵌入也被输入到配置代理中,该代理根据查询领域动态生成最佳MCTS配置。使用这个最佳配置,我们可以在策略树中搜索,以确定从检索到的潜在块中确定最佳块组合和顺序。最后,这个最佳块组合用于构建LLMs的最终提示。
4 块组合检索
如前所述,块组合的顺序对LLMs的提示构造效率有显著影响。由于潜在组合的数量巨大,特别是在涉及大量块的情况下,枚举所有可能的块组合顺序是不可行的。在本节中,我们提出了一种新颖的方法,在搜索最优块组合顺序问题上实现了效率与精度的良好平衡。在第4.1节中,我们将问题建模为在策略树中搜索最优节点(第4.1节)。然后,我们提出了一种基于MCTS的算法来解决这个节点搜索问题(第4.2节)。
4.1 策略树搜索建模
为了找到最优组合顺序,第一步是找到一个能够有效枚举所有可能组合顺序的数据结构。一个自然的选择是树,通过从根到叶节点的遍历,我们可以探索所有潜在的答案。
策略树。 如图3所示,我们构建了一个策略树来表示从向量数据库中提取的所有可能的块组合顺序。具体来说,根节点象征着没有任何块的初始状态,每个后续节点描绘了从潜在块中选出的一个块。因此,子节点通过从潜在块的队列中选择下一个可用块并将其合并到祖先节点建立的序列中,从而从其父节点生成。例如,如果一个节点代表块组合顺序{χ1},那么一个子节点可能包含后续组合顺序,如{χ1, χ2}、{χ1, χ3}或{χ1, χ4}。因此,我们正式定义策略树如下:
定义4.1(策略树)。 给定一个查询q和一组潜在块{χ1, χ2, ..., χn},我们构建一个策略树T。T的根节点代表初始状态,没有任何块。每个后续的非根节点包含一个块集,通过将新选择的块从剩余的潜在块中合并到其父节点的序列中。这个过程在每个非根节点中顺序地构造一个有序的块组合,我们的目的是找到具有最高效用分数的节点。
在策略树中,我们的目标是选择一个包含有序块的节点,以最低的成本提供最大的效益。为了实现这一点,我们需要设计一个效用计算函数来评估效益与成本之间的权衡。这个函数通过我们定义的“节点效用”来量化,如下所述。
节点效用。 效用指标包括两个组成部分:从选择块组合中获得的效益以及使用块作为LLMs提示的成本。具体来说,效益由LLMs量化,它可以衡量所选块与查询之间的相似性。特别是,我们将其表示为节点值V。接下来,我们进一步使用上限置信区间(UCB)[3]算法来平衡节点值V(vi)和搜索计数N(vi)的利用(节点值V(vi))和探索(搜索计数N(vi))之间的平衡。对于给定的节点vi,关于成本,我们考虑第2节中定义的标记成本,并通过当前块组合的成本相对于总分配预算B的比例来衡量。因此,节点效用定义如下:
定义4.2(节点效用)。 给定一个策略树和一个成本预算B,非根节点的效用定义为:
其中,V(vi)是节点vi处块组合的估计效益值,由训练模型确定,N(vi)是访问节点vi的次数,促进了较少访问节点的探索,N是策略树中所有节点的总访问次数,确保探索与利用之间的平衡。此外,cost(vi)表示节点vi的标记成本,B是总标记预算,c调节探索-利用权衡,λ作为成本惩罚因子以提高成本效率。
最优节点选择建模。 基于定义的节点效用,选择最优块组合顺序的任务,如第2节所述,被重新表述为在策略树T中选择最优节点。在给定的预算约束B下,目标是识别节点vi⊆T,以最大化效用U(vi),同时确保与vi相关的总成本不超过B。形式上,这表示为:
其中,V(vi)是节点vi处块组合的估计效益,而cost(vi)代表其相关成本。这个公式使得在给定预算内选择最大化效用的块成为可能。
4.2 基于MCTS的策略树搜索
动机。 枚举策略树中的所有节点将找到最优节点,但会导致高计算成本。为了解决这个问题,一个直接的方法是应用贪婪策略,从根节点开始迭代地导航树。在每次迭代中,选择效益最高的子节点,继续直到预算耗尽。然而,这种方法很可能会导致次优结果。例如,χ1的效益可能略高于χ2,但χ2+χ3的效益可能大大超过χ1+χ3。在这种情况下,贪婪的方法可能导致次优结果。因此,有必要重新访问高效益的父节点。同时,我们需要减少对低效益节点的探索。
为了实现我们的目标,我们提出了一种基于MCTS的策略树搜索方法,旨在有效地选择和排名块组合。这种方法迭代地探索潜在块顺序的空间,在指定的预算约束内优化给定的查询。
概述。 基于MCTS的策略概述在算法1中。我们首先使用输入查询初始化策略树的根节点。当计算预算没有耗尽时,我们迭代地执行两个关键步骤:节点选择和效用更新。一旦达到迭代限制或预算,我们就停止该过程,并在树中递归搜索以找到最高效用的节点。与通常只关注根节点的传统的MCTS策略不同,我们的方法还考虑了有希望的中间层节点,以最大化块组合效用。
关键功能的详细解释。 我们进一步解释这两个关键功能如下:
- 节点选择(算法2)。 我们递归地选择具有最高效用值的节点,该节点最有可能导致最优的块组合。具体来说:
- 选择: 我们识别出具有最大效用值的节点。如果𝑣尚未扩展,我们生成所有可能的子节点并将它们纳入策略树。如果𝑣已经扩展,我们选择具有最高效用的后代节点以进行进一步探索。
- 扩展: 在选择具有最高效用的节点后,我们通过生成所有潜在的子节点来扩展它。每个子节点代表一个新的可能的块组合顺序。我们的方法采用并行扩展,它可以同时计算和评估多个子节点。这种并行性利用了值网络处理多个组合的能力,其计算成本与单个节点相似,增强了搜索效率。
- 计算效用: 我们使用效用公式计算每个新子节点的效用值。重排名器模型R并行处理多个块组合,生成:
5 配置代理
在解决了用户预算内处理块关联的效率问题后,剩余的任务是设计一个系统以适应每个查询的领域。MCTS过程涉及几个关键配置,包括重排名器选择、迭代次数、探索系数和成本系数。在不同查询类型中优化这些配置尤其具有挑战性。为了解决这个问题,我们提出了一个配置代理框架,它预测每个查询的最佳重排名器和配置。在本节中,我们首先在第5.1节中介绍代理框架,然后在第5.2节中概述模型学习流程。
5.1 模型框架
动机: 为了应对挑战3,需要适应每个查询的领域并推荐最佳配置,一个直接解决方法是使用MLP分类器将每个查询分配给其最佳重排名器。然而,初步实验表明,MLP分类的性能不佳。经过进一步分析,我们观察到,相似类型的查询往往共享相同的最佳重排名器和配置。因此,利用具有对比学习的Siamese网络将同一类别的查询拉近,同时将不同类别的查询推开,是一种更可行的方法。
图4提供了我们配置代理的概览,它由两个主要模块组成,负责将输入转换为特征图。首先,输入嵌入模块生成输入查询的嵌入。随后,编码网络处理这些嵌入以产生特征图,然后用于导出MCTS设置的各种配置。
以下各节将详细介绍每个组件,并解释其设计原理。
(1) 输入查询嵌入: 为了有效捕捉各种查询的因素,鉴于查询类型的多样性,如显性事实、隐性事实、可解释的理性、隐藏的理性[47],我们采用BGE-M3[6]嵌入模型为每个查询生成嵌入。这些嵌入增强了学习框架的能力,通过将相似类型的查询映射到同一重排名器类别。表示在1024维空间中,嵌入捕获了基本的语义特征,使编码网络能够高效地进行比较和分类。这一步有助于提高不同查询类型的检索相关性。此外,使用相同的嵌入模型也生成最优重排名器注释嵌入,包括其独特特征和关联的元数据,使模型能够将查询与最优重排名器对齐。
(2) 使用编码网络生成特征图: 为了优化重排名器选择任务,并为不同查询类型推荐每个查询的最佳重排名器,我们利用编码网络来有效地学习对分类和配置预测都有用的表示。我们使用由三个全连接层组成的Siamese网络来完成这一任务。它处理维度d=1024的输入查询嵌入,并学习分类输出和MCTS配置预测(即迭代次数和λ)。编码网络的分支共享权重,每个分支应用线性变换,然后进行RELU激活。依次地,第一个隐藏层将维度减少到512,第二个减少到256,第三个减少到128。最终输出层提供了分类预测,指定每个查询的最佳重排名器,以及回归输出,用于预测最有效的MCTS配置,指导搜索过程。分类输出识别每个查询的最佳重排名器,而回归输出决定了最佳MCTS配置设置。
5.2 联合训练
在本节中,我们概述了开发配置代理的训练流程。如图4所示,我们实施了三个联合训练任务,以提高模型的训练效率。前两个任务涉及分类和回归,分别用于选择最佳重排名器和预测MCTS超参数的最佳值。此外,我们还结合了对比学习方法来进一步改进学习过程。
5.2.1 分类和回归损失。
给定预测的重排名器标签Y(pred)及其对应的实际最优重排名器Y(true),分类损失L(cla)计算如下:
其中,F(cla)代表预测和实际重排名器标签之间的交叉熵损失。这个损失函数有助于准确分类每个查询的最优重排名器。同样,回归损失L(reg)定义为:
其中,𝐹reg 是预测的 MCTS 参数 𝑝pred 与实际的 MCTS 参数 𝑝true 之间的均方误差 (MSE)。此指标确保了对 MCTS 配置的精确预测,包括迭代次数和 𝜆。
5.2.2 对比学习。
为了有效地区分不同的查询领域,并为每个查询推荐最佳配置,我们利用对比学习将同一领域的查询拉近,同时将不同重排名器类的嵌入推开。
对比对准备。为了准备训练数据集,我们必须确定每个查询的最优重排名器和配置。在本研究中,通过对各种设置进行广泛的实验,确定了每个查询最适合的重排名器和相应的配置。随后,根据这些最优重排名器注释生成查询对。正对由共享同一最优重排名器的查询组成,促进它们在特征空间中的嵌入最小化。相反,负对由具有不同重排名器的查询组成,其目标是最大化它们之间的嵌入距离。由于某些重排名器在某些查询上表现相似,我们只选择ROUGE-L差异超过10%的案例来形成我们的训练数据集。
对比损失。如图 4 所示,对于给定的正对 (𝑥𝑖, 𝑥+𝑖) 和负对 (𝑥𝑗, 𝑥−𝑗),我们首先使用编码模型生成它们对应的特征图。然后,这些特征图被用于计算对比损失 𝐿con。具体来说,这个过程可以表示如下:
其中,𝑓𝜃(𝑥) 表示嵌入函数,𝐹con 是应用于两种类型对的相似性函数:正对(具有相似的重排序器)和负对(具有不同的重排序器)。此损失函数旨在确保具有相同重排序器的查询在嵌入空间中更接近,而具有不同重排序器的查询则更远。
5.2.3 整个训练过程。
最后,总损失函数L(total)是对比、分类和回归损失的总和,如下所示:
特别是,对比损失Lcon(θ)鼓励具有相同最优重排名器的查询的嵌入更接近,同时将具有不同重排名器的查询的嵌入推开。分类损失Lcla(θ)帮助模型使用交叉熵正确识别重排名器,而回归损失Lreg(θ)最小化预测最优MCTS配置的错误。
备注。一旦总损失Ltotal计算完毕,网络参数θ使用梯度下降法以学习率η进行更新。这个优化过程在多个周期E和批次中重复进行,确保随着时间的推移,分类器和参数预测都得到改进。
6 实验
实验研究旨在回答以下问题:
• RQ1 我们的CORAG在成本受限的RAG管道中的效果与其他方法相比如何?
• RQ2 CORAG在不同的块大小下的效率如何?
• RQ3 当前RAG的瓶颈是什么?
• RAG的瓶颈是什么?
• RQ4 CORAG在不同的数据集大小下的可扩展性如何?
• RQ5 CORAG中每个设计的有效性如何?
6.1 实验设置
环境。我们将我们的系统与流行的RAG框架LlamaIndex集成。实验在配备Intel Core i7-13700K CPU(12核心,24线程,5.3 GHz),64GB RAM和1 TiB NVMe SSD的Linux服务器上运行。配置代理模块在配备24GB VRAM的NVIDIA RTX 4090 GPU上实现,使用PyTorch 2.0实现。
表格1:实验中使用的统计数据。
数据集 | #train | #dev | #test | #p |
MSMARCO | 502,939 | 6,980 | 6,837 | 8,841,823 |
Wiki | 3,332 | 417 | 416 | 244,136 |
数据集。为了评估CORAG在不同场景下的性能,我们在两个不同的数据集上进行了实验,具有不同的任务焦点:(1) WikiPassageQA[7]是一个问答基准,包含4,165个问题以及超过100,000个文本块,旨在评估段落级检索。(2) MARCO[24]是一个综合数据集,专门针对自然语言处理任务,主要强调问答和信息检索。如表1所示,WikiPassageQA和MARCO都提供高质量的问题和段落注释,使其适合评估检索有效性。在我们的实验中,我们提示LLMs为每个数据集生成地面真相答案。例如,如果我们使用Llama3来评估CORAG的性能,我们也在相同的实验设置中使用Llama3生成地面真相,以保持公平性和与LLMs特性的对齐。
基线。我们将CORAG的性能与两个典型的RAG基线进行比较:
• RAPTOR[29]: RAPTOR通过递归嵌入、聚类和总结文本块,构建了一个分层文档摘要树,实现了多层次抽象。这种方法与第1节中讨论的基于聚类的方法一致。我们在大约预算限制内完成树构建。
• NaiveRAG: 这是一种检索相关块的基本方法。首先,从向量数据库中基于向量相似性搜索检索候选块,然后使用重排名器模型对它们进行排名。这种方法是第1节中提到的AKNN方法。为了满足成本约束,我们采用贪婪的预算分配策略,检索块直到预算完全耗尽。
此外,我们移除我们方法中的配置代理作为基线,以评估其对CORAG性能的影响,将这个版本称为CORAG w/o Agent。最后,我们实施了一种称为CORAG Upper的方法,通过探索所有可能的块组合并选择最佳顺序来建立上界。由于潜在组合的数量巨大,在CORAG Upper的情况下,我们将探索限制在少于六个块的组合中。
备注。其他方法,如GraphRAG[22],严重依赖于频繁调用LLMs来总结块和构建索引,产生了巨大的成本(例如,数十亿个标记),超出了我们严格的成本约束。因此,这些方法对于解决我们的问题并不可行。为了公平比较,我们在实验中排除了这些类型的RAG方法。
超参数设置:CORAG的超参数由配置代理自动确定,而NaiveRAG不需要任何超参数。对于其他基线方法,我们确保一致性,使用相同的超参数进行公平比较。具体来说,我们设置探索系数为2.4,迭代次数为10,成本系数λ为0.1。初步实验表明,这个配置优化了基线性能。我们还将进行进一步的消融研究以验证这些设置。
学习参数设置。在我们的方法中,配置代理使用对比学习进行训练。在此过程中使用的超参数包括对比损失(margin=1.0),学习率(lr=0.001),批量大小(32),周期数(num_epochs=60)以及嵌入模型(即BAAI/bge-m3[6])。
评估指标。我们通过比较地面真相答案和生成响应的Rouge分数来评估有效性,使用Rouge-1、Rouge-2和Rouge-L作为评估指标。为了评估效率,我们测量使用不同方法回答查询所需的延迟。
6.2 性能比较
6.2.1 RQ1:Rouge比较。
如图2所示,我们比较了CORAG与几个基线在不同数据集上的性能,主要使用WikiPassageQA和MARCO。评估在三个不同的块大小下进行,使用Rouge-1、Rouge-2和Rouge-L指标来评估LLM由于我们的检索方法而生成的响应的改进。CORAG与主流RAG方法(如NaiveRAG和RAPTOR)相比有显著提高,大约25%。不出所料,CORAG没有超过上界,这代表了一个极端情况,即所有可能的组合顺序都被详尽枚举,这显然效率低下且不切实际。总之,CORAG优于基线,增强了检索相关性,同时有效地剪枝了搜索空间。
6.2.2 RQ2:效率评估。
如图5所示,由于CORAG基于树搜索算法,代理有助于预测给定查询的最佳重排名器和参数。因此,评估不同块大小和数据集对检索优化任务效率的影响至关重要。我们测试了使用不同数据集和块大小的效率,观察到使用传统检索方法的NaiveRAG实现了更短的检索时间,但Rouge分数较低。CORAG upper在Rouge方面表现良好,但效率显著降低,因为它探索了整个搜索空间。同样,RAPTOR利用外部LLM进行总结,表现出了较差的效率。相比之下,我们的CORAG方法在效率与检索相关性之间取得了平衡,实现了有效的权衡。
6.2.3 RQ3:性能分解。
我们展示了基线NaiveRAG的性能分解,以突出当前RAG系统的瓶颈。为了解决搜索最优块组合顺序的挑战,使用NaiveRAG实现它需要以下步骤:(a)获得查询嵌入,(b)检索潜在块组合,(c)重排名潜在块组合,以及(d)提示完善。图6报告了每个步骤的平均延迟,在之前的实验设置下。
6.2.4 RQ4:可扩展性评估
CORAG表现出卓越的可扩展性,特别是在处理大型数据集(如WikiPassageQA和MARCO)时,这些数据集包含大量段落。通过将每个段落分割成256个标记的块,块的数量可以轻松扩展到100k或更多。尽管数据量大幅增加,我们的检索时间仅比传统方法增加了10%,展示了我们的系统在管理大规模检索任务方面的效率。值得注意的是,我们的系统优于CORAG Upper方法,仅需十分之一的检索时间,同时仍能提供有竞争力的ROUGE分数。这种在性能和计算开销之间的有效平衡突显了系统有效剪枝搜索空间的能力,确保即使在庞大的数据集中也能快速检索。因此,我们的方法非常适合需要大规模数据处理和高检索准确性的场景。
6.3 RQ5:消融研究
6.3.1 不同预算的消融研究
如表3所示,我们评估了CORAG作为一个成本受限的系统,研究了不同预算对整体性能的影响。使用MARCO数据集,我们将预算限制设置为1024、2048和8192个标记,并使用ROUGE评估结果。CORAG在这些预算水平上始终优于所有基线。值得注意的是,CORAG的平均标记成本保持在每个预算限制以下,这表明块的效用不是单调的,正如挑战2中强调的那样。随着预算的增加,CORAG从扩大的搜索空间中受益,能够包含更多相关信息,而不仅仅是增加块的数量。
6.3.2 不同探索系数的消融研究
如图7所示,我们进行了消融研究,以评估不同探索系数对系统性能的影响,具体测试了C值为0、1、2和3。结果表明,探索系数约为2时提供了最佳性能,在搜索过程中实现了探索和利用之间的最佳平衡。这种平衡使系统能够有效地发现相关信息,同时保持对高潜力块的关注,最终导致RAG响应的改进。相比之下,较低的探索系数由于探索不足而导致了次优结果,而较高的探索系数则由于过度分散焦点而导致了次优结果。这些发现强调了探索系数在CORAG搜索过程中的性能中的关键作用,并突显了仔细参数调整的重要性。
6.3.3 不同成本系数的消融研究
如图8所示,我们进行了消融研究,以评估不同成本系数对系统性能的影响,具体测试了值0、0.1、0.2和0.3。结果显示,在效用中引入成本系数导致ROUGE分数略有下降。这种下降是因为在没有成本约束的情况下,CORAG倾向于产生更长的输出,尽管这是以成本效率为代价的。然而,尽管ROUGE分数略有下降,但下降幅度保持在5%以内,这是可以接受的。这些结果突显了有效调整成本系数的重要性,以平衡输出丰富度和成本约束,进一步强调了我们的配置代理在实现高效配置调整以优化CORAG性能方面的作用。
6.3.4 不同重排器的消融研究
为了评估不同重排器对检索性能的影响,我们使用六种广泛认可的重排器模型进行了消融研究:jina-reranker-v1-turbo-en、jina-reranker-v2-base-multilingual、bge-reranker-v2-m3、bge-reranker-large、bge-reranker-base和gte-multilingual-reranker-base。这些重排器在MARCO数据集上使用llama3-8B模型进行评估,配置了固定成本系数0.1,探索系数2.4和预算限制1024。
表4中的结果显示,不同重排器之间的性能存在差异,突显了在特定操作约束下仔细选择重排器以优化RAG系统性能的重要性。在重排器中,gte-multilingual-reranker-base和bge-reranker-large在QA任务上表现出始终如一的强大性能,这表明这些重排器模型在捕捉不同QA查询的相关信息方面具有很高的效率。我们观察到,随着消融研究中块大小的增加,每个单独的重排器在针对不同查询的推荐重排器方面的性能都低于代理。这表明配置代理有效地利用了重排器的多样性,动态调整配置以改进检索结果。配置代理推荐更好的重排器选择和参数配置的能力突显了它在最大化RAG系统性能方面的作用,特别是在预算有限等约束条件下。
6.4 案例研究
图9展示了三个示例,比较了CORAG和传统的NaiveRAG方法的检索质量,重点说明了我们的方法为何优于基线方法。由于其简单的top-k检索和重排序,NaiveRAG经常错过与查询意图相关的重要信息,因为它通常基于关键词匹配而不是与查询的相关性来检索块。例如,对于查询“叶子花是灌木吗?”,NaiveRAG检索包含匹配关键词的内容,但未能提供叶子花的实际分类。相比之下,CORAG的块组合策略检索到的上下文包括叶子花的类别,使LLM能够给出更准确的响应。在另一个案例中,NaiveRAG检索包含“oxyfluorfen”的术语和法律条款,但缺乏对查询意图的理解,而CORAG提供了将oxyfluorfen与在棉花中的使用案例联系起来的内容,这需要NaiveRAG的向量相似性搜索无法捕捉的块之间的逻辑关系。最后,对于查询“细菌从哪里来?”,NaiveRAG检索包含关键词“细菌”的块,但没有解决其起源问题,而CORAG提供了更完整的响应,包括细菌的来源及其繁殖条件。这些案例说明,CORAG在检索逻辑上连接的信息方面表现出色,使其比NaiveRAG更适合需要超过简单关键词匹配的查询。
7 对RAG的见解和未来设计选择
7.1 当前RAG的缺点
我们提供了对当前RAG系统的分析,揭示了在检索(S1)、增强(S2)和生成(S3)阶段面临的性能挑战。
S1:检索开销 当前RAG系统通常利用LLMs进行摘要和索引结构,忽视了与外部LLMs相关的计算成本,从而增加了计算开销。基于模型的重排器虽然在检索过程中提高了相关性,但引入了显著的延迟,这可能会阻碍对延迟敏感的上下文效率。成本效益高的索引构建和重排序优化对于平衡效率和性能至关重要。
S2:增强开销 检索后的技术,如优化的块组合排序,增强了上下文相关性,但需要额外的计算。最小化搜索空间和优化组合顺序的剪枝策略对于平衡计算成本和增强上下文相关性至关重要。高效的块组合优化,强调顺序和连贯性,对于降低成本和提高检索性能至关重要。
S3:生成开销 为最佳块组合进行有效的提示工程需要大量的计算资源。特定于查询的提示细化和压缩对于在保持输入相关性和简洁性的同时减少开销至关重要。处理不同查询类型和特定领域需求的适应性策略确保了提示效率,同时不会影响输出质量。
7.2 设计选择
为了应对上述挑战,以下设计选择旨在优化RAG系统的性能:
P1:检索和重排序过程的协同设计 在CORAG中,树搜索中的并行扩展通过实现并发检索和重排序,加速了查询处理,显著减少了延迟。未来的优化可以通过消除阶段特定的延迟,进一步提高排名效率,从而解决瓶颈。这种协同设计方法有效地管理块组合顺序,改善排名过程和相关评分。
P2:树结构和搜索迭代的优化 结果表明,较短的策略树高度提高了搜索效率,通过减少计算开销,这对大型数据集尤其有利。在基于树的搜索中最小化树高度可以提高上下文相关块的搜索速度,显著降低延迟和计算成本。这种优化方法提高了RAG系统在大规模数据集上的性能。
P3:动态提示工程 根据查询类型选择重排器并使用适应性提示模板可以提高LLM的检索相关性。动态提示结构与查询意图和特定领域上下文保持一致,可以在资源约束内保持输出质量。这种适应性提示工程方法在效率与检索质量之间实现了有效平衡,解决了RAG系统查询的动态性。
8 结论
考虑到块的效用的非单调性、块之间的相关性以及不同查询领域的多样性,我们提出了一个基于学习的检索优化系统CORAG。最初,我们将块组合顺序建模为策略树,并使用MCTS来探索这个策略树,旨在找到最优的块组合顺序。我们引入了一个配置代理,它准确地预测给定查询的最佳配置和重排器。此外,我们还设计了一种并行查询扩展策略,以在每次迭代中扩展多个节点。实验结果表明,我们的方法在成本限制内显著优于最先进的方法,并且在效率方面也表现出色。