DocReLM:自建數據集訓練檢索模型和重排序模型提高 RAG 效果
背景
根據 Crossref 的數據,目前學術界已經積累了 1.4 篇文獻,每年還以幾百萬的速度增加。這種驚人的增長速度,讓研究人員跟蹤業內動態變得越來越困難。因此,大家對高效的語義檢索系統需求越來越強烈。
這類系統的開發面臨非常多的困難:
? 首先,學術論文使用高度專業化的語言,需要深刻理解上下文才能提取關鍵信息。
? 其次,理解一篇論文不僅要閱讀其內容,還要理解引用及其相互關系,這要求系統能夠把握論文內引用間的聯系。
? 最后,用戶可能無法用精確關鍵詞描述信息需求,系統需能洞悉用戶意圖,提供最相關文獻。盡管已有嘗試采用通用嵌入模型構建檢索系統,但模型訓練與實際搜索之間存在差距,導致性能不佳。
基于以上背景,今天這篇論文的作者推出了 DocReLM,一個利用語言模型強化的文檔檢索系統。
在 DocReLM 中,LLM 以兩種獨特方式發揮作用:
? 首先,作為自動數據標注器,生成高質量的訓練數據以提升檢索模型性能。使得密集檢索器和重排模型在性能上超越了競爭對手。其他模型由于訓練于通用語義相似性任務,性能不及 DocReLM 的模型。通過 LLM 生成的數據訓練,DocReLM 更適合文檔檢索任務,同時領域專業知識也在這一過程中被融入模型。
? 另外,LLM 作為搜索 Agent,根據檢索文檔的文本信息和引用關系來優化論文選擇。類似于人類通過閱讀論文引用來尋找最相關文獻的推理過程。LLM 理解上下文和引用關系的能力使其能夠精準識別相關論文,從而提升系統性能。在自然科學領域,這一功能尤為有用,因為論文間高度互聯,理解引用對于找到相關文獻至關重要。這也是一種新穎的在檢索系統中應用大型語言模型的方法。
系統架構
圖片
DocReLM 系統整體結構如上圖所示,主要分為三個部分:檢索器、重排器和引用提取器。
- ? 用戶提出查詢時,檢索器利用嵌入技術迅速從文檔庫中篩選出候選段落。
- ? 接著,重排器進一步精準排序這些段落。
- ? 最后,引用提取器分析重排器提供的前 k 個結果,根據文檔內容生成恰當的引用。
檢索器
檢索器的作用是高效地從龐大的文檔庫中提取精選文檔集。
一般來說,檢索器會用到兩種模型:稀疏檢索器和密集檢索器。
例如,著名的 BM25 稀疏檢索器利用詞袋模型向量匹配候選項。而密集嵌入模型則將句子明確編碼為密集向量,并通過余弦相似度來評估查詢與段落嵌入之間的匹配可能性。
在 DocReLM 中,作者采用了密集檢索器,并基于開源的 jina-embedding-v2-base 模型進行定制訓練。這個檢索器充當編碼器角色,已在科學領域多種任務中得到應用。
通過對比學習,密集檢索器生成的嵌入能夠突出關鍵段落的細節,即便術語不同,也能連接語義相似的句子。然而,預訓練的嵌入模型在文檔檢索任務中常常缺乏特異性,與特定領域的論文存在知識鴻溝。為此,作者引入大型語言模型自動用學術論文庫標注偽查詢,以此優化嵌入模型,使其更貼合學術領域的任務需求。
構建這一檢索系統時,將長文檔切分為段落,并將它們轉換為嵌入向量,存儲于 Faiss 向量數據庫中。用戶輸入查詢后,DocReLM 利用檢索模型提取查詢的語義嵌入,并基于文檔和查詢嵌入之間的余弦距離挑選候選段落。
重排器
快速檢索之后,重排器被用來提升結果的準確性。雖然它通常更精確,但與檢索器相比,重排器的處理速度較慢。這決定了它通常在檢索后使用,檢索器首先通過排除大部分不相關段落來縮小選擇范圍。
在此任務中,作者采用了交叉編碼器,它處理由查詢字符串和檢索子系統提供的候選段落組成的串聯序列,以[SEP]標記分隔。
通過在輸入中整合查詢和段落,重排器增強了兩者的交互,從而提高了其表達力,并在效果上超越了檢索器。交叉編碼器的一個局限是它無法預處理文檔庫,需要對每個新查詢進行實時推理,因此不能單獨作為系統使用,它需要檢索器提供一小批相關候選。
在 DocReLM 中,使用了開源的 XLM-RoBERTa-large 模型作為基礎,并用作者的自定義數據進一步訓練,采用 LCE 方法。檢索器確定的前 200 個段落隨后輸入到重排器中,為每個串聯對輸出一個分數,然后根據這些分數對候選段落進行重新排序。
引用提取器
經過精確的段落檢索和交叉編碼器的重排,相關論文根據與查詢的相關性被依次展示。
但是,通過一個額外步驟可以顯著提升搜索結果的質量:DocReLM 系統中的引用提取器,它負責閱讀段落內容并從中抽取最佳引用文獻。
這種做法更貼近實際的搜索習慣。雖然檢索器和重排器挑選出的論文可能并非最佳答案,但其引用列表中很可能包含了正確答案。為了提高搜索的準確性,引導大語言模型來處理檢索結果,并從引用文獻中識別出更加合適的論文。
具體操作中,通過在論文段落中嵌入參考文獻的標識符來調整文本。當這段經過修改的文本被大型語言模型處理時,它會根據上下文信息抽取出這些標識符。如果模型認定段落本身足以回答問題,它將輸出引用文獻的標識符。系統將根據提供的標識符搜索相應的論文。
圖片
如上圖展示的,檢索器迅速從完整文獻庫中篩選出 n 個候選文段。緊接著,重排器對這些文段進行排序。繼而,引用提取器細讀排名最前的 10 篇論文,并為每篇提取三條引用,這些引用被巧妙地嵌入到對應原文之后。
模型訓練
作者分別對檢索器和重排序模型進行了訓練。
先利用對比學習在生成的訓練集上訓練檢索器。隨后利用這個檢索器篩選出難度較高的負樣本,用于對重排器進行更加精細的對比學習微調。
利用 LLM 生成訓練數據
為了通過對比學習訓練檢索器和重排器,首先需要搜集一系列的用戶查詢和文檔段落。盡管市面上有 MS MARCO 這類帶有人工標注的數據集,但人工標注的成本非常高,使得獲取領域特定的檢索數據集不是很容易。
在本項研究中,運用 LLM 來從文檔中生成偽查詢。選用了 unarXive 上的論文來構建數據集,對他們的原始數據進行處理,剔除了語義上不相關的文本,并加入了數學 LaTeX 格式和引用標識符。這些引用標識符按照 Ref.X of IDFp 的格式進行編排,其中 IDFp 代表的是原論文的標識符。將整篇論文切分成多個段落,并動用 vicuna-7b-v1.5-16k 根據段落、標題和摘要為每個句子生成查詢。借鑒思維鏈方法,首先采用零樣本方法引導大語言模型 Agent 應用基于標題、摘要和內容草擬大綱,繼而生成查詢。
作者為實驗建立了兩個子集。其一是量子物理領域,包含 56927 篇論文,搜集了 2.8??? 的訓練數據。其二是計算機視覺領域,包含 37390 篇論文,搜集了 1.1??? 的訓練數據。這些數據被用于訓練密集型段落檢索器與重排器。
在對比訓練過程中,檢索器和重排器的負樣本生成策略各有不同。檢索器的負樣本是通過隨機抽樣來獲得的。在每次訓練批次中,我們會抽取 N+1 組正樣本對,每組正樣本對將伴隨 N 組負樣本對。增加批次大小可以提升檢索器的性能,因為它增加了引入困難負樣本的可能性。為了在對比學習中獲得更佳表現,采用了梯度緩存技術來擴大訓練批次的規模。
至于重排器的訓練,我們采用已訓練好的檢索器,按照文獻[8]中提出的 LCE 方法來挑選負樣本。對于每一個查詢檢索器會返回與之相關性最高的 N 個段落。這些段落與查詢緊密相連,表明挑選出了除了正確答案之外最具迷惑性的錯誤選項。因此,這些誤導性強的段落成為了訓練重排器時使用的高質量困難負樣本。
效果如何
為了檢驗 DocReLM 的效果,作者根據真實的研究場景構建了一個基準測試,邀請各領域的研究人員根據他們的研究方向提出查詢,并附上論文作為標注的參考。為確保這些標注的準確性,研究人員必須提供調查研究中的論斷作為依據。接著,選用 unarXive 的一個子集作為檢索的文檔庫。最終的基準測試涵蓋兩個領域:量子物理領域包含 301 個查詢和 56927 篇論文,計算機視覺領域則包含 204 個查詢和 37390 篇論文。將 DocReLM 與 Google Scholar 進行了對比。此外,還進行了消融分析,以評估系統中三個獨立組件的有效性。
圖片
作者對比了三種基礎檢索模型:BM25、OpenAI 的 ada-002 以及 jina-base-v2。它們的準確度展示在商標中。稀疏型檢索器 BM25 相較其他模型表現欠佳,其基于詞頻的方法不足以捕捉任務中的語義關聯。DocReLM-retriever 在兩個領域中均超越了這些模型。特別是在計算機視覺領域,DocReLM-retriever 的表現顯著超越其他模型,其 top-1 準確度比其他最佳嵌入模型高出 5.89%,而 top-5 和 top-10 的準確度分別提高了 13.72%和 19.13%。值得一提的是,在計算機視覺領域,所有三種模型都展現出了較高的平均準確度。而在量子物理領域,嵌入模型的表現相近,top-1 結果并無太大差異;但在 top-5 結果上,作者的模型將準確度從 9.97%提升至 11.96%,超越了其他模型。
圖片
為測試重排器的性能,挑選了三種不同的重排器進行比較:bge-reranker-large、Cohere 以及 DocReLM-reranker。這些重排器以 DocReLM-retriever 篩選出的前 200 個段落作為輸入(如上表)。在計算機視覺領域,DocReLM-reranker 超越了其他模型,其 top-1 準確度比其他最佳模型高出 0.49%,而 top-5 和 top-10 的準確度分別提升了 1.97%和 2.94%。在量子物理領域,DocReLM-reranker 與 Cohere 表現相當,DocReLM-reranker 的模型在 top-20 的表現更佳,而 Cohere 在 top-1 和 top-10 的得分上更為突出。這些結果證明 DocReLM-reranker 在兩個領域的效率均高于其他模型,凸顯了重排器在提升檢索系統準確度中的關鍵作用。
最后一項實驗評估了引用提取模型的有效性。對比了開啟和未開啟引用提取功能的 DocReLM-reranker 的性能,結果列在上表中。在計算機視覺領域,引入了引用提取功能的模型 internLM 將 top-5 和 top-20 的準確度分別提升了 0.98%和 2.94%,證實了引用提取功能在提升檢索系統準確度方面的有效性。在量子物理領域,同樣的模型將 top-5 和 top-10 的準確度分別提升了 9.63%和 16.28%。
討論
LLM 的精華提煉
神經密集型檢索器和交叉編碼重排器能夠精準捕捉查詢與文檔間的語義聯系,使得相關文檔的檢索更為精確。使用 LLM 生成的數據訓練模型后,與基礎模型相比,性能有了顯著飛躍。這些數據由 LLM 生成,使得 LLM 的理解力得以凝練并注入到較小型的模型中,讓通用模型能夠適應特定任務。
LLM 與檢索系統的融合創新
作者提出了一種新穎的大語言模型與檢索系統集成方法。隨著 LLM 的流行,人們一直在嘗試將其與搜索引擎結合。這些嘗試主要有兩種方向:一是利用 LLM 擴展搜索查詢,二是用 LLM 總結搜索結果。
而 DocReLM 則開辟了新徑,讓 LLM 理解搜索結果并繼續搜索過程,推薦更佳的候選文檔。此概念有進一步擴展的空間。該過程可以迭代執行,讓 LLM 在引用網絡中穿梭多個節點,直至鎖定最相關的結果。這一任務要求 LLM 不僅要理解引用文獻間的邏輯聯系,還需具備足夠的領域知識以精確提取引用。通過整合能夠進行多次搜索并分析結果以優化搜索的 LLM,可以顯著提升通用檢索系統的性能。這種方法預示著檢索系統未來發展方向的希望之光。
圖片
圖片
? 論文原文: https://arxiv.org/abs/2405.11461
