Hugging Face TRL RLOO Trainer 發布

效能與效率提升

  • 記憶體效率: RLOO 使用大約 50-70% 更少的 vRAM,相較於 PPO,具體取決於模型大小。
  • 訓練速度: RLOO 在 1B 參數模型上比 PPO 快 2 倍,於 6.9B 模型上可快至 3 倍。
  • 模型品質: RLOO 在回應勝率(由 GPT-4 評估)上與 PPO 競爭,且持續優於離線方法如 Direct Preference Optimization (DPO)。

技術架構與動機

傳統的 PPO 需要大量資源,因為它必須在記憶體中載入四個模型副本:policy model、reference policy model、reward model 與 value model。RLOO 透過移除 value model,將架構簡化為僅需三個副本:policy model、reference policy model 與 reward model。

主要演算法差異

RLOO 在處理動作與獎勵的方式上與 PPO 不同:

  1. 完成層級的動作: PPO 將每個完成 token 視為單獨的動作,而 RLOO 將整個模型的完成視為單一動作。這解決了獎勵稀疏的問題,通常只有 EOS(句子結束) token 會得到真實獎勵;RLOO 會將 EOS 的獎勵歸屬於整個完成。
  2. REINFORCE 損失: RLOO 使用 REINFORCE 損失,將獎勵與基線的差異乘以動作的對數機率。此做法消除了 PPO 所需的 value model 與 Generalized Advantage Estimation (GAE)。
  3. Leave-One-Out 基線: RLOO 透過使用同一批次中其他樣本的獎勵作為特定樣本的基線來計算基線。對於給定的提示,一個完成的基線是該提示下所有其他完成的平均獎勵。

在 TRL 中的實作

RLOO Trainer 建立於實驗性的 PPOv2Trainer 之上。雖然 RLOO 概念上基於 REINFORCE 演算法,TRL 的實作仍使用 PPO 損失,因為 REINFORCE 損失是 PPO 損失的特例。

基準測試結果

  • Pythia 6.9B: 使用 GPT-4 作為評審,取得 78.7% 的偏好率 (k=2),超過原始論文報告的 77.9% (k=4) 與 74.2% (k=2)。
  • Pythia 1B: 取得 40.1% 的勝率,較 SFT(Supervised Fine-Tuning)檢查點的 21.3% 勝率有明顯提升。

使用 bf16 時的數值穩定性挑戰

RLOO 在使用 bf16 精度時會出現顯著的數值穩定性問題。生成過程中取得的對數機率可能與訓練前向傳播時取得的略有差異。

在 PPO 中,剪裁係數(通常為 0.2)會在比例超過 1.2 或低於 0.8 時將單個 token 的梯度置零。由於 RLOO 將整個完成視為單一動作,這些數值差異會被放大,導致整段序列的梯度更頻繁被置零。實驗觀察顯示,PPO 大約會置零 3% 的批次資料,而 RLOO 在 bf16 下會置零 20%~40% 的批次資料。

Sources