Consistent Diffusion Language Models

以下は、「Consistent Diffusion Language Models (CDLM)」に関する論文の詳細な技術的サマリーです。

1. 問題定義

拡散言語モデル（DLM）は、並列トークン生成と亜線形時間推論を可能にすることで、自己回帰（AR）モデルに代わる有望な代替手段を提供します。しかし、重大なボトルネックが残っています：高品質なテキストを生成するには通常、数百の反復ステップが必要であり、並列化が約束する効率性の向上を相殺してしまいます。

連続領域（例えば画像）では、Consistency Modelsが、Probability Flow Ordinary Differential Equation (PF-ODE) で定義される一意の決定論的軌道に沿ってモデルが不変となるように学習させることで、拡散を成功裏に加速しました。

• ギャップ: 離散領域（テキストなど）では、同様のサンプル空間 PF-ODE は存在しません。離散的な汚損プロセス（例えば、マスキングや一様ノイズ）は単一の決定論的経路に従うのではなく、確率的遷移を含みます。
• 結果: 連続的な一貫性手法を離散データに直接適用することは、定義が不明確です。既存の解決策は、しばしば教師モデルを用いた多段階蒸留パイプラインや連続緩和の代理モデルに依存しており、これらは複雑で計算コストが高く、モード崩壊や多様性の損失に悩まされることが多いです。

2. 手法：マルチパス離散一貫性（MPDC）

著者らは根本的な転換を提案します：存在しない決定論的経路を探す代わりに、離散拡散プロセスに内在する正確な事後ブリッジを活用します。

中核概念：事後ブリッジ

離散拡散プロセスにおける任意の 2 つのノイズレベル $s < t$ に対して、閉形式の正確な事後ブリッジ $q(x_s | x_t, x_0)$ が存在します。このブリッジは、真のデータ $x_0$ に条件付けられた、ノイズ状態 $x_t$ からよりクリーンな状態 $x_s$ へ接続する有効な確率的経路を定義します。

• 連続 PF-ODE とは異なり、このブリッジは確率的です。
• 重要なのは、これらのブリッジが半群（Chapman-Kolmogorov 性質）を形成し、経路が推移的に合成可能である点です。

MPDC の原理

著者らは、Multi-Path Discrete Consistency (MPDC) を導入し、ノイズ除去器を期待値における経路不変性を持つように学習させます。

• 定義: モデル $f_\theta(x_t, t)$ が一貫性を持つとは、ノイズ状態 $x_t$ における予測が、正確な事後ブリッジを介してよりクリーンな状態 $x_s$ へ中間的な「ホップ」を行った後の予測と（期待値において）一致することを意味します。
• 目的: モデルは、$x_t$ からの直接予測と、$x_s$（ブリッジからサンプリングされる）からの期待予測との間の発散を最小化します。
• 数式定式化:
  $$\mathcal{L}_{CDLM} = \mathbb{E}_{t, \delta, x_t} \left[ \sum_{i \in \mathcal{M}(x_t)} w(t, \delta) \cdot D\left( f_\theta(x_t, t)_i \parallel \tilde{f}_\theta(x_s, s)_i \right) \right]$$
  ここで、$x_s \sim q(x_s | x_t, x_0)$、$\tilde{f}_\theta$ はターゲットネットワーク（EMA）であり、$D$ は発散指標（通常は Jensen-Shannon 発散）です。

安定化のための主要な設計選択

モード崩壊（モデルが単純に一貫性を満たすために反復的なテキストを出力する）と一様ドリフトを防ぐため、著者らは 3 つの安定化メカニズムを導入します：

1. ステップサイズをカリキュラムとして: 特定の範囲（例：$1/8$ から $3/8$）からステップサイズ $\delta$ をサンプリングし、$1/\delta$ で重み付けすることで、短い経路と長い経路の学習信号をバランスさせます。
2. 最大ステップ拡散アンカー: 訓練損失の小さな割合を、標準的な拡散目的関数を用いて真のデータ（$s=0$）にアンカーします。これにより、モデルが退化した定数解に収束するのを防ぎます。
3. 最適化の非対称性と発散: ターゲットネットワークにストップグラディエントを使用し、Forward/Backward KL ではなく**Jensen-Shannon 発散（JSD）**を採用します。JSD は、ゼロから訓練する際の KL 発散の不安定性を回避しつつ、モード探索とモードカバリングの間の安定したトレードオフを提供します。

3. 主要な貢献

1. 離散生成モデリングのための新たな原理: 本論文は、正確な事後ブリッジが欠如している PF-ODE のための厳密な閉形式の代替手段として機能することを確立し、「離散加速」の障害を解析的な資産へと変換します。
2. 統合された訓練フレームワーク: CDLM は、単一段階かつ教師不要の目的関数を提供します。著者らは、標準的なマスク拡散、連続一貫性モデル、プログレッシブ蒸留、および最近の 2 段階蒸留手法（SDTT や DUO-DCD など）のすべてが、CDLM フレームワークの特殊な場合、極限、または近似結合であることを証明しています。
3. 最先端のパフォーマンス: CDLM は、無条件および条件付きテキスト生成の両方で優れた結果を達成し、強力なベースラインモデルを上回り、教師を必要としない複雑な多段階蒸留モデルに匹敵するか、それ以上の性能を発揮します。

4. 実験結果

著者らは、OpenWebText データセット上で、MDLM、DUO、SDTT、DUO-DCD などのベースラインに対して、CDLM（特にマスク拡散用のMCDLMと一様拡散用のUCDLM）を評価しました。

• 無条件生成:

  • 少数ステップ効率: CDLM は、単一段階モデルとして新しい最先端（SOTA）の結果を確立しました。すべてのサンプリング予算において、ベースモデル（MDLM、DUO）を大幅に上回ります。
  • 蒸留との比較: 驚くべきことに、単一段階の CDLM は、Perplexity（PPL）の面で多段階蒸留モデル（SDTT、DUO-DCD）を上回ることが多く、同時にはるかに高いエントロピー（多様性）を維持しています。
  • モード崩壊: 貪欲サンプリャーを使用する蒸留モデル（例：DUO-DCD）は、低いエントロピー（約 3.9）を示し、重度のモード崩壊（反復的な文字）を示しました。CDLM は健全なエントロピー（約 5.0 以上）を維持し、多様性を保持する能力を実証しました。
  • 速度: CDLM は、自己回帰ベースラインに対して32 倍〜64 倍、標準的な MDLM に対して64 倍〜128 倍の高速化を達成し、かつ同等の品質を維持しています。

• 条件付き生成:

  • CDLM は、複数のデータセット（OpenWebText、Lambada、Wikitext-103、PTB）において、50% 汚損された入力から元の文を復元する際、堅牢な性能を示しました。
  • 他のモデルが失敗する入力条件を保持する点において、SDTT や DUO よりも優れた BLEU スコアと低い Perplexity を達成しました。

5. 意義と影響

• 理論的統合: 本論文は、離散生成モデリングに対する「統合的なレンズ」を提供し、多様な加速技術が経路一貫性の概念を通じて本質的につながっていることを示しています。
• スケーラビリティ: 多段階蒸留と教師モデルの必要性を排除することで、CDLM は、次世代の高速離散生成モデルのための原理的、スケーラブル、かつ教師不要な基盤を提供します。
• 汎用性: テキスト上で実証されましたが、このフレームワークは、扱い可能な事後ブリッジを持つ任意の離散領域（生物学的配列、グラフ、プログラム合成など）に適用可能です。
• パラダイムシフト: 高速離散拡散を、連続 ODE の近似としてではなく、確率的ブリッジ上の経路非依存ノイズ除去器の学習として再定義し、離散空間における少数ステップ生成の問題を成功裏に解決しました。

結論として、CDLMは、連続一貫性モデルの効率性と離散データの現実の間のギャップを埋め、蒸留パイプラインの複雑さなしに、高忠実度かつ高速なテキスト生成を実現する堅牢な単一段階ソリューションを提供します。