技術要約: Forward-Free Diffusion Language Models (FReDA)

問題提起

拡散言語モデル（LLM）は、テキスト生成を反復的なデノイジング（除去）プロセスとして定式化することにより、自己回帰（AR）生成に代わる有望な選択肢を提供します。しかし、離散的な言語空間においては、根本的な課題が残っています。画像のような連続領域ではガウス摂動が自然なユークリッド近傍を利用できるのに対し、離散的なトークン空間には固有のメトリック（距離尺度）が存在しません。その結果、既存の拡散言語モデルは、ノイズ注入の軌跡を定義するために、手作業で設計されたフォワード・コラプション（前方的な破壊）プロセス（例：吸収状態のマスキングや構造化された遷移カーネル）に依存しています。

著者らは、これらの規定されたフォワードプロセスが、生成中に遭遇する実際のドラフトやエラーとは数学的には都合が良いものの、それらとは乖離した中間状態を生成していると主張しています。この乖離は、訓練と推論の間のギャップを生み出します。つまり、モデルは人工的なコラプションを逆転させるのではなく、現実的なモデル出力を洗練することを学習すべきところ、人工的なノイズを逆転させることを学習してしまい、結果としてサンプル品質の低下や自己回帰モデルに対する性能差を招いています。

手法: FReDA

これに対処するため、著者らは、定義済みのフォワード・コラプション・プロセスを必要としないフレームワークであるFReDA (Forward-Free Diffusion Language Model) を提案します。FReDAは、拡散言語モデリングを**再帰的な分布洗練（Recursive Distribution Refinement）**として定式化することで、この問題を解決します。

コアメカニズム: 再帰的周辺分布の洗練 (Recursive Marginal Refinement)

FReDAにおいて、「中間状態」は外部のノイズ注入スケジュールによって生成されるのではなく、モデル自身によって生成された暗黙的なドラフトです。プロセスは以下のように動作します：

1. 初期化: モデルは初期ドラフト分布 $p_{\theta, T}$（多くの場合、空白のシーケンスまたは粗い予測）から開始します。
2. 洗練の遷移: モデルは、現在のドラフト分布をターゲットデータ分布 $p_0$ へと写像する遷移カーネル $K_{\theta, t}$ を学習します。
3. 再帰的適用: モデルは自身の出力を反復的に洗練します。学習目的関数は、ターゲット分布と、カーネルによって誘導される洗練された周辺分布との間のカルバック・ライブラー（KL）ダイバージェンスを最小化するように導出されます。

決定的なのは、学習目的関数（式 11）が、ターゲット $x$ とモデル生成ドラフト $x'$ の負の対数尤度を最小化する扱いやすいサロゲート（代理）関数である点です：
$$\min_{\theta} \mathbb{E}_{x \sim p_0, x' \sim p_{\theta, t}} [-\log K_{\theta, t}(x | x')]$$
これにより、モデルは人工的なノイズを逆転させるのではなく、自身のドラフトを洗練することを学習するため、学習ダイナミクスが生成ダイナミクスと直接一致します。

パラメータ化設計

本フレームワークは柔軟な洗練パラメータ化をサポートしており、論文内では以下の2つの具体的な設計を用いてインスタンス化されています：

1. 自己洗練 (Self-Refinement): モデルは自身のドラフトシーケンスを直接修正します。現在のドラフトを条件とし、トークンごとの洗練確率を出力することで、実質的に確率的な不動点反復を実行します。
2. 確率的 Best-of-N 洗練 (Stochastic Best-of-N Refinement): 探索空間を拡大するため、モデルは現在のドラフトから $N$ 個の候補シーケンスを並列に生成します。学習されたスコアラーヘッド ($Q_\psi$) が、提案された尤度と学習された補正信号の組み合わせに基づいて、これらの候補を評価し、最も有望なものを選択します。この選択は、訓練中の微分可能性を確保するために Gumbel-Max トリックに従います。

実装の詳細

• アーキテクチャ: Qwen3-4B-Base バックボーンに基づいています。
• アテンション: ブロック単位のセミコーザル・アテンション（ブロックサイズ 4）を利用しています。これにより、効率性のためにブロック間で因果構造を維持しつつ、ブロック内での双方向アテンションを可能にし、グローバルな洗練を実現します。
• 訓練: 2段階で行われます：(1) 自己洗練目的関数を用いた継続的な事前学習、および (2) ランキング目的関数を用いた Best-of-N スコアラーヘッドのオプションの訓練。
• 推論: 柔軟な停止基準をサポートしています。ユーザーは、追加の洗練反復を行うか、信頼度閾値に基づく早期終了戦略を使用することで、計算量と品質をトレードオフできます。

主な貢献

1. フォワードフリー・パラダイム: 離散空間における手作業によるフォワード・コラプション・プロセスへの依存を排除する FReDA の提案。これは、モデル生成のドラフトを唯一の中間状態として扱うことで、近傍に依存せず（neighborhood-agnostic）、かつ**モデルの複雑さを考慮した（model-complexity-aware）**アプローチとなります。
2. 再帰的洗練の定式化: 拡散を再帰的な周辺分布の洗練として捉える理論的フレームワークを提供し、学習された洗練演算子が最適条件下でターゲット分布を保持し、単調な改善をもたらすことを形式的に正当化しました。
3. 柔軟なパラメータ化: 自己洗練と Best-of-N 洗練の両方の戦略を実証し、利用可能な計算資源に基づいてモデルが探索戦略を適応できることを示しました。
4. 実証的検証: 一般知識、数学的推論、およびコード生成のベンチマークにおいて、より大規模な拡散ベースラインを上回る性能を示す広範な評価を実施しました。

実験結果

著者らは、一般的な知識、数学的推論、およびコード生成のベンチマークにおいて、FReDA-4B をいくつかの最先端の拡散ベースライン（例：DiffuLLaMA-7B、LLaDA-8B、Dream-7B）および AR ベースラインと比較評価しました。

• 性能: FReDA-4B (Best-of-N) は、推論およびコーディングのベンチマーク（例：GSM8K、MATH-500、HumanEval）において、より大きな 7–8B の拡散ベースモデルに対して 5% から 15% の絶対的な利得を達成しました。
• AR との比較: FReDA-4B は、ほとんどのベンチマークにおいて Qwen3-4B AR ベースモデルと同等またはそれを上回る性能を示し、拡散モデルと自己回帰型アプローチの間のギャップを大幅に縮小しました。
• 効率性 (パレート境界): FReDA は優れた品質-速度のトレードオフを確立しました。同等の、あるいはより優れた精度を維持しながら、拡散ベースラインに対して 1.5–1.8倍の平均高速化を達成しました。他の並列デコーディング手法で見られる典型的な品質低下を起こすことなく、より多くのトークンをフォワードパスあたりに生成（TPF）します。
• スケーリング: 洗練の反復回数が増えるにつれて一貫した性能向上が示され、再帰的洗練の仮説が検証されました。

意義と主張

本論文は、FReDA が効率的かつ高品質な拡散ベースの言語生成に向けた有望な方向性を提示していると主張しています。拡散プロセスを人工的なフォワード摂動から切り離すことで、FReDA は学習目的関数を実際の生成ダイナミクスに一致させます。

著者らは、このアプローチが近傍に依存せず（neighborhood-agnostic）、離散空間における任意のトークンレベルの距離を定義する必要がないことを強調しています。さらに、これはモデルの複雑さを考慮しており（model-complexity-aware）、有用な洗練ステップの数が学習されたモデルの表現力に応じて自然に適応することを可能にします。結果は、フォワードフリーの再帰的洗練が、サンプル品質を犠牲にすることなく、スケーラブルで並列的なテキスト生成への実行可能な道を提供し、拡散モデルと自己回帰モデルの性能差を大幅に縮めることができることを示唆しています。