RAG

摘要¶

大型预训练语言模型已经被证明能够在其参数中存储事实性知识, 并在下游NLP任务的微调中取得了最先进的结果. 然而, 这类模型在访问和精确操作知识的能力上仍然有限, 因此在知识密集型任务上的表现仍然落后于特定任务架构. 此外, 如何为其决策提供可追溯性以及更新其世界知识, 仍然是未解决的研究问题. 迄今为止, 具备可微分访问机制的预训练模型(能够访问显式的非参数化记忆)仅被用于抽取式下游任务. 本文探索了一种通用的检索增强生成(Retrieval-Augmented Generation, RAG)微调方法——这种模型将预训练的参数化记忆与非参数化记忆结合, 用于语言生成. 我们提出的RAG模型中, 参数化记忆由预训练的seq2seq模型构成, 非参数化记忆则是一个基于稠密向量的Wikipedia索引, 并通过预训练的神经检索器访问. 我们比较了两种RAG形式: 一种是在整个生成序列中始终基于相同的检索片段进行条件建模, 另一种则允许在每个token生成时使用不同的检索片段. 我们在广泛的知识密集型NLP任务上对模型进行微调和评估, 在三个开放域问答任务中达到了最新的最优效果, 超越了仅参数化的seq2seq模型和特定任务的检索-抽取架构. 在语言生成任务中, 我们发现RAG模型相比最新的仅参数化seq2seq基线, 能够生成更加具体, 多样且符合事实的语言.

可微分访问机制用于抽样式下游任务/参数化记忆和非参数化记忆结合用于语言生成

参数化记忆是指模型参数本身(Transformer权重)存储的知识. 非参数化记忆是指显式存储库(如Wikipedia文档向量库), 可微分访问机制是指用Attention这样的方式去读取外部记忆, 整个过程可以端到端训练. 以往这种机制主要应用在抽样任务中, 比如, 抽取问答(QA), 文档匹配, 分类等等需要选中某些文本片段的服务, 这类服务只需要检索或者选择已有的文本, 不涉及复杂的生成. 本文不仅仅在抽取式任务上使用外部记忆, 而是把检索到的非参数化记忆和模型内部的参数化记忆结合, 用于生成任务, 比如开放域问答, 摘要生成, 对话生成, 这类任务需要模型在检索知识的基础上进行自然语言生成.

两种RAG的建模方式

第一种: 同一检索片段用于生成整个生成序列. 模型先根据输入去检索若干段落(比如top-k个文档), 在生成整个答案的过程中, 这些文档保持不变, 换句话说, 检索在生成开始前就完成了, 后续生成答案时, 每一步都是基于相同的检索结果. 例如, 问题式"爱因斯坦出生在哪一年", 检索到片段: "阿尔伯特·爱因斯坦出生于1879年3月14日…" 生成序列 "他出生于1879年" 时, 每个token都参考同一个检索片段. 这种方式简单高效, 但如果生成到一半发现不够信息, 就没办法动态换文档.
第二种: 不同token可以使用不同检索片段. 生成的时候不是一次性检索, 而是逐个token动态检索, 模型在生成每个单词的时候, 都可以决定再去查询一次外部知识库, 得到更加相关的片段, 所以不同位置的token, 背后的检索片段可能不同. 例如, 问题: "请介绍爱因斯坦的出生年份和主要贡献", 生成"出生于1879年"时, 模型检索出生信息的片段, 生成"相对论"时, 模型又检索贡献相关的片段, 整个答案过程中, 检索内容是随生成进度动态更新的. 这种方式更灵活, 信息覆盖更全面, 但计算代价更高.

动机贡献¶

预训练的神经语言模型已被证明能够从数据中学习到大量深入的知识. 它们即便不访问任何外部记忆, 也可以作为一种参数化的隐式知识库来工作. 尽管这一进展令人振奋, 这类模型仍存在不足: 它们难以轻松扩展或修订其记忆, 难以直接给出其预测的可解释依据, 并且可能产生"幻觉". 将参数化记忆与非参数化(即基于检索的)记忆相结合的混合模型能够缓解其中一些问题, 因为外部知识可以被直接修订与扩展, 且被访问的知识可被检查与解释. REALM与ORQA这两类将掩码语言模型与可微分检索器结合的最新方法已展示出有前景的结果, 但此前仅在开放域抽取式问答上进行了探索. 在此, 作者将参数化与非参数化记忆的混合范式引入NLP的"work horse", 即序列到序列(seq2seq)模型.

作者通过一种通用的微调方法, 为预训练的, 以参数化记忆为基础的生成模型赋予非参数化记忆, 将其称为检索增强生成(RAG). 作者构建的RAG模型中, 参数化记忆是预训练的seq2seq transformer, 非参数化记忆是Wikipedia的密集向量索引, 由预训练的神经检索器访问. 作者将这些组件组合成一个端到端训练的概率模型(见图1). 检索器(Dense Passage Retriever, 下文简称DPR)在给定输入的条件下提供潜在文档, 而seq2seq模型(BART)则在这些潜在文档与输入的联合条件下生成输出. 作者使用top-K近似对潜在文档进行边际化, 可以按两种粒度执行: 以输出为粒度(假设同一文档为所有token提供证据), 或以token为粒度(允许不同文档分别为不同token提供证据). 与T5或BART相似, RAG可以在任意seq2seq任务上进行微调, 并对生成器与检索器进行联合学习.

图1: 方法概览. 作者将一个预训练的检索器(Query Encoder + Document Index)与一个预训练的seq2seq模型(Generator)结合, 并进行端到端微调. 对于查询$x$, 使用最大内积搜索(MIPS)找到top-K个文档$z_i$. 在得到最终预测$y$时, 将$z$视为潜在变量, 并对不同文档条件下的seq2seq预测进行边际化.

已有大量研究提出用于为系统引入非参数化记忆的架构, 并针对特定任务从零开始训练, 例如memory networks, stack-augmented networks和memory layers. 与之相对, 作者探讨一种设置: 参数化记忆(指隐式存于模型参数中的知识)与非参数化记忆(指可从外部检索的显式存储)两个组件均已预训练并预先载入丰富知识. 关键在于, 通过使用预训练的访问机制(如检索器/索引器), 系统无需额外训练即可具备访问知识的能力.

实验结果表明, 将参数化与非参数化记忆与生成式建模结合, 对于知识密集型任务(即在人类无法在无外部知识源的情况下合理完成的任务)具有显著优势. 作者提出的RAG模型(检索增强生成)在开放域Natural Questions, WebQuestions和CuratedTREC上达到SOTA, 并在TriviaQA上显著超越使用特定预训练目标的最新方法. 尽管这些任务是抽取式任务, 作者发现不受约束的生成式方法相较以往的抽取式方法表现更优. 对于知识密集型生成, 作者在MS MARCO与Jeopardy问答题生成上进行实验, 结果显示相较BART基线, 该模型生成的回答在事实性, 具体性与多样性方面更佳. 用于FEVER事实验证时, 其成绩距离采用强检索监督的SOTA流水线模型仅差4.3%. 最后, 作者证明可通过替换非参数化记忆来随实时世界变化而更新模型的知识.

方法¶

作者探讨了RAG模型, 该模型使用输入序列$x$来检索文本文档$z$, 并在生成目标序列$y$时将这些文档用作额外上下文. 如图1所示, 作者的模型利用了两个组件: (i)一个带有参数$η$的检索器$p_η(z|x)$, 它在给定查询$x$的情况下返回文本段落上的(top-K截断)分布; (ii)一个由$θ$参数化的生成器$p_θ(y_i|x, z, y_{1:i−1})$, 它基于先前的$i-1$个词元(token)$y_{1:i−1}$, 原始输入$x$以及检索到的段落$z$的上下文来生成当前词元(token).

为了端到端地训练检索器和生成器, 作者将检索到的文档视为一个隐变量. 作者提出了两种模型, 它们以不同的方式对潜在文档进行边缘化, 以产生一个关于生成文本的分布. 在第一种方法RAG-Sequence中, 模型使用同一份文档来预测每个目标词元(token). 第二种方法RAG-Token, 可以基于不同的文档预测每个目标词元(token). 接下来, 作者将正式介绍这两种模型, 然后描述$p_η$和$p_θ$组件, 以及训练和解码过程.

模型¶

RAG-Sequence模型: RAG-Sequence模型使用同一份检索到的文档来生成完整的序列. 从技术上讲, 它将检索到的文档视为单个隐变量, 并通过top-K近似对其进行边缘化, 以获得seq2seq概率$p(y|x)$. 具体来说, 检索器会检索排名最高的K个文档, 而生成器则为每个文档生成输出序列概率, 之后再对这些概率进行边缘化,

\[p_{\text{RAG-Sequence}}(y|x) \approx \sum_{z \in \text{top-k}(p(\cdot|x))} p_{\eta}(z|x) p_{\theta}(y|x, z) = \sum_{z \in \text{top-k}(p(\cdot|x))} p_{\eta}(z|x) \prod_{i}^{N} p_{\theta}(y_i|x, z, y_{1:i-1})\]

RAG-Token模型: 在RAG-Token模型中, 他们可以为每个目标token采样一个不同的隐变量文档并进行相应的边缘化. 这使得生成器在生成答案时能够从多个文档中选择内容. 具体来说, 首先使用检索器检索出排名最高的K个文档, 然后生成器为每个文档生成下一个输出token的分布, 接着进行边缘化, 并对后续的输出token重复此过程. 形式上, 他们定义:

\[p_{\text{RAG-Token}}(y|x) \approx \prod_i^N \sum_{z \in \text{top-k}(p(\cdot|x))} p_{\eta}(z|x) p_{\theta}(y_i|x, z, y_{1:i-1})\]

最后, 他们指出RAG可用于序列分类任务, 只需将目标类别视为一个长度为一的目标序列, 在这种情况下RAG-Sequence和RAG-Token是等价的.

检索器¶

检索组件$p_{\eta}(z|x)$基于DPR. DPR遵循一种双编码器架构:

\[p_{\eta}(z|x) \propto \exp(\mathbf{d}(z)^\top \mathbf{q}(x)) \quad \mathbf{d}(z) = \text{BERT}_d(z), \mathbf{q}(x) = \text{BERT}_q(x)\]

其中$\mathbf{d}(z)$是由一个基于BERT$_{\text{BASE}}$的文档编码器生成的文档密集表示, $\mathbf{q}(x)$是由一个同样基于BERT$_{\text{BASE}}$的查询编码器生成的查询表示. 计算top-k$(p_{\eta}(\cdot|x))$, 即具有最高先验概率$p_{\eta}(z|x)$的$k$个文档$z$的列表, 是一个最大内积搜索(MIPS)问题, 可以在亚线性时间内近似求解. 他们使用一个来自DPR的预训练双编码器来初始化他们的检索器并构建文档索引. 该检索器经过训练, 用于检索包含TriviaQA问题和Natural Questions答案的文档. 他们将文档索引称为非参数化记忆.

生成器¶

生成器组件$p_{\theta}(y_i|x, z, y_{1:i-1})$可以使用任何编码器-解码器进行建模. 他们使用BART-large, 一个拥有4亿参数的预训练seq2seq transformer. 当使用BART生成时, 为了结合输入$x$和检索到的内容$z$, 他们只是简单地将两者进行拼接. BART使用去噪目标和多种不同的加噪函数进行了预训练. 它在多种多样的生成任务上取得了当前最佳结果, 并且性能优于规模相当的T5模型. 他们此后将BART生成器的参数$\theta$称为参数化记忆.

训练¶

他们联合训练检索器和生成器组件, 且不对应该检索哪个文档提供任何直接的监督信号. 给定一个由输入/输出对$(x_j, y_j)$组成的微调训练语料库, 他们使用带Adam的随机梯度下降法来最小化每个目标的负边际对数似然$\sum_j -\log p(y_j|x_j)$. 在训练期间更新文档编码器BERT$_d$的代价高昂, 因为这需要像REALM在预训练期间那样周期性地更新文档索引. 他们发现这一步骤对于获得强大性能并非必要, 因此保持文档编码器(及索引)固定, 仅微调查询编码器BERT$_q$和BART生成器.

解码¶

在测试时, RAG-Sequence和RAG-Token需要使用不同的方法来近似$\arg \max_y p(y|x)$.

RAG-Token - RAG-Token模型可以被视为一个标准的自回归seq2seq生成器, 其转移概率为$p'_{\theta}(y_i|x, y_{1:i-1}) = \sum_{z \in \text{top-k}(p(\cdot|x))} p_{\eta}(z|x) p_{\theta}(y_i|x, z, y_{1:i-1})$. 为了解码, 他们可以将$p'_{\theta}(y_i|x, y_{1:i-1})$代入一个标准的束搜索解码器.

RAG-Sequence - 对于RAG-Sequence, 似然$p(y|x)$无法分解为常规的逐token似然, 因此他们无法通过单次束搜索来求解. 作为替代, 他们为每个文档$z$运行束搜索, 并使用$p_{\theta}(y_i|x, z, y_{1:i-1})$对每个假设进行评分. 这会产生一个假设集$Y$, 其中某些假设可能并未出现在所有文档的束搜索结果中. 为了估计一个假设$y$的概率, 他们为每个$y$未出现在其束中的文档$z$运行一次额外的正向传播, 将生成器概率与$p_{\eta}(z|x)$相乘, 然后对所有束的概率求和以得到边缘概率. 他们将此解码过程称为"完全解码". 对于较长的输出序列, $|Y|$可能会变得很大, 需要多次正向传播. 为了更高效地解码, 他们可以作进一步的近似, 即对于在文档$z_i$的束搜索过程中未生成的$y$, 假设$p_{\theta}(y|x, z_i) \approx 0$. 这避免了在生成候选集$Y$后还需要运行额外正向传播的需求. 他们将此解码过程称为"快速解码".