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

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

圖片

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

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

結(jié)果發(fā)現(xiàn)，不僅數(shù)據(jù)效率提高了，模型性能也有顯著增強，相比今年5月剛提出的ChatQA-1.5系列有顯著優(yōu)勢。

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

圖片

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

圖片

檢索增強生成技術(shù)，簡稱為RAG（Retrieval-Augmented Generation），被廣泛適用于LLM的定制化，尤其是知識密集型的NLP任務(wù)?？梢詭椭Ｐ驮诓桓淖儥?quán)重的情況下掌握「長尾知識」和最新信息，并適應(yīng)到特定的領(lǐng)域。

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

圖片

來源：AWS

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

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

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

圖片

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

那把k值就定在這個區(qū)間不行嗎？

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

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

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

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

基于上述考慮，作者認為可以只使用LLM同時完成上下文檢索和內(nèi)容生成任務(wù)，通過設(shè)計RAG的指令調(diào)優(yōu)來實現(xiàn)，這種新穎的框架被命名為RankRAG。

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

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

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

RankRAG微調(diào)框架

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

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

圖片

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

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

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

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

圖片

實驗

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

圖片

- 優(yōu)于現(xiàn)有的RAG方法

可以看到，在8B的參數(shù)規(guī)模上，RankRAG始終優(yōu)于當前的開源SOTA——ChatQA-1.5 8B。

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

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

- 在更具挑戰(zhàn)性的數(shù)據(jù)集上表現(xiàn)出更大的改進

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

這些發(fā)現(xiàn)表明，在具有挑戰(zhàn)性的OpenQA數(shù)據(jù)集中，檢索器中的頂級文檔與答案的相關(guān)性較低，而上下文排名能有效提高性能。

消融研究

- 設(shè)計組件的效果

表3展示了以Llama3 8B為骨干的RankRAG，在九個通用領(lǐng)域數(shù)據(jù)集上的消融結(jié)果?？傮w來看，所有的新組件都對最終性能有所貢獻。

圖片

- 不同大語言模型的性能表現(xiàn)

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

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

圖片

- 不同檢索器的性能表現(xiàn)

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

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

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

圖片

特定領(lǐng)域的RAG基準

為了證明RankRAG在專業(yè)領(lǐng)域的適應(yīng)性，作者在最新的生物醫(yī)學(xué)RAG基準Mirage上進行了實驗。

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

從表5中可以看出，即使沒有進行過微調(diào)，RankRAG依然可以在醫(yī)學(xué)問答任務(wù)上有著出色的表現(xiàn)——

不僅在8B的規(guī)模下部分超越了醫(yī)學(xué)領(lǐng)域的開源SOTA Meditron 70B，而且在70B的規(guī)模下達到了GPT-4超過98%的性能。

圖片

排名模塊

- RankRAG具有高數(shù)據(jù)效率

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

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

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

圖片

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

這一發(fā)現(xiàn)證實了RankRAG的數(shù)據(jù)效率——只需少量排名數(shù)據(jù)就能實現(xiàn)有效性能，并能在各種任務(wù)中保持適應(yīng)性。

圖片

- RankRAG的性能與時間效率

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

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

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

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

圖片

案例研究

表7展示了NQ數(shù)據(jù)集上的一個案例。

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

相比之下，重新排名技術(shù)則發(fā)現(xiàn)了另外兩個相關(guān)的段落，從而幫助模型給出了正確的答案。

圖片

作者介紹

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

Yue Yu目前是GaTech CSE學(xué)院的五年級博士生，由Chao Zhang教授指導(dǎo)，研究興趣是LLM和以數(shù)據(jù)為中心的AI的交叉領(lǐng)域，博士期間曾在Meta、英偉達、谷歌研究院、微軟研究院等機構(gòu)實習(xí)。

在進入GaTech前，他在2019年從清華大學(xué)電子工程系獲得學(xué)士學(xué)位，曾在Yong Li教授指導(dǎo)下進行時空數(shù)據(jù)挖掘和推薦系統(tǒng)方向的研究。

另一位共同一作Wei Ping是英偉達應(yīng)用深度學(xué)習(xí)研究團隊的首席科學(xué)家，他的研究重點關(guān)注LLM和GenAI，致力于為文本、音頻等多模態(tài)數(shù)據(jù)構(gòu)建生成模型。

他2016年從加州大學(xué)歐文分校獲得機器學(xué)習(xí)專業(yè)的博士學(xué)位，本科和碩士分別畢業(yè)于哈爾濱工業(yè)大學(xué)和清華大學(xué)。博士畢業(yè)后曾在吳恩達創(chuàng)立的百度硅谷AI Lab領(lǐng)導(dǎo)文字轉(zhuǎn)語音團隊，之后加入英偉達。

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

圖片

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