【RAG 论文】RAG 模型的初代产品：Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

论文：Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks
⭐⭐⭐⭐
NeuraIPS 2020, Facebook
Code：https://huggingface.co/facebook/rag-token-nq

文章速读

这是早期提出通过结合文档检索 + 语言模型生成来实现 QA 等 NLP 任务的思路的文章，这篇文章提出了 RAG 模型。

这篇文章提出的模型叫做 RAG 模型，区别于现在经常泛指的 RAG 架构。

文章提出的整个模型架构如下：

最左边的文字是用户的自然语言问句 $x$，问句 $x$ 经过组件 Query Encoder 进行编码得到向量 $q(x)$，使用 $q(x)$ 利用最大内积搜索（MIPS）方式从文档库中查找到最相近的 K 个文档，之后把 K 个文档和用户问题合并后输入给 Generator 来生成最终的 answer。这其中：

• Query Encoder 是一个 BERT 模型，用来将自然语言问句 $x$ 编码为向量 $q(x)$
• Generator 是一个生成式预训练语言模型 BART-Large，采用的是 Encoder + Decoder 架构，共有 4 亿个参数
• 文档库中包含的是：预先将一堆文档输入给 Document Encoder 得到的一堆文档向量索引。其中“Document Encoder”也是一个 BERT 模型，用来将一个文档编码为文档向量。

论文指出，检索器中的文档采用维基百科截至 2018 年 12 月的全量数据，将每篇维基百科文档切分为互不重叠的包含 100 个单词的块，每个块作为一个文档，共有 2100 万个文档。

在以上过程中，RAG 模型实现了根据问题来检索相关文档，并将相关文档和问题一并输入给 Generator 来完成最终的 answer 生成。

RAG 模型的详细介绍

1. RAG 架构

论文指出，这个模型总共可以分成两个部分：

1. 由 Query Encoder 和 Document Encoder 组成的 Retriever，也被称为非参数化的部 memory，意思是这部分的知识来源于外部数据源而不是模型的参数
2. 由生成式预训练语言模型构成的 Generator，也称为参数化的 memory，意思是这部分的知识来源于模型的内部参数

这样来看，整体流程就是：用户问句经过 Retriver 得到相关文档，然后问句 + 相关文档一并输入给 Generator 来得到最终答案。

2. 模型的公式描述

• 问题的 token 序列为 $x$
• Retriver 是 $p_{\eta }(z|x)$，它根据 $x$ 返回 K 个最相关的文档 $z$
• Generator 是 $p_{\theta }(y_i|x,z,y_{1:i-1})$，采用自回归的方式，根据 $x$ 和 $z$ 以及回答的前 $i-1$ 个 token，来预测下一个 token，直至生成完整的回答。

论文提出了两种模型，

• 一种是 RAG-Sequence 模型，在整个生成序列的过程中使用同一个检索到的文档，这意味着，对于给定的输入 $x$，模型会从文档索引中检索出一个最相关的文档，并用这个文档来辅助生成整个目标序列 $y$。适用于那些单个文档就能提供足够信息来生成整个回答的任务，例如当问题的答案可以从一个单一的、信息丰富的文档中获得时。
• 一种是 RAG-Token 模型，它允许模型在生成每个 token 时使用不同的检索文档。这意味着模型可以为序列中的每个标记选择最合适的文档，从而更灵活地整合来自多个文档的信息。适用于需要从多个文档中提取和整合信息来生成回答的任务，例如当问题的答案分散在多个文档中，或者需要结合多个文档的信息来提供更全面的回答时。

论文给出了两种模型的“对于给定输入问题 x 而生成回答 y”的概率计算公式：

在第一个公式中，给定 x 生成 y 的概率，就等于给定一个 x 和检索的 z 来生成 y 的概率并对 z 做 marginalized，然后就有了上面那个公式。第二个公式也是类似。

从公式对这两种模型进行理解，RAG-Sequence 模型先计算每个文档条件下回答词元的概率分布，再连乘得到每个文档条件下回答的概率分布，最后再求和得到所有最相关文档条件下回答的概率分布，而 RAG-Token 模型先计算每个文档条件下回答词元的概率分布，再求和得到所有最相关文档条件下回答词元的概率分布，最后再连乘得到所有最相关文档条件下回答的概率分布。

3. Retriever

这里 Retriever 是一个 bi-encoder 架构，公式描述就是：

意思也就是，对于问句 $x$，经过 Query Encoder（BERT）转为向量 $q(x)$，对于一个文档 $z$，经过 Document Encoder（BERT）转为文档向量 $d(z)$，然后 $p_{\eta }(z|x)$ 就是指的对于 x 选择 z 的概率。从这些概率中选出最高概率的 K 个文档，也就是 MIPS 问题，这 K 个文档就是检索的结果。

4. Generator

Generator 是一个生成式语言模型，在这里选用了 BART-Large 模型，是一个带有 400M 参数的 seq2seq Transformer 模型。

在 RAG 架构中，论文是直接将问题和检索器返回的相关文档拼接在一起作为 Generator 的输入。

5. 解码

RAG model 的解码的目标就是生成一个目标序列 y，这个序列是对于给定输入 x 的响应。

在 RAG-Token Model 中，它允许在生成序列的每个 token 时考虑不同的检索文档，解码大致过程如下：

1. 使用 retriever 基于输入 $x$ 从文档库中检索出一组最相关的文档 $z_1,z_2,...,z_k$，这些文档将用于生成目标序列
2. 使用转移概率 $p_{\theta }^{\prime }(y_i|x,y_{1:i-1})$ 运用 beam search 来完成自回归的解码

在 RAG-Sequence Model 中，不能简单使用 beam search 来实现解码，具体原因涉及到一些关于 beam search 的知识，这里暂时先跳过。

6. 模型的训练

论文将 retriever 和 generator 联合在一起进行 end-to-end 的训练。

训练数据是问答对 $(x_j,y_j)$ 的集合，损失函数采用了负对数似然函数：

$$\sum _j-\log p(y_j|x_j)$$

使用 Adam 优化器进行梯度下降。

另外论文指出，因为文档库的所有文档向量都是预先使用 Document Encoder 计算构建出来的，训练时如果调整 Document Encoder 所使用的 BERT 模型的参数，那么就需要定期重新更新所有文档的向量并构建索引，成本较高。所以，训练时需要固定住 Document Encoder 的参数，只更新 Query Encoder 和 Generator 的参数。

总结

这篇论文是早期提出使用 Retrival 来增强语言模型的 generation 效果的思路的论文，这套架构中具体使用的组件模型可能不再适用（比如 BART），但整体架构的思想仍然值得学习。

示例

参考我的另一篇文章：【实践】LangChain + OpenAPI 实现基本的 RAG

参考文章：

• RAG论文阅读笔记