Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks
概
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}). \]
Generator: BART
- 生成器采用的是 BART-large.
