自我糾正(Self-correction)是大語言模型 (LLM) 非常重要的能力,但人們發現這種能力在現代 LLM 中基本上很少存在。現有的訓練自我糾正的方法要麼需要多個模型,要麼依賴於更強大的模型或其他形式的監督。
我們如何才能讓 LLM 具備自我糾正能力?之前的研究要麼依賴於提示工程,要麼依賴於專門用於自我糾正的微調模型。但前者通常無法有效地進行有意義的內在自我糾正,而後者基於微調的方法需要在推理時執行多個模型,例如需要 oracle「教師」來監督指導自我糾正過程。
在最近提交的一篇論文中,來自 Google DeepMind 的研究者開發了一種無需上述任何要求即可有效進行自我糾正的方法,即透過強化學習進行自我糾正(SCoRe,Self-Correction via Reinforcement Learning),只需訓練一個模型,該模型既可以對推理問題做出響應,也可以糾正錯誤,儘管沒有收到任何 oracle 反饋。更重要的是,SCoRe 完全透過在自生成資料上進行訓練來教模型具備這種能力,而無需任何 oracle。
- 論文標題:Training Language Models to Self-Correct via Reinforcement Learning
- 論文地址:https://arxiv.org/pdf/2409.12917
本文主要貢獻在於提出了一種多輪強化學習方法 ——SCoRe,用於教 LLM 如何糾正自己的錯誤。相對於基礎 Gemini 模型,SCoRe 在 MATH 推理問題的自我糾正方面獲得了 15.6% 的增益,在 HumanEval 編碼問題上獲得了 9.1% 的增益。
SCoRe 原理介紹
為了教 LLM 進行自我糾正,SCoRe 將標準單輪強化學習(公式 2)擴充套件到 Zhou 等人提出的分層框架下的多輪設定。
不過這樣做面臨諸多挑戰。首先,最佳化公式 1 解決了分佈偏移問題,但尚不清楚它是否也能滿足要求 [D2]。
這裡的 [D2] 如下圖所示,圖中展示了 SFT 方法失敗的兩個原因。而有效的解決方案必須滿足兩個要求:[D1] 模型應該直接在自生成軌跡上進行訓練,以緩解 SFT 的分佈不匹配(圖 4),[D2] 所採用的自生成軌跡應防止在學習過程中因進行微小編輯而崩潰。
作者開發了一種線上 RL 方法,透過仔細的初始化和獎勵塑造來解決這些挑戰。
其次用於微調的基礎模型初始化在編輯距離上呈現出高度傾斜的分佈(圖 3a),這使得它們容易受到模式崩潰的影響,這是深度強化學習中一個常見的問題。即使基礎模型可以在自我校正過程中產生編輯距離比傾斜度較小的分佈,但仍然需要強化學習訓練過程從訓練資料中學習一種可以推廣到測試提示的自我校正策略。
SCoRe 旨在解決上述關鍵挑戰,其分為兩個階段,這兩個階段都是透過適當初始化模型和控制後續 RL 使模型偏向學習自我糾正。
具體而言,這兩個階段包括:
階段 I:訓練模型初始化以防止崩潰
SCoRe 第一階段的目標是透過提高基礎模型對第二次嘗試響應的覆蓋率來獲得良好的模型初始化,以便後續自我糾正訓練不會出現 STaR/SFT 中觀察到的崩潰現象。
為了達到此目的,該研究不採用 SFT 來初始化 RL 訓練,而是開發了第一階段來產生不易崩潰的單獨初始化。
作者微調基礎模型,以便在第二次嘗試時產生高獎勵修正,同時透過使用 KL 散度將第一次嘗試的響應分佈限制為儘可能接近基礎模型的響應分佈,從而強制模型不改變其第一次嘗試的響應。雖然這看起來不是最優的 —— 但第一次嘗試的響應錯誤較少,可以糾正為更好的第二次嘗試響應。最佳化的目標可以表示為:
其中 𝛽_2 是一個超引數,旨在僅在第一次嘗試時強制執行嚴格的 KL 懲罰,以避免第一輪響應發生偏移(用藍色項表示)。請注意,作者仍然使用公式 2 中的預設 KL 散度懲罰,但該懲罰的權重要小得多,並且為了簡潔起見,公式 3 中省略了它。事實上,與簡單的多輪 RL 不同,階段 I 在分離兩個響應方面更有效(圖 5b)。
階段 II:帶有獎勵的多輪強化學習
藉助第一階段的模型初始化,該模型在耦合兩個響應時表現出更小的偏差,SCoRe 的第二階段現在可以訓練兩次嘗試的響應,並根據公式 1 最佳化獎勵。當然,作者還希望確保在此過程中不會降低第一次嘗試的響應。因此,對於兩輪自我糾正問題,作者針對以下目標訓練策略 𝜋_𝜃(⋅∣⋅):
圖 6 為階段 I、階段 II 流程說明。可以看出 SCoRe 以交錯方式應用階段 I 和 II 進行多次迭代。
實驗評估
該研究進行了一系列實驗,來驗證 SCoRe 在教 LLM 具備自我糾正能力方面的有效性,並透過消融實驗探索了 SCoRe 的每個元件的影響。
該研究主要關注數學和編碼任務,使用以下基準來評估方法的有效性:
- MATH;
- MBPP 和 HumanEval。
幾種方法在 MATH 基準上的實驗評估結果如下表 3 所示:
在程式碼生成方面,實驗結果如下表 4 所示:
消融研究
為了探究以下幾個問題,該研究進行了消融實驗:
- 多輪訓練的重要性
- 多階段訓練的重要性
- 獎勵函式設計(reward shaping)的影響
- on-policy 強化學習的重要性
消融實驗結果如下表 5 所示:
感興趣的讀者可以閱讀論文原文,瞭解更多研究內容。