在需要大量事實知識的文本生成任務中，RAG成為了常用的LLM部署技巧。

但佐治亞理工學院和英偉達最近發表的一篇論文提出——RAG可以不止停留在用于推理的pipeline中，類似的思路完全可以移植到微調階段，于是有了這個名為RankRAG的框架。

圖片

論文地址：https://arxiv.org/abs/2407.02485

他們的思路可以概括為：用微調拓展模型的能力，把原來RAG需要額外模型的檢索、排名任務全丟回給LLM自己。

結果發現，不僅數據效率提高了，模型性能也有顯著增強，相比今年5月剛提出的ChatQA-1.5系列有顯著優勢。

在9個通用基準和5個生物醫學的知識密集型基準上，RankRAG用Llama 3 8B/70B微調出的模型分別超過了同樣基座上ChatQA-1.5的兩個微調模型，Llama3-ChatQA-1.5-8B和Llama3-ChatQA-1.5-70B。

圖片

ChatQA-1.5項目地址：https://chatqa-project.github.io/

圖片

檢索增強生成技術，簡稱為RAG（Retrieval-Augmented Generation），被廣泛適用于LLM的定制化，尤其是知識密集型的NLP任務。可以幫助模型在不改變權重的情況下掌握「長尾知識」和最新信息，并適應到特定的領域。

通常情況下，RAG的工作流程大致是：對于給定問題，由一個基于文本編碼的稠密模型從外部數據庫中檢索到top-k個文本段，然后輸入給LLM進行讀取，以此為基礎進行生成。

圖片

來源：AWS

這個pipeline看起來非常符合直覺，也已經被廣泛使用，但作者在論文開篇指出了其中的固有局限，首先就是k值的選擇。

如果k值較大（比如top-100），即使是支持長上下文的窗口的LLM也很難快速讀取這么多文本塊。隨著k值的增大，性能會很快飽和。

除了效率原因，之前還有研究表明，k值在5或10這個量級時，生成結果的準確性更高。因為過多上下文會引入不相關內容，妨礙LLM生成準確答案，

圖片

《Retrieval meets Long Context Large Language Models》https://arxiv.org/abs/2310.03025

那把k值就定在這個區間不行嗎？

如果給定一個較小的k，我們需要一種機制來保證檢索結果的高召回率（recall）。

鑒于檢索器的表達能力有限（通常是稀疏檢索模型如BM25，或中等大小的編碼模型如BERT-based），通常無法捕獲所有相關信息，因此實際的應用過程還會加上一個交叉編碼（cross-encoding）的排名模型。

排名模型從數據庫中檢索到top-N個候選 (N ? k)，再經過一次排名得到最終top-k結果。

這種方案的缺陷在于，與通用的LLM本身相比，專家排名模型的零樣本泛化能力相對有限，上游檢索結果的質量很可能造成下游LLM生成任務的瓶頸。這在許多實證研究中都得到了驗證。

基于上述考慮，作者認為可以只使用LLM同時完成上下文檢索和內容生成任務，通過設計RAG的指令調優來實現，這種新穎的框架被命名為RankRAG。

OpenAI的GPT-4報告中就發現，檢索、排名過程中發展出的「確定文本塊與問題是否相關」的能力對答案的生成同樣有用，這兩者可以被視為「雙重能力」。

RankRAG在訓練過程中引入了一項帶指令的問答任務，讓模型能夠識別出與問題相關的上下文或段落，便于在推理時對檢索結果進行排名。

如果將一部分排名數據集成到指令微調中，還能大大增強LLM在RAG排名任務中的性能，甚至超過了單獨用LLM和10×排名數據進行微調的結果。

RankRAG微調框架

在推理階段，RankRAG的pipeline與上述的的「檢索-排名-生成」流程幾乎相同，首先檢索出帶有相關性分數的top-N結果，然后進行重新排名并保留top-k段落，將其與問題連接到一起進行生成。

主要的不同點在于模型訓練過程，使用了兩個階段的指令微調（圖2）直接增強LLM的相關能力，而不是在模型外部添加額外操作。

圖片

第一階段首先進行監督微調（SFT），128k個樣例來自多個數據集的混合，包括對話數據集SODA、Dolly、OpenAssistant，長格式QA數據集ELI5（需要詳細答案），LLM合成的指令，以及CoT數據集FLAN。

這個階段的SFT主要是為了提高LLM的指令跟隨能力，雖然與RAG關系不大，但可以為接下來的指令微調過程做好鋪墊。

為了提升LLM的檢索、排名性能，第二階段的微調數據集由以下幾個部分混合組成（表1）：

• 第一階段的SFT數據：用于維持指令跟隨能力
• 上下文豐富的QA數據：涵蓋了DROP、NarrativeQA、Quoref、ROPES、NewsQA、TAT-QA等數據集，每條數據包含問題、黃金上下文（golden context）和答案
• 會話QA數據集：如Synthetic Conversation和HumanAnnotatedConvQA，同時包括對話內容以及一份背景文檔
• 檢索增強的QA數據：不僅包括SQuAD和WebQuestions中的問題和答案，還用BM25將黃金上下文和檢索到的top結果組合起來，確保每條數據都有5個上下文，其中有些上下文可能不包括問題答案，甚至是hard-negative，這是為了重點提高LLM對不相關上下文的魯棒性
• 上下文排名數據：使用流行的MS Marco語義相關性數據集，將其中的黃金樣本視為相關的查詢-段落對 (??,??+)，BM25挖掘的hard negtive (??,???)則被視為不相關，讓LLM對這些樣本的相關性進行二元分類（True或False）
• 檢索增強的排名數據：同樣使用QA數據集SQuAD和WebQuestions，以及BM25檢索到的上下文，訓練LLM的對相關性進行排名的能力

圖片

實驗

表2展示了RankRAG與基線模型的對比結果：

圖片

- 優于現有的RAG方法

可以看到，在8B的參數規模上，RankRAG始終優于當前的開源SOTA——ChatQA-1.5 8B。

與更大的模型相比，RankRAG 8B依然顯著優于InstructRetro（5倍參數量）、RA-DIT 65B（8倍參數量），甚至在NQ和TriviaQA任務中超越了參數多達8倍的Llama3-instruct 70B。

在增加模型參數后，RankRAG 70B的表現不僅優于強大的ChatQA-1.5 70B模型，并且還顯著優于之前以InstructGPT為底層大語言模型的RAG基線。

- 在更具挑戰性的數據集上表現出更大的改進

相對于基線模型的性能提升，RankRAG在更具挑戰性的QA數據集上更加明顯。例如，在長尾QA（PopQA）和多跳QA（2WikimQA）任務中，RankRAG比ChatQA-1.5提高了10%以上。

這些發現表明，在具有挑戰性的OpenQA數據集中，檢索器中的頂級文檔與答案的相關性較低，而上下文排名能有效提高性能。

消融研究

- 設計組件的效果

表3展示了以Llama3 8B為骨干的RankRAG，在九個通用領域數據集上的消融結果。總體來看，所有的新組件都對最終性能有所貢獻。

圖片

- 不同大語言模型的性能表現

表4展示了RankRAG和最新基線ChatQA，在使用不同參數量的Llama2作為骨干時的性能表現。

可以看到，RankRAG在各種類型和規模下，性能都有提升——分別在7B/13B/70B變體上提高了7.8%/6.4%/6.3%。

圖片

- 不同檢索器的性能表現

圖3展示了RankRAG和ChatQA-1.5在三項代表性任務中，使用不同密集檢索器（DPR和Contriever-MS MARCO）的性能。

盡管初始檢索結果并不理想，但RankRAG在平均性能上仍然比ChatQA-1.5高出10%以上。

總的來說，RankRAG在選擇檢索器方面表現出很強的魯棒性。

圖片

特定領域的RAG基準

為了證明RankRAG在專業領域的適應性，作者在最新的生物醫學RAG基準Mirage上進行了實驗。

其中，MedCPT作為檢索器?，MedCorp2作為語料庫??。

從表5中可以看出，即使沒有進行過微調，RankRAG依然可以在醫學問答任務上有著出色的表現——

不僅在8B的規模下部分超越了醫學領域的開源SOTA Meditron 70B，而且在70B的規模下達到了GPT-4超過98%的性能。

圖片

排名模塊

- RankRAG具有高數據效率

如表6所示，相比于在10倍數據上訓練的專用排名模型，RankRAG在大部分情況下都取得了更好的召回率。

甚至，RankRAG還能在大多數任務中超越BGE-ranker，要知道后者曾在超過100萬對排名數據上進行了廣泛的訓練，其中還包括一些與評估任務重疊的排名對。

值得注意的是，直接使用ChatQA-1.5對段落進行排名會降低性能，這表明在指令微調中加入排名數據的很有必要的。

圖片

如圖4所示，僅使用5k排名數據（約占MS MARCO數據集的1%），RankRAG即可獲得極佳的結果，而將排名數據增加到50k更是帶來了顯著的增益。

這一發現證實了RankRAG的數據效率——只需少量排名數據就能實現有效性能，并能在各種任務中保持適應性。

圖片

- RankRAG的性能與時間效率

隨著模型規模的擴大，不僅延遲開銷會增加，逐樣本排名的時間開銷也會增大。

為了研究時間效率與性能之間的關系，作者改變了重新排名中使用的N，并在圖5中繪制了N與最終準確率的關系。

可以看到，即使N=20，RankRAG仍然比不進行重新排名的基準模型有所改進。

此外，當N從20逐漸增加100時，重新排名將精確匹配分數提高了5.9%到9.1%，并且只額外增加了0.9到6倍的時間——遠低于預期的20到100倍的預期增長。

圖片

案例研究

表7展示了NQ數據集上的一個案例。

可以看到，由于存在干擾因素，而且一些上下文（例如，ChatQA-1.5的段落4和段落5）對回答問題沒有幫助，因此只用檢索器的話會導致上下文中出現錯誤信息。

相比之下，重新排名技術則發現了另外兩個相關的段落，從而幫助模型給出了正確的答案。

圖片

作者介紹

論文的兩位共同一作分別是佐治亞理工學院的Yue Yu和英偉達的Wei Ping。

Yue Yu目前是GaTech CSE學院的五年級博士生，由Chao Zhang教授指導，研究興趣是LLM和以數據為中心的AI的交叉領域，博士期間曾在Meta、英偉達、谷歌研究院、微軟研究院等機構實習。

在進入GaTech前，他在2019年從清華大學電子工程系獲得學士學位，曾在Yong Li教授指導下進行時空數據挖掘和推薦系統方向的研究。

另一位共同一作Wei Ping是英偉達應用深度學習研究團隊的首席科學家，他的研究重點關注LLM和GenAI，致力于為文本、音頻等多模態數據構建生成模型。

他2016年從加州大學歐文分校獲得機器學習專業的博士學位，本科和碩士分別畢業于哈爾濱工業大學和清華大學。博士畢業后曾在吳恩達創立的百度硅谷AI Lab領導文字轉語音團隊，之后加入英偉達。

值得一提的是，本篇論文對標的基準方法ChatQA也是Wei Ping之前的研究。

圖片

參考資料：https://arxiv.org/abs/2407.02485