技術要約：離散拡散モデルは何を学習しているのか？

問題提起

離散拡散モデルは、テキストなどのカテゴリカルなデータに対する生成モデリングの強力なフレームワークとして台頭している。これらのモデルは、前方ノイズ付加プロセス（連続時間マルコフ連鎖、またはCTMC）を通じてデータを破損させ、そのプロセスを逆転させることを学習する。しかし、既存の文献では、学習された逆プロセスがいくつかの異なるように見えるパラメータ化によって提示されている：デノイザー（クリーンなトークンを予測するもの）、スコア比（ノイズを含んだ条件付き確率の比）、およびブリッジ・プラグインまたはキャビティ則（コンテキストは考慮するが、局所的なノイズを含むトークンは無視してクリーンなトークンを予測するもの）。

この分野には重大な曖昧さが存在する：これらは別々のオブジェクトなのか、それとも同じ基礎となるプロセスの異なる座標系に過ぎないのか？ 本論文は、ニューラルネットワークの出力を「誤った」座標系で読み取ると、最適化およびサンプリングされる連続時間マルコフ連鎖自体が変化してしまうことを主張している。具体的には、一様拡散（uniform diffusion）において、標準的なブリッジ・プラグイン目的関数は、通常のデノイザーではなくキャビティ則を訓練するが、マスク拡散（masked diffusion）ではこの区別が消失し、両者が一致する。さらに、訓練目的関数（ELBO）を単なる尤度の緩い代用物と解釈することは、その真の性質であるパス空間のダイバージェンスとしての側面を覆い隠してしまう。

手法

著者らは、CTMCに基づく連続時間離散拡散の厳密かつ自己完結した理論を展開している。彼らのアプローチは以下の通りである：

1. 一般化されたCTMC ELBOの導出： 変分オートエンコーダ（VAE）の枠組みから出発し、有限状態ジャンププロセスに対する連続時間ELBOを導出する。決定的なのは、文献ではしばしば破棄される境界項（$t_1$での再構成および$t_2$での終端事前分布）を保持することであり、これにより目的関数が有効な尤度境界となり、異なるノイズ付加プロセス間で比較可能になる。
2. オラクル距離恒等式： 負のELBO（NELBO）は単なる境界ではなく、以下の完全な分解であることを証明する：
   $$\text{NELBO} = H(q_0) + \text{KL}(P^\star_{[t_1, t_2]} \parallel P^\theta_{[t_1, t_2]}) + \text{Boundary Terms}$$
   ここで、$H(q_0)$はデータのエントロピーであり、KL項はオラクル逆プロセス（真の周辺逆ダイナミクス）と学習されたプロセスとの間のパス空間ダイバージェンスである。
3. 逆レート投影（Reverse-Rate Projection）： ELBOの集団最適化（population optimizer）は、周辺逆ジャンプレート、すなわち現在のノイズ状態が与えられた下でのクリーン条件付き逆レートの事後期待値であることを示す。これにより、ELBOが投影問題であることが確立される。
4. 座標辞書： トークン分解可能なノイズ付加プロセスに対して、最適な逆レートの3つの座標間の正確な閉形式の変換を導出する：
   • デノイザー： $\pi^\star_i(z^i_0 | z_t)$
   • キャビティ： $\mu^\star_i(z^i_0 | z^{\setminus i}_t)$
   • スコア： $s^\star_i(z^i_t, y | z^{\setminus i}_t)$
5. 情報理論的分析： ELBOの不可避なコストが情報の損失率 $-\frac{d}{dt}I(Z_0; Z_t)$ であることを特定し、ノイズ付加プロセスに関わらず、NELBOの普遍的な下限がデータエントロピー $H(q_0)$ であることを証明する。

主な貢献

1. オラクル距離定理

本論文は、NELBOを最小化することが、学習された逆ダイナミクスと真のオラクル逆ダイナミクスとの間のパス則KLダイバージェンスを最小化することと等価であることを証明している。これにより、ELBOは緩い変分境界から、軌跡の不一致を測定する正確な指標へと変貌する。一意の最適化対象は周辺逆レートであり、不可避なコストはデータエントロピーである。

2. パラメータ化の一致

著者らは、デノイザー、キャビティ則、およびスコアが、同一の最適な逆レートの3つの正確な座標であることを示す「辞書」を提供する。

• 変換： これら3つの座標間を変換するための閉形式の公式を提供している（例：デノイザーは、局所的な前方カーネルによって重み付けされたキャビティ則である）。
• 示唆： 変換を行わずに、キャビティ則のために設計された損失を用いてデノイザーヘッドを訓練することは、劣った学習につながる。

3. マスク拡散と一様拡散の区別の解消

論文は、なぜマスク拡散（MDM）と一様拡散（UDM）が異なる挙動を示すのかを説明している：

• マスク拡散： マスクされた位置では、局所的な観測は情報を与えない。その結果、デノイザーとキャビティ則は一致する（$\pi^\star = \mu^\star$）。これにより、標準的なデノイジング・クロスエントロピーが完璧に機能する。
• 一様拡散： 局所的な観測は情報を運ぶ。デノイザーとキャビティ則は異なる。論文は、一様拡散においてデノイザーのパラメータ化を使用すると、初期化時（$t \to 0$）にELBOが発散することを証明している。これは、無情報な事前分布がレート計算において特異な商を生じさせるためである。対照的に、キャビティ（ブリッジ・プラグイン）のパラメータ化は有限であり、適切に較正されている。

4. 実践的な較正と診断

• 初期化チェック： 著者らは較正恒等式を導出している：無情報なキャビティヘッド（一様分布）は、トークンあたりのNELBOが正確に $\log V$ （$V$は語彙サイズ）となる。これは実装の正当性を確認するためのデバッグツールとして機能する。
• 境界項： 境界項を落とすと、ELBOの値が異なるノイズ付加プロセス間（例：マスク型 vs 一様型）で比較不可能になることを示している。

5. エラー源の分離

本フレームワークは、2つのエラー源を明確に分離している：

1. 逆レート誤差： モデルのレートとオラレルレートの間のギャップであり、これをELBOが測定する。
2. 分解誤差： 結合逆ステップをトークンごとのステップの積として近似することによって導入される誤差（並列サンプリング）。この誤差はELBOには現れないが、たとえオラクルモデルであっても、生成パープレキシティに影響を与える。

結果と検証

理論的主張は、厳密に解けるトイモデル（8トークンの同次バイグラム・チェーン）を用いて数値的に検証されている。

• オラクル距離： 実験により、超過NELBO（データエントロキの床を超えた部分）が学習に伴ってゼロに減衰することが確認され、これはパスKLの予測と一致している。
• 座標ペナルティ： 誤った座標でレートを読み取る（例：キャビティの出力をデノイザーとして扱う）と、正しい変換を適用しない限り、NELBOと生成パープレキシティに大きなペナルティが生じることが数値テストで示された。
• 発散： 実験により、一様デノイザーヘッドは初期化時に発散する一方で、一様キャビティヘッドは安定していることが確認された。
• 較正： 無情報なキャビティヘッドに関する $\log V$ の恒等式が正確に成立している。

意義と主張

本論文は、バラバラに見えるパラメータ化を単一の幾何学的および情報理論的枠組みの下で統一する、連続時間離散拡散に関する初の厳密かつ一般的な理論を提供すると主張している。

• 理論的明晰さ： 離散拡散モデルが実際に何を学習しているのかについての長年の混乱を解消し、周辺逆レートが真の目的関数であり、ELBOがパスダイバージェンスであることを特定した。
• 実践的な指針： 実務家に対して具体的なルールを提示している：
  • 一様拡散およびハイブリッド拡散には、発散を避けるために**キャビティ（ブリッジ・プラグイン）**のパラメータ化を使用すること。
  • マスク拡散にはデノイザーのパラメータ化を使用すること。
  • 訓練中やサンプリング中にパラメータ化を切り替える場合は、必ず閉形式の変換を適用すること。
  • 手法間での公平な比較を保証するために、ELBOの計算には常に境界項を含めること。
• 限界： 著者らは、検証が小規模で厳密に解けるモデル上で行われたことを注記している。大規模なスケールでの一様拡散の有用性については、並行研究（[GJS+26]）が証拠を提供しているが、著者らは、これらの利点を大規模言語モデルのスケールで確認することが次の必要なステップであると述べている。また、彼らの分析は「ティーチャー・フォース（教師強制）」の状態（真の前方プロセスでの訓練）を想定しており、サンプリング・ドリフト（自己生成された状態への汎化）は、ELBO目的関数自体では対処されない別の課題であることを明確にしている。

要約すれば、本論文は、離散拡散はアドホックな手法の集合体ではなく、パラメータ化の選択は座標の選択であり、損失関数の選択は、最適な逆ダイナミクスへの収束を確実にするためにそれらの座標と一致していなければならない、統一されたプロセスであると主張している。