DRAGIN
Paper : DRAGIN: Dynamic Retrieval Augmented Generation based on the Information Needs of Large Language Models
Author : Weihang Su , Yichen Tang , Qingyao Ai , Zhijing Wu , Yiqun Liu
Affiliation : Department of Computer Science and Technology, Tsinghua University , School of Computer Science and Technology, Beijing Institute of Technology
Publication : ACL 2024
Code :GitHub - oneal2000/DRAGIN: Source code of DRAGIN, ACL 2024 main conference Long Paper
引入了一个新框架 DRAGIN,即基于 LLM 信息需求的动态检索增强生成。
确定激活检索模块的最佳时机(决定何时检索)和在触发检索后设计适当的查询(确定检索什么)
对于检索的时间,提出RIND:实时信息需求检测,考虑 LLM 对其自身生成内容的不确定性、每个 token 对后续 token 的影响以及每个 token 的语义重要性。
对于检索查询的制定,提出 QFS:基于自注意力的查询制定,通过利用 LLM 在整个上下文中的自注意力来创新查询制定
0 前景
LLM虽能力出众但是会产生一些连贯且合理但是不符合实际的内容。RAG是解决这个问题的有效方案,通过从外部数据库检索相关信息并将其整合到 LLM 的输入中来增强 LLM。传统的 RAG 方法通常依赖于单轮检索,使用 LLM 的初始输入从外部语料库中检索相关信息。虽然这种方法对于简单的任务很有效,但它往往不适用于复杂的多步骤任务和长格式生成任务。相比之下,动态 RAG在 LLM 的生成过程中执行多次检索。它包括两个步骤:确定激活检索模块的最佳时机(决定何时检索),以及在触发检索后设计适当的查询(确定要检索什么)。
1 动机
现有的动态 RAG 方法面临着几个关键挑战,主要是在确定检索的最佳时机和在触发检索时制定有效的查询方面。首先,现有方法通常依赖静态规则来决定何时进行检索,而忽略了对必要性和潜在风险的评估。一方面,根据输入查询和检索模型的质量,不必要的检索增强可能会向 LLM 引入不相关或嘈杂的数据,从而危及输出的质量。另一方面,进行检索增强将不可避免地增加 LLM 推理的时间和计算成本,如果 LLM 可以自己生成正确的输出,那么这种成本是值得的。此外,现有研究在确定检索内容方面的策略通常局限于 LLM 的最近一句或最后几个标记。这种方法可能无法捕捉模型的实时信息需求,因为 LLM 的信息需求实际上可能与涵盖整个上下文的术语相关。因此,在许多情况下,以这种方式检索文档并不是最理想的。
2 方法
DRAGIN 由两个组件组成:实时信息需求检测 (RIND) 和基于自注意力的查询公式 (QFS)
2.1 Real-time Information Need Detection
如上所述,大多数现有的动态 RAG 框架基于静态的预定义规则触发检索模块。据我们所知,唯一值得注意的例外是 FLARE (Jiang et al, 2023),当 LLM 对下一个 token 的置信度(即生成概率)低于某些阈值时,它会动态触发检索。
然而,检索增强的必要性不仅取决于生成置信度,还取决于 token 的重要性、token 的语义以及每个 token 对后续 token 的影响。为了解决现有方法的局限性,我们提出了一种在动态 RAG 框架内触发检索的增强方法,称为实时信息需求检测 (RIND)。该方法通过评估每个 token 的不确定性以及其语义贡献和对后续上下文的影响来改进检索激活过程。
RIND 首先量化 LLM 推理过程中生成的每个标记的不确定性。这是通过记录标记在词汇表中的概率分布的熵来实现的。考虑一个由LLM生成的输出序列,记作T = {t1, t2, …, tn},其中ti表示序列中位于位置i的单个词元。对于任意词元ti,其熵Hi计算如下:
其中,pi(v)表示在位置i生成词汇表V中所有词元v的概率。这种不确定性的测量构成了我们多维度词元评估的第一个维度。
此外,RIND 利用基于 Transformer 的 LLM 固有的自注意力机制为 token 分配权重,这些权重表示 token 对后续上下文的影响。具体来说,对于任何给定的 token ti ,我们通过记录最大注意力值 $ a_{max(i)} $ 来量化其影响力,该值记录了所有后续 token (选择 LLM 的最后一层 Transformer 的注意力分数) 的最大注意力。两个标记 ti 和 tj ( i < j ) 之间的 注意力值Ai,j 计算如下:
其中 Qi 表示 token ti 的查询向量,Kj 是 token tj 的键向量,dk 表示键向量的维数。softmax 函数应用于 Qi 和 Kj 的点积,并由 dk 的平方根归一化。随后,通过找到所有 j > i 中最高的 Ai,j 来确定 token ti 的最大注意力值 $ a_{max(i)} $ :
考虑每个标记的语义贡献,RIND 采用二进制语义指示符来滤除停用词,从而集中于具有重要语义价值的标记:
其中 S 是停用词集,si 是标记 ti 的语义贡献分数。此过程确保只有语义上有效的标记才能对检索决策过程做出贡献。
结合不确定性、重要性和语义,RIND 为每个标记 ti 计算出一个综合分数:
令 T = {t1, t2, . . . , tn} 表示 LLM 已生成的标记集。当任何标记的分数$ S_{RIND}(ti) $ 超过预定义阈值 θ 时,检索模块就会激活。
2.2 Query Formulation based on Self-attention
一旦确定了进行检索增强的位置,RAG 框架中的下一步就是制定查询,以从外部数据库检索必要的信息,以便继续生成 LLM。在现有的动态 RAG 框架中,所有查询制定方法都将重点限制在 LLM 的最新句子或最后几个标记上。这种狭窄的范围无法充分满足模型的实时信息需求,这些需求可能涵盖整个上下文。 为了克服这些方法的缺点,我们提出了一种新颖的策略,利用基于 Transformer 的 LLM 固有的自注意力机制。我们的方法称为“基于自注意力的查询制定”(QFS),旨在通过检查其对完整上下文的理解来更准确地确定 LLM 的信息需求。
考虑由 LLM 生成的输出序列,表示为 T = {t1, t2, . . . , tn},每个 ti 代表序列中的单个标记。假设 RIND 模块识别位置 i 处的标记,这需要外部知识并触发检索模块。然后,QFS 方法将重点关注此特定位置以制定查询。对于位置 i 处的标记,QFS 方法会评估前一个标记序列 {ti-1, ti-2, …, t1} 的注意权重。 由于 LLM 生成 ti 是基于其对整个前一个上下文的解释,因此注意权重反映了模型在生成 ti 时对每个标记的自我评估重要性。QFS 方法根据这些标记的注意力得分对其进行优先排序,选择前 n 个标记来构建查询。查询制定过程包括以下步骤:(1)为 T 中的每个 token ti 提取最后一个 Transformer 层 Ai = {ai,1, ai,2, …, ai,i−1} 的注意力分数,其中 ai,j 表示 ti 分配给 tj 的注意力分数;(2)按降序对 Ai 进行排序,以确定注意力分数最高的前 n 个 token;(3)从词汇表中找出与这些 token 相对应的单词,并按照它们在文本中的原始顺序排列;(4)使用这些前 n 个 token 中的单词构造查询 Qi,确保查询反映了 LLM 的自注意力机制确定的上下文中最相关的方面。
2.3 Continue Generation after Retrieval
一旦 RIND 模块检测到需要外部知识的位置 i,QFS 模块就会创建查询并利用现成的检索模型(例如 BM25)从外部知识库中检索相关信息。假设检索到的文档表示为 Di1、Di2 和 Di3。成功检索后,动态 RAG 框架的下一步是将这些外部知识集成到 LLM 的生成过程中。 此集成从截断已识别位置 i 处的 LLM 输出开始,以进行检索增强:
其中 T ′ 表示截断的输出,T 是 LLM 生成的原始序列,ti 是 RIND 识别出需要外部知识的标记。为了整合检索到的知识Di1、Di2、Di3,我们采用精心设计的提示模板3,其结构如下:
此时,LLM 继续根据检索到的外部知识和截断的输出 T ′ 生成内容。在整合检索到的知识之后,LLM 将从截断点继续生成内容,并通过附加信息进行增强。此过程允许 LLM 弥补先前确定的知识差距,从而更明智、更准确地继续输出。
假设在后续位置 j,RIND 再次检测到 LLM 需要外部知识。在这种情况下,QFS 模块在位置 j 再次触发以生成新查询,检索一组新文档 Dj1、Dj2 和 Dj3 来替换 Di1、Di2 和 Di3。然后,LLM 将根据新检索到的文档继续从位置 j 生成,并遵循相同的过程。
3 实验设置
3.1 数据集
我们选择了两个 MultihopQA 数据集 2WikiMultihopQA (Ho et al, 2020) 和 HotpotQA (Yang et al, 2018) 来评估 RAG 框架回答需要多跳推理的复杂问题的能力。我们选择了 IIRC (Ferguson et al, 2020) 数据集来评估 RAG 框架在阅读理解任务中的能力。此外,我们利用 StrategyQA (Geva et al, 2021) 数据集来评估 DRAGIN 和其他基线的常识推理能力。
3.2 Settings for each Dataset
- 2WikiMultihopQA (Ho et al, 2020)
我们遵循 (Wang et al, 2022) 的设置来生成思路链 (CoT) 推理过程以及最终答案。我们遵循 (Trivedi et al, 2022) 和 (Jiang et al, 2023) 的提示模板。对于评估指标,我们使用模式匹配技术从生成的输出中提取最终答案。然后将提取的答案与参考答案进行比较,使用答案级别的精确匹配等方法,以及 F1 分数和精度的标记级别测量。
- HotpotQA (Yang et al, 2018)
我们遵循 (Trivedi et al, 2022) 的设置和提示模板来生成思路链 (CoT) 推理过程以及最终答案。我们在这个数据集上的评估指标与 2WikiMultihopQA 相同。
- StrategyQA (Geva et al, 2021)
我们遵循 (Wei et al, 2022) 的设置来生成 CoT 推理过程以及最终答案。 我们遵循 (Wei et al, 2022) 和 (Jiang et al, 2023) 的提示模板。对于评估指标,提取获得的是/否响应并使用精确匹配方法与标准正确答案进行比较。
- IIRC (Ferguson et al., 2020).
我们遵循 (Trivedi et al, 2022) 的设置和提示模板来生成最终答案。我们在这个数据集上的评估指标与 2WikiMultihopQA 相同。
每个数据集对应的具体提示模板在附录F中给出。附录A对每个数据集的设置提供了更详细的描述。
3.3 Baselines
我们选择了以下文本生成基线进行比较。遵循 FLARE(Jiang et al,2023)的设置,我们使用相同的设置实现了现有的多轮 RAG 框架,唯一的变化是触发检索的时间(何时检索)和触发检索时的查询制定方法(检索什么)。
• wo-RAG。LLM 无需 RAG 即可直接回答问题。
• SR-RAG(单轮 RAG)。根据初始问题从外部语料库中检索相关段落。然后将检索到的段落添加到 LLM 的输入中。
• FL-RAG(固定长度 RAG)(Khandelwal 等人,2019;Borgeaud 等人,2022;Ram 等人,2023)。一种多轮检索增强方法,每 n 个 token 触发一次检索模块。在前一个 token 窗口中生成的 token 用作查询。
• FS-RAG(Fix Sentence RAG)(Trivedi 等人,2022 年)。一种多轮检索增强方法,每句话都会触发检索模块。最后生成的句子用作查询。
• FLARE (Jiang 等,2023)。一种多轮检索增强方法,每次遇到不确定的 token 时都会触发检索。 当触发检索模块时,最后生成的不含不确定 token 的句子被定义为查询。
为了直接说明 DRAGIN 与其他动态 RAG 基线之间的差异,我们在表 1 中对每个动态 RAG 框架的检索时间和查询形成方法进行了比较。
3.4 Selected LLMs
为了验证 DRAGIN 和其他 RAG 基线的有效性,我们使用以下 LLM 进行了实验:
• LLaMA-2-Chat(Touvron 等人,2023b)。 LLaMA2 是一组经过预训练和微调的 LLM。该系列包括经过微调的 LLM,称为 Llama2-Chat,专为在基于对话的应用程序中实现最佳性能而设计。我们选择了 LLaMA-2-Chat7B 和 LLaMA-2-Chat-13B。
• Vicuna-13B-v1.5(Chiang 等人,2023)是一组开源聊天机器人,使用从 ShareGPT 收集的用户共享对话从 LLaMA 进行微调。我们选择了最新版本的 Vicuna,即 Vicuna-13B-v1.5。
3.5 Implementation Details
超参数:所有超参数均在附录 E 中列出。
检索器:我们根据 (Ram et al, 2023) 的发现采用 BM25 作为检索模型,该模型在 RetrievalAugmented Generation 中表现出色,甚至优于一些密集检索模型。我们还探讨了用 SOTA 密集检索方法 SGPT (Muennighoff, 2022) 替换 BM25 的影响,详见第 5.5 节。
停用词:为了识别 RIND 模块中的停用词,我们使用了 Spacy 库中的 en_core_web_sm 语言模型,该工具因其在自然语言处理任务中的有效性和效率而得到认可,这已被先前的研究证明 (Shelar et al, 2020)。
外部知识语料库:我们采用维基百科作为外部知识语料库。每篇文章都分为 100 个标记段落。
LLM 配置:对于选定的 LLM,我们直接从 Hugging Face 官方仓库下载每个模型的模型参数,并使用 Hugging Face 提供的代码进行文本生成。对于生成配置,我们选择了贪婪解码作为 LLM 推理的解码策略,以确保实验结果的可重复性。然而,对于实际应用,我们建议使用每个模型提供的官方默认生成配置,因为这将产生更好的性能。
4 实验结果
4.1 Overall Results of DRAGIN and Baselines
我们的实验全面评估了 DRAGIN 在四个基准数据集上相对于各种基线的性能:2WikiMultihopQA、HotpotQA、StrategyQA 和 IIRC。表 2 中总结的结果强调了几个关键见解:
(1) 与 DRAGIN 相比,单轮检索增强的集成来指导问答可以持续提高 LLM 在所有数据集上的性能,证实了 RAG 的有效性。
(2)尽管检索增强总体上有积极影响,但我们观察到基于固定规则的检索方法(例如 FL-RAG 和 FS-RAG)并不总是优于单轮检索。这一观察证实了我们的假设,即如果在错误的位置进行检索增强,可能无助于提高 LLM 输出的质量。这强调了激活检索过程的时间安排的重要性,检索过程应根据大型语言模型 (LLM) 的信息需求进行量身定制,仅当 LLM 需要外部知识时才激活检索。
(3)**FLARE 的性能在不同数据集之间差异很大。**有趣的是,如我们的消融研究(§ 5.4)所示,查询制定策略明显优于其他基线使用的策略,但其整体性能并非如此。这表明 FLARE 中检索增强的时机远非完美。
(4) 我们提出的 DRAGIN 方法在大多数 LLM 和数据集上都表现出优异的性能。这表明 DRAGIN 在增强 LLM 功能方面的稳健性和有效性。
(5) DRAGIN 在 MultihopQA 任务(例如 2WikiMultihopQA 和 HotpotQA)中表现出的性能改进比在需要常识推理的任务(例如 StrategyQA 数据集中的任务)中表现出的性能改进更为显著。这一差异凸显了 DRAGIN 在管理复杂、多步骤推理任务方面的专业能力。
4.2 效率
在本节中,我们研究了各种动态 RAG 框架在多个数据集中的效率。我们根据进行的检索调用次数来衡量效率,如表 3 所示。由于 FS-RAG 的特殊设计,FS-RAG 的 #NUM 还表示 LLM 响应此数据集上的查询时生成的平均句子数。在评估的框架中,FLARE 以其效率脱颖而出,在所有数据集中需要的检索调用最少。 DRAGIN 紧随其后,检索次数较少。
4.3 检索时间
在本小节中,我们研究了检索时间对动态 RAG 框架性能的影响。具体来说,我们固定了查询制定方法,使用 LLM 生成的最后一个完整句子作为查询,并将检索时间作为唯一变量进行更改。我们在 IIRC 数据集上对 DRAGIN 以及三个现有框架:FLARE、FL-RAG 和 FS-RAG 进行了测试。如表 4 所示,DRAGIN 始终优于所有其他动态 RAG 方法。这凸显了我们确定最佳检索时刻的新方法的有效性。
DRAGIN 的卓越性能表明,其检测 LLM 的实时信息需求并相应地触发检索的方法特别擅长提高生成文本的质量。 我们还评估了 RIND 模块中不同阈值对 DRAGIN 框架性能的影响。我们在表 5 中展示了 HotpotQA 数据集上的实验结果。我们的实验结果表明,DRAGIN 的性能在阈值设置上保持稳定,表明对该超参数的变化的敏感度较低。阈值是确定检索模块激活频率的关键。随着阈值的增加,检索模块的激活频率会降低,这表明在实际应用中,调整阈值可以在系统的效率和输出的准确性之间取得平衡。
4.4 查询定制
本小节深入探讨了查询制定技术对动态 RAG 框架性能的影响。我们将触发器检索的时间标准化为 RIND,根据第 5.3 节详述的实验结果,这被证明是最有效的时间。我们重点关注 DRAGIN 与三个现有框架的比较:FLARE、FL-RAG 和 FS-RAG。 FLARE 的查询制定方法是最后生成的句子排除低概率标记。FL-RAG 选择最接近此位置的 25 个标记作为查询。FS-RAG 选择此位置之前的句子作为查询。我们还评估了使用完整上下文作为查询的有效性。如表 6 所示,DRAGIN 的查询制定方法在所有动态 RAG 框架中表现最佳。FLARE 成为第二有效的查询制定方法,优于 FS-RAG 和 FL-RAG 方法。 此外,利用整个上下文作为查询并没有产生最佳结果,这表明整个上下文中存在潜在的冗余。这一发现验证了我们提出的 QFS 方法的有效性,该方法旨在从上下文中选择能够代表 LLM 的实时信息需求的标记作为查询。
4.5 检索器的影响
在动态 RAG 范式中,检索器的选择对于根据给定查询从语料库中检索相关段落起着重要作用。在信息检索领域,两种流行的检索方法是词汇匹配(Zhai,2008;Robertson 等,2009)和密集检索(Su 等,2023b;Gao and Callan,2021;Su 等,2023a;Muennighoff,2022;Li 等,2023b;Ma 等,2023;Ye 等,2024;Su 等,2023c;Li 等,2023a;Chen 等,2023,2022;Li 等,2023c;Fang 等,2024)。在词汇匹配技术中,BM25 因其广泛采用和有效性而脱颖而出(Robertson 等人,2009 年)。相反,在现有的密集检索方法中,没有一种能像 BM25 那样广泛流行。我们选择了 SGPT,它最近在各种数据集上都取得了最先进的性能。(Muennighoff,2022 年)。 在表 7 中显示的实验分析中,我们发现,尽管 SGPT 在许多信息检索任务中通常表现更好,但 BM25 在动态 RAG 框架内各种数据集上的性能始终优于 SGPT。这一结果与先前研究的结果一致,例如(Ram 等人,2023 年)的研究,该研究强调了 BM25 在 RAG 任务中的有效性。这些结果表明,尽管 SGPT 等密集检索技术取得了进展,但更简单的基于词典的 BM25 算法仍然是增强 LLM 在 RAG 任务中性能的强大基础。
5 限制
我们承认本文存在局限性。主要限制之一是实时信息需求检测 (RIND) 和基于自注意力的查询公式 (QFS) 都依赖于基于 Transformer 的 LLM 的自注意力机制。虽然所有开源 LLM 都可以访问自注意力分数,但需要注意的是,我们的方法不适用于某些不提供自注意力分数访问权限的 API。因此,我们未来的工作旨在开发更多方法来克服这一限制。
6 结论和未来工作
制
我们承认本文存在局限性。主要限制之一是实时信息需求检测 (RIND) 和基于自注意力的查询公式 (QFS) 都依赖于基于 Transformer 的 LLM 的自注意力机制。虽然所有开源 LLM 都可以访问自注意力分数,但需要注意的是,我们的方法不适用于某些不提供自注意力分数访问权限的 API。因此,我们未来的工作旨在开发更多方法来克服这一限制。
6 结论和未来工作
在这项工作中,我们提出了 DRAGIN,这是一个动态 RAG 框架,专门用于满足文本生成过程中 LLM 的实时信息需求。 通过集成 RIND 以实现及时检索激活和 QFS 以实现精确查询制定,DRAGIN 在各种知识密集型基准测试中的表现明显优于现有的动态 RAG 方法。