💬 NLP

d-TreeRPO: Towards More Reliable Policy Optimization for Diffusion Language Models

本論文は、木構造ロールアウト、検証可能なステップごとの報酬、時間スケジュール付き自己蒸留を通じて、報酬の希薄性と確率推定のギャップに対処し、複数の推論ベンチマークで顕著な性能向上を実現する拡散言語モデル向けの信頼性の高い強化学習フレームワークである d-TreeRPO を紹介する。

原著者： Leyi Pan, Shuchang Tao, Yunpeng Zhai, Zheyu Fu, Liancheng Fang, Minghua He, Lingzhe Zhang, Zhaoyang Liu, Bolin Ding, Aiwei Liu, Lijie Wen

公開日 2026-05-14

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

CC BY 4.0

原著者： Leyi Pan, Shuchang Tao, Yunpeng Zhai, Zheyu Fu, Liancheng Fang, Minghua He, Lingzhe Zhang, Zhaoyang Liu, Bolin Ding, Aiwei Liu, Lijie Wen

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

複雑なパズル、例えば数独や数学の問題を解くロボットを教えることを想像してください。このロボットは、**拡散型大規模言語モデル（dLLM）**と呼ばれる特別な種類の脳を使用します。答えを一文ずつ（文章をタイプするように）書き出す従来のロボットとは異なり、このロボットは真っ白でかき混ぜられたページから始め、徐々に「ノイズ除去」を行い、完全な解答が現れるまで、カオス的で非線形な順序で正しい単語を明らかにしていきます。

この論文は、このロボットをより賢く、より信頼性の高いものにするための新しい訓練手法、d-TreeRPOを紹介しています。その仕組みを簡単な概念に分解して説明します。

1. 問題：「目隠し」されたロボット

著者らは、これらのロボットを訓練する既存の手法には、2 つの重大な欠陥があると述べています。

「全か無か」の報酬: 現在、ロボットがパズルを解ければ高得点を与えられ、失敗すればゼロ点を与えられます。どの特定のステップが良かったのか、悪かったのかはわかりません。まるで、ゲームオーバー画面が最後の最後まで表示されず、どの動きが失敗の原因だったのかというヒントも一切ないビデオゲームをプレイしているようなものです。これにより、学習は遅く、不正確になります。
「混乱した」確率: ロボットは任意の順序で単語を明らかにできるため、特定の単語に対する確信度を正確に計算することが困難です。既存の手法はこの確信度を推測しますが、その推測はしばしば誤っており、ロボットが不適切な決定を下す原因となります。

2. 解決策：「木探索者」(d-TreeRPO)

これを修正するため、著者らはd-TreeRPOと呼ばれるフレームワークを構築しました。これは、ロボットに地図と虫眼鏡を与えるようなものです。

A. 木構造（地図）

ロボットが単一の答えへの経路を推測するのではなく、d-TreeRPO はロボットに複数の経路を同時に探索させます。まるで木の枝のようですね。

幹: 出発となる質問。
枝: ロボットがパズルを埋めるためのさまざまな方法を試すこと。
葉: 最終的な答え。

もしある枝が行き止まり（誤った答え）につながれば、ロボットはその枝のどこで間違ったのかを正確に知ることができます。その後、木を「登り返って」、その特定のステップが悪かったと判断できます。これにより、最終結果だけでなく、すべての単一のステップに対して詳細なフィードバックが得られます。

B. 自己蒸留損失（「確信度コーチ」）

これが 2 番目の主要な革新です。著者らは、厄介なトレードオフに気づきました。

ロボットが好奇心が強すぎる（確信度が低い）場合、多くのアイデアを探索しますが、いい加減な推測をしてしまいます。
ロボットが頑固すぎる（確信度が高い）場合、正確に推測しますが、新しいことを試すのをやめてしまいます。

d-TreeRPO は、これを管理するために時間スケジュール付き自己蒸留損失を使用します。トレーニングキャンプの日に応じてロボットと異なるように話すコーチを想像してください。

初期の日々: コーチは、「好奇心を持て！すべてを試せ！完璧である必要はない」と言います。これはロボットに探索を促します。
後期の日々: コーチは、「さて、選択肢を見た今、決断力を持て！最善の動きに固執し、直感を信じるのだ」と言います。これはロボットをより確信に満ち、精密なものにします。

ロボットを「好奇心旺盛な探検家」から「確信に満ちた専門家」へとゆっくりと移行させることで、この手法はロボットの内なる数学（確率推定）が時間とともにはるかに正確になることを保証します。

3. 結果：より賢い解決

著者らは、この新しい手法を 4 種類の異なるパズルでテストしました。

数独（論理グリッド）
カウントダウン（数学を用いた数字作り）
GSM8K（小学校レベルの数学の文章題）
Math500（より難しい数学の問題）

結果:
d-TreeRPO で訓練されたロボットは、以前のバージョンと比較して劇的な改善を見せました。

数独では、**86%**改善し、成功率がほぼ倍増しました。
カウントダウンでは、**51%**改善しました。
数学ベンチマークにおいても確実な gains（改善）が見られました。

結論

この論文は、ロボットが学習するプロセスを木（各ステップに対するより良いフィードバックを得るため）に整理し、時間ベースのコーチングシステム（好奇心と確信度のバランスを取るため）を使用することで、拡散型言語モデルに推論を教えるはるかに信頼性の高い方法を生み出したと主張しています。その結果、以前よりも論理パズルや数学パズルを大幅にうまく解くロボットが実現しました。

技術的サマリー：d-TreeRPO

問題定義

拡散型大規模言語モデル（dLLM）は、自己回帰（AR）モデルに対する並列デコーディングの代替手段を提供し、より高速な推論を可能にします。しかし、dLLM の推論能力を強化するために強化学習（RL）を適用する際、既存の手法では対処されていない 2 つの重要な信頼性のボトルネックに直面します。

報酬の疎性と検証性の欠如: 既存の dLLM ポリシー最適化手法は、しばしば粗い結果報酬（例：最終回答の二値的な正誤）や検証されていないプロセス信号に依存しています。この疎性は正確なアドバンテージ計算を妨げ、検証されていない信号は報酬ハッキングを招く可能性があります。
信頼性の低い確率推定: 連鎖則を用いてシーケンスを分解し、単一のフォワードパスで正確な確率計算を行う AR モデルとは異なり、dLLM は「任意の順序」のノイズ除去プロセスを通じてテキストを生成します。真の行動対数確率は、すべての可能なデコーディング順序にわたる期待値を必要とし、これは計算的に困難です。既存の近似手法（例：ELBO ベースの下限または単一ステップのフォワードパス）は、分析されていない近似誤差に悩まされ、ポリシー勾配更新の信頼性を損なっています。

手法：d-TreeRPO

著者は、dLLM 向けに設計された信頼性の高い RL フレームワークであるd-TreeRPOを提案します。この手法は、以下の 3 つの中核コンポーネントを通じて前述のボトルネックに対処します。

1. 微細な報酬のための木構造ロールアウト

報酬の疎性を克服するため、d-TreeRPO はロールアウトプロセスを、独立したサンプルの平坦なシーケンスではなく、木構造として構成します。

構造: $N$ 個のノイズ除去ステップを $H$ 個の木ステップにグループ化します。各非葉ノードにおいて $B$ 個の独立サンプリングを実行し、終端状態を表す $B^H$ 個の葉ノードを持つ木を生成します。
ボトムアップのアドバンテージ計算: 葉ノードには検証可能な結果報酬（例：パズル解の正誤）が割り当てられます。親ノードの状態価値は、子ノードの価値の平均としてボトムアップに計算されます。
微細なアドバンテージ: 親ノード $p$ から子ノード $c$ への遷移に対するアドバンテージは、 $A_p^c = V_c - V_p$ と定義されます。これにより、均一な結果報酬よりも微細で、検証可能なステップごとの報酬信号が得られます。

2. 確率推定の理論的解析

本論文は、dLLM における不偏確率推定の計算困難性に対処します。

理論的証明: 著者は（定理 1 において）、単一フォワードパスの確率推定値と不偏期待確率との間のギャップが、モデルの予測信頼度が高まるにつれて減少することを証明しました。具体的には、モデルが決定的になるにつれて、推定誤差の上限は厳密になります。
含意: 高い信頼度は推定精度を向上させますが、探索と利用のトレードオフを生み出します。トレーニングの初期段階で過度に決定的なポリシーを採用すると、探索が阻害されるためです。

3. 時間スケジュール付き自己蒸留損失

トレーニング初期の探索と、トレーニング後期の正確な確率推定の間のトレードオフをバランスさせるため、d-TreeRPO は時間スケジュール付きの自己蒸留目的関数を導入します。

メカニズム: この損失関数は、正のアドバンテージを持つ子ノードから構築されたターゲット分布とポリシーの整合を促します。
スケジュール:
- 温度 ( $\tau(t)$ ): 多様なサンプリングを可能にするため高く始まり、時間の経過とともに低下してターゲット分布を鋭くします。
- 重み ( $\lambda(t)$ ): 探索フェーズ中の損失を最小化するため 0 に近く始まり、後期段階で決定性を強制し推定誤差を低減するため段階的に増加します。
目的: この損失は、標準的な GRPO スタイルのポリシー勾配目的関数に追加され、安全に行える場合のみモデルの決定性を高めることで、RL 更新で使用される単一パス確率推定の精度を向上させます。

主要な貢献

d-TreeRPO フレームワーク: 木構造ロールアウトを活用して、微細で検証可能なステップごとの報酬とアドバンテージを提供する、dLLM 向けの新しいポリシー最適化アルゴリズム。
理論的洞察と自己蒸留: 予測信頼度と推定誤差の低減を結びつける理論的証明、および dLLM 固有の探索と利用のトレードオフを動的に管理する時間スケジュール付き自己蒸留損失の設計。
実証的検証: Sudoku、Countdown、GSM8K、Math500 といった複数の推論ベンチマークにおいて、d-TreeRPO が既存のベースライン（Diffu-GRPO、VRPO、wd1 など）を上回ることを示す広範な実験。

実験結果

著者は、4 つの推論ベンチマーク（Sudoku、Countdown、GSM8K、Math500）において、LLaDA-8B-InstructおよびLLaDA-MoE-7BA1B-Instructに対して d-TreeRPO を評価しました。

性能向上: ベースモデルと比較して、d-TreeRPO は顕著な改善を達成しました。
- Sudoku: +86.2% (LLaDA-8B)
- Countdown: +51.6% (LLaDA-8B)
- GSM8K: +4.5% (LLaDA-8B)
- Math500: +5.3% (LLaDA-8B)
アブレーション研究:
- 疎な報酬: 最終報酬をすべてのトークンにブロードキャストする疎な結果報酬を使用する変種は著しく性能が低下し、木構造による微細なクレジット割り当ての必要性を確認しました。
- 自己蒸留: 自己蒸留損失を除去すると、性能の低下と確率推定誤差の増加を招きました。一方、高いエントロピーを維持することを意図した多様性促進損失は、最悪の性能と不安定なトレーニングをもたらしました。
- スケジュール: 逆スケジュール（初期に強い蒸留、後期に弱い蒸留）は、初期の獲得後に性能低下を引き起こし、前方スケジュール設計の有効性を検証しました。
効率性: 木ロールアウトにより、d-TreeRPO は一部のベースライン（例：Diffu-GRPO）よりも計算コストが高くなりますが、報酬の観点ではより速く収束し、優れた最終性能を達成するため、実用的な選択肢となります。

意義と主張

本論文は、d-TreeRPO が報酬の疎性と確率推定のバイアスという二重の問題を解決することで、dLLM における RL のより信頼性の高い基盤を確立すると主張しています。木構造ロールアウトを活用することで、アドバンテージ信号が検証可能かつ微細であることを保証します。さらに、確率推定における信頼度の必要性を理論的に裏付け、時間スケジュール付き自己蒸留メカニズムを実装することで、dLLM 固有の探索と利用のトレードオフを効果的に navigates（航行）します。結果は、拡散モデルに対する信頼性の高いポリシー最適化には、単純な結果報酬や単一パス近似を超え、構造化された検索と動的な信頼度制御を採用する必要があることを示唆しています。