用於大型語言模型強化學習的預測分歧遮罩

摘要

大型語言模型（LLMs）的強化學習通常依賴信任區域掩碼來穩定離策略更新。主流的PPO風格方法使用取樣詞元的重要性比率作為兩項準則的依據：一為鄰近準則，用以檢查策略是否已偏離行為策略過遠；二為方向準則，用以評估更新是否使策略進一步背離。近期提出的DPPO方法透過以行為策略與訓練策略之間的概率散度取代PPO基於比率的測試，改進了鄰近準則。然而，其方向準則仍沿用PPO的設計——只有當取樣詞元的重要性比率偏離1時，該詞元才會被掩蓋。我們觀察到，這種基於比率的方向準實質上是一種單樣本代理指標，其正負號可能與定義鄰近準則之散度的變化方向不一致。為此，我們提出預測性散度掩碼，該方法能判定下一個策略梯度步驟將增大還是減小與信任區域相同的散度。針對LLM強化學習中使用的離散Softmax策略，我們推導出此預測的封閉形式解。由於生產環境中的推論引擎僅暴露詞彙表的截斷（Top-K）視圖，我們開發了兩種輕量化的Top-K估計器來進行此預測。詳細分析表明，相較於取樣比率，基於散度的方向準則更能與散度的實際變化保持一致；由此產生的掩碼能在不同模型規模與精度設定下改善強化學習訓練。

English

Reinforcement learning for large language models (LLMs) typically relies on trust-region masks to stabilize off-policy updates. The dominant PPO-style approach uses the sampled-token importance ratio for two criteria: a proximity criterion, which asks whether the policy has moved too far from the behavior policy, and a direction criterion, which asks whether the update pushes it farther away. Recent work DPPO improves the proximity criterion by replacing PPO's ratio-based test with a probability divergence between the behavior and training policies. However, its direction criterion is still inherited from PPO. A token can be masked only when the sampled-token importance ratio moves away from one. We observe that this ratio-based direction criterion is a single-sample proxy that can disagree in sign with the change of the divergence that defines the proximity criterion. We therefore propose the predictive divergence mask, which asks whether the next policy-gradient step will increase or decrease the same divergence used by the trust region. For the discrete softmax policies used in LLM RL, we derive this prediction in closed form. Because production rollout engines expose only a truncated (top-K) view of the vocabulary, we develop two lightweight top-K estimators for this prediction. Detailed analysis shows the divergence-based direction is better aligned with the realized change of the divergence than the sampled ratio, and the resulting masks improve RL training across model scales and precision settings.