技術的概要：マスク拡散言語モデルのための拡散状態方策最適化

問題定義

マスク拡散言語モデル（MDLM）は、複数のノイズ除去ステップにわたってマスクされた位置を反復的に埋めることでテキストを生成する。方策最適化（PO）手法は言語モデルにおける推論やアライメントの向上に成功してきたが、既存の大部分のアプローチは最終結果のみの報酬に依存している。これらの手法では、最終的な完成に対してスカラー報酬が割り当てられ、それがすべての中間生成決定に対して逆伝播される。

このアプローチは粗いクレジット割当をもたらす。それは、局所的な埋め込み決定の全シーケンスを単一の単位として扱い、ロールアウト中にすでに生成されている豊富な中間状態 - 行動情報を活用できない。具体的には、任意の中間ノイズ除去ステップにおいて、モデルは現在マスクされているすべての位置に対するトークンレベルの分布を生成するが、標準的な最終結果フィードバック PO はこの粒度を無視し、最終結果への特定の貢献度に関係なく、すべての決定に同じ結果シグナルを割り当てる。

手法：拡散状態方策最適化（DISPO）

著者らは、追加の多ステップ拡散ロールアウトを必要とせずに中間の埋め込み決定を直接最適化するように設計されたプラグイン型のクレジット割当層である**拡散状態方策最適化（DISPO）**を提案する。

中核概念

DISPO は、MDLM の生成プロセスを状態 - 行動フレームワークとして再定義する：

• 状態（$s_{k,t}$）： ノイズ除去ステップ $t$ における中間的な部分的にマスクされたシーケンス。
• 行動（$a_{k,t}$）： その状態における現在マスクされている位置の同時埋め込み。

アルゴリズム的メカニズム

1. ロールアウトとキャッシング： 固定された挙動パラメータ（$\theta_{old}$）を用いた標準的なロールアウト中に、モデルは部分的にマスクされたシーケンスの軌跡を生成する。重要なのは、訪問されたすべての状態におけるすべてのマスク位置のログオッズがキャッシュされる点である。
2. 状態条件付き分岐： 選択された中間状態において、DISPO はロールアウトからキャッシュされたログオッズを使用して、マスクされた位置の代替埋め込みを再サンプリングすることで「分岐」する。これにより、同じ中間状態から複数の「反事実的」完成（$Z$ 分岐）が生成される。
3. スコアリングとアドバンテージ： 各分岐された完成は、標準 PO で使用されるのと同じ最終報酬関数（$R$）を用いてスコアリングされる。特定の状態からの $Z$ 個の完成の報酬に基づき、各分岐に対してグループ相対アドバンテージが計算される。
4. トークン局所更新： 方策勾配は、現在の状態における新たに埋められたトークン（行動）のみを更新するように計算され、以前に埋められたトークンは固定状態の一部として扱われる。
5. 代理尤度： MDLM における正確なシーケンス尤度の扱いにくさを処理するため、DISPO は扱いやすい**1 ステップマスクトークン代理（$\tilde{\pi}_\theta$）**を利用する。これにより、更新対象のマスク位置に特化した尤度比の計算が可能となる。

訓練目的

DISPO は固定状態期待リターン目的を定義する。これは、重み付き和を介してベースの最終結果フィードバック最適化器（例：diffu-GRPO または SPG）と組み合わせることができる：
$$\mathcal{L}(\theta) = \alpha_{base} \mathcal{L}_{base}(\theta) + \alpha_{step} \mathcal{L}_{step}(\theta)$$
ここで、$\mathcal{L}_{step}$ は分岐された完成から導出された状態ごとの損失である。理論的解析により、この組み合わせ目的が、最終目標と時間ステップ平均化された固定状態目標からなる混合目的に対する有効な方策勾配シグナルを提供することが確認されている。

主要な貢献

1. 状態条件付きクレジット割当： 本論文は、中間マスクシーケンスを状態として、マスク埋め込みを行動として扱う手法を導入し、標準的な最終結果フィードバック PO が欠いている微細なクレジット割当を可能にする。
2. キャッシュログオッズによる効率的な分岐： 重要な革新は、同じ状態からの代替完成を生成するためにロールアウトでキャッシュされたログオッズを再利用することである。これにより、追加の多ステップ拡散フォワードパスなしで同じ状態での比較が可能となり、支配的な訓練コスト（ロールアウトと最適化器ステップ）はベースラインと一致したままとなる。
3. 理論的定式化： 著者らは DISPO を固定状態目的として定式化し、方策勾配推定量を導出した。彼らは、ステップごとの損失が固定状態目的に対する不偏（スケーリング済み）推定量であることを証明し、それを最終損失と組み合わせることで明確に定義された混合目的を最適化することを示した。
4. 分散低減戦略： この手法は分散分析に基づいた 2 つの設計選択を採用している：
   • トークン局所更新： 実行可能な（マスクされた）トークンのみに勾配を制限することで、すべてのトークンを更新する場合と比較して分散が低減される。
   • 同一状態平均： $Z$ 分岐全体で平均化することで、分散が $O(1/Z)$ の因子で低減される。

実験結果

著者らは、LLaDA-8B-Instruct上で DISPO を評価し、2 つのベース最適化器であるdiffu-GRPOとSPGと統合した。

• ベンチマーク： 数学的推論（GSM8K, MATH500）と記号計画（Sudoku, Countdown）において実験が行われた。
• 性能： 計算予算が一致した場合（ロールアウト数、グループサイズ、最適化器ステップ数が同じ）、DISPO はすべてのベンチマークにおいて両方のベース最適化器を一貫して改善した。
  • 計画タスク（例：数独の精度が diffu-GRPO を用いて約 14% から約 31% に向上）において特に顕著な改善が見られた。
  • 改善は異なる生成予算（$N_{gen} \in \{128, 256, 512\}$）においても持続した。
• アブレーション研究：
  • 状態ごと対最終結果： 状態ごとの目的のみが最終結果のみのベースラインを上回り、それらの組み合わせが最良の結果をもたらした。
  • タイミング： 完了に近い状態に焦点を当てた後期バイアス付き時間ステップサンプリングは、一様または早期バイアス付きサンプリングよりも良好なパフォーマンスを示し、最終完了に近い中間状態がより信頼性の高いクレジットシグナルを提供することを示唆している。
  • 勾配分散： 実行可能なトークンのみを更新し、分岐サイズ $Z$ を増加させることが、勾配分散を大幅に低減することを実験が確認した。
• エラー分析： 数独における「早急なコミットメント」の分析により、DISPO はベースラインと比較して制約違反の最初の発生をより後のノイズ除去ステップにシフトさせたことが示され、中間行動の最適化が早期に修正困難なエラーを回避するのに役立つことを示している。

意義と主張

本論文は、DISPO がマスク拡散方策最適化のための汎用プラグインとして機能すると主張している。その主な意義は以下の点にある：

• 中間情報の活用： MDLM 生成プロセスに内在する中間状態 - 行動情報が、完全に新しい軌跡を生成する計算オーバーヘッドなしに、微細なクレジット割当のために活用できることを示している。
• 相補性： DISPO は、中間状態に基づいて報酬を整形する SAPO などの既存手法や、軌跡レベルの目的と相補的であると提示されている。SAPO が中間状態に基づいて報酬を整形するのに対し、DISPO は固定状態の下で埋め込み行動を直接最適化する。
• 効率性： キャッシュされたログオッズを再利用することで、DISPO は支配的な訓練コスト（多ステップロールアウト）を増加させることなくこれらの改善を達成し、既存の MDLM 訓練パイプラインへの実用的な追加となっている。

著者らは、現在の実装は代理評価と報酬評価の繰り返しによりウォールクロックオーバーヘッドを伴うが、アルゴリズム的コストはベースラインと一致しており、性能向上は堅牢であると指摘している。彼らは、DISPO を拡散ベース言語モデルにおけるより効率的かつ効果的な推論への一歩として位置づけている。