- 完全双方向Attentionを採用した8Bパラメータの拡散型言語モデル「iLLaDA」が、12兆トークンの大規模事前学習で先行モデルLLaDAをBBH+21.6点・HumanEval+16.5点で大幅に上回った
- 複数のベンチマークで自己回帰型のQwen2.5 7Bと同等水準を達成し、拡散型が大規模学習で真の競合になり得ることを実証した
- コードとモデル重みはGitHubで公開済みであり、研究コミュニティが再現・発展させやすい環境が整っている
研究の背景
テキスト生成AIの世界では長い間、GPTシリーズやLLaMAに代表される「自己回帰型」モデルが主流でした。自己回帰型は文章を左から右へ1トークンずつ生成し、各ステップで過去のトークンのみを参照する「因果Attention(Causal Attention)」を使います。生成の自由度は高いものの、一度出力したトークンを見直すことが難しく、後半の文脈から前半を修正するような柔軟な生成は苦手です。
これに対して注目されているのが「拡散型言語モデル(Diffusion Language Model)」です。画像生成でStable DiffusionやFluxが成果を上げたように、言語にも拡散プロセスを適用する研究が進んでいます。PerceptionDLMのように拡散型言語モデルを視覚言語タスクに応用する研究もあり、テキスト以外の領域でも応用範囲が広がっています。
言語向け拡散モデルの中でも「マスク拡散(Masked Diffusion)」は、トークンをランダムにマスクして復元するタスクで学習し、全トークンに対して双方向Attentionを使える点が特徴です。ノイズ除去のプロセスを繰り返しながら文章を生成するため、途中のトークンを後からまとめて修正できる柔軟性があります。ただし、従来の拡散型言語モデルは自己回帰型と比べてベンチマーク性能で大きく見劣りしていました。
ByteDance Seed(北京字節跳動)の研究チームが発表した「iLLaDA(Improved Large Language Diffusion Attention)」は、この差を大規模なデータと設計改善によって縮めることを目指した研究です。
iLLaDAの提案手法
iLLaDAは8Bパラメータの Transformerで、完全双方向Attention(Fully Bidirectional Attention)を採用しています。自己回帰型が左から右にしか文脈を参照できないのに対し、iLLaDAは入力の全トークンを双方向に参照できるため、文全体の文脈を踏まえた予測が可能です。
学習の仕組みはマスク拡散です。入力トークンを一定の確率でランダムにマスクし、モデルがマスクされた全ての位置を同時に予測するよう訓練します。損失はマスク位置のみで計算されます。BERTのマスク言語モデリングに似た考え方ですが、拡散プロセスとして定式化されており、生成時には段階的なノイズ除去としてテキストを出力します。
アーキテクチャの工夫として、計算効率を高める「グループ化クエリAttention(Grouped-Query Attention)」と、入出力の埋め込み行列を共有してパラメータ数を抑える「埋め込み重みの共有(Tied Embeddings)」を採用しています。また、推論時の効率を上げる「可変長生成」と、多肢選択問題の評価精度を高める「信頼度ベースのスコアリング」も実装されています。
事前学習は最大シーケンス長8192トークンで、12兆トークンという大規模なデータを使用しました。学習率スケジュールは線形ウォームアップ後に一定を保ち、途中から余弦減衰に切り替えています。事前学習後は25Bトークンの指示データで12エポックの微調整(SFT)を行い、この段階でも拡散目的関数を維持したまま訓練を続けています。
実験結果
iLLaDAの性能を先行モデルLLaDA 8Bと同規模の自己回帰型モデルQwen2.5 7Bとベンチマーク上で比較した結果は以下のとおりです。
ベンチマーク | iLLaDA 8B | LLaDA 8B | Qwen2.5 7B |
|---|---|---|---|
MMLU | 74.8 | 65.9 | 71.9 |
BBH | 71.3 | 49.7 | 63.9 |
GSM8K | 81.9 | 70.3 | 78.9 |
HumanEval | 50.0 | 35.4 | 56.7 |
先行モデルLLaDAとの比較では、BBHで21.6点、ARC-Challengeで14.9点、HumanEvalで16.5点の大幅な改善を達成しています。微調整後の数学ベンチマーク(MATH)でも14.5点の向上が確認されました。大規模なデータと設計改善の効果が明確に現れた結果といえます。
自己回帰型のQwen2.5 7Bとの比較では、MMLU・BBH・GSM8KでiLLaDAが上回りました。一方、HumanEval(コード生成)ではQwen2.5 7Bに約6.7点及びませんでした。コード生成は逐次的な論理展開が重要なタスクであり、双方向型モデルの相対的な課題として残っています。また、同じ拡散型モデルのDream 7Bとの比較でも、一般知識・数学タスクでiLLaDAの優位性が示されています。
まとめと今後の展望
iLLaDAは拡散型言語モデルが大規模訓練によって自己回帰型と競合できることを、具体的な数値で示した研究です。MMLU・BBH・GSM8Kなどの主要ベンチマークでQwen2.5 7Bを上回り、LLaDAとの差を大幅に縮めたことで、拡散型のスケールアップに一定の有効性があることが裏付けられました。
ただし、研究チーム自身が認める課題も残っています。強化学習(RLHF/RLAIF)による指示追従性能の向上は未実施であり、8Bスケールを超えた実験も行われていません。コード生成での遅れも今後の改善点です。現時点での比較対象が同規模の7〜8Bモデルに限られているため、より大きなモデルとの差がどうなるかはまだ不明です。
コードとモデル重みはGitHubで公開されており、追試や発展研究がしやすい状況が整っています。スケールアップや強化学習との組み合わせによってコード生成での差がさらに縮まるか、今後の動向が注目されます。
