[論文閱讀] Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection

前言

RAG-based LLM 是當前使用大型語言模型（Large Language Model, LLM）的一種知名架構，也就是通過『檢索』，來給模型提供訓練時所沒有的先驗知識，讓模型能夠在得知特定資訊的情況下進行問題的回答。

對很多企業或商家來說，『大型語言模型處理文字任務』都是實實在在可以利用的技術。好比說 2023 年這一年來，無數廠商都希望透過 ChatGPT、Llama、Mistral… 等等各種模型建立屬於自己的、懂得企業內部資料的客服機器人一樣。

在許多場景中，重新訓練一個專門的客服 AI 是資源開銷極大的。所以訓練一個懂得看企業內部文件、並以此作為根據回答的模型是很有幫助的。而通過使用者的提問 Q，進而讓檢索系統 R 檢索出相關 top-K 個文件，再讓大型語言模型看著 top-K 文件回答問題 Q 的這一系統，就被稱作 RAG（Retrieval-Augmented Generation）架構。

但問題是，大型語言模型的『幻覺問題』在過去的一年裡仍然沒有得到解決。就算是透過 RAG 的方法，檢索出了相關的資訊讓模型去回答使用者的問題，也依然存在著以下可能的異常情境：

檢索出的資訊（context）相當依賴 embedding model，可能檢索到與使用者問題（query）無關的資訊，反而導致模型輸入了噪音（noise）
對所有的文本都進行檢索，是低效率的行為，提供給模型的 context 越多，反而讓自回歸模型（auto-regressive model）計算時的負擔成指數型成長

而 Self-RAG 在一定程度上減少了模型的幻覺，透過 reflection token（反思標記），模型允許在認為需要檢索時提出檢索、並對檢索到的段落進行評分與接續回答，最後再替自己的回應生成評分。檢索到的不同段落的續寫是平行的，所以最後需要使用 reflection token 的評分來替自己生成的段落打分，選擇最優的解碼結果。這一行為就是論文中提到的 tree decoding（樹狀解碼）。

Self-RAG 的流程與一般 RAG 的比較可以參考上圖。

Self-RAG – 訓練

訓練階段和推理階段的流程是不同的。首先先紀錄訓練流程。
在訓練時，模型分成兩個部份：評論模型（Critic Model, 簡稱 C）和生成模型（Generator Model，簡稱 M）。

評論模型是為了生成反思標記而特意訓練的，它的工作就是專門生成對於檢索結果、生成回答的評估。而生成模型就是我們最後用於推理回應的模型，它會基於評論模型生成的資料繼續往下訓練。

1. 評論模型（C）的訓練

C 的任務是生成反思標記，用於評估檢索到的段落以及給定生成回應的分數
C 在訓練時所使用的資料集是經過 GPT-4 標註好的，其中大部分都是由檢索系統（R）所檢索到的段落以及 C 所要輸出的反思標記
最後再使用 C 模型去更新訓練資料，由 C 所標記的訓練資料後續用於生成模型（M）的訓練 —— 這是為了在真正現實世界任務推理時，生成模型 M 有能力自己去生成這些反思標記。

根據上圖，反思標記有以下種類：

Retrieve: yes/no/continue
IsREL: 相關/不相關
IsSUP: 完全支持/部份支持/不支持
IsUSE: 5/4/3/2/1

2. 生成模型（M）的訓練

M 從 C 模型標註好的資料開始往下訓練，能夠自行生成文本和反思標記。

可以想像成先訓練一個專門的模型學會這些反思任務，然後再讓最後生成的模型去學習這個過程。也可以理解成用 GPT-4 標註全部反思標記成本實在太高了，於是把 GPT-4 的在這個特有任務上的知識蒸餾到一個小的語言模型上，再讓這個小的語言模型完成所有任務標註。

Self-RAG – 推理

Self-RAG 跟一般的 RAG 最大的不同之處在於，是由以下流程組成：

模型主動提出檢索 token，進行資訊的檢索
平行化地根據檢索到的資訊生成片段的結果
檢查生成的結果是否相關，選擇最需要的片段
重複檢索…
依此類推，直到生成結束

當然，模型可能從頭到尾都不進行檢索，僅靠自己本來的知識回答。

更詳細地說，就是可以分成以下數個階段：

檢索階段：

檢索系統 R 根據當前的輸入 x 和先前的生成 y (時序 t 以前) 進行檢索，尋找 K 個相關的文本段落。
這些段落將被用作生成模型 M 進行下一步生成的基礎。

並行生成階段：

生成模型 M 處理每一個檢索到的段落 d ，為每一個 d 生成不同的續寫候選。
此過程產生 K 個不同的生成候選，每一個都基於一個不同的檢索到的段落。

決定最佳生成段落階段：

使用解碼演算法中經典的 beam search 在每個時間點 t 進行段落的評估，從 K 個候選中選出排名前B 的續寫。

這些候選是基於生成段落中的反思標記：相關性（ISREL）、支持性（ISSUP）和使用性（ISUSE）等評價指標進行評分和排名的。

最終，選擇得分最高的續寫作為該時間戳的最佳輸出。如果剛好同分，則應依賴一些事前設定，如隨機、選取特定長度、融合同分回答、embedding 計算… 等等方式去選擇。

這個過程確保了生成的回答不僅與檢索到的資訊相關，而且考慮到了生成模型自身對這些檢索段落的評價。這樣，Self-RAG G模型能夠在保持事實準確性和資訊豐富性的同時，生成品質較高的回應。

結果

在整體性能上，Self-RAG 在多種任務上顯著優於更多 LLM 模型與一般 RAG 模型。在 PubHeath、PopQA 等任務上甚至贏過 ChatGPT。

感想

這種使用 special token 調用工具的方法，讓我想到了 Toolformer，只不過在 Self-RAG 的任務中被換成了『檢索系統』跟『評估』兩種工具，其中評估甚至還是自己去做的。

Self-RAG 雖然在論文中寫道其效果超過許多重要的模型，但考量到多次檢索的 tree decoding，其實本質上確實就是拿推理時間換取了參考的性能。

當然，隨著各種量化、推理工具（如 FlashDecoding-V4，我一直想去看卻沒時間）的不斷發展，或許現在我們所擔心的推理時間，會在某一天被徹底解決吧。

[論文閱讀] Self-RAG: Learning to Retrieve, Generate, and Critique through Self-Reflection

前言