拡散言語モデルの構築方法

拡散言語モデルの構築方法：拡散言語モデルへの入門と、今日の拡散LLMの基盤となる研究の進歩について説明します。最近のオープンソースモデルの構成要素を、単純なマスキング拡散から、反復的洗練、ポストトレーニング、可変長生成のテクニックまで含めて記述します。この資料は、ICLR 2026およびMLSS 2026でのワークショップ講演や講義から抜粋したものです。 はじめに：自己回帰モデルと拡散言語モデル 今日広く使われている生成AIアルゴリズムには2つのファミリーがあります。画像やビデオのような連続データに対しては、最先端のアプローチは拡散モデルに基づいています。テキストやコードのような離散データに対しては、標準的なアプローチは自己回帰モデルです。この記事では、離散データのための代替案、すなわち現代の拡散パラダイムに基づいたものを探求します。 主流の言語モデルは自己回帰型です。これらはトークンを左から右へ、一度に1つずつ生成し、それぞれが先行するトークンに条件付けられます。このアプローチは強力ですが、固有の限界もあります。 エラー訂正なし：一度トークンが出力されると、修正できないため、初期の間違いが累積します。 生成が遅い：シーケンスを生成するには、トークンの数と同じステップが必要であり、高速な並列生成には自然には適しません。 因果的注意：生成は常に過去のみを参照し、未来のコンテキストを見ることはありません。 拡散モデルは異なるアプローチを取ります。テキストを一度に1トークンずつ生成するのではなく、初期の推測から始めて、数ステップにわたって反復的に洗練させることで、シーケンス全体を一度に生成します。これにより、いくつかの利点が生まれます。 生成は、より少ないまたはより多くのステップを使用することで、速度と品質をトレードオフできます。 間違いは途中で修正できます。 各ステップは双方向のコンテキストに注意を払います。 自己回帰LLMは、左から右へ一度に1トークンを生成し、トークンの数と同じステップ数を要します（上）。拡散LLM（ここではGemma Diffusionを示す）は、代わりにラフな全長ドラフトから開始し、各ステップでシーケンス全体を書き換えることで、数ラウンドにわたってすべての位置を並列に洗練させます（下）。図のクレジット：M. Grootendorst & Gemma Diffusion。 言語への拡散の適用は、長い間未解決の問題でした。2024年に、拡散モデルが品質において自己回帰モデルに匹敵するようになり、この分野は転換点を迎えました。2026年までに、拡散LLMは現実のものとなり、主要な産業ラボ（Mercury 2 (Inception Labs)、Gemma Diffusion (Google)、Nemotron Diffusion (NVIDIA)）からリリースされています。この記事は、これらの現代モデルの基盤となるアイデアと論文をたどります。 背景：ガウシアン拡散 言語のための拡散を導入する前に、画像生成のためのガウシアン拡散の簡単な概要から始めます。その後、類推によって離散拡散を構築していきます。 反復的デノイズングによる生成 拡散モデルの根底にある中心的な概念はデノイズングです。拡散モデルは、画像を一度に描画するのではなく、ステップバイステップで画像を生成します。純粋なランダムノイズから開始し、各ステップで少しずつノイズを除去していき、一貫した画像が現れるまで続けます。多くの小さなステップを経て画像を生成することは、一度に生成することよりもはるかに簡単であることが判明しており、これが拡散モデルを非常に効果的にしている理由です。 モデルはどのようにデノイズングを学習するのでしょうか？トリックは、ノイズが徐々に画像に変換される例を示すことによって、それを教え込むことです。拡散は、2つの相補的なプロセスを通じてこれを達成します。まず、フォワードプロセスはクリーンなソース画像を取り、それをステップバイステップで純粋なノイズに変換します。次に、リバースプロセスは、この変換を逆転させることを学習し、純粋なノイズを画像に戻します。これは、フォワードプロセスによって生成された画像からノイズへの軌跡でトレーニングされます。 フォワードプロセス フォワードプロセスは、クリーンなトレーニング画像を取り、クリーンなデータから純粋なノイズへのパスをたどる、ますますノイズの多い画像のシーケンスを生成します。これは、各ステップで増加する量のランダムガウシアンノイズを混合することによって行われ、画像が純粋な静止画に溶解するまで続きます。このステップは学習を全く必要としません。私たちは単にノイズを追加しているだけですが、それは無限のトレーニングデータ（画像がノイズに変換され、その逆の例）を製造するため、非常に役立ちます。 画像上のガウシアン拡散軌跡。左から右へ、フォワードプロセスは徐々にノイズを追加し、犬のクリーンな写真が純粋な静止画に溶解します。この軌跡は、リバースプロセスのトレーニングデータとして使用されます。 リバースプロセス リバースプロセスは、実際の学習が行われる場所です。フォワードプロセスによって生成されたステップを逆方向にたどることで、ノイズを画像に変換するようにモデルをトレーニングします。 ガウシアン拡散のリバースパス。再び左から右へ、生成リバースプロセスは、ノイズから開始し、元の画像を再構築して、フォワードプロセスの軌跡を逆方向に再現するようにトレーニングされます。具体的には、ノイズの多い画像が与えられた場合、ノイズを基になる画像から分離する、または同等に、ノイズの多い入力が与えられた場合、一方を知れば他方が決定されるため、追加されたノイズまたはクリーンな画像自体を予測するように機械学習モデルをトレーニングします。モデルがこれを行うことができるようになれば、生成は簡単です。純粋なノイズから開始し、モデルに少しずつノイズを推定して除去するように依頼し、繰り返します。各パスは、サンプルを実際のデータのように見えるものに少しずつ近づけ、クリーンな画像が残るまで続きます。 拡散を定義する2つのプロセス。フォワードプロセス（上、左から右）は、各ステップで少量のガウシアンノイズを追加することにより、クリーンなデータ $x_0$ を完全なノイズ $x_T$ に変換します。生成リバースプロセス（下、右から左）は、このデータでトレーニングされ、同じステップシーケンスを使用して $x_T$ を $x_0$ にデノイズします。十分な大きさの軌跡セットでトレーニングした後、モデルは一般化を学習し、ランダムなホワイトノイズサンプルから開始して新しい画像を生成します。このフォワード/リバースのレシピ—データをノイズで破損させ、次に破損を1ステップずつ逆転させることを学習する—は、すべての拡散モデルの設計図です。 単純なマスキング拡散モデル 言語に拡散をもたらす際の主な障害は、離散トークンにとって「ノイズ」が何を意味するかを決定することです。たとえば、古典的な拡散で使用されるノイズはガウシアンですが、カテゴリカル変数に連続的なガウシアンノイズを追加することは適切に定義されていません。以下では、マスキングによってノイズを定義する、シンプルでありながら効果的なアプローチを紹介します。私たちのグループはこのアプローチを普及させ、現在ではほとんどのオープンソース拡散言語モデルの基盤となっています。 マスキング拡散の概要 マスキング拡散を理解する最も簡単な方法は、アンマスキングトランスフォーマーとして考えることです。クリーンなシーケンスを取り、トークンのランダムな割合をマスクし、双方向トランスフォーマーに空白を埋めるように依頼することでモデルをトレーニングします。BERTを知っていれば、これは基本的にランダム化されたマスキング率を持つBERTですが、BERTとは異なり、結果のモデルは生成型です。マスキング拡散を生成型BERTと考えることができます。 アンマスキングトランスフォーマーとしてのマスキング拡散のトレーニング。フォワードマスキングプロセスは、ランダムなノイズレベル $0 < t < 1$ をサンプリングし、クリーンなデータポイント $x$ のその割合のトークンを隠し、部分的にマスクされた $z_t$ を生成します。次に、モデル $x_ heta$ は元のトークンを再構築するようにトレーニングされます。図のクレジット：Sasha Rush。 アンマスキングトランスフォーマーをトレーニングしたら、完全にマスクされたシーケンスから開始し、次の2つのステップを何度も繰り返すことでテキストを生成できます。 インフィリング：現在のシーケンスのすべての空白をモデルに埋め込ませ、クリーンなトークンのラフな推測を生成します。 リマスキング：インフィリングされたシーケンスをランダムに再ノイズ化します。トークンをマスクに置き換えますが、前のラウンドよりも多くのトークをマスクせずに保持します。 1つのサンプリングステップ。デノイズングモデルは、現在のシーケンス $z_t$ のマスクされた位置を埋め込みます。次に、それらの新しい予測のランダムなサブセットが再マスクされて $z_s$ 、次のステップのわずかにマスクされていないシーケンスが形成されます。これを繰り返すことで