技術要約：グラフ上のランダムウォークによる、マスク付き拡散モデルにおける並列サンプラーの理解

問題提起

マスク付き拡散モデル（MDM）は、マスクされたトークンを反復的にデノイジングすることで離散データを生成する。MDMの主要な利点は、その「任意の順序」での推論能力であり、これによりトークンを逐次的、適応的、あるいは並列に明らかにすることができる。並列アンマスキングは、関数評価数（NFE）を削減することで大幅な高速化を実現するが、統計的なリスクを伴う。すなわち、並列更新は、真のブロック条件付き分布を、各座標の周辺分布の積へと置き換えてしまう。この近似は、選択された座標が現在のコンテキストに対して条件付き独立である場合にのみ有効である。

中心となる課題は、いつ並列アンマスキングが統計的に正確になるのか、そして、ターゲットとなる分布の条件付き依存構造を尊重するようにサンプラーをいかに設計するかである。自然言語では、モデルのエラー、サンプラーのエラー、および評価ノイズが絡み合っているため、これらの要因を分離することは困難である。本論文は、グラフ上のランダムウォークを用いた制御されたベンチマークを提案することで、この問題に対処する。

手法

グラフ・ランダムウォーク・ベンチマーク

著者らは、固定されたグラフ $G=(V, E)$ からのランダムウォーク・サンプルを用いてMDMを学習させる手法を提案している。決定的なのは、モデルにはグラフの構造や遷移カーネルが明示的に示されないことである。モデルは単にノードIDのシーケンスのみを見る。このセットアップは、「数独のような」妥当性チェックを提供する：

1. 潜在構造: グラフのトポロジー（例：木構造、ボトルネック、コミュニティ構造）が、制御可能な潜在構造として機能する。
2. 検証可能性: 生成されたシーケンスが有効（一貫している）であるための必要十分条件は、連続するすべてのトークンペアがグラフ内のエッジに対応していることである。
3. 評価: 指標には、一貫性（Coherence）（有効なウォークの割合）と、遷移全変動（Transition Total Variation, TV）距離（分布の忠実度を測定するもの）が含まれる。

このフレームワークは、無条件生成、エンドポイント条件付きブリッジ、および高次マルコフ・ウォーク（遷移が過去$k$個の状態に依存する場合）をサポートしている。

並列アンマスキングの理論的分析

本論文は、なぜ単一の並列アンマスキング・ヘリスティック（例：最小エントロピー、ランダム選択）が常に最適とは限らないのかを理論的に分析している。特定のグラフ構成を用いることで、以下を示している：

• Tree-Line-DAG: エントロピー誘導型サンプラーが、決定的な「セパレーター（分離子）」（ルート）を早期に特定することで、残りのチェインの依存関係を解決し、ランダムサンプラーよりも優れた性能を示すことを示す。
• Bottleneck DAG: ランダムサンプラーがエントロピー誘導型よりも優れた性能を示すことを示す。この構造において、エントロピーサンプラーは更新を依存関係のあるコンポーネント（コリドー／通路）内に集中させやすく、独立したトークンが同じコリドーからサンプリングされる際にエラーを引き起こす。一方、ランダムサンプラーは更新を分散させ、これらの結合されたエラーを回避する。

分析の結論として、サンプラーの有効性は、単なる局所的な不確実性のスコアではなく、その更新ポリシーが残されたマスクされた座標間の条件付き依存関係をどのように再形成するかによって決まる。

二分サンプラー（Bisection Sampler）

ランダムウォークのマルコフ分離特性（サイズ$k$の明らかになったブロックが、過去と未来を分離する性質）に着想を得て、著者らは二分サンプラーを導入する：

1. メカニズム: サンプラーは、各マスクされた区間内のサイズ$k$（または$\min\{k, \ell\}$）の中央ブロックを開示する。このブロックはセパレーターとして機能し、左側の部分区間と右側の部分区間を条件付き独立にする。
2. 再帰: プロセスは、新たに作成された部分区間上で並列に再帰する。
3. 厳密性: 完全な条件付き周辺分布という仮定の下で、この手法はオーダー$k$のマルコフ・ウォーク（無条件およびエンドポイント条件付きの両方）に対して、証明可能なほど厳密である。
4. 複雑性: これは $O(k \log(L/k))$ の並列深度を達成し、逐次的なサンプリングと比較してNFEを対数的に削減する。
5. バリアント: また、区間の中央半分の中で最も不確実性が低いピボットを選択することで、セパレーターの特性を維持しながら分割のバランスをとるスコア誘導型二分法も提案している。

主な貢献

1. 制御されたベンチマーク: グラフ・ランダムウォークを、潜在構造がモデルには隠されているが、厳密な評価には利用可能であるような、マスク付き拡散タスクとして定式化した。これにより、サポートの妥当性と分布の忠実度を直接測定することが可能となった。
2. アンマスキング・ポリシーの分離: 並列アンマスキングはオーダー依存であり、単純なヘリスティック（最小エントロピーなど）が一律に最適ではないことを理論的に証明した。性能は、基礎となるグラフの条件付き依存構造に決定的に依存する。
3. 二分サンプラー: マルコフ・セパレーターを利用する新しいサンプリング戦略の導入。これは完全な条件付き分布の下で厳密であり、対数的な並列深度を実現する。
4. 言語への転用: 二分法スタイルのスケジュールが、事前学習済みマスク付き言語モデル（OpenWebText）において速度と品質のトレードオフを改善することを実証し、グラフ・ウォークからの知見が自然言語生成に一般化できることを示唆している。

実験結果

グラフ・ウォーク実験

• ST-ER グラフ（全域木 + ランダムエッジ）: 木のような構造では、エントロピー誘導型サンプラーは相関のある選択に早くコミットしてしまうため、並列設定において失敗することが多い。ランダムエッジが追加されたグラフでは、エントロピーがアンカーを固定するのに役立つ。二分サンプラーは、攻撃的な並列ベースラインと比較して、NFEを大幅に削減しながら高い一貫性を一貫して維持した。
• ボトルネック・グラフ: コミュニティ間の移動を必要とするタスク（エンドポイント条件付きブリッジ）において、エントロピーベースの並列サンプラーはランダムサンプラーよりも性能が悪かった。エントロピー誘導型二分サンプラーは、計算効率を維持しつつ、すべてのベースラインを上回る一貫性を達成した。
• 高次ウォーク: 非局所的な依存関係タスク（オーダー$K$ウォーク）において、サンプラーのメモリ・オーダーが真の依存スケール（$K=4$）に近づくにつれ、二分サンプラーは一貫性が向上した。エントロピー誘導型二分法は、ランダムなベースラインよりも大幅に高い一貫性を達成した。

言語生成（OpenWebText）

OpenWebTextを用いた事前学習済みMDLMを用いた実験：

• 速度と品質のトレードオフ: 二分法ベースのサンプラー（特にエントロピー誘導型）は、最高の非自己回帰的な速度・品質トレードオフを実現した。
• 指標: エントロピー誘導型二分法は、約145 NFEでMAUVEスコア $\approx 0.16$ に達した。これに対し、より優れた性能を得るための自己回帰型ベースラインは約1023 NFEを必要とした。
• 多様性: 二分サンプラーは、他の非自己回帰的ベースラインと同等、あるいはそれ以上の多様性（エントロピー、反復率）を維持した。
• 観察: 低い生成パープレキシティだけでは高品質な生成を予測できない（例：一定エントロピー・サンプラーはパープレキシティは低いがMAUVEは低い）。これは、構造的なサンプラーの必要性を裏付けている。

重要性と主張

本論文は、グラフ・ランダムウォークを、マスク付き拡散モデルのための並列サンプラーを診断・設計するためのメカニスティックなベンチマークとして位置づけている。主要な主張は、並列デコーディングは局所的な不確実性だけでなく、条件付き依存構造によって支配されるということである。

二分サンプラーを導入することで、著者らは（完全な学習が行われているという仮定の下で）厳密性と高い一貫性を維持しながら、対数的な並列深度を達成できることを示している。これらの結果は、グラフ理論から導き出された粗から密へのアンマスキング・パターンが、実用的な言語生成タスクにおける速度と品質のトレードオフを効果的に改善できることを示唆しており、ヒューリスティックな並列デコーディング戦略に対する原理的な代替案を提供している。

著者らは、自らの理論が正確な条件付き周辺分布を想定していることを認め、謙虚な姿勢を保っており、今後の課題として、推定エラーがどのように伝播するか、および、自然言語に固有のより長期的な非マルコフ的依存関係にどのように適応するかを挙げている。