技術的サマリー：拡散の双対性、第 II 章：Ψ-サンプラー

問題定義

離散拡散モデルは、言語および画像生成における自己回帰モデルに対する強力な代替手段として登場し、並列デコーディングと双方向コンテキストを提供する。これらのうち、**一様状態拡散モデル（USDMs）**は、トークンを複数回修正することで自己修正する能力により、少数ステップでの生成やガイド付きタスクにおいて優れている。しかし、USDMs は 2 つの決定的な限界に直面している：

1. サンプリング品質の頭打ち： USDMs の標準的な祖先的サンプリングは、少数ステップ領域での性能とは異なり、サンプリングステップ数（NFEs）の増加に伴い生成品質が向上しない。
2. トレーニングの非効率性： 曲線学習（curriculum learning）（例：Duo）を通じて USDMs とマスク拡散モデル（MDMs）間の尤度ギャップを狭めようとする以前の試みは、トレーニング段階で高次元の潜在ベクトルを具体化する必要があるため、計算コストが高く、高いメモリ使用量と長いトレーニング時間を必要とする。

一方、**マスク拡散モデル（MDMs）**は一般的に高い尤度を達成するが、推論中に累積誤差に悩まされる。なぜなら、一度マスクが解除されたトークンは再マスクできないからである。**予測子 - 修正子（PC）**手法（例：ReMDM）に関する最近の研究は、再マスクを可能にすることで MDMs を改善したが、USDMs に対する同様の進展は未だ十分に探求されておらず、既存の USDM 用 PC 手法は祖先的サンプリングよりも性能が劣ることが多い。

手法

本論文は、Ψ-サンプラーによる推論とスケーラブルな曲線学習によるトレーニングという 2 つの主要な革新を通じてこれらの課題に対処する**Duo++**というフレームワークを導入する。

1. Ψ-事後分布とΨ-サンプラー

著者は、Ψ-事後分布と呼ばれる非マルコフ的事後分布の一族を提案する。これは、任意のノイズ事前分布（マスク型および一様型の両方）に対して従来の PC 手法を一般化するものである。

• 定義： Ψ-事後分布は、標準的な逆プロセス事後分布（$q_{s|t}$）と前方プロセス周辺分布（$q_s$）の線形重ね合わせとして定義される：
  $$\Psi_{s|t}(\cdot|x^\ell, z^\ell_t) = \kappa_t q_{s|t}(\cdot|z^\ell_t, x^\ell) + (1-\kappa_t) q_s(\cdot|x^\ell)$$
  ここで、$\kappa_t \in [0, 1]$は時間依存係数である。
• メカニズム： この定式化により、サンプリングプロセス中に$q_s$項を介してノイズを注入することが可能になる。
  • MDMsの場合、これにより再マスクが可能となる。デノイザーが誤りを犯した場合、以前にデコードされたトークンをマスク状態に戻すことで、累積誤差を修正できる。
  • USDMsの場合、デノイザーがほぼゼロの確率を割り当てた場合でも、すべてのトークンがサンプリングされる非ゼロの確率を保持することを保証し、サンプリングが局所最適解に陥るのを防ぐ。
• 理論的保証： Ψ-事後分布は、標準的なマルコフ的離散拡散プロセスと正確に同じ周辺分布を共有するように構築されている。これにより、十分なサンプルが与えられれば、プロセスが正しいデータ分布に収束することが保証される。
• 一般化： このフレームワークは、特定の$\kappa_t$スケジュールを選択した場合の特殊ケースとして、既存の PC 手法（例：ReMDM、Campbell ら、Gat ら）を包含する。

2. 高速トレーニングのためのスケーラブルな曲線学習

Duo（Sahoo ら、2025a）で提案された曲線学習戦略の計算ボトルネックに対処するため、著者はメモリ効率の良い実装を導入する。

• 疎性の活用： 従来の曲線学習は、離散トークン入力をガウス潜在変数の低温ソフトマックス緩和に置き換える。大規模語彙（$K > 100,000$）の場合、完全な$K$次元重みベクトルを具体化することは不可能である。
• Top-$k$近似： 著者は、低温においてソフトマックス分布が疎となり、質量がわずか数個の項に集中することを観察する。これに基づき、完全な線形結合をTop-$k$項（$k \ll K$）のみを用いて近似することを提案する。
• 効率的サンプリング： 完全なベクトルを生成してソートする代わりに、この手法は順序統計量を用いて、完全なベクトルを具体化することなく、基礎となるガウス分布から直接 Top-$k$値とその対応するインデックスをサンプリングする。
• 正規化： サンプリングされた Top-$k$項と、サンプリングされなかった項の期待寄与を考慮することで、ソフトマックス正規化項（$\tilde{Z}$）を近似する閉形式の式が導出される。
• 効率性： このアプローチにより、ピークメモリ使用量が**33%削減され、トレーニング時間が25%**短縮される一方、尤度性能は維持される。

主要な貢献

1. Ψ-事後分布： 非マルコフ的重畳事後分布を導入する離散拡散の統一フレームワーク。周辺分布の一貫性を保ちながら、任意のノイズ事前分布（マスク型および一様型）に対する予測子 - 修正子サンプリングを可能にする。
2. Ψ-サンプラー： 従来の PC 手法（ReMDM など）を一般化するサンプラーの一族。実証的に、USDMs が NFEs の増加に伴ってより良好にスケーリングし、テキスト生成において MDMs との性能ギャップを埋め、画像生成においては MDMs を凌駕することを可能にする。
3. 効率的な曲線学習： 低温ソフトマックス緩和の疎性を活用し、モデル品質を犠牲にすることなくメモリと計算コストを大幅に削減するスケーラブルなトレーニング戦略。

実験結果

言語モデリング（OpenWebText）

• 性能： Ψ-サンプラーは、祖先的サンプリングと比較して、USDMs の**生成パープレキシティ（Gen. PPL）**を大幅に改善する。
• スケーリング： NFEs の増加に伴って頭打ちになる祖先的サンプリングとは異なり、Ψ-サンプラーはステップ数の増加に伴い改善を続ける。一様エントロピー（多様性）が一致する条件下では、Ψ-サンプラーを備えた Duo++ は、ReMDM を用いた MDLM よりも低い Gen. PPL を達成する。
• 下流タスク： Duo++ は、25% 少ない FLOPs を使用しながら、多肢選択式質問応答タスク（例：ARC、HellaSwag）において元の Duo モデルと同等の性能を達成する。

画像モデリング（CIFAR-10）

• 品質： CIFAR-10 において、Ψ-サンプラーを備えた Duo++ は、祖先的サンプリングおよび ReMDM と比較して、優れたFIDおよび**Inception Score（IS）**を達成する。
• スケーリング： 画像品質の向上は、サンプリングステップ数の変化に対して一貫しており、Duo++ は MDLM のバリエーションを上回る。

トレーニング効率

• メモリと速度： 提案された効率的な曲線学習は、元の Duo 曲線学習と比較して、ピークメモリ使用量を**33%削減し、総トレーニング時間を25%**短縮する。
• 尤度： 効率的な曲線学習でトレーニングされたモデルは、LM1B および OpenWebText において、元の Duo モデルと同等のパープレキシティを達成する。

意義と主張

本論文は、マスク拡散が拡散ベースの言語モデリングの避けられない未来であるという仮定は疑問であると主張する。Ψ-サンプラーを導入することで、著者は**一様状態拡散モデル（USDMs）**が以下の達成が可能であることを実証する：

1. 優れたスケーリング： 祖先的サンプリングとは異なり、サンプリングステップ数の増加に伴い性能が継続して向上する。
2. 競争力のある、あるいは優れた品質： （PC サンプリングと組み合わせることで）テキスト生成において MDMs と同等かそれ以上の性能を発揮し、画像生成においては MDMs を凌駕する。
3. 実用的な効率性： 従来の曲線学習アプローチよりもメモリおよび時間効率に大幅に優れたトレーニングパイプライン。

この研究は、ガウス拡散と離散拡散の間の「双対性」と、非マルコフ的サンプリング戦略の組み合わせが、離散生成モデリングのための堅牢でスケーラブルな道筋を提供し、マスクアプローチの優位性に挑戦することを示唆している。