Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

概
符号说明
RAG
- Retriever: DPR
- Generator: BART

Lewis P. and Perez E., et al. Retrieval-augmented generation for knowledge-intensive nlp tasks. NIPS, 2020.

概

RAG: 赋予模型检索的能力.

符号说明

$\bm{x}$ , 输入序列;
$\bm{y}$ , 输出序列, 长度为 $N$ ;

RAG

RAG 的思路其实很简单, 我们知道, 生成模型通过建模条件分布:

$p(\bm{y}|\bm{x}),$
并从中不断地采样序列.
不依赖检索的生成模型可以理解为 $\bm{x} \rightarrow \bm{y}$ 的过程, 现在我们希望赋予模型检索的能力, 即希望通过如下方式进行生成

$\bm{x} \rightarrow (\bm{x}, \bm{z}) \rightarrow \bm{y},$
即模型需要先通过 $\bm{x}$ 检索得到 $\bm{z}$ 并一起生成最后的 $\bm{y}$ .
实际上, 就是 (此处的积分是勒贝格积分)

$\tag{1} \bm{y} \sim p(\bm{y}|\bm{x}) = \int p_{\theta}(\bm{y}|\bm{z}, \bm{x}) p_{\eta}(\bm{z}|\bm{x}) \mathrm{d} \bm{z},$
其中 $p_{\theta}, p_{\eta}$ 是我们构建的两个条件模型:
让我们接下来考虑两种实际的任务.
RAG-Sequence: 即通过检索后的文档生成完整的序列, 假设我们检索出 (相同的) Tok-k 最相关的文档 (注意这些文档), 我们可以用如下方式近似 (2)

$p(\bm{y}|\bm{x}) \approx \sum_{\bm{z} \in \text{top-}k (p(\cdot|\bm{x}))} p_{\eta}(\bm{z}|\bm{x}) p_{\theta}(\bm{y}|\bm{x}, \bm{z}) = \sum_{\bm{z} \in \text{top-}k (p(\cdot|\bm{x}))} p_{\eta}(\bm{z}|\bm{x}) \prod_{i=1}^N p_{\theta}(y_i|\bm{x}, \bm{z}, y_{1:i-1}).$
RAG-Token: 采用的是一种迭代的方式, 对于第 $i$ 个需要预测的 Token, 它

$p(y_i|\bm{x}; \bm{y}_{1:i-1}) \approx \sum_{\bm{z} \in \text{top-}k (p(\cdot|\bm{x}, \bm{y}_{1:i-1}))} p_{\eta}(\bm{z}|\bm{x}, \bm{y}_{1:i-1}) p_{\theta}(y_i|\bm{x}, \bm{z}, \bm{y}_{1:i-1}),$
于是

$p(\bm{y}|\bm{x}) \approx \prod_{i=1}^N \sum_{\bm{z} \in \text{top-}k (p(\cdot|\bm{x}, \bm{y}_{1:i-1}))} p_{\eta}(\bm{z}|\bm{x}, \bm{y}_{1:i-1}) p_{\theta}(y_i|\bm{x}, \bm{z}, \bm{y}_{1:i-1}).$
这相当于, 每一个 Token 的生成我们进行一次文档的检索 (而 RAG-Sequence: 则是只检索一次).

Retriever: DPR

检索模型, 作者用的是 DPR, 其是一个 bi-encoder 的架构:
$p_{\eta}(\bm{z}|\bm{x}) \propto \exp(d^T(\bm{x}) q(\bm{x})), d(\bm{x}) = \text{BERT}_d(\bm{z}), q(\bm{x}) = \text{BERT}_q (\bm{x}).$