RAG都可以在哪进一步优化？

        RAG（Retrieval-Augmented Generation，检索增强生成） 指利用外部知识库来提高模型生成质量，目前已被广泛应用于工业界LLM产品中。本文通过综合对比几篇有代表性的文章，尝试摸索现有的RAG都还可以在哪些点上进行优化，希望能给大家的优化之路上提供一点点参考。

一、开篇

• Reading Wikipedia to Answer Open-Domain Questions，FaceBook17

        RAG并非LLM时代的产物，其思路最早出现于2017年的开放域问答任务，在FaceBook《Reading Wikipedia to Answer Open-Domain Questions》文中，将“外挂外部知识库”这种方法应用到了开放域问答任务上，其基本流程如上图所示。与早期的RAG方法类似，该文使用的检索器基于Sparse retrieval（稀疏检索），虽然比较简单，但在当下的搜索工程中仍占有重要位置，近来仍有不少研究如FLARE、In-Context RALM在LLM推理过程中也使用了稀疏检索作为检索器。

        RAG名称的真正由来是在21年Facebook《Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks》文中，这篇后面会进行介绍。但Facebook早期在开放域问答上的上述工作，给后面RAG的思路奠定了基础，且后续方法流程大同小异，基本都由“检索器”+“PLM生成器”组成，都遵循“Retrieval then Read”这一过程，即：

1. 根据用户query从预先构建的知识库中检索相关段落。

2. 将检索到的信息与query以某种方式结合，以提供给模型进行生成。

        之后的优化方法也都大多围绕检索器、生成器两项展开，接下来将简单介绍几个主要的优化方向。

二、优化生成器

• FiD：Leveraging Passage Retrieval with Generative Models for Open Domain Question Answering，Facebook21

• RETRO：Improving language models by retrieving from trillions of tokens ，deepmind22

• kNN-LM：Generalization through Memorization: Nearest Neighbor Language Models，斯坦福19

• REPLUG: Retrieval-Augmented Black-Box Language Models，Meta23

                在检索器提供的额外知识的加持下，如何有效利用这些信息进行更好的生成，是一大热门研究方向。相关的工作对应生成器的两个过程：训练和推理。

1、训练过程优化

        其中在训练过程优化生成器的代表性工作有FiD和RETRO。

        FiD（Fusion-in-Decoder）的想法简单直接，将检索回来的每个passage都与question通过encoder分别编码，然后concat在一起输入decoder生成最终的回复，整个过程只有生成器encoder-decoder的优化。

        相比之下，RETRO并不是直接将检索回的知识和原query进行拼接，而是传入模型的中间层参与训练，已起到更好的“指导生成”的效果。

        RETRO的整体流程图如上左图所示，其基本架构同样是transformer的encoder-decoder架构，在decoder部分的块交叉注意力层CCA上融合了检索到的知识。简单描述一下训练的基本流程：

1. 把训练数据集的每个句子分成大小一样的块chunk，这里作者取的chunk大小是64。

2. 每一个块Ci都进行向量编码，为了减少重复的编码，作者这里设置了一个缓存，存储了所有chunk的编码结果。注意这里的编码器和RETRO的base模型是两个东西，这里就是用的BERT。

3. 每一个块在检索数据集中检索，召回最相似的文本。这里的检索模型是kNN。

4. 检索回的文本进行编码，注意这里的输入不单单是检索回的文本，还有输入块Ci在decoder阶段进行CCA前的向量表示。作者这里解释，因为这里的encoder也是要进行fine-tune的，所以也将原输入块Ci以condition作为输入。

5. 每一个块Ci用encoder进行编码，然后进入改造后的decoder部分，也就是图中的RETRO block。

6. 分别进行ATTN、CCA和FFW，和常见的transformer decoder不同的就是这里的CCA部分，其示意图如上右图所示。右图乍一看不是很好理解，但是其实描述的也很清楚。在这里作者没有按原来的块Ci为单位进行计算，而是丢掉最后一个token，加入上一个块的最后一个token，这里个人猜测是为了增强各个块之间的信息交互。然后这个新块就和检索回的文本的编码一起进行coss attention的计算，最后把各个块的结果拼接。

7. 继续后续流程。

        在实际应用时，我们也可以借鉴这种RAG的方法，一般来说，把信息有效地放在中间层和输出层，效果是要优于直接放在输入层的。但是这类方法需要对模型本身动刀，实现上可能比较复杂一些。

2、推理过程优化

        这里同样介绍两篇文章kNN-LM和REPLUG。

        我们知道，语言模型在推理时都是根据前n个token预测下一个要出现的token，kNN-LM在每一个token的预测过程中都请教一下外部知识库，然后和本身预测的结果进行加权求和，来得到最终的结果。整个过程如上图所示，就以图中的示例为例，模型在推理时遇到“Obama's birthplace is”这个问题，

1. 一方面自己会给出答案可能是“Illinois”、“Hawaii”；

2. 另一方面，模型会把问题“Obama's birthplace is”作为key，去训练语料库里做相似度检索，获取和他最相像的句子的next token作为答案。

        最后模型综合上面两部分结果，得出最终的正确答案“Hawaii”。

        整个过程思路比较简单，知识就是（left context，token）这种形式，检索的key就是left context，value就是token的概率分布。在进行实验的时候，作者为了证明该部分知识库是有效的，并没有对模型进行fine-tune，而是对知识库的表示进行了fine-tune，最终在几个数据集上表明确实对ppl值有帮助。

        而Meta23年的RePlug的推理过程如下所示，即在推理过程中将检索到的文档并行与输入结合送给大模型，然后综合多个结果得分得到最终答案。RePlug的优势在于即插即用，也正是因为这点，RePlug在网上还火了一把，当时记得不少网友要借鉴进行落地，不知道最后落地结果咋样了。

三、端到端联合训练

• ORQA：Latent Retrieval for Weakly Supervised Open Domain Question Answering，谷歌19

• REALM: Retrieval-Augmented Language Model Pre-Training，谷歌20

• RAG：Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks，Facebook21

• Atlas: Few-shot Learning with Retrieval Augmented Language Models，Meta22

        在RAG中应用稀疏检索虽然简单，但还是存有一定的缺点。而后续大多数研究虽然将稀疏检索器升级为向量检索器，但其和生成器间还是存在一定的gap。因此不少研究将两者进行联合训练，以端到端的方式进行检索增强。

        最早的工作来自谷歌19年ORQA（Open Retrieval Question Answering），作者阵容非常豪华。要实现端到端训练，关键在于检索这个过程必须是可训练的。在训练过程中，ORQA通过逆完形填空（ICT）任务思想，将一段话中的某一句子作为query、其上下文作为候选正样本doc+、batch内的其他句子作为候选负样本doc-，摆脱了训练过程对于检索系统或者黄金检索数据集的依赖，证明了在没有任何信息检索系统的情况下，仅从Q-A句子对中联合学习检索器和阅读器是有可能的。其模型训练过程如上图所示，训练得分为检索得分Sretr与阅读理解得分Sread的加和。

        基于ORQA做baseline，20年谷歌又发布了REALM（Retrieval-Augmented Language Model）一文对其进行了改进。和ORQA一样，REALM也采用了ICT构建样本的思想，不同的是，REALM使用了更好的预训练方法来优化PLM的语言能力，添加了一个新颖的语言模型预训练步骤，并将loss反向传播到检索的向量表示中。直观上理解，就是使用query信息，以及检索到的doc信息，共同预测被mask掉的token。

        类似地，21年Facebook的RAG相当于在REALM上做了进一步拓展，将模型任务从REALM地开放域问答拓展至了所有Seq2Seq任务，这个思路跟之前T5的思路就很像了，RAG的名称的起源也是来自这里，之后被用作Seq2Seq下检索增强生成一类方法的统称而广泛使用。

        大模型时代端到端联合训练的代表性工作是Atlas，Atlas通过联合预训练检索器模块和语言模型，在广泛的知识密集型任务上具有强大的小样本学习能力。作者在论文中给出了四种联合LLM和检索器进行训练的方法，分别是：Attention Distillation (ADist)、End-to-end training of Multi-Document Reader and Retriever (EMDR2)、Perplexity Distillation (PDist)、Leave-one-out Perplexity Distillation (LOOP)，感兴趣的可以去看一下原文。

四、检索前处理

        为了缓解一些query太短或者意图模糊的问题，不少研究引入了改写器，在检索前对query进行改写。这一思路和之前搜推的思路很像，这里介绍两个研究

• Rewrite-Retrieve-Read：Query Rewriting for Retrieval-Augmented Large Language Models，上交23

• STEP-BACK PROMPTING：Take a Step Back: Evoking Reasoning via Abstraction in Large Language Models，DeepMind24

        Rewrite-Retrieve-Read这篇文章的思路比较清晰，和先前的流程对比如上所示，从左到右分别为：

1. 通用的RAG流程框架，即检索-生成。

2. 基于通用RAG框架加入重写模块，使用LLM先对question（用户query）进行重写，然后结合重写的结果一起送入检索器进行检索。

3. 引入强化学习，将模型的生成结果作为Reward反馈给改写器，相当于对改写器进行了适应性微调，使得改写器、检索器、生成器三者间更好协同。

        24年DeepMind的工作后退提示（STEP-BACK PROMPTING）思路非常巧妙。一般的改写都是希望将简短query拓展为更详尽的query，而后退提示如其名字，则是退一步思考，对query进行抽象（abstraction），获取该query的“前提query”作为其改写结果，其流程图如上所示，包含两个步骤：

1. 抽象（Abstraction）：这一步的目的是提出问题的基础原理或概念问题，从而让模型先确认这部分知识。

2. 推理（Reasoning）：结合基础原理再进行原始问题的解答。

        从图中例子可以很好理解，当模型被问到“Estella Leopold在1954年8月至1954 年11月期间去了哪所学校？”，大模型对于如此详细的时间范围限制很难准确把握，不出所料输出了错误的答案。而后退提示会对该问题进行抽象，询问“Estella Leopold的教育经历是什么？”，然后再进行常规RAG流程，即可获取该问题的相关必要信息，得到正确答案。

        后退提示的思路很巧妙，利用类似后退CoT的方法，进一步有效解决了大模型幻觉问题，个人直观感觉应该是有用的，且不难落地。

五、检索后处理

• Re2G: Retrieve, Rerank, Generate，IBM22

• MemWalker：Walking Down the Memory Maze: Beyond Context Limit through Interactive Reading，Meta23

        在搜索工作中，如果说改写对应召回前QU阶段的前处理，那么重排就对应QU阶段的后处理了。在RAG中也是如此，检索到的内容往往存在过长或噪声问题，对生成器最终的生成质量产生影响。那么利用“后处理”的思想，对检索到的结果进行进一步筛选或者提炼，能一定程度上缓解上述问题。

        最早加入重排器的相关工作来自Re2G，该工作的思想就是既然重排器对检索结果是有益的，那对RAG中的检索结果也应该是有益的。Re2G通过在检索阶段后加入一个基于BERT的重排器，并进行端到端训练及推理。这里重排器的模型结构分为交互式和双塔式两种，分别如上左右两图所示。两者的优劣也是老生常谈了，主要区别在于效果和线上推理速度的权衡。

        类似于重排器，总结器能够对检索出的内容进行进一步提炼，来达到优化LLM生成结果的目的，去年较火的MemWalker正是在RAG后加入了这样类似的“总结器”，分为以下两个主要步骤：

1. MemWalker首先对长文本进行分割成多个节点，然后递归地对这些节点的内容进行提炼总结，直到所有节点汇聚到一个根节点上为止。最终原始的文本就被表示成一棵树的数据结构。至于分割的具体长度，也要根据具体场景决定，毕竟每段长度越长，层级就会越少，虽然利于后续检索，但会导致准确度下降。

2. 在推理的检索阶段，模型会从这棵树的根节点出发，然后结合query来决定该进入哪个叶子节点，如此一层层往下进行查找。要注意的是这里并不是一定要找到最底层的叶子节点，如果中间某一结点适合生成答案则会直接返回，甚至如果发现进入了错误的路径还可以返回查找。

        MemWalker在长文本问答任务上取得了不错的效果，其思路和上面提到的STEP-BACK PROMPTING有异曲同工之妙，都是在LLM训练或推理时显式模拟人类的思考过程。但是，因为要构建记忆树，MemWalker在上述两阶段的复杂度是高于其他方法的，尤其在文本很长时，复杂度呈指数上涨。

六、pipeline路线

• SKR：Self-Knowledge Guided Retrieval Augmentation for Large Language Models ，清华23

• SuRe: Summarizing Retrievals using Answer Candidates for Open-domain QA of LLMs，CMU24

        除了上述在“Retrieval-Read-Generate”框架内的工作，还有不少工作在RAG的这条pipeline上进行决策调整，使得整个RAG过程更加智能化。比如，在检索前决策是否要进行RAG，要不要进行多轮RAG；在检索时分多路进行，最终综合多路结果等等。

        SKR在检索过程前加入了一个分类器，用于判断这个query需不需要进行RAG，也就是说，会不会这个query不需要RAG就已经能得到不错的答案了，整个过程如上图所示。SKR中，这个分类器的训练过程大体如下：

1. 找一个问答数据集，其中每一条数据是一个问答对<问题qi，答案ai>。

2. 把问题qi分两路，分别输入LLM和LLM+RAG，得到两个结果。其中这里外挂的RAG是一个预训练的检索器。

3. 分别给<标准答案，LLM答案>、<标准答案，LLM+RAG答案>打分，如果前者的得分优于后者，则表明模型不依靠RAG就能得到接近标准答案的结果，即模型知道这部分知识，标记为正样本，相反则标记为负样本。这里的打分作者并没有写具体用的啥。

4. 根据上面的正负样本训练一个”分类器“，这里的”分类器“作者给出了四种，分别是：直接prompt、上下文学习、基于BERT-BASE的分类器、基于kNN的最近邻搜索。其中前两者的基座模型可以是LLM模型本身，考虑到推理速度，也可以是其他size小一些的模型。

5. 至此，完成判断，根据上述分类器的结果选择进不进行RAG，对应图中最右边部分。

        在检索时分多路进行检索，最终综合多路检索结果，这一方案的具有代表性的工作有RePlug和SuRe，RePlug上面简单介绍了，这里再介绍一下SuRe，其思路也非常简单清晰，流程图如上图所示。从图中可以看出，SuRe的区别主要在于在检索回文本后加入了候选答案生成、条件总结和验证三个步骤，包括：

1. 多个候选检索结果生成多个候选答案。

2. 让LLM对每个候选答案进行总结。

3. 根据每个候选答案及其总结，选出最好的答案。这里作者使用的验证方法还是基于prompt进行的，包括单个答案的评估以及pair-wise的对比评估。

        初次之外，还有不少优化RAG pipeline路线的方法，比如在CoT过程中进行RAG（IRCoT）、让LLM进行自我提问（Self-Ask）等等，篇幅原因这里就不介绍了。

        这类方法大多对RAG的基本流程进行分支、重复等，思路比较清晰，方法也不算难。但是在提升效果的同时，多次LLM推理或检索往往带来额外的时间消耗，需要根据实际应用情况进行权衡。

小结

        综上，本文首先介绍了RAG的开篇，后又分五个部分从生成器、联合训练、前处理、后处理、pipeline路线五个角度介绍了一些相关的RAG优化工作。有的工作虽然效果出色，但是会带来更高的推理成本、消耗更多计算资源，需要大家根据自身情况斟酌选择。