技術的概要：非自己回帰生成のための離散確率的局所化

問題定義

連続拡散は、画像、動画、音声における高次元生成の標準的な枠組みとなってきたが、言語生成においては、マスク付き離散拡散モデル（MDM）に一貫して劣ってきた。既存の連続拡散言語モデルは、通常、分散保存（VP）または分散爆発（VE）ガウスノイズを用いて、ノイズの付いたトークン埋め込み上で動作する。著者らは、ボトルネックは連続性そのものではなく、表現の選択にあると主張する。

標準的な連続拡散において、ノイズ除去プロセスは時間ステップでインデックス付けされたノイズ領域に依存する。モデルは、ノイズの付いた埋め込みがどの程度の証拠を含んでいるかを較正するために、時間ステップまたはノイズレベルのラベルを受け取る必要がある。その結果、この表現におけるノイズ除去器は、マスクされた（ゼロ情報）状態と不確実な（有限の SNR）状態を区別するために、時間条件付けに極めて敏感である。この明示的な時間条件付けへの依存は、単一のネットワークが、個別のサンプリャ固有のモデルや再トレーニングなしに、ランダム順序の自己回帰サンプリングやハイブリッド連続・離散復号化のような多様な推論パスを自然にサポートすることを妨げている。

手法：離散確率的局所化（DSL）

本論文は、ベイズ最適ノイズ除去器を名义信号対雑音比（SNR）に対して不変にするように設計された連続状態枠組みである**離散確率的局所化（DSL）**を導入する。

1. 単位球トークン幾何学

中核的な革新は、すべてのクリーンなトークン埋め込みを単位球上（$\|x_i\|_2 = 1$）に制約することである。ノイズの付いた観測が $z_t = t x + \sqrt{t}\epsilon$ である確率的局所化チャネルの下では、ベイズ事後確率の因数分解には $t \|x\|_2^2$ に比例する項が含まれる。埋め込みが単位ノルムであるため、$\|x\|_2^2$ は所定の長さの任意の文に対して一定である。この項は正規化中に相殺され、事後確率、したがって最適ノイズ除去器を、名义 SNR $t$ に依存しないものとする。

最適ノイズ除去器は、ノイズの付いた状態 $z$ だけの関数となる：
$$\hat{x}(z) = \mathbb{E}[x | z] = \frac{\mathbb{E}_P[x e^{x \cdot z}]}{\mathbb{E}_P[e^{x \cdot z}]}$$
これにより、単一の時間非依存的なネットワークが、トークンごとの SNR パスの全ファミリーにサービスを提供できるようになる。

2. トークンごとの SNR パス

ノイズ除去器がグローバルな時間ステップを必要としないため、DSL は各トークン位置が独自の任意の SNR パス $\gamma_i(t)$ の下で進化することを可能にする。これにより、単一のモデルの下でいくつかの推論戦略が統合される：

• ランダム順序自己回帰（ROAR）： トークンを 1 つずつ明らかにすることは、他のトークンを 0 に保ったまま、特定の $\gamma_i$ を 0 から $\infty$ へ駆動することに相当する。
• マスク再精製： トークンを再マスクすることは、それらの $\gamma_i$ を 0 にリセットすることに相当する。
• 連続拡散： すべての $\gamma_i$ を同時に増加させることは、標準的な連続ノイズ除去に相当する。

3. 訓練目的とアーキテクチャ

これらの多様なパスに対して単一のモデルを訓練するために、DSL は混合サポート SN サンプリング戦略を採用する：

• 連続ブランチ： 有限 SNR 状態をカバーするために、対数正規分布から SNR をサンプリングする。
• エンドポイントブランチ： マスクされた（$\gamma=0$）およびクリーンな（$\gamma=\infty$）極端な状態の近くに集中した、ROAR 風の分布からサンプリングする。

両方のブランチは、単一の**トークンレベルの交差エントロピー（CE）**損失を使用して訓練される。著者らは、カテゴリカルなトークン事後確率上の CE を最小化することが、正確なパス積分平均二乗誤差（MSE）NLL の代理である一方、埋め込み空間回帰に関連する病理（埋め込みの崩壊や記憶など）を回避することを示すことで、これを正当化する。

アーキテクチャには、ノイズの付いた埋め込み $z_i$ を有界なトークン混合表現 $m_i$ にマッピングする事後ビュー変換器が含まれる。これにより、SNR レベルに関わらず、Transformer のバックボーンへの入力が安定し、語彙上の確率分布として解釈可能であることが保証される。

主要な貢献

1. 離散生成のための連続状態枠組み： DSL は、単位球埋め込み上の確率的局所化チャネルを導入し、ベイズ最適ノイズ除去器がノイズの付いた状態だけの関数となるようにすることで、明示的な時間条件付けの必要性を排除する。
2. 実用的なレシピ： 本論文は、エンドポイント（マスク/明らかにする）と有限 SNR 状態の両方に対して単一のノイズ除去器を訓練する混合サポート SN サンプリング戦略を提案し、連続事後幾何を標準的な Transformer バックボーンに橋渡しする変換器と組み合わせる。
3. パス全体にわたる証拠： 著者らは、単一の DSL チェックポイントが、再トレーニングなしでマスク再精製、ランダム順序自己回帰サンプリング、およびハイブリッド連続・離散サンプリングをサポートできることを実証する。

実験結果

著者らは、OpenWebText（OWT）および Text8 上で事前学習された MDLM チェックポイントを微調整することで DSL を評価した。

• OpenWebText 生成： DSL は、ベースライン MDLM と比較して、すべてのステップ予算（$T=128$ から $T=1024$）において、分布忠実度（MAUVE で測定）を大幅に改善する。
  • マスク再精製デコーダ（ReMDM）を使用すると、DSL は $T=1024$ で 0.722 の MAUVE を達成し、PRISM-loop（0.527）や ADLM（0.699）などの以前の最先端手法を上回る。
  • 反復的再精製なし（ナイーブな ROAR またはバニラ MDLM デコーダを使用）であっても、DSL 微調整済みチェックポイントはベースラインを大幅に上回り、この改善がサンプリングスケジュール単独ではなく、学習された事後確率に由来することを示している。
• ハイブリッドサンプリング： 短い連続ノイズ除去ステージを実行し、その後離散マスク再精製を行うハイブリッドサンプリャは、わずか 48 ステップ（MAUVE 0.662）で整合性のある生成を実現し、はるかに高いステップ数での非再精製ベースラインを上回る。
• Text8 尤度： Text8 ベンチマークにおいて、DSL は $\le 1.45$ の文字あたりビット数（BPC）を達成する。これは上限（ELBO）であり、自己回帰モデルの正確な NLL と直接比較できるものではないが、マスク付き離散拡散モデル（MDLM $\le 1.40$）の範囲に近づき、他の連続拡散ベースライン（Plaid $\le 1.48$）を改善する。

意義と主張

本論文は、言語における連続拡散と離散拡散の間のギャップは、連続性の性質によるものではなく、ノイズ領域の表現によるものであると主張する。ノイズの付いた状態のみがクリーンなトークンの事後確率を決定するように汚染をパラメータ化することで、DSL は、マスク拡散の「エンドポイントを明らかにする」振る舞いと、拡散モデルの「連続的再精製」能力を統合する。

中心的な主張は、エンドポイントを明らかにすること、マスク再精製、および連続/ハイブリッドノイズ除去は、個別のサンプリャ固有のモデルを必要とするのではなく、単一の事後推定器によってサービスを提供できるというものである。これは、モデルが時間非依存的な事後マップを学習し、推論パス（AR、マスク、またはハイブリッド）がサンプリング時に動的に決定される、非自己回帰生成の新たなパラダイムを示唆している。

著者らは限界に関して謙虚であり、評価が OWT と Text8 に限定されていること、報告されたサンプリャ（ReMDM ファミリーおよびハイブリッド）が網羅的に最適化された復号ポリシーではなく、DSL 事後確率の柔軟性を示す概念実証に過ぎないことを指摘している。今後の研究方向には、最適化された再マスクスケジュール、大規模モデル、条件付き生成が含まれる。