Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks

发布时间 2023-07-18 14:42:29作者: 馒头and花卷

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}). \]

Generator: BART

  • 生成器采用的是 BART-large.