技術要約: Set Diffusion

問題提起

離散拡散モデルは、自己回帰（AR）言語モデルに代わる競争力のある選択肢として台頭しており、生成速度の向上と制御性の改善を実現しています。しかし、既存の拡散アプローチには以下のような重大な制限があります：

1. 固定長および双方向の制約: ほとんどの離散拡散モデルは、完全な双方向コンテキストを必要とするアーキテクチャに依存しており、その結果、KVキャッシュの使用を妨げ、生成を固定長に制限しています。
2. ブロック拡散の硬直性: Block Diffusion (BD3LM) は、左から右へ固定サイズのシーケンシャルなブロックを生成することでこのギャップを部分的に埋めていますが、柔軟性に欠けています。これは、任意の位置への挿入（インフィリング）やブロック内での可変長生成をサポートできず、KVキャッシュの更新はブロック全体がデコードされるまで遅延されます。
3. トレードオフの限界: 現在の手法は、主に固定のブロックサイズを変化させることで自己回帰と拡散の間を補間していますが、これはデコーディングの柔軟性と並列性を制限します。

核心となる課題は、ブロックサイズによってではなく、トークンセットが生成される順序を変化させることで、自己回帰と拡散の間を補間し、高いサンプル品質を維持しながら、柔軟な長さの生成、KVキャッシュ、および任意の順序によるデコーディングを可能にする言語モデルを開発することです。

メソドロジー

著者らは、固定されたブロックではなく、柔軟な位置および柔軟な長さのトークンセット上で生成プロセスを分解するフレームワークである Set Diffusion を導入します。

1. 確率モデル

Set diffusion は、一連のトークンセット $\sigma = (\sigma_1, \dots, \sigma_N)$ に対する自己回帰的な確率分布を定義します。ここで、各 $\sigma_n$ はトークン位置のサブセットです。尤度は次のように分解されます：
$$p_\theta(x) = \sum_{\sigma} \prod_{n=1}^N \pi(\sigma_n | x_{<\sigma_n}) \cdot p_\theta(x_{\sigma_n} | x_{<\sigma_n})$$
ここで、$\pi$ は次の位置のセットを選択する順序ポリシーであり、$p_\theta$ は拡散を用いてそのセット内のトークンの条件付き尤度をモデル化します。モデルは、これらのセットにわたる拡散損失を合計する負の証拠下限（NELBO）を最適化します。

2. Set-Causal アーキテクチャと KV キャッシュ

主要なアーキテクチャ上の革新は、set-causal transformer です。標準的な双方向拡散モデルとは異なり、このアーキテクチャは毎回の推論ステップ後に KV キャッシュをサポートします。

• メカニズム: トークンセット $x_{\sigma_n}$ が生成されるとき、そのキーおよび値のベクトルはキャッシュされ、後続の予測の条件付けとして再利用されます。
• 学習効率: 著者らは、特殊なアテンションマスクを使用して、クリーンなコンテキストとノイズを含むクエリを単一のフォワードパスで同時に処理する、ベクトル化された学習アルゴリズムを提案しています。このマスクは、セット内のノイズを含むトークンが、同じセット内の他のノイズを含むトークンおよび先行するセット内のすべてのクリーンなトークンにのみ注目することを保証します。

3. トークン順序の補間

サンプル品質（境界の緊密さ）と推論効率（並列性）のトレードオフを制御するために、著者らは**位置依存の公開タイミング・スケジュール（reveal-time schedules）**を導入しています。

• 位置オフセット・スケジュール: 各トークン位置 $\ell$ には、累積分布関数 $\alpha^\ell_\tau$ によって定義される分布から抽出されたランダムな公開時刻 $R_\ell$ が割り当てられます。
• スライディングウィンドウ実装 (SW-SetDLM): 著者らは、トークンがアクティブな生成間隔 $[a_\ell, a_\ell + w]$ を持つ特定のインスタンス化を提案しています。ウィンドウサイズ $w$ は「左から右へ」のバイアスを制御します：
  • $w$ が小さい場合（$1/L$ に近づく場合）、順序が左から右への生成に集中し、AR挙動を回復して ELBO をタイトにします。
  • $w$ が大きい場合（1 に近づく場合）、順序に依存しない生成が可能になり、標準的なマスク拡散を回復します。
  • 中間値は、スライディングウィンドウ・デコーディングを可能にし、任意の位置における柔軟な長さのトークンセットの予測、およびインフィリングのようなタスクをサポートします。

4. 推論

推論中、モデルは現在の順序タイミングと位置オフセット・スケジュールに基づいて候補となる位置をサンプリングします。その後、KV キャッシュを使用して、これらの位置を並列にデノイズします。受理されたトークンはコミットされ、キャッシュは即座に更新されます。これにより、可変長生成と効率的なデコーディングが可能になります。

主な貢献

1. Set Diffusion フレームワーク: 柔軟な位置、柔軟な長さのトークンセット上で生成を分解する新しいクラスの言語モデルであり、自己回帰、順序に依存しない拡散、およびブロック拡散を一般化します。
2. Set-Causal アーキテクチャ: 毎回の推論ステップ後に KV キャッシュの更新をサポートするトランスフォーマー・アーキテクチャであり、効率的な可変長生成を可能にします。
3. 順序の補間: 位置依存の順序付けスケジュール（具体的にはスライディングウィンドウ・スケジュール）を調整することにより、AR と拡散の間を補間する、原理に基づいた手法を提供します。
4. 最先端のパフォーマンス: 数学的推論、要約、および無条件生成において、Set Diffusion が従来の拡散モデルよりも優れた速度と品質のトレードオフを実現することを実証しました。また、インフィリング・タスクにおいて Block Diffusion を大幅に上回っています。

実験結果

著者らは、いくつかのベンチマークで Set Diffusion (特に SW-SetDLM) を評価しました：

• 数学的推論 (GSM8K): Set diffusion は、Block Diffusion よりも厳密に優れた速度と精度のパレート・フロンティアを達成しています。例えば、デコーディング・ウィンドウサイズ $S \le 8$ の場合、$\ge 60.42$ tokens/sec のスループットで 66.41% の pass@1 精度を達成し、$\ge 55.39$ tokens/sec で 63.53% の精度を達成した Block Diffusion ($S=8$) を上回りました。
• インフィリング (ROCStories): Set diffusion は、インフィリング・タスク（中間の1文または3文を埋める）において Block Diffusion を大幅に上回り、最大で約 25% 高速なデコーディングでより高い ROUGE スコアを達成しました。Block diffusion は、その硬直した左から右へのブロック構造のために、このタスクに苦戦しています。
• 要約 (CNN/DailyMail): Set diffusion は、拡散ベースラインと比較して競争力のある ROUGE スコアを達成し、Block Diffusion よりも最大 10% 高速なデコーディングを実現しました。また、著者らが過学習の傾向があると指摘している AR ベースラインをも上回っています。
• 無条件生成および尤度: OpenWebText および LM1B において、Set diffusion は Block Diffusion よりも高い MAUVE スコアを達成しつつ、より高速にデコードします。LM1B では、高いスループットを維持しながら、拡散モデルの中で最高のパープレキシティ（$S \le 8$ で 27.83 PPL）を達成しました。

意義と主張

本論文は、Set Diffusion が、硬直したブロック構造から自己回帰と拡散の間の補間を切り離すことにより、重要な進歩を遂げたものであると主張しています。トークンの順序付けを通じて補間することにより、このフレームワークは以下を提供します：

• より大きな柔軟性: 任意の位置への挿入および可変長生成のサポート。これらはインフィリングのようなタスクにおいて極めて重要です。
• 効率性: 毎ステップごとに KV キャッシュを更新できる能力。これは従来、拡散モデルでは利用できなかった機能であり、高速な推論につながります。
• パフォーマンス: 既存の拡散言語モデルと比較して、サンプル品質と推論速度の優れたトレードオフを提供します。

著者らは、Set Diffusion を、AR、順序に依存しない拡散、および Block Diffusion を特殊なケースとして回収する、統一的なフレームワークとして位置付けています。これにより、推論タスクの特定の要件に基づいて、並列性とサンプル品質のバランスを調整できる、調整可能なメカニズムを提供しています。