Diffusion Language Models

拡散言語モデルは、通常はすべてマスクトークンで構成された、完全に破損したシーケンスからテキストを生成し、複数回のノイズ除去ステップを通じて多くの位置を並列に更新します。従来の大規模言語モデル（LLM）が左から右へ1つずつトークンを生成するのとは異なります。この手法は、画像の拡散モデルの考え方を離散的なテキストに応用します。順方向の処理でシーケンスを徐々に破損させ、Transformerでそれを元に戻す方法を学習します。このアプローチは2025年に本番規模に達しました。スクラッチからトレーニングされた80億パラメーターの拡散モデルLLaDAは、コンテキスト内学習でLLaMA3 8Bと競合する性能を示しました。Inception LabsのMercury Coderモデルは、独立したテストでNVIDIA H100 GPU上で毎秒1,109トークンと737トークンを生成しました。また、Google DeepMindの実験的なGemini Diffusionは、平均サンプリング速度として毎秒1,479トークンを公表しています。

マスク付き拡散の仕組み#

LLaDAをはじめとする近年のほとんどのモデルは、マスク付き拡散、または吸収状態拡散を使います。順方向の処理では、0から1までのノイズレベルtを選び、確率tで各トークンをマスク記号に置き換えます。

import torch

torch.manual_seed(6)
tokens = ["the", "cat", "sat", "on", "the", "warm", "mat", "by", "the", "door"]

# Forward process of a masked diffusion LM: at noise level t, each token is masked with probability t
for t in [0.25, 0.5, 0.75, 1.0]:
    masked = torch.rand(len(tokens)) < t
    noisy = ["[M]" if m else token for token, m in zip(tokens, masked)]
    print(f"t={t:.2f}: {' '.join(noisy)}")

実行すると、2語、5語、8語、最後には10語すべてをマスクした文が出力されます。トレーニング時には、因果マスクを使わないTransformerがこのようなシーケンスを受け取り、すべてのマスクトークンを一度に予測します。LLaDAはマスク位置のクロスエントロピー損失関数を1/tで重み付けします。この目的関数は、常にトークンの15％を予測するBERTのマスク言語モデリングを、あらゆるマスキング率に一般化し、モデルの対数尤度に上限を設けることで、完全な生成モデルを実現します。

生成では、t = 1から逆方向に処理を進めます。各ステップですべてのマスク位置を予測し、最も信頼度の高い予測を確定して、残りを再びマスクします。LLaDAはこの戦略をlow-confidence remaskingと呼びます。長い出力はブロック単位で生成されることが多く、この手法はブロック拡散として定式化されています。これにより、前のブロックをキャッシュでき、出力長を事前に固定する必要がありません。

D3PMから本番モデルへ#

GoogleのD3PMは2021年、離散状態上の拡散を定義し、現在多くのモデルが使う吸収マスク状態も導入しました。StanfordのDiffusion-LMは2022年、連続的な単語埋め込みに拡散を適用し、生成テキストを構文やトピックの制約に沿うよう誘導しました。ICML 2024のBest Paper賞を受賞したAaron Lou、Chenlin Meng、Stefano ErmonによるSEDDは、従来のテキスト拡散モデルのパープレキシティを25～75％低減し、GPT-2を上回りました。その後、LLaDAはこの手法が数十億パラメーター規模に拡張できることを示しました。また、文脈を双方向から読み取るため、左から右に処理するモデルの反転の呪いを検証する逆順詩の補完タスクでGPT-4oを上回りました。InceptionはOpenAI互換のChat Completions APIを通じてMercuryを提供しているため、Ultralytics LLMインターフェースからapi="chat.completions"とInceptionのbase_urlを使って呼び出せます。たとえば、YOLO26の検出結果を文章のレポートに変換できます。

強みとトレードオフ#

並列マスク解除では、トークン数より少ない順方向パスで済むため高速ですが、ステップ数を減らして一度に確定するトークン数を増やすと、品質が低下します。同じステップで確定するトークンは独立してサンプリングされるため、一度に多くを確定すると、互いに整合しない組み合わせが生じる可能性があります。また、各ステップではブロック全体を双方向に処理するため、KVキャッシュを使う自己回帰ステップよりもコストが高くなります。双方向の文脈により、穴埋めや編集が自然にできます。たとえば、関数や文書の途中にある欠落箇所を、前後のテキストを使って補完できます。そのため、コード補完や既存ファイルを編集するエージェント型コーディングツールに適しています。

拡散LMと関連用語#

• 自己回帰型テキスト生成： 左側のトークンから次のトークンを予測し、トークンごとに順方向パスを1回実行します。
• 画像拡散： Stable Diffusionなどの潜在拡散モデルは、連続値のピクセルまたは潜在表現にガウスノイズを加え、ノイズ除去器として拡散Transformer（DiT）を使うこともあります。
• 拡散フォーシング： 各トークンに個別のノイズレベルを設定し、次トークン予測とシーケンス全体の拡散を組み合わせます。主に動画や計画立案に使われます。
• 投機的デコーディング： 小さなモデルで複数のトークンを下書きし、大きなモデルの1回の処理で検証することで自己回帰モデルを高速化します。出力分布は変わりません。