Qwen3 低成本手撕Search-R1的強化學習訓練框架
OpenAI-o1和DeepSeek-R1,通過獎勵驅動的學習,而無需依賴于明確的逐步監督,在邏輯推理和迭代推理方面取得了顯著的成果。
使用RL訓練能夠更有效地搜索相關信息的策略模型。例如,Search-R1、R1-Searcher和ReSearch等模型都嘗試通過強化學習來提升模型的搜索能力。
這些方法通常使用靜態的本地文本語料庫,如維基百科,無法捕捉到現實世界交互的復雜性。
DeepResearcher引入了與商業搜索引擎(如谷歌)的實時交互,使得模型能夠在接近真實世界網絡搜索的環境中進行訓練。
但它面臨著重大挑戰:
- 一是搜索引擎返回的文檔質量不可預測,這給訓練過程帶來了噪聲和不穩定性;
- 二是強化學習訓練需要頻繁的rollout,可能涉及數十萬次搜索請求,這導致了巨大的API費用,嚴重限制了系統的可擴展性。
一、ZEROSEARCH核心思想
- ZEROSEARCH利用LLMs在大規模預訓練過程中所積累的豐富世界知識,使其能夠根據搜索查詢生成相關文檔。通過輕量級的監督式微調,即使是相對較小的LLMs也能夠有效地模擬真實搜索引擎的行為。
- 通過模擬搜索引擎的交互過程來訓練策略模型,使其能夠在沒有真實搜索引擎的情況下有效地進行信息檢索和推理。
- 其創新點在于,通過控制文檔質量,解決搜索引擎返回文檔質量不可預測的問題,并且消除了API費用,極大地提高了訓練的可擴展性。
二、架構介紹
- ZEROSEARCH的整體架構圍繞著一個強化學習框架展開,其核心組件包括:
模擬LLMs在訓練過程中扮演著搜索引擎的角色,它接收策略模型生成的搜索查詢,并返回相應的文檔,這些文檔可以是有用的,也可以是噪聲的,具體取決于訓練階段和課程策略的設置。
獎勵模型根據策略模型生成的答案與真實答案之間的匹配程度給出獎勵信號,以驅動策略模型的學習過程。
參考模型則用于計算策略模型輸出與參考輸出之間的差異,從而為策略模型的更新提供指導。
策略模型是需要優化的對象,負責根據輸入的查詢生成搜索請求,并基于檢索到的信息生成最終答案。
一個策略模型(policy model)
一個參考模型(reference model)
一個獎勵模型(reward model)
一個模擬LLMs(simulation LLM)
三、技術細節
- 強化學習目標:
ZEROSEARCH的優化目標是最大化策略模型在給定查詢下生成正確答案的期望獎勵,同時最小化策略模型輸出與參考模型輸出之間的差異。
具體來說,優化目標可以表示為:
- 訓練模板:
內部推理階段(??<think>...</think>?
?)
搜索查詢階段(??<search>...</search>?
?)
答案生成階段(??<answer>...</answer>?
?)。
- ZEROSEARCH采用了一個多輪交互模板,將策略模型的輸出分為三個階段:
- 這種結構化的輸出格式有助于模型更好地組織推理過程,并使訓練過程更加透明和可控。
- 課程學習策略,是ZEROSEARCH的一個關鍵創新點:
通過逐漸增加生成文檔的噪聲比例,模擬越來越復雜的檢索場景,從而逐步提升策略模型的推理能力。
這種策略類似于人類學習過程中的逐步引導,有助于模型更好地適應不同難度的任務。
- 損失掩碼機制:
在訓練過程中,ZEROSEARCH對檢索到的文檔token應用損失掩碼,避免了這些外部生成的token對策略模型更新過程的干擾。
這一機制有助于穩定訓練過程,并提高模型的最終性能。
四、整體流程
- 數據準備:首先,收集與真實搜索引擎交互的軌跡數據,并將其標記為正樣本或負樣本,用于后續的輕量級監督式微調。
- 輕量級監督式微調:使用標記好的軌跡數據對LLMs進行微調,使其能夠根據搜索查詢生成有用或噪聲文檔。
- 策略模型初始化:初始化策略模型,并將其與模擬LLMs和獎勵模型連接起來,形成完整的強化學習框架。
- 訓練:
在訓練過程中,策略模型根據輸入的查詢生成搜索請求,模擬LLMs根據搜索請求生成文檔,策略模型根據生成的文檔進行推理并生成答案。
獎勵模型根據答案的準確性給出獎勵信號,策略模型根據獎勵信號進行更新。
隨著訓練的進行,逐漸增加生成文檔的噪聲比例,以提升策略模型的推理能力。
- 模型評估:在訓練完成后,使用各種問答數據集對模型的性能進行評估,比較其與真實搜索引擎和其他基線方法的性能差異。
舉個例子:
假設問題是“誰是《冰與火之歌》系列小說的作者?”
- 策略模型首先在?
?<think>...</think>?
?階段進行內部推理,可能會思考“我需要找到《冰與火之歌》系列小說的作者是誰”。 - 然后,在?
?<search>...</search>?
?階段,策略模型生成一個搜索查詢,如“《冰與火之歌》作者”。模擬LLMs接收到這個查詢后,生成一些文檔,這些文檔可能是有用的,也可能是噪聲的,具體取決于當前的課程階段。 - 策略模型根據這些文檔在?
?<answer>...</answer>?
?階段生成最終答案,如“喬治·R·馬丁”。 - 獎勵模型根據生成的答案與真實答案之間的匹配程度給出獎勵信號,策略模型根據這個獎勵信號進行更新,以提升其在后續任務中的性能。
4.1 輕量級監督式微調的詳細步驟
- 數據收集:
通過讓LLMs與真實搜索引擎進行多輪交互,收集一系列的交互軌跡。
這些軌跡包括查詢、搜索引擎返回的文檔以及LLMs基于這些文檔生成的答案。
- 數據標注:
根據LLMs生成的答案是否正確,將這些軌跡標注為正樣本或負樣本。
如果答案正確,則將該軌跡標記為正樣本,表示檢索到的文檔是有用的;
如果答案錯誤,則將該軌跡標記為負樣本,表示檢索到的文檔是噪聲。
- 提取查詢-文檔對:
從正負樣本軌跡中提取查詢-文檔對。
這些查詢-文檔對將用于后續的微調過程。
- 設計微調提示:
為了使LLMs能夠生成有用或噪聲文檔,設計微調提示。
在提示中加入“有用”或“噪聲”等關鍵詞,并將輸入問題及其對應答案納入提示中,以拓寬LLMs的知識邊界。
- 微調過程:
使用提取的查詢-文檔對和設計好的微調提示對LLMs進行輕量級監督式微調。
在微調過程中,調整LLMs的參數,使其能夠根據搜索查詢生成符合提示要求的文檔。
- 評估微調效果:
通過一些驗證數據集評估微調后的LLMs的性能,確保其能夠有效地模擬真實搜索引擎的行為,并且能夠根據提示生成有用或噪聲文檔。
4.2 基于課程的rollout策略的詳細步驟
- 初始化課程參數:
設置初始噪聲概率ps、最終噪聲概率pe以及指數基數b等課程參數。
這些參數將決定在訓練過程中生成文檔的噪聲比例如何變化。
- 策略模型生成搜索查詢:
在每次rollout過程中,策略模型根據輸入的查詢生成一個或多個搜索查詢。
這些查詢將被發送到模擬LLMs。
- 模擬LLMs生成文檔:
模擬LLMs接收到搜索查詢后,根據當前的課程階段生成相應的文檔。
在訓練初期,生成的文檔主要是有用的,隨著訓練的進行,逐漸增加噪聲文檔的比例。
- 策略模型進行推理:
策略模型根據模擬LLMs返回的文檔進行推理,并生成最終答案。
在推理過程中,策略模型可能會根據需要進行多輪搜索。
- 獎勵模型給出獎勵信號:
獎勵模型根據策略模型生成的答案與真實答案之間的匹配程度給出獎勵信號。
如果答案正確,獎勵信號較高;如果答案錯誤,獎勵信號較低。
- 更新策略模型:
根據獎勵信號和課程參數,更新策略模型的參數。
在更新過程中,策略模型將學習如何在不同難度的檢索場景中生成更準確的答案。
- 調整課程階段:
隨著訓練的進行,逐步調整課程參數,增加生成文檔的噪聲比例。
這將使策略模型逐漸適應更具挑戰性的檢索場景,提升其推理能力。
4.3 訓練算法的細節
- 選擇合適的強化學習算法:
根據具體任務和模型的特點,選擇合適的強化學習算法,如PPO、GRPO或Reinforce++等。
- rollout過程:
在rollout過程中,策略模型根據輸入的查詢生成搜索查詢,模擬LLMs根據搜索查詢生成文檔,策略模型根據生成的文檔進行推理并生成答案。
這個過程將產生一系列的交互數據,包括查詢、文檔、答案以及對應的獎勵信號。
- 計算損失函數:
策略損失:衡量策略模型輸出與參考模型輸出之間的差異,
價值損失:衡量價值模型預測與真實獎勵之間的差異,
熵損失:用于鼓勵策略模型的探索行為。
根據策略模型的輸出和獎勵信號,計算損失函數。
損失函數通常包括策略損失、價值損失以及熵損失等。
- 更新模型參數:
根據計算得到的損失函數,使用反向傳播算法更新策略模型、參考模型和獎勵模型的參數。
在更新過程中,應用梯度裁剪等技術,以避免梯度爆炸或梯度消失的問題。
- 重復訓練過程:
我
ZEROSEARCH優點
- ZEROSEARCH框架通過輕量級監督式微調將LLMs轉變為一個檢索模塊,并采用基于課程的rollout策略逐步提升策略模型的推理能力,有效地解決了現有基于強化學習的LLMs搜索能力訓練方法中存在的搜索引擎返回文檔質量不可預測和API費用高昂的問題。
- ZEROSEARCH在多個問答數據集上均取得了優異的性能,超越了真實搜索引擎和其他基線方法。
- ZEROSEARCH還具有良好的可擴展性,能夠適應不同大小和類型的LLMs,并且與多種強化學習算法兼容。
- ZEROSEARCH在部署模擬LLMs時需要一定的GPU資源,但其成本遠低于商業搜索引擎API費用,具有較高的經濟可行性。
- ZEROSEARCH可以應用于其他類型的自然語言處理任務,如文本摘要、機器翻譯等,以充分發揮其在提升LLMs性能方面的潛力。
本文轉載自???CourseAI???,作者:CourseAI
