技術的概要：証明可能な自己修正のためのマスクド拡散モデルのファインチューニング

問題定義

マスクド拡散モデル（MDM）は、離散空間（テキスト、コード、数独など）における生成モデリングにおいて、自己回帰モデルに代わる有望な代替手段として登場しました。トークンを逐次的に生成する自己回帰モデルとは異なり、MDM は完全にマスクされたシーケンスから開始し、任意の順序でトークンを徐々にマスク解除します。この並列かつ少数ステップの推論は効率性と柔軟性を提供しますが、決定的な限界をもたらします。すなわち、標準的な MDM には自己修正の能力が欠如していることです。

標準的な MDM 推論では、一度トークンがマスク解除されると、それは固定されたままになります。この硬直性は、特に推論ステップ数が少ない場合に、位置間の依存関係エラーに対して MDM を脆弱にします。これを解決しようとする以前の試みは、以下のような重大な欠点に直面しました：

1. 不正確な代理指標： 学習不要な手法（ReMDM、DFM など）は、ランダムな再マスクや、トークンの品質を表すように学習されていないロジットから導出されたスコアに依存しており、非効率的または不正確な修正をもたらします。
2. アーキテクチャの大規模改変： 自己修正のための有効なターゲットを定義する手法は、しばしば MDM アーキテクチャ（例：ホロウ・トランスフォーマー）や事前学習の目的関数の変更を必要とし、既存の事前学習済みモデルへの適用性を制限します。

既存の事前学習済み MDM のコアアーキテクチャを変更することなく、任意の事前学習済み MDMに自己修正能力を付与し、かつ「トークン品質」の理論的に根拠のある定義を提供する手法の必要性があります。

手法：PRISM

著者らは、マスクド拡散の推論時自己修正のためのプラグイン型再マスク（Plug-in Remasking for Inference-time Self-correction of Masked Diffusions）であるPRISMを提案します。これは、軽量かつモデル非依存のファインチューニングフレームワークです。

1. 理論的基盤：トークンごとの品質

PRISM は、自己修正のための厳密なターゲットを定義します。それはトークンごとの品質スコアです。部分的にマスクされたシーケンス $y$ に対して、位置 $i$ におけるトークン $y_i$ の品質は、$y_i$ がマスクされた残りのシーケンスが与えられた場合のそのトークンの尤度として定義されます：
$$g^*_i(y) := p(x_i = y_i \mid y \oplus m_i)$$
ここで、$y \oplus m_i$ は、$i$ 番目のトークンをマスクトークン $m$ に置き換えたシーケンス $y$ を表します。$y_i$ が文脈と整合している場合、この確率は高くなります。もしそれがエラーである場合、確率は低くなります。

2. PRISM 損失

$g^*$ を近似するモデル $g_\phi$ を学習するために、PRISM は新しい損失関数を導入します。課題は、$g_\phi$ が入力 $y$ においてすでにマスク解除されているトークン $y_i$ の品質を予測しなければならない一方で、ターゲットはマスクされた文脈 $y \oplus m_i$ 上で定義された事後確率であるという点です。

PRISM は、特定のデータ構築と損失定式化を通じてこれを解決します：

1. データ構築： 真のシーケンス $x$ とそのマスク版 $z$ をサンプリングします。事前学習済み MDM $f_\theta$ を用いて、$z$ のトークンのサブセットをマスク解除し、新しいシーケンス $y$ を作成します。
2. 損失関数： モデル $g_\phi$ は、バイナリ交差エントロピー（BCE）を用いて、マスク解除されたトークンが真値と一致するかどうかを示すバイナリラベル $1[x_i = y_i]$ を予測するように学習されます：
   $$\mathcal{L}(\phi) := \mathbb{E}_{x,z,(y,i)} \left[ \text{BCE}(1[x_i = y_i], g_\phi^i(y)) \right]$$
   理論的に、この損失の唯一の最小化点は、$y$ を生成するために使用された事前学習済み MDM $f_\theta$ の品質に関係なく、まさに望ましいトークンごとの品質 $p(x_i = y_i \mid y \oplus m_i)$ と一致します。

3. 実装の詳細

• プラグ＆プレイ： PRISM は、任意の事前学習済み MDM に軽量な補助アダプター（「品質ヘッド」）を接続します。バックボーンは共有され、一つのヘッドがマスク解除の事後確率（元の MDM 機能）を出力し、新しいヘッドがスカラー品質スコアを出力します。
• 正則化： MDM の元のマスク解除能力を維持するため、標準的な MDM 訓練損失がファインチューニング中に正則化項として追加されます。
• 効率性： このフレームワークでは、複数のトークンを同時にマスク解除することで、$f_\theta$ の単一のフォワードパスから複数の訓練ペア $(y, i)$ を生成でき、データ効率を向上させます。

4. 推論手順

推論中、ファインチューニングされた MDM は各ステップで 2 つの操作を同時に実行します：

1. マスク解除： 標準的な MDM 事後確率に基づいて、マスク解除するマスクトークンのサブセットを選択します。
2. 再マスク： 最も低い品質スコア（$g_\theta$ によって予測されたもの）を持つすでにマスク解除されているトークンのサブセットを選択し、それらを再マスクします。
   これらのステップにより、モデルはシーケンスが確定する前にエラーを反復的に修正することができます。

主要な貢献

1. 理論的保証： PRISM は、完全なクリーンシーケンスを生成したり、検証器を使用したりすることなく、MDM におけるトークンごとの品質のための最初の証明可能な学習目的を提供します。この損失の最小化点が、マスクされた文脈が与えられた場合のトークンの真の条件付き尤度に収束することを証明しています。
2. モデル非依存設計： 先行研究とは異なり、PRISM はアーキテクチャの変更やゼロからの再訓練を必要としません。既存の任意の事前学習済み MDM に対する「プラグイン」型のファインチューニング層として機能します。
3. 効率性： この手法は、事前学習に比べて桁違いに少ない訓練サンプルと計算リソースで自己修正能力を学習します。

実験結果

著者らは、PRISM を 3 つのドメインとスケールで評価しました：

1. 数独（30M スケールの MDM）：

   • PRISM はファインチューニング中に順調に成功率を向上させ、数エポック以内に ReMDM や ReMDM-conf などのベースラインを上回りました。
   • 可視化により、PRISM がグリッド内の誤ったセルに正しく低い品質スコアを割り当て、効果的なエラー検出を可能にしていることが確認されました。

2. 無条件テキスト生成（170M スケールの MDM）：

   • OpenWebText において、生成ペレプレキシティと MAUVE で評価されました。
   • PRISM はベースラインを上回り、特に依存関係エラーが最も顕著な、サンプリングステップ数が少ない領域（例：64〜256 ステップ）において顕著でした。
   • 較正研究により、予測された品質スコアがトークンの経験的尤度と密接に追従することが示されました。

3. コード生成（8B スケールの MDM - LLaDA）：

   • 500 GPU 時間未満のファインチューニングで LLaDA-8B-Instruct に適用されました。
   • MBPP および HumanEval ベンチマークにおいて、PRISM は様々なサンプリングステップにおいて、バニラ MDM や他の再マスクベースラインを上回りました。
   • 機械的な分析により、PRISM が生成されたコードにおける構文エラー（例：不完全な elif 文）や論理エラー（例：区切り文字の配置）を正常に修正することが示されました。
   • 転移性： 1 つの MDM バックボーンで訓練された PRISM 調整モデルは、同じ分布上で動作する異なる事前学習済み MDM バックボーンに対する自己修正器として機能でき、大幅な性能向上（例：HumanEval で +10.9%）をもたらします。

意義と主張

本論文は、PRISM を修正と再順序付けを通じて離散シーケンスを構成する生成モデルの構築への一歩として位置づけています。これは人間のような生成プロセスを模倣するものです。

• 証明可能な自己修正： 著者らは、PRISM が Reinforcement Learning（完全なシーケンス生成が必要）や外部検証器に依存することなく、トークンごとの品質のための正確でテスト可能なターゲットを定義し、学習する点で唯一無二であると強調しています。
• スケーラビリティ： この手法は、最小限の計算リソースでさえも、大規模モデル（80 億パラメータ）においてさえ、ファインチューニングを通じて効率的に自己修正を獲得できることを実証しています。
• 限界： 著者らは控えめに、品質スコアが位置ごとの事後周辺分布に基づいているため、PRISM は複数のトークンにまたがる構造的論理エラーなどのグローバルな推論エラーには対処できない可能性があると認めています。しかし、彼らはこの研究が、離散拡散の効率的な訓練パラダイム内でグローバルなエラーに対処するための将来のフレームワークへの扉を開くと主張しています。