拡散言語モデルのためのマスク解除ポリシーの学習
拡散大規模言語モデル（dLLM）は、多くのタスクで自己回帰モデルと同等の性能を達成しつつ、推論効率が向上しています。その重要な設計要素の一つが、各拡散ステップでどのトークンをマスク解除するかを選択するサンプリング手順です。現在のヒューリスティック手法は手動調整が必要で、ブロックサイズが大きくなると性能が低下します。本研究では、強化学習を用いてサンプリングポリシーを訓練することを提案し、マスク付き拡散サンプリングをマルコフ決定過程として形式化し、軽量な単層トランスフォーマーポリシーを採用します。実験では、訓練されたポリシーが半自己回帰（ブロック）生成で最先端のヒューリスティックと同等の性能を示し、全拡散設定ではそれを上回ります。
拡散大規模言語モデル（dLLM）は、近年自然言語処理の分野で急速に発展しており、特に生成タスクにおいて優れた性能を発揮しています。これらのモデルは複数のトークンを並列にデコードすることで、自己回帰モデルと比較して推論効率を大幅に向上させることができます。しかし、dLLMの重要な設計上の課題の一つはサンプリングプロセスにあります。各拡散ステップでどのトークンをマスク解除（生成）するかを決定する必要があり、従来は信頼度しきい値などのヒューリスティック手法が用いられてきました。これらの手法はランダムなマスク解除と比較してサンプル品質とトークンスループットを改善しますが、手動調整が必要であり、ブロックサイズが大きくなると性能が顕著に低下することが観察されています。
この問題に対処するため、アムステルダム大学、マサチューセッツ工科大学、およびAppleの研究者らは、強化学習（RL）を活用した新しいアプローチを提案しました。彼らはマスク付き拡散サンプリングをマルコフ決定過程（MDP）として形式化し、拡散言語モデル自体を環境と見なします。そして、軽量な単層トランスフォーマーポリシーを設計し、モデルが出力する各トークンの信頼度をマスク解除の決定にマッピングします。このポリシーは強化学習によって訓練され、生成されるサンプルの品質と効率を最適化します。手動調整は不要です。
実験では、訓練されたポリシーを最先端のヒューリスティック手法と比較しました。半自己回帰（ブロック）生成の設定では、訓練ポリシーはヒューリスティック手法と同等の性能を達成しました。さらに、全拡散（各ステップで1トークンのみデコード）の設定では、訓練ポリシーはヒューリスティック手法を上回る結果を示しました。この結果は、強化学習によるサンプリングポリシーの学習が手動調整の必要性を排除するだけでなく、より広範な設定で性能を向上させることを示しています。
本研究は、拡散言語モデルのサンプリング戦略に新たな視点を提供し、コード生成やテキスト要約などのタスクへの応用が期待されます。関連成果は2026年7月にICML会議で発表されました。