🤖 machine learning

Predictive Divergence Masks for LLM RL

本論文は、PPOの比率に基づく方向基準を、近接基準に使用される確率ダイバージェンスが次の勾配ステップで増加するか減少するかに関する閉形式の予測に置き換えることで、モデルのスケールを問わず学習の安定性と性能を向上させる、LLMの強化学習のための新しい手法であるPredictive Divergence Masksを提案する。

原著者： Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang

公開日 2026-07-14

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Xiangxin Zhou, Jiarui Yao, Penghui Qi, Bowen Ping, Jiaqi Tang, Haonan Wang, Tianyu Pang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、超スマートなロボットに物語の書き方を教えていると考えてください。あなたは、ロボットに間違いから学んでほしいと考えていますが、一つの厳格なルールがあります。「一度にあまりに激しく書き方を変えてはいけない」というものです。もしロボットが興奮しすぎて、全く異なる声（スタイル）で書き始めたら、あなたはブレーキをかけなければなりません。これが、大規模言語モデル（LLM）における**強化学習（RL）**の仕事です。

長い間、そのブレーキを踏むための標準的な方法は、PPOと呼ばれる手法でした。それは、単一の単語に基づいた単純な信号機のように機能しました。ロボットが単語を選び出すたびに、システムはこうチェックしました。「この単語は、ロボットが通常選ぶものよりも大幅に確率が高いか？そして、報酬はそれをさらに高い確率にすることを求めているか？」もし答えが「はい」で、かつ両方の条件を満たす場合、システムは「止まれ！それはやりすぎだ！」と言い、その単語からの教訓を無視しました。

しかし、ここには問題があります。ロボットは単に一つの単語を選ぶだけでなく、その単語を選ぶために自身の「人格（確率分布）」全体を再構成しているのです。古い手法は、ロボットが選んだ特定の単語だけを見ていました。それはまるで、一つのバイオリンの音を聞いてオーケストラ全体を判断するようなものでした。他の楽器がどのように変化したかを見落としていたのです。

新しいアイデア：予測的ダイバージェンス・マスク（Predictive Divergence Mask）

Tencent、UIUC、NUSの研究チームは、この「単一単語の視点」が誤解を招くものであることに気づきました。彼らは、予測的ダイバージェンス・マスクと呼ばれる新しいルールを提案しました。

単にロボットが選んだ一つの単語を見るのではなく、この新しい手法はより大きな問いを投げかけます。「もしこのステップを踏んだとしたら、ロボットの『人格全体』は元の自分からさらに遠ざかってしまうだろうか？」

次のように考えてみてください：

旧来の方法（PPO）： 子供が崖に向かって走っているのが見えます。あなたは「その特定のステップ」が危険かどうかをチェックします。もし危険なら、「止まれ！」と叫びます。
新しい方法（予測的ダイバージェンス）： あなたは子供が走っているのを見ます。あなたは「経路全体」を見ます。たとえその特定のステップ自体が安全に見えたとしても、子供の勢いによって回転がかかり、別の崖へと放り出される可能性があるかもしれません。新しい手法は、ステップが行われる前に、その「回転全体」を予測します。

その仕組み（魔法の数学）

研究者たちは、実際にシミュレーションを実行することなく、この「回転全体」を予測する巧妙な数学的トリックを発見しました。

ローカル項（Local Term）： これは、古い手法がすでに知っていた部分です。ロボットが選んだ特定の単語の変化です。
グローバル項（Global Term）： これが新しい秘伝のソースです。一つの単語を変更することが、他のすべての単語（選ばれなかった単語）の確率をどのように強制的に調整させるかを考慮に入れます（合計確率が常に100%にならなければならないため）。

古い手法は、このグローバル項を無視していました。新しい手法は、これらを合算します。もし数学的に、ロボットの人格全体が遠ざかりすぎていることが示された場合、マスクが更新をブロックします。もし数学的に、ロボットが実は安全な方向へと戻ろうとしていることが示された場合、更新は許可されます。

「Top-K」問題

現実の世界では、ロボットのエンジン（ロールアウト・エンジン）は、時間とメモリを節約するために、最も可能性の高い上位20個の単語（「Top-K」）のみを表示します。それ以外の何千もの単語は「テイル（裾）」バケットに隠されています。

この限定的な視点でも予測を機能させるために、研究者たちは2つの軽量な推定器を考案しました。

集約テイル推定器（Aggregated-Tail Estimator）： 隠されたすべての単語を一つの巨大な塊として扱います。
一様テイル推定器（Uniform-Tail Estimator）： 隠された単語がすべて等確率であると仮定します。

彼らは、上位20個の単語が通常、ほぼすべての確率（99%以上）を保持しているため、どちらの手法もほぼ同じ結果になることを発見しました。それは、空を見て天気を予想するようなものです。雨が降るかどうかを知るために、雲を一つ一つ数える必要はありません。

実験が示したこと

チームは、さまざまなサイズのロボット（40億から300億パラメータの範囲のモデル）でテストを行い、さらに「FP8」精度（トレーニングや生成エンジンを高速化し、精度をわずかに下げた数学モード）で動作させた場合もテストしました。

結果： 新しいマスクは、トレーニングをより安定させました。
証拠： 彼らは61種類の異なる実験（シード）を実行しました。その結果、古い手法は、実際にはロボットを遠ざけてしまう更新を保持してしまうことがあった（「不適切な保持」率が36.9%）ことが分かりました。新しい手法はこの数値を**34.2%**に減少させました。
改善点： 数字自体は小さく見えるかもしれませんが、巨大なAIモデルのトレーニングの世界では、たとえわずかな悪いステップを防ぐだけでも、ロボットがより良く、より速く学習する助けになります。新しい手法は、すべてのモデルサイズと精度設定において一貫して機能しました。

明確に否定されたこと

論文では、何がうまく機能しないかについても明確に述べています。

彼らは、更新の方向を決定するために、サンプリングされたトークンの重要度比（sampled-token importance ratio）（古い「単一のバイオリン」によるチェック）を使用することに反対しています。彼らは、この単一単語のチェックが、ロボットの人格全体の真の変化としばしば食い違うことを示しました。
また、単に「安全圏（信頼領域）」を狭めること（閾値を0.15から0.05にするなど）は、すべてを悪化させることも示しました。これは、単に制限を厳しくすることが答えではなく、更新の「方向」についてより賢くなることが重要であることを示唆しています。

どの程度確実なのか？

著者らは、実験から得られた測定データに基づき、自らの発見に自信を持っています。彼らは単に真空状態でシミュレーションを行ったのではなく、実際の数学問題（DAPO-Math-17kデータセットを使用）を用いてモデルを訓練し、AIME24やAIME25といったベンチマークで精度を測定しました。

彼らは、このアプローチがより優れたルールの整合性をもたらすと示唆しています。もし、安全圏を「人格全体（ダイバージェンス）」を見て定義するのであれば、更新の方向を確認する際も、単一の単語ではなく、人格全体を見るべきです。結果は、これがより安定したトレーニングにつながることを示唆していますが、これは依然として「局所的な一次近似（local first-order approximation）」であることも認めています。つまり、これは次のステップに基づいた非常に優れた推測であり、未来全体を見通す完璧な水晶玉ではないということです。

要するに、新しい手法は、今踏み出そうとしているステップだけを見るのではなく、地図全体を見るGPSをロボットに与えるようなものです。これにより、ロボットは世界の端から外れることなく、より速く学ぶことができるのです。

技術要約：LLMの強化学習における予測的ダイバージェンス・マスク

1. 問題提起

大規模言語モデル（LLM）における強化学習（RL）は、ロールアウトに使用される推論エンジンと学習スタックの間の不一致、およびポリシーの陳腐化により、本質的にオフポリシー（off-policy）となります。これらのオフポリシー更新を安定させるため、PPO（Proximal Policy Optimization）のような標準的な手法は、信頼領域（trust-region）マスクを採用しています。これらのマスクは、トークンの勾配をマスク（ゼロにする）するかどうかを決定するために、以下の2つの基準に基づいています。

近接性基準（Proximity Criterion）： ポリシーが行動ポリシーからどの程度離れたかを判断する。
方向性基準（Direction Criterion）： 更新がポリシーを行動ポリシーから遠ざける方向（外向き）に向かっているか、あるいは引き戻す方向（修正的）に向かっているかを判断する。

最近の研究、特にDPPOは、単一サンプルの重要度比（importance ratio）を、トップK語彙全体で推定された分布的ダイバージェンス（例：KLダイバージェンス）に置き換えることで、近接性基準を改善しました。しかし、DPPOは依然としてPPO本来の方向性基準を保持しており、これはサンプリングされたトークンの重要度比（ $r_t - 1$ ）のみに依存しています。

著者らは、このハイブリッドなアプローチにおける根本的な不整合を指摘しています。近接性基準は現在、分布的なダイバージェンスによって定義されている一方で、方向性基準は単一サンプルのプロキシ（代理指標）のままです。その結果、サンプリングされた比率の変化の符号と、実際のダイバージェンスの変化の符号が一致しない現象が観察されます。したがって、比率が「修正的更新」を示しているために（＝更新を維持する）、実際にはグローバルなダイバージェンスを増大させてしまう、あるいはその逆のケースが発生します。

2. 手法

本論文では、方向性基準を近接性基準の分布的な性質に適合させる手法として、**予測的ダイバージェンス・マスク（Predictive Divergence Mask）**を提案しています。これは、「サンプリングされた比率が一つから離れるか？」と問うのではなく、「次の勾配ステップによってダイバージェンスが増大するか？」と問う手法です。

理論的導出

自己回帰型LLMで使用されるソフトマックス・ポリシーに対して、著者らは、ポリシー勾配ステップに沿ったカルバック・ライブラー（KL）ダイバージェンスの方向微分 $\dot{D}$ の閉形式（closed-form）の式を導出しています。
$\dot{D} = \underbrace{(\pi_k - \mu_k)}_{\text{局所項 (Local Term)}} + \underbrace{\sum_{i} \pi_i (\mu_i - \pi_i)}_{\text{グローバル項 (Global Term)}}$
ここで：

$\pi_k$ および $\mu_k$ は、学習中のポリシーおよび行動ポリシーにおけるサンプリングされたトークン $k$ の確率です。
**局所項（Local Term）**は、標準的な比率ベースの方向性基準で使用される量（ $r - 1$ ）と正確に一致します。
**グローバル項（Global Term）**は、語彙全体にわたる正規化による結合（normalization coupling）を捉えています。この項は単一のサンプリングされた比率からは見えませんが、真のダイバージェンスの変化を決定する上で極めて重要です。

予測的マスクは、トークンがすでに信頼領域の外にある場合（ $D_t > \delta$ ）、 $\hat{A}_t \cdot \dot{D}_t$ の符号が非正（すなわち、修正的または中立的な動きを示す）である場合にのみ、その更新を保持します。

実用的な推定

プロダクション環境のロールアウトエンジンは、通常、全語彙の分布ではなく、トップ $K$ のトークン確率のみを公開します。この制限されたビューの下で $\dot{D}$ を計算するために、著者らは、分布の未観測な裾（テール）に対する2つの軽量な推定器を提案しています。

集約テール推定器（Aggregated-Tail Estimator）： 残りの確率質量全体を単一の集約されたトークンとしてまとめる。
一様テール推定器（Uniform-Tail Estimator）： 残りの質量を未観測の語彙トークンに均等に分配する。

どちらの推定器も、保持されたトップ $K$ の確率とテールの質量に対する単純なリダクション（集約操作）のみを必要とするため、計算オーバーヘッドは無視できるほど小さいです。

3. 主な貢献

不整合の特定： 分布的な近接性基準（ダイバージェンス）と、単一サンプルの方向性基準（比率）を混在させると、方向性マスクが真のダイバージェンスの変化を追跡できなくなるという不整合が生じることを明らかにしました。
予測的ダイバージェンス・マスク： ダイバージェンス自体の一次変化を用いて更新の方向を決定する、新しいマスキング規則を提案しました。
閉形式の導出： ソフトマックス・ポリシーにおける方向微分の導出を行い、局所的（サンプリングされた）成分とグローバル（正規化）成分を明示的に分離しました。
切り詰められた語彙の推定器： 集約テールおよび一様テールという、低オーバーヘッドの2つの実用的な推定器を提案し、標準的なプロダクション・ロールアウトの制約下（トップ $K$ ビュー）での動作を可能にしました。

4. 実験結果

著者らは、Qwen3モデル（4B, 8B, 30B）を用い、様々な精度設定（FP8）においてDAPO-Math-17kデータセットで評価を行いました。

学習の安定性： すべてのモデルスケールおよび精度設定（FP8を含むストレス・テストを含む）において、予測的ダイバージェンス・マスクは、ベースラインであるDPPO-TopK-KLと比較して学習の安定性を向上させました。対照的に、clip-higherテクニックを用いた標準的なGRPOは、より大きなモデル設定において不安定性や崩壊を示しました。
方向性基準の整合性： トークンレベルの分析により、ダイバージェンスベースの方向性基準は、比率ベースの基準よりも実現されたダイバージェンスの変化とよく一致していることが明らかになりました。
- 両方の基準が一致しなかったトークンのうち、ダイバージェンスベースの基準は「不安全な保持（unsafe-keep）」率が低かった（34.2% vs 36.9%）ため、実際にダイバージェンスを増大させてしまう更新を保持してしまう可能性が低いことが示されました。
- 逆に、ダイバージェンスを減少させることに成功した更新を保持する割合は高くなりました（65.1% vs 62.9%）。
堅牢性： 本手法は、信頼領域の閾値（ $\delta$ ）に対するロバスト性を示し、より厳しい制約下（ $\delta = 0.05$ ）でもベースラインより優れた性能を発揮しました。
推定器の等価性： 集約テール推定器と一様テール推定器はほぼ同一の結果を生み出し、本手法がテール・モデリングの仮定に敏感ではないことを示唆しました。

5. 意義と主張

本論文は、予測的ダイバージェンス・マスクが、現在のLLM向けダイバージェンスベースの強化学習手法における構造的な不整合を解決すると主張しています。近接性と方向性の両方の基準を同じ分布的ダイバージェンスから導出することで、本手法はより忠実な信頼領域制御を提供します。

著者らは、不整合が発生するトークンの集合が小さいことから「不安全な保持」率の改善は緩やかではあるものの、その一貫性は理論的に妥当であり、あらゆるスケールにおいて経験的に有益であると強調しています。結論として、ダイバージェンスベースの信頼領域の利点を最大限に活用するためには、方向性基準も近接性基準と同様に分布的な扱いを受けるべきであるとしています。本手法は、既存のダイバージェンス閾値以外の新しいハイパーパラメータを必要としない、既存のDPPO実装に対するドロップイン（そのまま置き換え可能な）リプレースメントとして提示されています。