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

1. 問題定義

連続拡散モデルは、画像、動画、音声における高次元生成の標準となっています。しかし、トークン埋め込みに対する連続拡散を介した言語生成への適用は、**マスク付き離散拡散モデル（MDM）**と比較して一貫して性能が劣っています。既存の連続拡散言語モデルは、分布の忠実度（MAUVE などの指標で測定）において、MDM よりも遅れをとっています。

著者らは、ボトルネックは連続性そのものではなく、使用される表現にあると主張しています。標準的なアプローチ（分散保存型または分散発散型）は、ゼロ情報限界を等方性ガウスノイズではなく、マスクのようなトークン状態として表現します。さらに、これらのフレームワークでは、ノイズの埋め込みにおける証拠の量を較正するために、デノイザがタイムステップまたはノイズレベルのラベルを受け取る必要があります。これにより、デノイザは時間条件付けに対して極めて敏感となり、マスクされた状態と不確実な状態に対して異なる処理が必要となります。その結果、連続拡散モデルは、有限の信号対雑音比（SNR）状態の連続体が提供する豊富な監視情報を活用することに失敗してきました。

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

本論文は、デノイジング過程を名义 SNR に対して不変にするように設計された連続状態フレームワークである**離散確率局所化（DSL）**を導入します。

中核的な設計選択

1. 単位球トークン埋め込み: DSL は、すべてのクリーンなトークン埋め込みが単位球上（$\|x_i\|_2 = 1$）にあるように制約します。これは損失正則化子ではなく、幾何学に対する構造的制約です。
2. 確率局所化チャネル: 汚染プロセスは $z_t = t x + \sqrt{t} \epsilon$ として定義されます。ここで、$t$ は SNR を表します。単位球の制約下では、ベイズ事後確率における $x$ のノルムに関する項が定数となり、正規化中に打ち消し合います。
3. SNR 不変デノイザ: この幾何学の下での重要な理論的結果として、ベイズ最適デノイザ（事後平均）は、名义 SNR $t$ ではなく、ノイズのついたトークン状態 $z$ のみに依存することです。
   $$\hat{x}(z, t) = \hat{x}(z) = \frac{\mathbb{E}_P[x e^{x \cdot z}]}{\mathbb{E}_P[e^{x \cdot z}]}$$
   これにより、単一の時間非依存ネットワークが、完全にマスクされた状態からほぼクリーンな状態までのすべての SNR 領域におけるデノイザとして機能できます。

学習とアーキテクチャ

• 混合サポート学習: 連続パス積分尤度とマスク付き拡散のエンドポイント尤度の両方を活用するために、DSL は混合サポート学習目的を採用します。SNR を 2 つのブランチからサンプリングします。
  • 連続パス: 対数正規分布からサンプリングされた SNR（有限 SNR 状態）。
  • ROAR（ランダム順序自己回帰）エンドポイント: 極端な値からサンプリングされた SNR（マスクされた場合は 0 に近く、露出された場合は $\infty$ に近い）。
    両方のブランチは、トークン事後確率に対する単一の交差エントロピー損失を使用して学習されます。これは、埋め込み空間の最適化の病理（例：埋め込みの崩壊）を回避しつつ、パス積分の平均二乗誤差（MSE）の代理として機能します。
• 事後ビュー変換器: ノイズのついた状態 $z$ は連続的な埋め込み空間に存在しますが、バックボーン（Transformer/DiT）はトークン様の入力を期待するため、変換器が $z$ を「トークンの混合」表現にマップします。これにより、ノイズのついた状態をノルム 1 の埋め込みの有界な組み合わせにマップし、事後幾何（方向がトークンの好みを制御し、ノルムが信頼度を制御する）を保持しつつ、安定した自己注意を確保します。

推論の柔軟性

デノイザが時間非依存であるため、DSL は単一の生成実行内で任意のトークンごとの SNR パスをサポートします。

• マスク付き洗練: 標準的な反復的リマスキング（ReMDM）。
• ランダム順序自己回帰（ROAR）: リマスキングなしで、ランダムな順序でトークンを 1 つずつ露出させる。
• ハイブリッドサンプリング: 連続デノイジング段階に続き、離散マスク付き洗練段階を行う。

3. 主要な貢献

1. 離散生成のための連続状態フレームワーク: DSL は、単位球埋め込み上の確率局所化チャネルを導入し、ベイズ最適デノイザがノイズのついた状態のみの関数となるようにします。これにより、単一の時間非依存ネットワークが、トークンごとの SNR パスの全体族をサポートできます。
2. 実用的な学習レシピ: 著者らは、同じデノイザをエンドポイント（マスク/露出）と有限 SNR 状態の両方で学習させる混合サポート SNR サンプリング戦略を提案しています。また、標準的な Transformer アーキテクチャと連続事後ビューを整合させるために、トークンの混合変換器を利用しています。
3. 経路全体での証拠: 単一の DSL チェックポイントが、マスク付き洗練における OpenWebText での少数ステップ生成の改善を示し、再学習なしでランダム順序 AR サンプリングをサポートし、Text8 において競争力のある尤度を達成しています。

4. 実験結果

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

• OpenWebText 生成（MAUVE）:
  • 非洗練: 反復的リマスキングなしでも、DSL 微調整済みチェックポイントはベースライン MDLM を大幅に上回りました。$T=1024$において、MAUVE は 0.042（MDLM）から 0.495（DSL-FT + MDLM）および 0.551（DSL-FT + ROAR）に向上しました。
  • 反復的洗練（ReMDM）: ReMDM ファミリのデコーダと組み合わせると、DSL は非参照ベースラインの中で最先端の結果を達成しました。最良の設定（原理的な $\eta_{cap}$ を用いた ReMDM ループ）は、$T=1024$で MAUVE 0.722 に達し、PRISM ループ（0.527）や ADLM（0.699）などの以前の手法を超えました。
  • 少数ステップ効率: DSL は低ステップ予算において強力な性能を示しました。$T=128$において、DSL-FT + ReMDM ループは 0.639 の MAUVE を達成し、ベースライン MDLM の 0.015 と比較されました。
• ハイブリッドサンプリング: ハイブリッドサンプリャ（連続デノイジングに続き離散洗練）は、わずか48 ステップで MAUVE 0.662 を達成し、$T=1024$における非洗練 MDLM の結果を超えました。
• Text8 尤度: DSL は文字あたりのビット数（BPC）で$\le$ 1.45を達成し、マスク付き離散拡散モデル（MDLM $\le$ 1.40、MD4 $\le$ 1.37）の範囲に近づき、以前の連続拡散ベースライン（Plaid $\le$ 1.48）を改善しました。

5. 意義と主張

本論文は、DSL の主な利点が推論経路の統合にあると主張しています。連続拡散を、タイムステップインデックス付きデノイジングではなく、ノイズのついたトークン上の事後状態デノイジングとしてパラメータ化することで、DSL は以下を可能にします。

• エンドポイント露出、マスク付き洗練、連続/ハイブリッドデノイジングを単一の事後推定器がサービスすること。
• チェックポイントの互換性: 単一のモデルを、再学習や蒸留なしで複数の SNR パス（例：ランダム順序 AR、マスク付き洗練、ハイブリッド）に沿ってクエリ可能にすること。
• 改善された分布の忠実度: 連続有限 SNR 状態が豊富な監視を提供し、標準的な連続拡散またはマスク付き拡散ベースラインと比較して、すべてのステップ予算において大幅に優れた MAUVE スコアをもたらすこと。

著者らは、言語における連続拡散の欠落していた要素は連続性ではなく、表現の選択であったと結論付けています。単位球埋め込みと確率局所化を採用することで、彼らは連続拡散と離散拡散の間のギャップを埋め、単一のモデルが高効率かつ高品質で多様な生成戦略をサポートできるようにしました。

限界: 評価は現在 OpenWebText と Text8 に限定されています。報告されたサンプリャ（ReMDM ファミリ、ハイブリッド）は網羅的に最適化されたデコーディング方針ではありませんが、基盤となる DSL 事後の柔軟性を示すために機能します。今後の研究の方向性には、リマスキングスケジュールの最適化、ハイブリッドハンドオフルールの最適化、条件付き生成の探求が含まれます。