論文信息

論文題目：Query Rewriting for Retrieval-Augmented Large Language Models
論文原文：arxiv.org/pdf/2305.14283.pdf

摘要

大型語言模型（llm）借助retrieve-then-read的流程，發(fā)揮了強大的black-box閱讀器的作用，在知識密集型任務方面取得了顯著進展。這項工作引入了一個新的框架，用Rewrite-Retrieve-Read 替代之前的retrieve-then-read。我們的方法關注的是搜索查詢本身的適應性，因為輸入文本和檢索過程中所需的知識之間不可避免地存在差距。我們首先提示一個LLM來生成query，然后使用一個web搜索引擎來檢索上下文。此外，為了更好地將查詢與凍結(jié)的模塊對齊，我們?yōu)槲覀兊墓艿捞岢隽艘粋€可訓練的方案。采用小型語言模型作為可訓練的重寫器，以滿足黑盒LLM閱讀器。通過強化學習，利用LLM Reader 的反饋對改寫器進行訓練。對下游任務、開放域QA和多項選擇QA進行評估。實驗結(jié)果表明 表明我們的框架被證明是有效的和可擴展的，并為檢索增強的LLM帶來了一個新的框架。

簡介

• LLM在NLP方面顯示出了非凡的能力，并在few-shot或zero-shot的設置中顯示出了非凡的可伸縮性和適應性。然而，訓練過程依賴于大規(guī)模的高質(zhì)量的語料庫，但沒有感知現(xiàn)實世界。因此，llm仍然必須面對幻覺的問題和時間失調(diào)。這影響了LLM的可靠性，并阻礙了更廣泛的實際應用，因為LLM響應與現(xiàn)實世界之間的一致性需要進一步的驗證?，F(xiàn)有的工作已經(jīng)證明，將外部知識（即非參數(shù)知識）與內(nèi)部知識（即參數(shù)知識）相結(jié)合可以有效地緩解幻覺，特別是對于知識密集型的任務。事實上，檢索增強的LLM已經(jīng)被證明是如此有效，以至于它們被認為是解決原始LLM代中事實缺陷的標準解決方案。檢索增強用于選擇相對段落作為語言模型的外部上下文，這是retrieve-then-read 的框架。以開放域問答任務（開放域QA）為例，檢索器首先搜索相關文檔。然后LLM接收到問題和文檔，然后預測出一個答案。
• 由于大多數(shù)LLM只能通過推理api來訪問，所以它們在管道中扮演了black-box frozen Reader的角色。這使得以前需要完全訪問的檢索增強方法不再可行。最近對檢索增強語言模型的研究更傾向于面向llm的適應。一個想法是訓練一個密集的檢索模型，以迎合凍結(jié)的語言模型。其想法是通過精心設計的提示或復雜的提示管道來觸發(fā)突發(fā)能力。與外部知識的多重交互允許LLM逐步接近正確的答案。
• 然而，仍有一些問題有待解決?，F(xiàn)有的方法忽略了查詢的適應性，即Retrieve-than-Read 的輸入。檢索查詢要么是來自數(shù)據(jù)集的原始查詢，要么是直接由黑盒生成確定的，因此總是固定的。然而，在輸入文本和真正需要查詢的知識之間不可避免地存在差距。這限制了性能，并給檢索能力的增強和快速工程帶來了負擔。
• 考慮到這個問題，本文提出了一種新的檢索增強框架Rewrite-Retrieve-Read，可以進一步調(diào)整以適應llm。在檢索器的前面，添加了一個重寫輸入的步驟，填補了給定的輸入和檢索需求之間的空白，如圖1所示。我們采用了現(xiàn)成的工具，一個互聯(lián)網(wǎng)搜索引擎，作為檢索器，它避免了搜索索引的維護，并可以訪問最新的知識。與之前的研究不同，他們需要對每個樣本的檢索犬和LLM之間的多次交互輪進行記憶，我們重寫步驟的動機是澄清從輸入文本中檢索的需要。

• 我們還為我們Rewrite-Retrieve-Read框架提出了一個可訓練的方案（圖1 (c)）。black-box retriever和Reader形成了一個凍結(jié)的系統(tǒng)。為了進一步平滑管道中的步驟，我們應用了一個小的、可訓練的語言模型來執(zhí)行重寫步驟，表示為重寫器。通過使用LLM性能作為強化學習的獎勵，學習適應檢索查詢，以提高讀者的下游任務。

  
  

方法

• 我們提出了 Rewrite-Retrieve-Read，這是一個從查詢重寫的角度改進了檢索-增強的LLM的管道。

Rewrite-Retrieve-Read

• 一個具有檢索增強的任務可以表示如下。給定一個知識密集型任務的數(shù)據(jù)集（例如，開放域QA），D={(x，y)_i}，i=0,1,2…N，x（例如，一個問題）是管道的輸入，y是預期的輸出（例如，正確答案）。我們的管道包括三個步驟。(i)查詢重寫：根據(jù)原始輸入x生成所需知識的查詢?x。（ii）檢索：搜索相關上下文，文檔。（iii）閱讀：理解輸入和上下文[doc，x]，并預測輸出y?。
• 一種直接但有效的方法是要求LLM重寫查詢，以搜索潛在需要的信息。我們使用few-shot提示來鼓勵LLM思考，輸出可以是無，一個或多個查詢來搜索。

Trainable Scheme

• 此外，完全依賴于frozen的LLM也存在一些缺點。推理錯誤或無效搜索會阻礙性能。另一方面，檢索到的知識有時可能會誤導和損害語言模型。為了更好地與凍結(jié)的模塊對齊，添加一個可訓練的模型，并通過以LLM讀者的反饋作為獎勵來進行調(diào)整是可行的。
• 基于我們的框架，我們進一步建議使用一個可訓練的小語言模型來接管重寫步驟，如圖1的右側(cè)部分所示。可訓練模型用預先訓練的T5-large（770M）初始化初始化，記為可訓練重寫器，Gθ。首先根據(jù)偽數(shù)據(jù)對重寫器進行預熱訓練，然后通過強化學習進行持續(xù)訓練。

Rewriter Warm-up

• 查詢重寫任務與像T5這樣的序列到序列生成模型的訓練目標有很大的不同。首先，我們構(gòu)造一個用于查詢重寫任務的偽數(shù)據(jù)集。受最近的蒸餾方法啟發(fā)，我們提示LLM重寫訓練集中的原始問題x，并收集生成的查詢?x作為偽標簽。然后對收集到的樣本進行過濾：那些從LLM閱讀器中獲得正確預測的樣本被選擇到預熱數(shù)據(jù)集中，表示為DT rain = {（x，?x）|y?= y}。以標準對數(shù)似然為訓練目標，對DT雨進行微調(diào)改寫Gθ，記為

  
  

• 預熱后的重寫模型表現(xiàn)出適度的性能，這取決于偽數(shù)據(jù)質(zhì)量和重寫能力。高度依賴于人工編寫的提示行，?x可能是次優(yōu)的。相對較小的規(guī)模也是熱身后性能的限制。然后，我們轉(zhuǎn)向強化學習，將重寫器與以下的檢索器和LLM閱讀器對齊。

Reinforcement Learning

• 為了進一步調(diào)整重寫器，以迎合LLM的讀者，我們采用了一個策略梯度強化學習框架。
• 在強化學習的背景下，重寫器優(yōu)化被表述為一個馬爾可夫決策過程的5元組<S，a，P，R，γ>。(i)狀態(tài)空間S是一個受詞匯表和序列長度限制的有限集。（ii）動作空間A等于詞匯表。（iii）轉(zhuǎn)移概率P由策略網(wǎng)絡決定，該策略網(wǎng)絡是轉(zhuǎn)發(fā)器模型Gθ。（iv）獎勵函數(shù)R給出的獎勵價值取決于當前的狀態(tài)。政策梯度來源于獎勵，作為訓練目標。(v) γ為貼現(xiàn)因子。更具體地說，熱身后的重寫器Gθ是初始策略模型π0。在每一步t中，at的動作是基于對當前狀態(tài)的觀察，st =生成下一個標記??xt[x，??x<t]。當生成被句子結(jié)束標記停止時，有一個集結(jié)束。在完成檢索和閱讀后，通過評估最終輸出來計算獎勵，即LLM閱讀器預測的分數(shù)。

• 我們采用近端政策優(yōu)化（PPO。最大化獎勵的期望的R被表示為

  
  
  
  

• 獎勵函數(shù)R反映了生成的查詢的質(zhì)量，這需要與任務的最終評估相一致。??x被提供給Retrieval和Reader，以進行最終的預測y?。獎勵函數(shù)的一部分是將y?與黃金標簽y進行比較的測量方法（例如，預測答案的精確匹配和F1），記為Rlm。此外，還增加了kl發(fā)散正則化，以防止模型偏離初始化

  
  
  
  

實驗