【RAG 论文】Selfmem：使用 LLM 自己的输出来作为下一轮的 context 从而提升自己的生成效果

论文：Lift Yourself Up: Retrieval-augmented Text Generation with Self Memory
⭐⭐⭐⭐
NeurIPS 2023，北大

文章目录

- 一、论文速读
- 二、实现细节
- - 2.1 检索增强的 Generator
  - 2.2 Memory Selector
  - 2.3 Generator 的两种 mode
- 总结

一、论文速读

在以为 RAG 中，LLM 使用 user query 和检索到的文档（本文中称之为 memory）作为输入来生成回复。然后，这样检索到的 memory 是来自于一个固定的 corpus 并受限于 corpus 的质量，这样有界的 memory 明显限制了 memory-augmented generation model 的效果，因此本文提出：better generation also prompts better memory，意思也就是说，LLM 所生成的好的 generation 也可以作为一种好的 memory，从而将有界的 memory 推广到了无界中。

因此在这篇论文扩展了 LLM 输入的上下文，提出 LLM Generator 的输入应该是 user query + memory，其中 memory 既包含检索到的文档，也包含 LLM 自己的输出。这种将 generation 也作为 memory 的思路，是本文重点强调的创新点。

另外，本论文还抽象了 RAG 中各个概念，将 RAG 用了在诸如机器翻译等多个场景中。

本论文提出 Selfmem 框架：

Selfmem 框架

上图以机器翻译的场景为例，训练数据包含 source $X$ 和 target $Y$ 。Selfmem 框架包含两个部分：

检索增强的 Generator：它拿到 source $X$ ，调用 retriever 进行检索，形成 memory $M$ 。然后把 memory $M$ 和 source $X$ 输入给 LM 得到生成多个候选回复 candidates。
Memory Selector：Selector 接收到前面 LM 生成的多个 candidates，将其与 source $X$ 做计算，根据某个评价指标 metric（比如 BLEU）从 candidates 中选出一个 candidate 作为选择结果，然后把这个选择结果输送给 Generator，这样下一次 Generator 部分在形成 memory 时，会把 candidate 选择结果和 retrieval 结果一同用于形成 memory 作为下一轮 LM 的输入。

可以看到，Selfmem 的关键创新就是，改造了 RAG 框架中输入给 LLM 的上下文，这个上下文不仅仅包含了检索到的文档，还包含了 LLM 在上一轮中生成的回复，这个回复是由 Memory Selector 从多个候选回复中根据 metric 选出来的一个。

二、实现细节

2.1 检索增强的 Generator

文中提到，Generator 既可以是 frozen LLM，也可以是 trainable LM。对于一个任务，输入是 $x$ ，输出是 $y$ ，在不同的下游任务中，输入和输出的含义不同，比如 summarization 任务中，就是 (document, summary)，在对话任务中，就是 (context, response)，在机器翻译中，就是 (source, target)。

Generator 的输入包含任务的输入 $x$ 和 memory $m$ ，为了连接 $x$ 和 $m$ ，其架构可以有两种：

Joint-Encoder：将 $x$ 和 $m$ join 成一个 sequence 做 encode： $e n co d er (x [SEP] m)$ ，然后再解码。
Dual-Encoder：使用两个 encoder —— SourceEncoder 和 MemoryEncoder —— 分别对 $x$ 和 $m$ 进行编码，并将两个编码做 Cross Attention 后再解码。