AutoRAG-HP: 自动在线超参数调整用于检索增强生成
摘要
最近大型语言模型(LLMs)的进步改变了ML/AI开发,需要重新评估AutoML原则,以适应检索增强生成(RAG)系统。为了解决RAG中超参数优化和在线适应的挑战,我们提出了AutoRAG-HP框架,将超参数调整形式化为在线多臂老虎机(MAB)问题,并引入了一种新颖的两级层次化MAB(Hier-MAB)方法,以有效探索大型搜索空间。我们在ALCE-ASQA和自然问题数据集上进行了广泛的实验,调整了诸如检索的top-k文档、提示压缩比率和嵌入方法等超参数。我们的评估表明,基于MAB的在线学习方法可以在搜索空间中有显著梯度的场景中实现 Recall@5 ≈ 0.8 \text{Recall@5} \approx 0.8 Recall@5≈0.8,仅使用网格搜索方法所需的约20%的LLM API调用。此外,提出的Hier-MAB方法在更具挑战性的优化场景中优于其他基线。代码将提供在https://aka.ms/autorag。
1 引言
最近大型语言模型(LLMs)的进步代表了ML/AI解决方案开发的重大转变。传统上,大多数ML/AI应用都训练了特定场景的模型。然而,在LLM时代,基础模型作为基础,为实际应用添加了补充模块。这种转变显著影响了以前被称为AutoML的ML/AI解决方案开发的自动化。检索增强生成(RAG)已成为使用LLMs构建ML/AI解决方案的一个突出框架。虽然标准的RAG框架包括一个信息检索组件,以将LLM的输出基于相关数据,但现在许多变体集成了额外的模块,如查询重写、提示压缩和查询路由,以提高性能。RAG系统的复杂性增加带来了两个主要挑战。首先,模块和模块内的超参数的多样性使得确定最佳设置变得复杂。其次,我们经常从用户那里获得在线反馈(例如,通过拇指向上/向下功能),有效地利用这些反馈来持续调整系统也是至关重要的。为了解决这些挑战,我们提出了一个自主和自我优化的RAG系统,称为AutoRAG,符合AutoML的原则。作为第一步,本研究专注于RAG中的超参数调整(AutoRAG-HP)。虽然以前的工作讨论了RAG中的超参数调整,但它们倾向于关注LLM API调用中的可调超参数,或通过手动调整评估RAG超参数的性能和影响。在这项工作中,我们专注于可以在线应用的优化方法。具体来说,我们将超参数选择形式化为在线多臂老虎机(MAB)问题,并探索了MAB的几种变体。此外,为了在同时调整多个超参数时有效探索大型搜索空间,我们引入了一种新颖的两级层次化MAB(Hier-MAB)方法,其中高级MAB指导模块的优化,而几个低级MAB在每个模块内搜索最佳设置。我们的评估表明,基于MAB的在线学习方法对于搜索空间中有显著梯度的场景是有效的,提出的Hier-MAB方法在更具挑战性的优化场景中优于其他基线。我们的贡献可以总结如下:
- 我们引入了AutoRAG-HP框架,以解决RAG中最佳超参数调整的迫切需求。据我们所知,我们是第一个讨论RAG中自动在线超参数调整的。
- 我们将RAG中的在线超参数搜索形式化为多臂老虎机问题,并提出了一种新颖的两级层次化多臂老虎机方法,以有效探索大型搜索空间。
- 使用公共数据集在多个场景下验证了我们方法的有效性。
2 相关工作
2.1 AutoML和LLMs
在开发ML/AI解决方案的过程中,AutoML已经简化了特征工程、模型构建和超参数优化三个关键领域的自动化。在过去的十年中,AutoML取得了显著的成功,拥有大量使用的开源框架,如AutoSklearn、TPOT、Auto-Keras、Auto-PyTorch和FLAML。然而,随着LLMs的出现,发生了显著的转变,LLMs经常被选为基础模型,绕过了传统的模型设计和训练阶段。研究社区已经开始研究将AutoML应用于优化LLMs的预训练、微调和推理过程的机会和挑战。最近的一篇论文提出了一个及时的调查,讨论了AutoML和LLMs之间的潜在共生关系,同时也提供了一个面向未来的愿景。特别是在利用AutoML进行LLMs的预训练和微调阶段的超参数优化(Liu and Wang, 2021; Treviso et al., 2022; Tornede et al., 2024)。LLaMA-NAS也探索了LLMs的有效神经架构搜索。对于推理阶段,EcoOptiGen代表了将AutoML应用于优化LLM推理以进行文本生成的开创性步骤。这项工作针对在OpenAI完成标题中调整超参数,如温度和最大令牌。另一项工作探索了梯度下降和MAB在自动优化提示中的应用。
2.2 超参数优化
常见的超参数优化技术包括网格搜索、随机搜索、贝叶斯优化和手动调整,以确定最佳超参数。网格搜索涉及在预定义的超参数网格内进行详尽搜索,测试每种可能的组合以找到最佳选择。虽然直接,但这种方法可能会产生巨大的计算成本,特别是当超参数空间扩展时。随机搜索通过随机抽样选择超参数,但其结果可能缺乏稳定性。手动调整则基于领域知识或经验调整超参数。虽然灵活,但这种方法耗时且难以标准化。这些超参数优化方法通常忽略了评估成本,特别是在基于LLMs评估解决方案时。BlendSearch引入了经济预算以提高成本效率。然而,这些方法不是基于学习的,可能不适合需要随时间适应性优化的场景。
2.3 RAG中的超参数调整
大量注意力已经集中在独立地完善索引、检索和生成等各个模块的模型上。然而,就超参数而言,这些研究通常只报告导致最佳结果的那些,通常是在实验期间通过专家手动调整选择的。因此,很少有探索旨在调整每个模块内的超参数,更不用说集体调整RAG模块的各种超参数。
CRUD-RAG深入研究了RAG超参数的手动调整,并评估了RAG系统不同组件(如检索器和上下文长度)的性能和影响。虽然这些研究为优化RAG技术提供了宝贵的见解,但它们在多样化场景或现实世界应用中的适用性有限。此外,一个项目提到了通过贪婪方法进行优化,最初在每个节点生成所有可能的模块和超参数组合。
3 方法论
3.1 问题形式化
我们将RAG中的超参数调整问题形式化为多臂老虎机(MAB)问题,将其比作玩家在赌场中从多臂老虎机中选择的情景。玩家的目标是选择提供最高预期收益的臂。每次玩家拉动一个臂并获得收益或不获得收益时,他们更新对该臂潜在收益的估计。MAB问题涉及做出顺序决策,要求代理在探索不同臂以了解它们的奖励概率和利用基于过去观察预期会产生更高奖励的臂之间进行平衡。鉴于MAB是在线学习方法,它非常适合AutoRAG-HP的在线设置。
作为说明,我们在图1中展示了RAG系统的示例,该系统包括检索模块、提示压缩模块和提示构建模块(未显示),后者组装最终发送给LLMs以生成答案的提示。在检索模块中,我们引入了两个可调超参数:从外部知识库检索的top-k(K)文档块和用于对这些检索块进行排名的嵌入模型(E)。检索到top-k块后,提示压缩模块然后压缩每个块中的令牌,以消除无关信息并节省令牌成本。由于过度压缩也可能移除相关信息,导致性能下降,因此找到最佳压缩比率(C)至关重要。下面,我们介绍MAB中的术语在AutoRAG-HP的上下文中的含义。
-
臂:在此上下文中,臂指的是我们旨在优化的特定超参数组合。例如,如果我们正在优化top-k参数,一个臂可以代表top-k的候选值(例如, K = 3 K = 3 K=3)。当同时优化多个超参数时,一个臂对应于这些超参数的组合,如MAB问题的标准公式中定义的。注意,由于臂是离散的,搜索空间必须首先被离散化。
-
试验:试验是算法选择一个臂、观察相关奖励并更新其估计的单次迭代。在RAG系统中,一个试验可以涉及评估当前选择的超参数组合(即,臂)的一组查询,批量大小为B。在预定数量的迭代T之后或满足特定停止标准后,可以确定最佳设置。
-
奖励:奖励函数代表用户的目标,并在优化过程中指导臂的选择。对于AutoRAG-HP,常见目标包括在不太关注LLM API调用成本(由输入令牌数量化)的情况下最大化响应准确性,或平衡这些目标。为简单起见,我们将奖励函数形式化为响应准确性和输入令牌长度的线性组合:
Reward = w ⋅ acc − ( 1 − w ) ⋅ t t max \text{Reward} = w \cdot \text{acc} - (1 - w) \cdot \frac{t}{t_{\text{max}}} Reward=w⋅acc−(1−w)⋅tmaxt
其中 w w w是平衡权重, t t t表示输入令牌长度并由最大输入令牌长度 t max t_{\text{max}} tmax标准化, acc \text{acc} acc代表LLM的响应准确性。优化算法:MAB中的几种优化算法可以基于给定的奖励函数指导臂的选择。一个常见的选择是上置信界(UCB)算法,它通过选择基于它们的上置信界来有效地平衡探索和利用。这些界限是从代表臂值估计范围的置信区间推导出来的。UCB选择臂的公式如下:
A t = arg max a ∈ A ( Q t ( a ) + α ln ( t ) N a ( t ) ) A_t = \arg\max_{a \in A} \left( Q_t(a) + \alpha \sqrt{\frac{\ln(t)}{N_a(t)}} \right) At=arga∈Amax(Qt(a)+αNa(t)ln(t))
其中 A t A_t At是在时间步( t )选择的臂,即选择的超参数或其组合, Q t ( a ) Q_t(a) Qt(a)是在 t t t时臂 a a a的估计值。平方根项量化了不确定性。 N a ( t ) N_a(t) Na(t)表示臂 a a a被选择的次数, α \alpha α是调整探索和利用之间平衡的超参数。在迭代过程中,每个臂的 Q Q Q和上置信界(平方根项)都会更新,以指导臂的选择。汤普森采样(TS)是MAB中另一种流行的优化算法。它通过从每个臂的奖励的后验分布中采样来平衡探索和利用。基于最高采样奖励选择臂。总之,MAB的目标是在一系列选择中最大化总奖励,即使每个臂的奖励概率分布是未知的。经过一系列试验后,累积奖励最高的臂成为所需的RAG超参数。这种方法特别适合冷启动问题,在没有先前用户数据估计的情况下,利用MAB框架可以快速调整RAG中的超参数。
3.2 两级层次化MAB
将上述MAB的标准公式应用于RAG中的超参数调整可能导致在联合优化多个超参数时臂的数量过多,导致搜索空间过大,因为它需要展平搜索空间以获得离散臂。为缓解此问题,我们提出了两级层次化MAB(Hier-MAB),我们首先选择要调整的超参数,然后选择其可能值之一。在图2中,我们展示了在联合调整top-k(K)、嵌入模型(E)和压缩比率(C)超参数的上下文中两级Hier-MAB的示例。高级臂负责选择要调整的超参数,而低级臂控制每个超参数搜索空间内的超参数选择。因此,我们不是只有一个MAB,而是有四个MAB:一个用于高级臂选择,另外三个用于各个超参数。这种层次结构确保每个MAB都有合理数量的臂可供选择,而所有MAB组合可以覆盖大型搜索空间。这与单MAB方法形成对比,后者在调整多个超参数时需要枚举所有可能的组合。Hier-MAB的优化过程可以通过图2中显示的试验来演示。拉动高级臂(top-k),在此超参数内,拉动K = 3臂(其他超参数与上一试验保持相同)。拉动两级臂并观察相关奖励后,算法使用新信息更新所选臂的奖励分布估计,即,根据观察到的奖励更新平均奖励估计和置信区间。对于图2中的示例,正面奖励更新了被拉动臂的奖励分布,以反映更高的估计奖励。同时,之前迭代中拉动的其他高级和低级臂的奖励分布也得到更新。这个过程重复进行预定数量的迭代或直到满足停止标准。
4 评估
4.1 实验设置
数据集:我们使用ALCE-ASQA和自然问题(NQ)数据集进行实验。两个数据集都采用问答格式,并为每个问题包含候选文档块。我们使用它们的评估器来评估生成响应的准确性。为确保LLMs没有对基准问题的先验知识,我们排除了无需上下文即可正确回答的问题(即零样本)。从剩余问题中,我们从每个基准中选取350个问题作为我们的实验数据集。基础LLMs:我们采用GPT3.5-Turbo和GPT-4模型作为基础LLMs。尽管API参数是可调的,我们选择通过将温度设置为零并使用所有其他参数的默认设置来固定它们。搜索空间:我们检查了图1中所示的RAG设置。在检索模块中,我们通过考虑三种不同的选择来评估top-k超参数(K)和嵌入模型(E)的影响,即“mpnet”、“ada_002”和“contriever”。压缩模块使用LLMLingua-2工作中的方法实现,压缩比率表示为C。具体来说,我们根据超参数数量考虑了两个优化任务:
- 联合优化(K, C):它们从$$K \in [1,
通过检查图3和NQ数据集的其他网格搜索结果(在附录中显示),我们做出了以下观察:
- top-k和压缩比率对响应准确性有显著影响,而嵌入模型的效果在ASQA中比在NQ中更为明显。这突出了调整这些超参数的必要性,特别是对于“准确性中心”场景(即, w = 0.9 w = 0.9 w=0.9)。
- 通过更多地考虑代币长度,奖励函数的整体景观发生了变化。对于“成本中心”场景( w = 0.1 w = 0.1 w=0.1),首选的最优设置将是K和C的小值。由于对响应准确性的依赖减少,样本方差变得小得多,奖励函数在(K, C)上表现出更陡的梯度。这表明在 w = 0.1 w = 0.1 w=0.1时,基于MAB的在线学习方法可以更有效地找到最优超参数,因为搜索空间中的梯度更明显。
4.2 实验结果
我们的实验结果表明,在不同的权重设置下,提出的Hier-MAB方法在Recall@x指标上普遍优于其他基线方法。特别是在“成本中心”场景下,Hier-MAB能够快速收敛到最优或近似最优的超参数组合,这得益于其在探索和利用之间的有效平衡。此外,Hier-MAB在处理具有挑战性的优化场景时,如联合优化(K, C, E),也显示出了其优越性,这表明了其在处理更复杂超参数空间时的鲁棒性和有效性。
在具体的实验设置中,我们观察到Hier-MAB在不同的数据集和LLM模型上都能保持一致的性能优势。这验证了我们方法的泛化能力,表明其可以应用于不同的RAG系统配置和应用场景。此外,通过与网格搜索结果的比较,我们进一步证实了Hier-MAB在寻找最优超参数组合方面的有效性,尤其是在资源有限的情况下。
5 结论
本文提出了AutoRAG-HP框架,用于自动在线调整RAG系统中的超参数。通过将超参数调整问题形式化为多臂老虎机问题,并引入两级层次化MAB方法,我们的方法能够在大型搜索空间中有效地找到最优超参数组合。实验结果表明,与传统的网格搜索和其他在线学习方法相比,我们的方法在不同的数据集和LLM模型上都能取得更好的性能。未来的工作将探索将AutoRAG-HP应用于更广泛的ML/AI解决方案开发中,以及进一步优化Hier-MAB方法以处理更复杂的超参数空间。