LLM強化学習のための予測ダイバージェンスマスク

要旨

大規模言語モデル（LLM）の強化学習では、通常、信頼領域マスクを用いてオフポリシー更新を安定化させる。主流のPPO手法では、サンプリングされたトークンの重要度比を2つの基準に用いる。すなわち、近接性基準（ポリシーが行動ポリシーから大きく乖離していないかを問う）と方向性基準（更新がポリシーをさらに遠ざけるかどうかを問う）である。最近の研究であるDPPOは、PPOの比率ベースのテストを行動ポリシーと訓練ポリシー間の確率ダイバージェンスに置き換えることで近接性基準を改善した。しかし、その方向性基準は依然としてPPOから引き継がれている。トークンがマスクされるのは、サンプリングされたトークンの重要度比が1から乖離する場合のみである。我々は、この比率ベースの方向性基準が単一サンプルの代理指標であり、近接性基準を定義するダイバージェンスの変化と符号が一致しない場合があることを観察した。そこで我々は、次回のポリシー勾配ステップが信頼領域で使用されるのと同じダイバージェンスを増加させるか減少させるかを問う、予測的ダイバージェンスマスクを提案する。LLM強化学習で使用される離散ソフトマックスポリシーに対し、この予測を閉形式で導出する。プロダクションのロールアウトエンジンは語彙の切り詰められた（トップK）ビューしか公開しないため、この予測のために2つの軽量なトップ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.