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

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

圖片

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

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

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

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

圖片

ChatQA-1.5項(xiàng)目地址：https://chatqa-project.github.io/

圖片

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

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

圖片

來(lái)源：AWS

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

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

除了效率原因，之前還有研究表明，k值在5或10這個(gè)量級(jí)時(shí)，生成結(jié)果的準(zhǔn)確性更高。因?yàn)檫^(guò)多上下文會(huì)引入不相關(guān)內(nèi)容，妨礙LLM生成準(zhǔn)確答案，

圖片

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

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

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

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

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

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

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

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

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

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

RankRAG微調(diào)框架

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

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

圖片

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

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

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

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

圖片

實(shí)驗(yàn)

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

圖片

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

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

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

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

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

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

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

消融研究

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

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

圖片

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

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

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

圖片

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

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

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

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

圖片

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

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

其中，MedCPT作為檢索器?，MedCorp2作為語(yǔ)料庫(kù)??。

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

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

圖片

排名模塊

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

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

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

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

圖片

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

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

圖片

- RankRAG的性能與時(shí)間效率

隨著模型規(guī)模的擴(kuò)大，不僅延遲開(kāi)銷會(huì)增加，逐樣本排名的時(shí)間開(kāi)銷也會(huì)增大。

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

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

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

圖片

案例研究

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

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

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

圖片

作者介紹

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

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

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

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

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

值得一提的是，本篇論文對(duì)標(biāo)的基準(zhǔn)方法ChatQA也是Wei Ping之前的研究。

圖片

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