拡散言語モデルに強化学習の新手法、GSM8Kで87%達成

a close up of a sheet of paper with numbers on it AI

UT Austinなどの研究チームが、拡散言語モデル(文章全体の空欄を少しずつ埋めていく方式で文章を生成するAIモデル)向けに、強化学習の新手法「Mask-Aware Policy Gradients」を提案しました。生成する単語だけでなく「どの順番で空欄を埋めるか」も学習対象に含めたのが特徴です。算数の文章題ベンチマークGSM8Kで87.1%、Pythonコード生成ベンチマークMBPPで53.4%という高い成績を達成し、言語モデル研究の国際会議COLM 2026に採択されています。

背景と文脈

ChatGPTやClaudeなど、現在主流の大規模言語モデル(LLM)は、文章の先頭から1単語ずつ順番に生成する「自己回帰型」という方式を採っています。これに対し拡散言語モデルは、画像生成AIで使われる拡散モデルの発想を文章に応用したもので、最初は文章全体を空欄(マスク)で埋めておき、それを少しずつ具体的な単語に置き換えていきます。理屈の上では複数の空欄を同時並行で埋められるため、生成速度を上げられる可能性を持つ方式として近年注目されてきました。

一方で、自己回帰型モデルの性能を大きく伸ばしてきた強化学習(RLHF: 人間のフィードバックに基づく強化学習、などが代表例で、モデルの出力を評価して報酬を与えることで賢くする学習法)を拡散言語モデルにそのまま適用しようとすると、うまくいかないことが分かっていました。既存の手法は「どの単語を出力したか」だけを評価対象にしており、拡散言語モデル特有の「どの順番で空欄を埋めていくか」という生成プロセスの重要な部分を見落としていたためです。論文の著者らは、この見落としこそが拡散言語モデルの強化学習が伸び悩んできた一因だと指摘しています。空欄を埋める順序自体にも「先に決めておくべき単語から埋める」といった巧拙が存在するはずですが、従来の学習方法ではその巧拙を報酬に反映できていませんでした。

技術/ビジネス面

man and woman standing near laser light
Photo by Matthieu Joannon on Unsplash

研究チームが提案したのは、拡散言語モデルの生成過程を「トークン(単語の最小単位)を決める判断」と「次にどの位置の空欄を埋めるか決める判断」という2種類の意思決定からなる過程として捉え直すアプローチです。この2段階構造をモデル化し直すことで、学習時に計算する勾配(モデルをどちらの方向に調整すべきかを示す指標)が、トークンに関する項と、空欄を埋める順序に関する項へ自然に分解できることを示しました。

この2つの項を同時に最適化することで、数学推論のGSM8Kで87.1%、コード生成のMBPPで53.4%という結果を達成しています。生成順序という、これまで見過ごされがちだった要素まで学習対象に含めたことが、性能向上の主な要因だと研究チームは説明しています。単に正しい単語を選ぶだけでなく、どの順番で組み立てるかという「書き方」まで強化学習で磨き上げられる、という点が今回の技術的な肝と言えます。

論文は今年開催される言語モデル研究の国際会議COLM 2026に採択されており、査読を経た成果として研究コミュニティでの検証が進む見込みです。著者にはUT Austinのコンピュータサイエンス教授Adam Klivans氏、Philipp Krähenbühl氏らが名を連ねており、理論と実装の両面から裏付けを与えている点も評価のポイントになりそうです。

これからどうなるか

拡散言語モデルは、複数の空欄を並行して埋められる特性上、将来的には自己回帰型より高速に文章を生成できる可能性を秘めています。ただしこれまでは、数学やコーディングのような複雑な推論を要するタスクで、自己回帰型モデルに一歩譲るケースが目立っていました。今回のような強化学習手法の改良が積み重なれば、この差は縮んでいきそうです。

開発者にとっても他人事ではありません。もし拡散言語モデルが推論性能で自己回帰型に追いつけば、並列生成による低レイテンシーを活かしたコード補完ツールや、手元のCI環境で高速に動く軽量アシスタントが現実味を帯びてきます。特にMBPPのようなコード生成タスクでの性能向上は、日々の開発ワークフローに直結しやすい成果です。現時点ではまだ研究段階の成果であり、商用サービスに組み込まれるまでには追加の検証が必要ですが、拡散言語モデル陣営のキャッチアップが進んでいる動きとして押さえておく価値がありそうです。

まとめ

UT Austinなどの研究チームが、拡散言語モデル向けの強化学習手法「Mask-Aware Policy Gradients」を提案しました。生成する単語だけでなく空欄を埋める順序まで学習対象に含めることで、GSM8Kで87.1%、MBPPで53.4%という成績を達成しています。自己回帰型モデルに対する拡散言語モデルの弱点を埋める一歩として、今後の追試や実装への波及が注目されます。

参考リンク

アイキャッチ画像: Photo by Bozhin Karaivanov on Unsplash

タイトルとURLをコピーしました