拡散言語モデルの構築:アーキテクチャ、サンプリング、およびスケーリング

エグゼクティブ・サマリー

拡散言語モデル(dLLM)は、テキスト生成における標準的な自己回帰(AR)アプローチからのパラダイムシフトを表しています。ARモデルがトークンを左から右へと逐次的に生成するのに対し、拡散モデルはシーケンス全体を並列に生成し、粗いドラフトを洗練された出力へと反復的に洗練させていきます。このアプローチは、ARモデルの3つの核心的な限界、すなわち、エラー訂正の欠如、遅い逐次生成、および因果的(過去を振り返る)アテンションへの制限に対処します。

離散拡散の核となるメカニズム

テキストのような離散データに拡散を適用するには、「ノイズ」の再定義が必要です。ガウスノイズを使用する画像拡散とは異なり、言語拡散は通常、マスキングまたは状態置換戦略を採用します。

マスクド・ディフュージョン・モデル(MDLM)

マスクド・ディフュージョンは、「生成的なBERT」として機能します。プロセスは2つのフェーズで構成されます:

  1. 前方プロセス: クリーンなシーケンスは、特定のスケジュールに従ってトークンを [MASK] トークンに置き換えることで、シーケンスが完全にマスクされるまで徐々に汚染されます。
  2. 逆プロセス: 双方向トランスフォーマーが、マスクされたシーケンスから元のトークンを予測するように訓練されます。生成時、モデルは完全にマスクされたシーケンスから開始し、空白を反復的に埋めていきます。その際、洗練を可能にするために、時折トークンを再マスクすることもあります。

均一状態拡散(UDLM)

マスク・トークンを使用する代わりに、UDLMはトークンを語彙全体からランダムに選択されたものに置き換えます。これにより、中間シーケンスがマスクフリーの状態を維持し、モデルが任意のステップで任意のトークンを修正することを可能にし、より優れたエラー訂正と制御可能性を促進します。

高度なアーキテクチャとサンプリング

理論的なモデルからプロダクション環境向けのLLMへと移行するためには、いくつかのアーキテクチャおよびアルゴリズムの拡張が必要です。

ブロック拡散と可変長

標準的なMDLMは固定長シーケンスに限定されています。Block Diffusion は、以前に生成されたブロックを条件としてテキストをブロック単位(例:256トークン)で生成することで、この問題を解決します。この手法は、自己回帰モデルと同様にKV キャッシングをサポートしており、効率的な可変長生成を可能にします。

エンコーダー・デコーダー・フレームワーク

Gemma DiffusionNemotron Diffusion のような現代的な dLLM は、エンコーダー・デコーダー・アーキテクチャを利用しています。強力なエンコーダーがクリーンなコンテキストを一度処理し、軽量なデコーダーがトークンの「キャンバス」を反復的にデノイズ(ノイズ除去)します。この分離により、デノイズ・ステップの計算コストが削減され、トレーニングが加速されます。

反復的洗練とエラー訂正

標準的なマスキングでは、一度アンマスクされたトークンは二度と再マスクされないため、初期のミスは永続的になります。これには2つの主要な解決策が存在します:

  • Remasking: 予測されたトークンのサブセットをランダムに再マスクするプラグイン・サンプラーであり、これにより、より多くのコンテキストが埋まった際に、モデルが文法的なエラーや論理的なエラーを修正することを可能にします。
  • Uniform State Diffusion: すべてのトークンを潜在的にノイズを含んでいるものとして扱うことで、モデルはシーケンス内の任意のポジションを自然に修正することができます。

蒸留によるサンプリング加速

拡散モデルは、並列生成により、ARモデルよりも5〜10倍高速です。速度をさらに向上させるために、progressive distillation が使用されます。これは、生徒モデルが教師モデルの2つのデノイズ・ステップを単一のステップで模倣するように訓練される手法であり、再帰的に必要なサンプリング・ステップ数を半分に減らしていきます。

制御可能な生成とガイダンス

拡散モデルは、局所的で不可逆的な編集ではなく、グローバルなシーケンスを洗練させていくため、本質的に制御可能な生成において優れています。これは、サンプルの自然さと特性の充足度の間のパレート最適(Pareto trade-off)によって管理されます。

  • Classifier-Based Guidance (CBG): 別の予測モデルを使用して、デノイズ・プロセスをターゲットとなる特性(例:分子における結合親和性)に向けて誘導します。
  • Classifier-Free Guidance (CFG): 条件付け信号をランダムにドロップすることで、単一のモデルで条件付き生成と無条件生成の両方を扱えるように訓練し、2つの出力を組み合わせることで結果を誘導します。

実世界での応用とスケーリング

生物学的および科学的ドメイン

拡散は、生物学において初期の成功を収めています。なぜなら、生物学的シーケンスは人間の言語のような左から右へのバイアスがなく、高い制御可能性を要求するためです。

  • ESM3: 100B パラメータの MDLM であり、タンパク質シーケンス、構造、および機能のモデリングに使用されます。
  • Nucleotide Transformer v3 (NT-v3): DNAシーケンス生成のための数千億パラメータのモデルであり、CFG と remasking を利用して、特定の遺伝子発現レベルを持つ調節的 DNA シーケンスを設計するために使用されます。

汎用拡散 LLM

いくつかの最前線ラボが、スケーリングされた拡散モデルをリリースしています:

  • LLaDA: 8B パラメータのオープンウェイトモデルであり、GSM8K や MMLU のようなベンチマークにおいて AR 形式に近いスケーリングを示しています。
  • Mercury 2: 標準的な GPU 上で約 1,200 トークン/秒の速度を達成し、Claude Haiku のような速度最適化された AR モデルを大幅に上回る商用モデルです。
  • Gemma Diffusion: Google のオープンウェイトモデルであり、UDLM バックボーンとエンコーダー・デコーダー・アーキテクチャを利用しています。
  • Nemotron Diffusion: NVIDIA のモデルファミリー(最大 35B パラメータ)であり、結合 AR-diffusion 目標関数を使用し、従来の AR モデルと比較して 2〜8倍の throughput を提供します。

批判的分析と今後の展望

拡散モデルは、推論速度とハードウェア利用率において莫大な利点を提供しますが、いくつかの技術的課題が残っています。コミュニティのフィードバックによれば、トークンの調整(coordination)における特定の弱点が見られます。つまり、2つのポジションが共に2つの有効な状態の間で揺れ動いている場合、拡散モデルはそれらを一致するペアに収束させることができず、不整合が生じる可能性があります。

スケーリングの観点からは、拡散モデルは、トランスフォーマーが事前学習計算量(pre-training compute)に対して持っていた役割を、**推論時計算量(inference-time compute)**に対して担う可能性があります。AR 推論の逐次的なボトルネックを並列プロセスへと変換することで、拡散モデルは、秒間あたりのより高い FLOPs を実現し、スケーリングされた事後学習(post-training)および推論時探索(inference-time search)を通じて、新たな知能のレベルを解き放つ可能性があるでしょう。

Sources

関連

  • Dispatch
  • Dispatch
  • Dispatch
  • Dispatch
  • プロジェクト