💬 NLP

Mask-Aware Policy Gradients for Diffusion Language Models

本論文は、Masked Diffusion Language Modelによる生成を、トークンの選択と位置のマスキングを共同で最適化する二段階の意思決定プロセスとして定式化することで、対数尤度の計算困難性を克服し、数学的推論およびコーディングのベンチマークにおいて最先端の性能を達成する強化学習フレームワークである、Mask-Aware Policy Gradientsを導入するものである。

原著者： Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

公開日 2026-07-17

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Haran Raajesh, Kulin Shah, Adam Klivans, Philipp Krähenbühl

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

ロボットに物語を書かせたり、数学の問題を解かせたりする方法を教えようとしている場面を想像してみてください。長い間、最も優れた方法は、人間がキーボードでタイピングするように、ロボットに一単語ずつ書いていく方法でした。これは「自己回帰的（autoregressive）」な生成と呼ばれます。これは信頼性は高いのですが、高速道路を横切るカタツムリのように、非常に低速です。最近、科学者たちは「拡散（Diffusion）」と呼ばれるより速い方法を発見しました。ゼロから書き始める代わりに、ロボットは空白のマスク（[MASK]のようなもの）で埋め尽くされたページからスタートし、文章が完成するまで、一つずつ空白を埋めていくのです。これは、覆われた絵から、カバーを少しずつ持ち上げて画像を見せていくパズルのようなものです。

しかし、この速い「拡散」ロボットに対し、報酬を使って賢くする方法（「強化学習」と呼ばれる手法）を教えることは、悪夢のような作業でした。従来の「一単語ずつ」という手法では、ロボットが特定の選択をした確率を正確に計算することが容易でした。しかし、「拡散」法では、ロボットはあらゆるステップにおいて2つの選択を行います。つまり、「どの単語を空白に入れるか」と、「次にどの空白を明らかにするか」の2つです。これまでの手法は、ロボットに教える際に、最初の選択（単語）だけに着目し、二番目の選択（明らかにする順番）を無視しようとしてきました。これは、シェフに完璧な料理を作る方法を教える際に、選んだ食材についてだけ批評し、それらをどの順番で刻み、混ぜたのかについては完全に無視して教えるようなものです。

COLM 2026カンファレンスで発表されたこの論文は、私たちがパズルの極めて重要な部分を見落としていたことを示唆しています。著者であるHaran Raajesh氏、Kulin Shah氏、そしてテキサス大学オースティン校のチームは、拡散の技術を真にマスターするためには、正しい単語を選ぶことと、それを明らかにする正しい順番を選ぶことの両方に対して報酬を与えなければならないと主張しています。彼らは「マスク認識型ポリシー勾配（Mask-Aware Policy Gradients）」と呼ばれる新しい手法を開発しました。このシステムは、単に次の単語を推測するのではなく、「次にどのマスクを持ち上げるか」という決定を、ロボットの戦略の極めて重要な一部として扱います。学習プロセスをこれら2つの明確な部分に数学的に分解することで、ロボットがより速く学習し、ミスが減ることを発見しました。

その結果は非常に印象的です。彼らがこの新しい手法を難解な数学問題やコーディングの課題でテストしたところ、ロボットの解決能力は大幅に向上しました。有名な数学テストであるGSM8Kにおいて、彼らの手法は87.1%の精度に達し、コーディングテストであるMBPPでは53.4%を記録しました。これらの数値は、報酬を用いて拡散モデルを教えようとしたこれまでのどの手法よりも高いものです。この論文は、「マスキング」の決定、つまりロボットが自分の思考を明らかにする順番に注意を払うことで、これらの速く柔軟なAIモデルの新たなレベルの知能を解き放つことができることを示しています。これにより、モデルは単に速くなるだけでなく、より賢くなれるのです。

技術要約：拡散言語モデルのためのマスク認識型方策勾配（Mask-Aware Policy Gradients）

問題提起

強化学習（RL）は、実行可能な対数尤度に基づいて方策勾配を最適化することにより、自己回帰型の大規模言語モデル（LLM）における推論能力を向上させることに成功してきました。しかし、これらの手法を**マスク拡散言語モデル（MDLM）**に拡張することは依然として困難です。左から右へと逐次的にトークンを生成する自己回帰モデルとは異なり、MDLMは、完全にマスクされたシーケンスから反復的にトークンをアンマスク（非マスク化）することでテキストを生成します。

核心的な困難は、MDLMにおける対数尤度推定の計算不可能性（intractability）にあります。既存のアプローチは、証拠下限（ELBO）の変種を用いるか、あるいは生成軌道のトークン予測コンポーネントのみをモデリングすることで、この尤度を近似しています。これらの手法は、MDLMの生成における重要な側面、すなわちどの位置をアンマスクするかという順序を無視しています。MDLMにおいて、生成の軌道には2つの明確な決定が含まれます。(1) マスクされた位置にどのトークンを配置するか、および (2) それらの位置のうち、どれを公開し（アンマスクし）、どれを再びマスク状態に戻すか（リマスクするか）という決定です。この第2の決定を無視する先行研究は、方策最適化のために利用可能な生成プロセス内の完全な信号を活用できていません。

手法

著者らは、MDLMの生成を**2段階のアクション・マルコフ決定過程（MDP）**として定式化しています。各デノイジング・ステップ $t$ において、モデルは以下を行います：

トークン予測（Token Prediction）: 現在マスクされているすべての位置に対してトークンを予測する。
位置選択（Position Selection / Remasking）: 予測された位置のサブセットのうち、どれを明らかにするか、およびどれを再び [MASK] 状態に戻すかを決定する。

このプロセスに対して方策勾配最適化を可能にするため、著者らは**確率的リマスキング（Probabilistic Remasking）**を導入しています。標準的なMDLMの推論では、多くの場合、信頼度スコアに基づく決定論的な貪欲なトップ- $K$ 戦略を使用して、アンマスクする位置を選択します。この貪欲な選択は微分不可能であり、アンマスク順序の勾配ベースの最適化を妨げます。提案手法は、これを以下の確率的バリアントに置き換えます：

決定論的にトップ- $K$ の位置を選択する代わりに、モデルはプラケット・ルーチェ（Plackett–Luce）分布から位置のサブセット $U_t$ をサンプリングします。
この分布の確率は、モデル自身のトークン対数尤度（ロジット）から直接導出され、温度パラメータ $\tau$ によってスケーリングされます。
この定式化により、位置のサブセットに対する微分可能な分布が作成され、アンマスク決定の対数確率を生成軌道の尤度に含めることが可能になります。

方策勾配の分解

全軌道の尤度 $\pi_\theta(\hat{z}|c)$ を、トークン予測確率とアンマスク確率の積としてモデリングすることで、方策勾配は自然に以下の2つの項に分解されます：
$\nabla_\theta J(\theta) = \mathbb{E}_{\hat{z}} \left[ R(c, z_T) \sum_{t=1}^T \left( \underbrace{\nabla_\theta \log \pi_\theta(\hat{z}_t | c, z_{t-1})}_{\text{トークン勾配}} + \underbrace{\nabla_\theta \log p_{\text{unmask}}(U_t | \hat{z}_t, z_{t-1}, c)}_{\text{マスキング勾配}} \right) \right]$

トークン項（Token Term）: 予測されたトークンの正しさ（正確性）を強化します。
マスキング項（Masking Term）: アンマスクの順序の選択を強化します。
効率性: マスキング項は、トークン予測に使用される既存のモデルのロジットから計算されるため、追加のパラメータ、アーキテクチャの変更、または追加のフォワードパスを必要としません。

著者らはこれをグループ・シーケンス方策最適化（GSPO）アルゴリズムを用いて実装しており、トークン項と位置項の両方に対して個別に重要度比をクリッピングしています。全軌道の評価に伴う計算コストに対処するため、デノイジング・ステップをセグメントにグループ化して、軌道の構造を維持しながらフォワードパスの回数を削減するStepMerge近似を採用しています。

主な貢献

MDLMを2段階MDPとして定式化: 本論文は、MDLMの生成にはトークン選択と位置選択の両方が含まれ、方策勾配がそれに応じて分解されることを示す理論的枠組みを提供しています。
確率的リマスキング: モデル自身のロジットに基づく微分可能なアンマスク分布（プラケット・ルーチェ）を導出し、アーキテクチャの変更なしにアンマスク順序の最適化を可能にしました。
理論的洞察: 位置コンポーネントの方策勾配を無視すると、たとえトークン確率が変化しなくても、期待報酬を向上させる方向を見逃す可能性があることを実証しました。
最先端の性能: 提案手法が、既存の軌道ベースおよびELBOベースの手法を一貫して上回る、新しいSOTA（State-of-the-Art）の結果をもたらすことを実証しました。

実験結果

本手法は、数学的推論（GSM8K, MATH500）およびコード生成（HumanEval, MBPP）のベンチマークにおいて、LLaDA-8B-Instructをベースモデルとして評価されました。

性能向上: 提案手法は、すべてのベンチマークにおいて新しいSOTAの結果を達成しました。
- GSM8K: 正解率 87.1%（生成長 128において）、前述の最高値（StepMerge）に対し +2.5% の改善。
- MATH500: 正解率 53.4%、+4.0% の改善。
- HumanEval: 正解率 47.1%、+2.2% の改善。
- MBPP: 正解率 53.4%、+2.5% の改善。
アブレーション研究:
- ブロックサイズ: 提案手法とベースラインの性能差は、推論ブロックサイズが大きくなるにつれて（32からフルシーケンスへ）拡大しました。これは、位置選択の最適化による恩恵が、アンマスク決定の複雑さが増すにつれて大きくなることを示唆しています。
- 効率性: 提案手法は、強力なELBOベースのベースラインであるSPGと比較して、実時間（wall-clock time）でより速く収束しました。SPGが最終的な精度に達するのに約18時間を要したのに対し、提案手法は約15時間で到達し、かつ最終的な精度も高くなりました（GSM8Kにおいて 81.0% 対 78.5%）。
- 汎用性: 本アプローチは、異なるベースモデル（Dream-7B）やプランニングタスク（Sudoku, Countdown）にも適用可能であり、一貫してベースラインを上回りました。
DCoLTとの比較: 位置選択のための別個のヘッドを使用するLLaDOUと比較して、提案手法は大幅に低いトレーニングコスト（約160 GPU時間対約800 GPU時間）で、同等またはそれ以上の精度を達成しました。これは、追加のモジュールを必要としないためです。

意義と主張

本論文は、MDLMの柔軟性は、自己回帰モデルよりも豊かな軌道構造（具体的には、いつ位置をアンマスクするかという決定）を伴うものであると主張しています。この決定を方策の一部として扱い、微分可能な確率的フレームワークを通じて最適化することで、著者らは、トークン予測と位置選択の共同最適化が、強化学習設定における拡散言語モデルの性能を最大化するために極めて重要であることを示しました。

著者らは、本手法が実用的に効率的であることを強調しています。なぜなら、新しいパラメータやアーキテクチャの変更を必要とせず、既存のモデルのロジットを活用しているため、StepMergeのような標準的なMDLMトレーニング・インフラストラクチャと互換性があるからです。実験結果は、「マスキング勾配」が、トークンの尤度のみに焦点を当ててきた従来の手法では捉えられなかった、意味のある学習信号を提供していることを示唆しています。