Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection

背景

RAG通过外挂知识库解决事实错误，往往引入大量的上下文 token，有些时候可能不是必要的。

Self-RAG 从大模型自身的能力出发，自我决策整个过程，生成更高质量的问答。

下面是一个简单的对比图，后面有详细的解释：

推理

整个推理过程以 Segment 为基本单位，从论文后面给的例子来看这个 Segment 就是指每一句话。在拿到问题时，区别于传统的 RAG 直接进行检索，Self-RAG 的生成模型 Generator 会首先生成特殊的 token [Retrieve=Yes/No] ，这一步决定是否要通过检索来生成对应的回答，可以自适应的调整阈值来控制结果。

当 [Retrieve=Yes] 确定触发检索后，Generator 会首先对召回的文本做相关性评估，引入特殊 token [IsREL=relevant/irrelevant]，同时并行的对每个召回结果生成相应的回答，为每个回答再生成两个特殊token [IsSup=fully supported/partially supported,no support] 和 [IsUse=5/4/3/2/1]，分别代表回答是否能被召回的文本支持和回答是否有意义。

显然，对于每个 Segment，如果 [Retrieve=Yes] 我们会得到与召回文本数量 K 一致的回答数，每个回答都包含了下面四种特殊 token，蓝色的token基本就代表了回答的质量，可以根据属性值进行排序，选择最终的结果输出。

下面给出得分的计算规则：本身的概率得分加上蓝色 token 的评估得分。W 为超参数需要自己调节，看当前的应用场景更在乎相关性、意义还是完整的逻辑链，赋予不同 token 不同的权重。

训练

训练过程涉及到两个模型：Critic 和 Generator

Critic

Critic model 只需要为推理过程中出现的 4 种特殊 token 生成对应的数值，在实际推理过程中用不到它，是为了下一阶段更好的训练 Generator。

直接人工标注类似的数据比较困难，对数据的质量有比较高的要求，论文中选择使用 GPT-4 进行训练数据的生成。考虑分别为不同的 token 构造不同的 prompt，为每个 token 生成 4k - 20k 左右的数据量，人工简单的进行评估筛选掉明显不合理的数据即可。

模型上选择任何一个 LLM 都行，如果和 Generator 选择同一个 base model 明显会更加方便。训练过程和预训练过程一致，构造类似的 next token prediction loss

Generator

首先我们需要构造一批原始的对话数据，这部分仍然可以用 GPT-4 生成。再此基础上，我们可以用已经训练好的 Critic 对这批数据进行增强。

首先用 Critic 评估原始对话的输出，在 [Retrieve=Yes] 的地方插入相应的 token，然后走一遍和传统 RAG 一样的流程，召回 K 个相关文本，用 Critic 评判文本和问题的相关性 [IsREL]，相关性文本的后面插入[IsREL=relevant/irrelevant]，用 LLM 生成后续的回答，用 Critic 预测 [IsSup] 并插入结果 token；在整个回答结束后用 Critic 预测 [IsUse]，构成整个训练数据。这里和推理部分有一些不同，推理的时候是每个 Segment 都要生成 [IsUse]，是为了方便算得分。训练的时候没有这个要求，全部回答完算得分就行。

训练过程都是一样的，注意把召回的文档内容 mask 即可

效果

RAG的效果评价一直是个问题，这里列了一些评价指标看看就行

总结

整个训练和推理过程都不算复杂，借助了 LLM 本身的能力对外挂的知识库进行评估和反思，数据的准备也不算困难，有不错的借鉴意义，可以考虑内嵌进一些常见的 RAG 框架进行尝试。

这种分阶段的训练和 PPO 类似，应该比较考验 Critic 模型的能力，可能有一些训练的坑需要发掘。

比较注重推理的质量，没有考虑推理的时间成本，单个请求可能需要多次检索知识库。并行生成结果时看不到其他召回结果，也可能导致回答不全面，反而降低质量。期待后续的优化