【RAG 论文】Iter-RetGen：多轮迭代检索增强生成

论文：Enhancing Retrieval-Augmented Large Language Models with Iterative Retrieval-Generation Synergy

⭐⭐⭐

arXiv:2305.15294

论文速读

给定 question $q$ 和用于检索的 corpus $D$，ITER-RETGEN 重复 T 轮迭代检索，再每一次迭代 $t$ 中：利用上一轮生成的文本 $y_{t-1}$ 与 question $q$ 拼接，交给 retriever 去检索到 top-k 的段落，然后去生成一个输入给 LLM $M$ 的 prompt 从而得到本轮迭代的 LLM 生成 $y_t$，生成 prompt 的方式是：将本轮检索到的段落 $D_{y_{t-1}}$ 与 $q$ 填入 prompt template 中。

形式化上，每一轮迭代都可以写成下面的式子：

直至迭代 T 轮后得到 $y_T$ 作为最终的 response。

如下是一个使用 ITER-RETGEN 的图示：

每一轮迭代过程中，体现出了这个模型的一个特点：轮流地做 Generation-Augmented Retrieval 和 Retrieval-Augmented Generation。

对 retriever 的优化

这是本工作的一个特殊之处。

LLM 的 generations 不仅仅用于增强下一轮的 retrieval，还用于优化 retriever。

具体来说，本文采用了一个 dense retriever，其结构包含一个 query encoder 和 doc encoder，这个模型被训练根据 query vector 和 doc vector 的 inner product 相似度来完成检索。这个 retriever 在用于 ITER-RetGen 后，会固定住 doc encoder，微调 query encoder。

在 RAG 中，re-ranker 会将检索到的 docs 根据与 question 的相关性重新排序，论文认为，re-ranker 通常比 retriever 更擅长捕捉查询和段落之间的相关性，所以这里将 re-ranker 的知识蒸馏到 retriever 中，训练目标就是最小化 re-ranker 和 retriever 的输出的概率分布的 KL 散度：

• 第一个公式是训练目标，即找到能够最小化两个概率分布的 KL 散度的参数值
• 第二个公式是 re-ranker 输出的在给定 y1 和 q 下检索出文档 d 的概率
• 第三个公式是 retriever 输出的在给定 q 下检索出文档 d 的概率

它想让 retriever 根据 q 检索到的文档与 re-ranker 认为的在给了 y1 后与 q 最相关的文档是尽可能相似的，这样 re-ranker 可以教会 retriever 什么才是与 q 相关的文档，同时可以训练 retriever 不要偏离原问题 q，因为随着迭代，会引入很多其他信息，但不要让 retriever 偏离了原来的问题 q 而专注于子问题。