🤖 machine learning

Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models

本論文は、報酬最大化アプローチの「軌道固定」という失敗様式を克服し、方策を報酬傾斜目標分布と整合させることで、数学的推論およびコード生成ベンチマークにおいて一貫した性能向上を実現する拡散言語モデル向けの軌道バランス事後学習手法であるTraFLを導入する。

原著者： Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

公開日 2026-05-15

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

CC BY 4.0

原著者： Saba Ahmadi, Prasanna Parthasarathi, Yufei Cui

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

以下は、論文「Beyond Mode-Seeking RL: Trajectory-Balance Post-Training for Diffusion Language Models」について、平易な言葉と創造的な比喩を用いて解説したものです。

全体像：AI に教える新しい方法

あなたが、ノイズや静電気の混じったキャンバスから始めて、徐々にそれを清浄化し、次第に鮮明な画像が現れるまで描き上げる、非常に才能のある芸術家（AI モデル）を持っていると想像してください。これが拡散言語モデルの仕組みです。これらは古い AI モデルがそうであったように、タイプライターのように単語を一つずつ書くのではなく、文字の羅列から始めて、徐々に「ノイズ除去」を行い、一貫した文章へと変換していきます。

この論文は、これらの芸術家に数学やコーディングといった難しい問題を解決する能力を向上させる現在の教え方が欠陥があると主張しています。著者らは、**「軌道固定（Trajectory Locking）」**と呼ばれる特定の欠陥を修正する新しい手法、**TraFL（Trajectory Flow baLancing）**を提案しています。

問題点：「一本道」の罠（軌道固定）

この問題を理解するために、迷路を想像してください。

目標： この迷路には、数学問題の異なる有効な解答に対応する、多くの異なる正しい出口があります。
従来の方法（報酬最大化 RL）： 犬を出口を見つけるように訓練すると想像してください。犬がいずれかの出口を見つけると、ご褒美を与えます。
- 欠陥： 犬はご褒美を得るためにどの経路を通ったかは気にせず、ご褒美を得たことだけを気にします。
- 結果： 犬がたまたま早期に、ある特定の経路を辿ってご褒美にたどり着くと、ご褒美を得ます。その経路を記憶します。次に、同じ経路を試します。再びご褒美を得ます。その一本の経路に対して、より自信を持ちます。
- 軌道固定： 最終的に、犬は迷路全体を探検することをやめます。たとえ他にも 50 個の有効な出口があったとしても、その一本の狭い経路だけを走り続けます。それは単一の解答に「固定」され、他のものを見つける方法を忘れ去ってしまいます。

この論文において、著者らはこれを軌道固定と呼びます。AI は最終的な解答に対してのみ「報酬（スコア）」を得るため、そこに至るまでの多くの異なる方法（経路）が存在したという事実を無視します。その結果、すべての創造性が一本の狭い経路に収束し、再度試すよう求められた際に、別の正しい解答を見つける能力が低下してしまいます。

解決策：「ガイドブック」アプローチ（TraFL）

著者らは、TraFLを提案し、訓練のルールを変更します。いずれかの出口に対してご褒美を与えるのではなく、AI にガイドブック（凍結された参照モデル）と地図を与えます。

ガイドブック（参照モデル）： これは、特定の報酬に基づいてまだ訓練されていない AI のバージョンです。これは「健全で多様な」思考のあり方を表しています。問題には多くの解決方法があることを知っています。
地図（報酬偏向ターゲット）： 私たちは AI に次のように伝えます。「報酬（正解）を得る出口を見つけたいが、しかし、あなたの動きのパターンはガイドブックと似ているように保たなければならない」と。

比喩：
あなたが学生に数学の問題を解くことを教えると想像してください。

従来の方法： 「どんな方法でも、正しい答えを出せ！」と言います。学生は一つうまくいったトリックを見つけ、それを暗記し、他の方法を学ぶことを拒みます。
TraFL の方法： 「正しい答えを見つけなさい。しかし、多くの異なる戦略を知っているトップの学生のように、思考プロセスが多様で柔軟であるように保ちなさい」と言います。

TraFL は AI に二つのバランスを取ることを強制します。

報酬を得ること： 正しい答えを見つけること。
多様性を保つこと： 単一の反復的な経路に収束しないこと。それは「確率質量（経路を選ぶ可能性）」を、ガイドブックによって支えられながら、多くの異なる有効な解答全体に広げて保持します。

どのように機能させたか

この論文は、拡散モデルが古いモデルのように思考プロセスを段階的に示さないため、扱いが難しいと指摘しています。これを修正するために、著者らは二つのツールを考案しました。

代理スコア： 各ステップの正確な数学を見ることができないため、完全な解答がどれほど優れているかを推定する「代理」スコアを作成し、各微小ステップの完全な可視性を必要とせずにモデルを訓練できるようにしました。
学習された正規化器： 特定の質問（プロンプト）に基づいてタスクの難易度を調整する特別な「計算機」を AI に学習させ、訓練が公平かつバランスよく行われるようにしました。

結果：実際に役立つか

著者らは、この新しい手法を数学（文章題の解決）とコード（コンピュータプログラムの作成）でテストしました。

「一本道」の罠は破られた： 単一の解答を見つける能力は向上する一方で、異なる解答を見つける能力が低下する他の手法とは異なり、TraFL はすべてのテストで改善しました。
サンプル数増加＝結果向上： 1 つではなく 16 個の異なる解答を生成するように AI に求めたところ、TraFL は劇的に向上しました。これは、単なる幸運な推測ではなく、実際により多くの異なる正解を見つけたことを証明しています。
新しいものにも機能する： AI は訓練質問を単に暗記したわけではありません。Minerva Math や LiveCodeBench といった、これまで見たことのない新しい数学問題やコーディング課題でテストされた際、TraFL は最も強力なパフォーマンスを発揮しました。
多様性の確認： 彼らは「ジャッジ」（別の AI）を使って解答を確認しました。ジャッジは、TraFL が単に異なる言葉で同じ答えを返しているのではなく、同じ問題を解決するために異なる推論戦略と異なるアルゴリズムを実際に使用していることを確認しました。

まとめ

この論文は、AI モデルが問題解決の一つの方法に「閉じ込められ」、他の有効な解答を無視してしまう欠陥を特定しています。彼らはこれをTraFLで修正しました。これは、AI に正解を求める一方で、参照モデルによって導かれる異なる経路の多様な「探索」を維持することを教える手法です。その結果、AI はより賢くなるだけでなく、複数の解答を生成するよう求められた際にも、より創造的で信頼性の高いものになります。

技術的サマリー：拡散言語モデルのための軌道バランス事後学習

問題：拡散 RL における軌道の固定化

拡散言語モデル（dLLMs）は、推論やコード生成において自己回帰モデルに代わる有望な代替手段を提供する。しかし、dLLMs に対する既存の事後学習手法の多くは、自己回帰モデル向けに設計された報酬最大化の強化学習（RL）目的関数（例：PPO、GRPO）をそのまま適用している。著者らは、この適用における中心的な失敗モードとして「軌道の固定化（trajectory locking）」を特定した。

拡散モデルにおいて、報酬は最終的なノイズ除去完了（ $y$ ）のみに割り当てられるが、モデル分布はノイズ除去の全軌道（ $\tau$ ）にわたって誘導される。終端報酬は「経路非依存」であるため、すなわち同じ正しい $y$ に到達する任意の軌道が同じ報酬を受け取るため、サンプリングされた方策勾配更新は、学習の初期段階で偶然サンプリングされた特定の軌道を強化する傾向がある。これにより、確率質量が狭いノイズ除去経路の集合に収束する自己強化フィードバックループが生じる。その結果、モデルが正しい答えを生成することを学習したとしても、繰り返しサンプリング下で多様な有効な解（モード）を生成する能力を失う。著者らは定理 1を通じてこれを定式化し、終端モードのカバレッジは軌道の多様性によって上限が決まることを証明した。したがって、軌道の固定化はモデルの解カバレッジを直接的に制限する。

手法：TraFL（軌道フローバランス）

軌道の固定化に対処するため、著者らは生成フローネットワーク（GFlowNets）の軌道バランスの原理に基づいた事後学習目的関数TraFLを提案する。報酬を直接最大化するのではなく、TraFL は、凍結された参照モデルにアンカーされた報酬傾斜の目標分布と方策を一致させるように訓練する。

目的関数

目標分布 $p^*(\tau | x)$ は以下のように定義される：
$p^*(\tau | x) = \frac{1}{Z(x)} p_{\text{ref}}(\tau | x) \exp(\beta r(x, y))$
ここで：

$p_{\text{ref}}$ は凍結された参照モデルである。
$r(x, y)$ は終端報酬である。
$\beta$ は報酬の鋭さを制御する。
$Z(x)$ はプロンプト依存の正規化項（分配関数）である。

訓練目的関数は、現在のモデルの対数確率と目標との間の二乗残差を最小化する：
$\delta(\tau, x) = \log p_\theta(\tau | x) - \log p_{\text{ref}}(\tau | x) - \beta r(x, y) + \log Z(x)$
$L_{\text{TraFL}} = \mathbb{E}_{\tau \sim p_\theta} [\delta(\tau, x)^2]$

拡散モデルへの実用的実装

dLLMs は正確な軌道対数確率を露出しないため、著者らは TraFL を実用的にするために 2 つの主要なコンポーネントを導入した：

拡散互換シーケンスレベルの代理変数：マスク再構成の下限に基づき、扱いやすい対数確率の代理変数（ $\log \hat{p}_\theta(y|x)$ ）を導出した。これには、マスクレベル $l$ をサンプリングし、 $l$ 個のトークンを汚染し、再構成されたトークンの対数確率を平均化することが含まれる。 $1/l$ の因子は、マスクされたトークンの数に依存しないスコアの正規化に用いられる。
学習されたプロンプト依存正規化：バッチ全体での softmax 推定を通じて $Z(x)$ を推定する手法（濃縮を引き起こす可能性がある）とは異なり、TraFL は小さな MLP を通じて分配関数 $Z_\phi(x)$ を学習する。この項は、どの完了がサンプリングされるかに関わらず、すべてのステップで勾配信号を受け取るため、他の分布一致アプローチで見られる「バッチローカル」な正規化の失敗を防ぐ。

主要な貢献

軌道の固定化の特定：dLLMs における報酬最大化目的関数が軌道多様性の収束を引き起こし、それが結果として異なる正しい解モードのカバレッジを制限することを、形式的に特定し証明した。
TraFL 目的関数：方策を参照モデルにアンカーする軌道バランス目的関数を提案し、報酬の役割（終端完了の傾斜）と参照モデルの役割（妥当な軌道への確率質量の配分）を分離した。
実用的な代理変数と正規化：正確な軌道確率なしで安定した訓練を可能にする、拡散互換のシーケンスレベル代理変数と学習された正規化項を開発した。
実証的な優位性：TraFL は、すべてのベンチマーク長設定（GSM8K、MATH-500、HumanEval、MBPP）およびサンプリング予算においてベースモデルを一貫して上回る、評価された唯一の事後学習手法である。

結果

著者らは、数学的推論とコード生成のベンチマークにおいて、ベースモデル（LLaDA-8B-Instruct）および強力なベースライン（ESPO、JUSTGRPO）に対して TraFL を評価した。

分布内性能：TraFL は、すべてのサンプリング予算（ $k \in \{1, 2, 4, 8, 16\}$ ）および温度において Pass@k を改善した。特に、単一サンプル精度である Pass@1 を改善しつつ、サンプリング予算の増加に伴ってより良くスケーリングする。これは、単一サンプル精度を犠牲にしてマルチサンプルの利益を得る傾向があるベースラインとは対照的である。
保持済みデータへの一般化：利益は分布外ベンチマークへも転移する。Minerva Mathにおいて、TraFL はベースモデルを上回った。LiveCodeBenchにおいて、TraFL はすべての難易度分割（Easy、Medium、Hard）で最強の手法であり、中級および上級の問題においてベースラインを大幅に上回った。
多様性分析：LLM-as-a-judge プロトコルを用いて、TraFL の利益は単一のモードのスコアを鋭くするだけでなく、異なる証明アプローチやアルゴリズム的分解などの多様な解戦略のより広範なカバレッジに関連していることを示した。ESPO などのベースラインは精度を改善するものの、ベースモデルに対する多様性の有意な向上を示さず、軌道固定化仮説と一致する。

意義と主張

本論文は、TraFL が拡散モデルに対する標準的な RL 事後学習の「モード探索」失敗に対する原理的な解決策を提供すると主張する。報酬最大化から軌道バランスへ移行することで、この手法は複数の有効な解を網羅するために必要なノイズ除去経路の多様性を維持する。

著者らは、TraFL がテストされたすべてのベンチマーク長設定においてベースモデルを上回る唯一の評価済み手法であると述べている。改善はサンプリング予算の増加に対して頑健であり、保持済みデータセットへも一般化することを強調しており、軌道バランス目的関数が報酬信号を経路配分メカニズムから効果的に分離していることを示唆している。この研究は、dLLM 事後学習において純粋な報酬最大化よりも分布一致目的関数を推奨し、標準的な方策勾配アプローチに内在する濃縮失敗モードを回避する実用的な枠組みを提供する。

論文は限界に関して控えめであり、この手法が正確な軌道確率ではなく近似されたシーケンスレベルの代理変数に依存し、厳密な一致またはテスト通過といった粗い二値報酬を使用することを指摘している。これは、稀な成功事例における学習効率を制限する可能性がある。