大語言模型具有顯著的能力，但它們常常因為僅依賴于其參數化知識而產生包含事實錯誤的響應。傳統的檢索增強生成（Retrieval-Augmented Generation, RAG）方法雖然能減少這些問題，但是存在無差別地檢索和結合固定數量的段落，沒有考慮檢索是否必要或檢索結果的相關性，都會削弱語言模型的靈活性，或者導致無益的響應生成。

現在LLM+RAG（檢索增強）的技術方案已成為LLM在眾多應用場景的主流。在LLM+RAG技術框架下，存在的問題/技術難點主要有：

1）如何保證檢索內容是有效，或有用的？

2）如何驗證檢索的內容對輸出的結果是支持的？

3）如何驗證輸出的結果是來自檢索還是模型的生成？

今天來看一篇華盛頓大學和艾倫AI Lab發表在2024 ICLR上的工作，提出Self-RAG，使單一的語言模型能夠根據需求自適應地檢索段落，并使用特殊的反思Token（reflection tokens）來生成和反思檢索到的段落及自身生成的內容。

下面來詳細看下Self-RAG是如何解決這些問題的：

1、方法介紹

從Figure 1可以看出Self-RAG和Vanilla RAG的區別，

傳統RAG方法的工作流程通常包括：

• Step 1: 檢索K個文檔：無論是否有必要，都會從大型文檔集合中無差別地檢索出固定數量的相關段落。
• Step 2: 生成輸出：根據檢索到的段落和原始任務輸入，生成響應文本。這種方法可能無法確保所有生成的內容都得到了檢索段落的支持。

相比之下，SELF-RAG方法引入了一個更靈活且智能的流程：

• Step 1: 決定是否檢索：給定一個輸入提示和先前的生成，SELF-RAG首先判斷繼續生成時是否需要通過檢索來增強信息。如果確定需要檢索，它會輸出一個檢索Token（retrieval token），該Token會按需調用檢索模型。
• Step 2: 并行處理多個檢索段落：當接收到檢索Token后，SELF-RAG并行處理多個檢索到的段落，評估它們的相關性，并基于這些段落生成相應的任務輸出。
• Step 3: 自我評價與選擇最佳片段：SELF-RAG接著生成評價Token（critique tokens），用于評價自身生成的輸出，并從中挑選出最符合事實性和整體質量的輸出。

簡單來說，對比RAG，self-RAG框架的不同之處就是：在生成過程中利用特殊的token達到更精細的控制——要不要檢索、檢索內容相關性怎樣、利用檢索內容生成的質量怎樣。達到這些目的，就會讓RAG+LLM生成的內容在質量、事實性、驗證性上得到提升。

下面來詳細看下怎么利用特殊Token實現Self-RAG：

問題概述

給定一個輸入x，訓練目標是讓模型M順序生成由多個片段組成的文本輸出，其中代表第t個片段的一系列Token。每個片段可以是一個句子或者更小的文本單位，這取決于具體的實現。生成的Token不僅包含來自原始詞匯表中的文本，還可能包含特殊的反思Token（reflection tokens），這些Token用于指示是否需要進行檢索，或者確認輸出的相關性、支持性或完整性。

對于每一個輸入x及其之前生成的內容，模型首先解碼一個檢索Token來評估檢索的必要性。如果不需要檢索，模型會像標準的語言模型一樣預測下一個輸出片段。如果確定需要檢索，那么模型將：

1. 生成評價Token以評估相關性：模型生成一個評價Token來評價檢索到的段落是否與任務相關。
2. 生成下一個響應片段：根據檢索到的信息生成接下來的響應部分。
3. 生成評價Token以評估支持度：評價新生成的響應片段是否得到了檢索到的段落的支持。
4. 生成評價Token以評估整體效用：最后生成一個新的評價Token，對整個響應的總體效用進行評價。

在生成每個片段的過程中，SELF-RAG并行處理多個檢索到的段落，并使用其自身生成的反思Token來施加軟約束或硬控制。例如，在決定是否插入某個段落時，模型會考慮該段落的相關性和支持程度。如果存在多個滿足條件的段落，則選擇具有最高檢索分數的那個。

訓練概述

SELF-RAG通過將反思Token統一為從擴展模型詞匯表中預測的下一個Token，訓練一個任意的語言模型生成帶有反思Token的文本。

具體來說，使用由檢索器R檢索的段落和由評價模型C預測的反思Token交織在一起的精選語料庫訓練生成器模型M。使用評價模型，通過將反思Token插入任務輸出中來更新訓練語料庫。然后，使用傳統的語言模型目標訓練最終的生成器模型M，使其能夠在推理時不依賴評價模型自行生成反思Token。

訓練

在訓練過程中，需要兩個模型：一個評價模型C，一個生成模型M；利用評價模型C來生成M模型所需的監督數據，在推理過程中只依賴M模型，不用C模型。

訓練評價模型（Critic Model）

• 數據收集：
  手動為每個段落標注反思Token成本很高，因此研究人員采用GPT-4來生成這些反饋。通過向GPT-4發出指令提示，使其為特定的任務輸入-輸出對預測適當的反思Token，從而創建監督數據。這種方法需要對多個檢索到的段落以及任務輸入-輸出實例進行細粒度評估，以生成用于訓練SELF-RAG的數據，并將這些知識蒸餾到一個內部的評價模型C中。
  對于每組反思Token，從原始訓練數據中隨機采樣實例。然后使用不同的指令提示（表1）來生成不同類型的反思Token（如Retrieve、ISREL、ISSUP、ISUSE）。例如，對于Retrieve Token，使用原始任務輸入x和輸出y的少量示例提示GPT-4判斷是否需要從外部文檔中獲取信息以生成更好的響應。收集4k到20k的監督訓練數據用于每種類型的反思Token，并將它們組合成C的訓練數據。

• 評價模型學習：
  在收集到的訓練數據上，使用預訓練的語言模型初始化C，并進行訓練。最大化似然：
  
  其中r是反思Token。使用與生成器模型相同的預訓練語言模型（Llama 2-7B）來初始化C。訓練后評價模型的表現如下：
  
  
  

訓練生成器模型（Generator Model）

• 數據收集：
  給定輸入-輸出對，運行 C評估是否需要額外的段落來增強生成。
  如果需要檢索，添加檢索特殊Token Retrieve=Yes，并使用R檢索頂部K個段落D。對于每個段落，C 進一步評估段落的相關性并預測 ISREL。如果段落相關，C 進一步評估段落是否支持模型生成并預測 ISSUP。在輸出的末尾，C 預測整體效用Token ISUSE，并將帶有反思Token的增強輸出和原始輸入對添加到中。

• 生成器模型學習：
  在增強的語料庫上，使用標準的下一個Token目標訓練生成器模型 M進行訓練：
  
  與C的訓練不同，M 學習預測目標輸出以及反思Token。在訓練期間，對于損失計算，屏蔽檢索的文本塊，并將原始詞匯表 V 擴展為一組反思Token {Critique,Retrieve}。

以上就是訓練的過程，其核心就是生成模型M的訓練，而M模型的不同之處在于，在生成文本序列的同時，還會輔助生成多種標簽信息，以達到更細粒度的檢索控制能力。

推理

生成反思Token以自我評估其輸出使得Self-RAG在推理階段可控，使其能夠根據不同的任務需求調整其行為。對于需要事實準確性的任務，希望模型更頻繁地檢索段落，以確保輸出與可用證據緊密一致。相反，在更開放的任務中，如撰寫個人經歷的文章，重點則轉向減少檢索并優先考慮整體創造力或效用評分。

主要思路：1）在自回歸方式生成過程中，若生成有[Retrieve]標簽，且值為yes時，便使用檢索器（檢索模型）R召回d個相關片段；2）接著訓練利用召回的相關片段去預測是否相關[IsREL]標簽，以及評價標簽[IsSUP]和[IsUSE]標簽，跟進這些標簽對生成的序列進行排序，選top1的結果作為最終生成結果；3）若生成有[Retrieve]==NO時，就直接利用x生成。

具體實現如下：

自適應檢索（Adaptive Retrieval）

動態檢索決策：SELF-RAG在推理過程中動態決定是否需要檢索文本段落。這是通過預測“Retrieve”Token來實現的。如果模型預測“Retrieve = Yes”，則觸發檢索過程。否則，模型將繼續生成文本，而不進行檢索。為了更靈活地控制檢索頻率，作者引入了一個閾值機制。如果生成“Retrieve = Yes”Token的概率超過指定的閾值，則觸發檢索。這種機制允許在不同任務中靈活調整檢索頻率，以適應不同的任務需求。

基于評價Token的樹解碼（Tree-decoding with Critique Tokens）

• 并行處理和候選生成：
  當需要檢索時，SELF-RAG會并行處理多個檢索到的段落，并為每個段落生成不同的候選續寫。
  具體來說，對于每個輸入x 和之前的生成，模型會預測是否需要檢索。如果需要檢索，模型會檢索K個相關段落D，并為每個段落生成一個候選續寫。這些候選續寫會通過段落級別的束搜索（beam search=B）進行評估和選擇，并在生成結束時返回最佳序列。
• 評價Token的評分機制：
  每個候選續寫的評分由兩部分組成：生成概率和評價分數。
  生成概率表示在給定輸入x、段落d 和之前的生成的情況下，生成的概率。
  評價分數S(Critique) 是評價Token的加權線性組合。具體來說，對于每種評價Token類型G（如 ISREL、ISSUP、ISUSE），計算其在時間戳t的評分，然后使用加權和來計算段落的總評分：
  其中，，，表示評價Token類型G 的最理想反思Token （例如，[IsRel]=Relevant）的生成概率， 是在推理時調整的超參數，以在測試時實現自定義行為。

推理過程的靈活性和定制性

• 任務需求的適應性：SELF-RAG通過生成反思Token，能夠在推理階段自我評估其輸出的質量，并根據任務需求調整其行為。在需要高事實準確性的任務中，模型會更頻繁地進行檢索，以確保輸出與可用證據緊密對齊。在更開放的任務中，如創作個人經歷作文，模型會減少檢索頻率，優先考慮整體的創造力或效用得分。
• 推理時的定制化：通過調整評價Token的權重，可以在推理時對模型的行為進行定制，而無需額外的訓練。例如，可以通過增加 ISSUP的權重來確保輸出主要由證據支持，同時相對降低其他方面的權重。

2、實驗結果

無檢索的基線模型比較

Baseline：公開可用的預訓練模型（如Llama2-7B、Alpaca7B和Alpaca13B），以及使用私有數據訓練和強化的模型（如ChatGPT和Llama2-chat）。還包括CoVE65B，這是一種同時引入迭代提示工程來提高生成事實性的方法。

SELF-RAG在所有任務中均顯著優于沒有檢索的基線模型，包括監督微調的大型語言模型（LLMs）。在公共健康（PubHealth）、開放域問答（PopQA）、傳記生成和長文本問答（ASQA）任務中，SELF-RAG甚至超過了ChatGPT。與使用復雜提示工程的CoVE方法相比，SELF-RAG在生物傳記生成任務中表現更好。

帶檢索的基線模型比較

Baseline：

• 標準RAG基線：基于Llama2或Alpaca的語言模型，在生成輸出時會將檢索到的文檔附加到查詢之前。
• 微調后的Llama2（Llama2-FT）：在所有訓練數據上進行了微調，不包含反思Token或檢索到的段落。
• 私有數據檢索增強的模型：如Ret-ChatGPT和Ret-Llama2-chat，以及perplexity.ai，它們部署了類似的增強技術。
• 并行方法：如SAIL和Toolformer，這些方法分別在指令調整期間插入檢索到的文檔或預先訓練帶有API調用的語言模型。

SELF-RAG在許多任務中也優于現有的檢索增強模型（RAG），在所有任務中均獲得了非專有語言模型基礎模型中的最佳性能。盡管指令調優的LLMs（如Llama2-chat和Alpaca）在檢索后有顯著提升，但在需要復雜推理或生成的任務中，這些基線模型的表現并不如SELF-RAG。在引用準確性方面，大多數帶檢索的基線模型表現不佳，而SELF-RAG在ASQA任務中顯示出顯著更高的引用精度和召回率。

消融研究（Ablation Studies）

研究人員構建了幾個SELF-RAG的變體以識別哪些因素發揮了重要作用：

• No Retriever：該變體僅使用標準指令跟隨方法訓練語言模型，不利用檢索到的段落。
• No Critic：這個變體總是用最上面的一個檢索文檔來增強輸入輸出對，但沒有反思Token，類似于SAIL的方法，不過這里使用的是自己的指令輸出數據而不是Alpaca數據集。
• No Retrieval：推理時禁用檢索功能。
• Hard Constraints：當且僅當檢索Token為“是”時才進行檢索，而不是采用自適應閾值。
• Retrieve Top 1：始終檢索并使用最上面的一個文檔，類似于傳統的RAG方法。
• Remove ISSUP：在批判導向的束搜索過程中移除ISSUP評分。

與沒有檢索器或評價器的模型相比，SELF-RAG在所有任務中表現出顯著的性能提升，表明這些組件對性能的貢獻很大。

使用Top段落而不考慮其相關性（如標準RAG方法）會導致性能大幅下降，而僅依賴相關性分數也不足以提高性能。

檢索機制：檢索機制顯著提升了模型在需要外部信息的任務上的表現，特別是在事實驗證和引用準確性方面。

自我評價機制：反思Token（如ISREL, ISSUP, ISUSE）有助于提高輸出的相關性和支持度，從而改善模型的整體效用。

自適應閾值：相較于固定規則，自適應閾值使得模型能夠在適當的時候動態決定是否檢索，提高了靈活性和效率。

這些結果表明，SELF-RAG通過仔細選擇生成內容并基于多個細粒度標準進行評估，而不是簡單地使用檢索模型的所有段落或僅依賴相關性分數，從而提高了生成質量。

推理時定制化的效果

SELF-RAG的一個關鍵優勢是能夠在推理時控制每種評價類型對最終生成采樣的影響（圖3 b）。例如，在ASQA任務中，通過增加ISSUP的權重，可以提高模型的引用精度，因為這強調了模型生成是否由證據支持。然而，較大的權重會導致MAUVE分數降低，因為生成的文本更長、更流暢時，常常會有更多未被引用完全支持的主張。這種框架允許實踐者在測試時通過調整這些參數來選擇和定制模型的行為，而無需額外的訓練。

檢索頻率和歸一化準確率

隨著檢索閾值的變化，模型在不同任務上的檢索頻率和準確率也有所不同。較低的閾值可能會增加不必要的檢索次數，而較高的閾值則可能導致錯過重要的外部信息。因此，找到一個合適的閾值對于平衡這兩者至關重要。作者評估了這種適應性閾值對整體準確性和檢索頻率的影響，并在PubHealth和PopQA數據集上進行了測試。結果表明，模型的檢索頻率隨著閾值的變化而顯著變化。在PubHealth上，檢索較少導致的性能下降較小，而在PopQA上則較大。

3、總結

Self-RAG框架給LLM+檢索增強任務提供了一種新的結合方式——在生成過程中加入可以多維、更細粒度的控制與評價標簽，讓LLM對檢索的內容的利用，以及利用效果有了更直接的操作。

其有個缺陷時，在生成的時候要多次生成和判斷標簽，這會增加推理成本。此外，在此框架下，還是有很多優化空間，1）比如標簽的優化（用更少的標簽，或者代表其他含義的標簽）；2）在召回相關文檔后，用一個小模型來判斷，選擇top1作為最終結果，減少循環計算。

該方法目前也沒有找到落地方案，并且回答可能會很慢，還是有很大優化空間。