详解 SELF-RAG 框架

2023/12/27 21:53 下午 posted in  大模型

当前先进的大语言模型在回答问题时依然会存在胡说八道的现象,而检索增强生成(RAG)方法通过将相关的检索段落与语言模型的输入结合,可以减少在知识密集型任务中的事实错误。然而,这些方法可能会影响语言模型的通用性,引入不必要或离题的段落,导致生成的回答质量较低。此外,由于 RAG 检索段落时不考虑事实基础是否有帮助,生成的结果也不能保证与检索到的相关段落一致。

自我反思检索增强生成(Self-Reflective Retrieval-Augmented Generation,SELF-RAG)是一种通过检索和自我反思提高 LLM 质量和事实准确性的框架,而不损害 LLM 的原始创造力和多功能性。本文将详细介绍 SELF-RAG 框架。

概述

SELF-RAG 允许语言模型 $LM$ 根据检索到的段落生成信息,并且通过自我批判生成的内容来生成特殊 token。这些特殊的 token 称之为 reflection token(反思 token),表示是否需要检索或确认输出的相关性、或完整性。相比之下,常规的 RAG 方法会无差别地检索段落,且不确保引用来源的是否完全支持输入。

具体而言,SELF-RAG 首先确定是否通过在继续生成时使用检索到的段落来帮助生成,如果是,它会输出一个检索标记,调用一个检索模型(第一步)。接下来,SELF-RAG 同时处理多个检索到的段落,评估它们的相关性,然后生成相应的任务输出(第二步)。然后生成评论标记,对自己的输出进行批评并选择最佳输出(第三步),评价标准是事实准确性和整体质量。接下来将进一步介绍 SELF-RAG 的几个重要概念与算法。

反思 token

给定输入 \(x\),SELF-RAG 会训练语言模型 \(M\) 顺序生成文本 \(y\),且 \(y\) 由多个段落组成,记为 \(y=[y_1,...,y_t]\),其中,\(y_t\) 表示第 \(t\) 个段落的 token 序列,\(y_t\) 中的生成 token 包括原始文本和反思 token。

而反思 token 主要有 4 种,分别为 RetrieveIsRelIsSupIsUse,其含义分别如下:

按需检索(Retrieve)

表示对于给定输入,判断是否需要额外检索信息。

比如:

  • 问题 \(x\) :提供保持健康的三个秘诀
  • 输出:yes
  • 原因:因为可能有一些可靠的资源来解释不同行为对健康的影响。因此检索文档有助于提高对此查询的响应。

再比如:

  • 问题 \(x\) :描述一次你不得不做出艰难决定的经历
  • 输出:no
  • 原因:因为这个问题询问的是一些个人经验,因此不需要寻找一些外部文档。

相关性(IsRel)

表示对于给定输入,检索出的信息是否提供了有用信息来解决输入问题。

比如:

  • 问题 \(x\) :竞选美国众议院的年龄
  • 检索信息 \(d\) :宪法规定了在美国参议院任职的三项资格:年龄(至少三十岁);美国公民身份(至少九年);以及参议员在选举时所代表的州的居住权。
  • 输出:Irrelevant
  • 原因:这些检索信息仅讨论了竞选美国参议院议员的年龄,而不是讨论众议院议员的年龄。

支撑性(IsSup)

评估检索信息中提供的信息是否完全支持输出,输出为 “Fully supported, partially supported, no support”。

比如:

  • 问题 \(x\) :解释 embedding 在自然语言处理中的使用
  • 检索信息 \(d\) :embedding 是自然语言处理 (NLP) 中一组语言建模和特征学习技术的统称,其中词汇表中的单词或短语被映射到实数向量。单词和短语嵌入用作底层输入表示时,已被证明可以提高 NLP 任务的性能,例如句法解析、情感分析、下一个标记预测以及类比检测。
  • 回答 \(y\) :词嵌入对于情感分析、文本分类、预测序列中的下一个词以及理解同义词和类比等任务非常有用。
  • 输出:Fully supported
  • 原因:输出句子讨论了词嵌入的应用,证据提到了所有应用句法解析、情感分析、下一个标记预测以及类比检测作为应用。因此,评分应为 “Fully supported”。

有用性(IsUse)

表示回答是否对问题有用,输出为 “5,4,3,2,1”,5 为非常有用,1 为几乎不切题或完全不相关。

比如:

  • 问题 \(x\) :“2023 年英国现任首相是谁?”
  • 回答 \(y\) :“鲍里斯·约翰逊 (Boris Johnson) 于 2019 年至 2022 年担任英国首相。”
  • 输出:2
  • 原因:虽然输出提供了关于 2019 年至 2022 年英国首相的事实正确的陈述,但该指令询问 2023 年的首相是谁,因此它没有回答该指令。因此,评分为2。

SELF-RAG 训练

给定一组输入输出数据 \(D=\{X, Y\}\),Generator 模型 \(M\),Critic 模型 \(C\)。

  1. 用预训练的语言模型 LM 对 \(C\) 进行初始化;
  2. 对 \({X, Y}\) 进行采样得到训练数据 \(\{X^{sample}, Y^{sample}\}\)
  3. 对采样数据中每一对 \((x, y)\):
    1. 通过 GPT-4 收集 reflection token \(r\);
    2. 将 \(\{(x, y, r)\}\) 添加到 \(D_{critic}\) 中;
  4. 用下一个 token 预测损失更新 \(C\);
  5. 用预训练的语言模型 LM 对 \(M\) 进行初始化;
  6. 对 \((X, Y)\) 中的每一对 \((x, y)\):
    1. 运行 \(C\) 得到 reflection token \(r\);
    2. 将 \((x, y, r)\) 添加到 \(D_{gen}\) 中;
  7. 基于 \(D_{gen}\) 用下一个 token 预测损失更新 \(M\);

训练评论者模型

评论者模型数据生成

手动标记每个段落的反思 token 是不现实的,而我们可以使用像 GPT-4 这样的最先进的大语言模型来生成反思 token。通过引导 GPT-4 生成反思 token,可以将其知识提炼到内部的评论者模型 \(\mathcal{C}\) 中,从而创建了监督数据。如下图所示:

对每组反思 token,从原始训练数据中随机采样 \(\{X^{sample},Y^{sample}\}\sim \{X,Y\}\)。由于每组反思 token 有自己的定义和输入,我们会针对性使用不同的 prompt。

这里以 Retrieve 为例,通过使用类型特定的指令来引导 GPT-4,比如给定一条指令,在原始任务输入 \(x\) 和输出 \(y\) 上进行少量示范,判断从网络中找到一些外部文档是否有助于生成更好的响应,以生成适当的反思 token:\(p(r|I,x,y)\) 。

评论者模型训练

生成数据 \(\mathcal{D}_{critic}\) 后,使用预训练语言模型 \(LM\) 初始化评论者模型 \(\mathcal{C}\),并用 \(\mathcal{D}_{critic}\) 对其进行训练。其目标函数为(对每对 reflection token 来说):

\[\max_{c} \mathbb{E}_{((x,y),r)\sim D_{critic} } \log {p_c}(r|x,y), r \]

初始模型可以是任意的预训练语言模型 \(LM\),评论者模型在大多数 reflection token 类别上都与基于 GPT-4 的预测达成了超过 90% 的一致性。

训练生成器模型

生成器模型的数据生成

给定一个输入输出对 \((x,y)\),使用检索和评论者模型来扩充原始输出 \(y\),从而创建监督数据,精确地模拟 SELF-RAG 推理时的过程。整个过程如下:

对 \(y_t\in y\),运行批评者模型 \(\mathcal C\) 来评估需要额外的检索信息来帮助增强生成。如果需要,则加上 \(Retrieve=Yes\) token,并且使用 \(\mathcal R\) 来获取前 \(K\) 个信息段落 \(D\)。对每个段落来说,\(\mathcal C\) 会进一步评估相关性并预测 \(IsRel\)。如果某个段落是相关的,则 \(\mathcal C\) 会进一步评估该段落是否支持模型的输出,并预测 \(IsSup\)。评论 token \(IsRel\) 和 \(IsSup\) 会被附加到检索的段落或输出后面。在最后的输出 \(y\) 中,\(\mathcal C\) 会预测整体效用 token \(IsUSE\),并将带有反思 token 和原始输入对的扩充输出添加到 \(D_{gen}\)。

生成器模型训练

通过使用精选的增强语料库,以及 reflection token \(D_{gen}\) 来训练生成器模型 \(M\)。目标函数为:

\[\max_{\mathcal{M} } \mathbb{E}_{(x,y,r)\sim D_{gen} } \log {p_\mathcal{M} }(y,r|x). \]

与评判模型 \(\mathcal{C}\) 训练不同,生成器 \(\mathcal{M}\) 学习预测目标输出以及 reflection tokens。训练期间,将检索到的文本块(由 <p></p> 标记)进行遮挡以进行损失计算,这意味着模型在计算损失时不考虑这些检索到的文本块。原始词汇 \(\mathcal{V}\) 通过一组 reflection tokens(如 <Critique><Retrieve>)进行扩展,这表示这些 tokens 被加入到词汇中,使模型能够使用这些特定的 tokens 来生成输出。

SELF-RAG 推理

最后再来介绍一下 SELF-RAG 的推理过程。如下图所示:

对于每个输入 \(x\) 和前一代生成的 \(y_{<t}\),模型解码检索 token 以评估检索的效用。如果不需要检索,模型将直接预测下一段输出,这与标准的语言模型行为一致。如果需要检索,模型会生成:一个评估检索段落的相关性的反思 token、下一个回答段落、以及评估回答段落是否被检索信息支持的反思 token。最后,一个评估整体效用的新的反思 token。每生成一个回答,SELF-RAG 都会并行处理多个段落,并且使用其自动生成的反思 token 来控制生成的输出。

生成反思 token 以自我评估输出使得在推理阶段 SELF-RAG 更加可控,能够调整其行为以满足多样的任务要求。对于要求事实准确性的任务,目标是使模型更频繁地检索段落,以确保输出与现有证据紧密对齐。相反,在更为开放的任务中,例如撰写个人经历文章,重点转向更少的检索,优先考虑整体创造力或效用。接下来,将介绍在推理过程中如何实施控制以满足这些不同目标的方法。

基于阈值的自适应检索

SELF-RAG 可以动态决定何时检索文本段落,这是通过预测 Retrieve token 来完成的。此外,框架还允许设定一个阈值。具体而言,如果生成的 token 是 Retrieve=Yes,且在所有输出 token 中的标准化值超过了指定的阈值,则触发检索。

基于评判 tokens 的树解码

在每个段落步骤 \(t\) 中,当需要检索时,基于硬性或软性条件,\(\mathcal R\) 检索 \(K\) 个段落,并且生成器模型 \(\mathcal M\) 并行处理每个段落并输出 \(K\) 个不同的候选值。我们进行段落级的 Beam Search(使用 Beam 大小为 \(B\))以获取每个时间戳 \(t\) 的前 \(B\) 个段落,并在生成结束时返回最佳序列。

每个段落 \(y_t\) 相对于段落 \(d\) 的分数通过评论者模型的评分 \(\mathcal S\) 进行更新,该评分是每个评论 token 类型的标准化概率的线性加权和。对于每个评论 token 组 \(G\)(例如 \(IsREL\)),我们将其在时间戳 \(t\) 的分数表示为 \(s^G_t\),然后按以下方式计算段落分数:

\[\displaylines{ f(y_t,d,Critique)=p(y_t|x,d,y_{<t})+\mathcal S(Critique), \text{where} \\ \mathcal S(\text{Critique})=\sum_{G\in \mathcal G} \omega ^Gs^G_t \space \text{for} \space \mathcal{G} =\{\text{IsREL, IsSUP, IsUSE}\} } \]

其中,

\[\displaylines{s^G_t=\frac{p_t(\widehat{r} )}{ {\textstyle \sum_{i=1}^{N^G}}(P_t(r_i)) } } \]

代表最理想的反射 token 的生成概率

\[\displaylines{ \widehat{r}(e.g.,\text{IsREL}=\text{Relevant}) } \]

,其中 \(N_G\) 个不同的令牌表示 \(G\) 的不同可能值。权重 \(\omega ^G\) 为可以调整的超参,以自定义模型在推理期间的行为。另外,通过调整这些权重,可以强调某些期望的行为并降低其他行为。