💬 NLP

SPG: Sandwiched Policy Gradient for Masked Diffusion Language Models

本論文は、拡散言語モデル（dLLMs）の対数尤度の非計算可能性という課題を解決するため、真の対数尤度の上限と下限の両方を利用した「サンドイッチ方策勾配（SPG）」を提案し、GSM8K や MATH500 などのベンチマークにおいて既存の手法を大幅に上回る性能向上を実現したことを示しています。

原著者： Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

公開日 2026-04-16

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

CC BY 4.0

原著者： Chenyu Wang, Paria Rashidinejad, DiJia Su, Song Jiang, Sid Wang, Siyan Zhao, Cai Zhou, Shannon Zejiang Shen, Feiyu Chen, Tommi Jaakkola, Yuandong Tian, Bo Liu

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

この論文は、**「新しいタイプの AI（拡散モデル）」を、より賢く、より正確に答えられるように教えるための新しい勉強法（SPG）**について書かれています。

少し専門的な話を、日常の例え話を使って解説しますね。

1. 背景：2 種類の「AI の勉強方法」

まず、AI が文章を作る（生成する）方法には、大きく分けて 2 つの流派があります。

従来の方法（自動回帰型）：
- 例え： 「一文字ずつ、順番に書く」方法。
- イメージ： 手紙を書くとき、1 文字目を書いて、次に 2 文字目、3 文字目…と順番に書いていく感じ。
- 特徴： 確実だが、時間がかかる（一列に並んで待つ必要がある）。
新しい方法（拡散モデル）：
- 例え： 「真っ黒な紙から、少しずつ文字を浮かび上がらせる」方法。
- イメージ： 霧がかかった部屋で、最初は何も見えないが、時間をかけて霧を晴らしていくと、最終的にきれいな文章が見えてくる感じ。
- 特徴： 一度に複数の文字を同時に「晴らす」ことができるので、ものすごく速い！ これが今回の主人公です。

2. 問題点：「正解」を教えるのが難しい

この「新しい AI（拡散モデル）」は速いのですが、**「正解かどうかを評価して、もっと上手にさせたい（強化学習）」**という段階になると、大きな壁にぶつかります。

壁：「この答えが本当に正しいかどうか」を数値で正確に測る計算が、**「難しすぎて計算できない（計算不可能）」**のです。
これまでの対策： 過去の研究では、「計算できないなら、近似値（だいたいの値）を使おう」という方法をとっていました。
- 例え： 「真実の山の高さ」が測れないので、「山の高さの『最低保証値』だけ」を使って勉強させていました。
- 欠点： 「最低保証値」を使うと、「悪い答え（マイナスの評価）」を教えるときに、AI が混乱してしまいます。 「もっと悪くしなさい」と言っても、AI は「最低値」しか見えていないので、どうすればいいか分からず、学習が歪んでしまうのです。

3. 解決策：「サンドイッチ」作戦（SPG）

そこで、この論文の著者たちは**「サンドイッチ・ポリシー・グラディエント（SPG）」という新しい勉強法を提案しました。名前の通り、「サンドイッチ」**のように挟み込むのがポイントです。

アイデア：
- 良い答え（正解に近いもの）が出たとき： 「最低保証値（ELBO）」を使って、**「もっと良くしなさい！」**と褒めて伸ばす。
- 悪い答え（間違っているもの）が出たとき： 「最高保証値（EUBO）」という、逆の指標を使って、**「もっと悪くしないように（正解から遠ざからないように）」**と厳しく指導する。
例え：
- 生徒が良いテストを取ったら、「最低でもこの点数は取れるよ」と言って自信を持たせる（下限を上げる）。
- 生徒が悪いテストを取ったら、「この点数は最高でもこれだよ（もっと下がる可能性がある）」と示して、**「このレベルに落ちないように頑張れ！」**と厳しく指導する（上限を下げる）。
- この「良い時は下限、悪い時は上限」というサンドイッチ構造にすることで、AI は「正解」に近づき、「誤り」から遠ざかる方向へ、より正確に学習できるようになります。

4. 工夫：ブロックごとの学習（ブロック・マスキング）

さらに、この勉強法を安定させるために、**「ブロックごとの学習」**という工夫も加えました。

問題： 文章全体をバラバラに隠して学習させると、AI が混乱しやすい。
解決： 文章を「ブロック（塊）」に分ける。
- 例え： 長文を勉強する際、いきなり全文を隠すのではなく、「前半は読めるように、真ん中を隠す、後半は隠す」というように、**「一部はクリアで、一部は隠す」**という順序で学習させる。
- これにより、AI は「次はここを推測すればいい」という文脈を掴みやすくなり、学習が安定して、より早く上達します。

5. 結果：驚異的な成績向上

この新しい「サンドイッチ勉強法（SPG）」を試した結果、数学や論理パズルのテストで、従来の方法よりも劇的に成績が向上しました。

GSM8K（算数）： 3.6% 向上
MATH500（数学）： 2.6% 向上
Countdown（数字パズル）： 18.4% 向上（大幅！）
Sudoku（数独）： 27.0% 向上（爆発的！）

特に、論理的な思考が必要なパズルや数学の問題で、他のどんな方法よりもはるかに高い精度を達成しました。

まとめ

この論文は、「速いけど勉強の仕方が難しかった新しい AI（拡散モデル）」に対して、「良い時は褒め、悪い時は厳しく指導する『サンドイッチ』作戦」を導入したところ、驚くほど賢くなったという話です。

これにより、AI はより速く、かつより正確に、複雑な問題（数学やパズル）を解けるようになることが期待されています。

SPG（Sandwiched Policy Gradient）の技術的サマリー

本論文は、拡散大規模言語モデル（dLLMs）に対する強化学習（RL）の課題を解決し、より効率的かつ高精度な方策勾配法「SPG（Sandwiched Policy Gradient）」を提案するものです。ICLR 2026 で発表されたこの研究は、dLLM の推論における並列処理の利点を活かしつつ、人間の嗜好やタスク固有の報酬への整合性を高めるための新たなアプローチを示しています。

以下に、問題定義、手法、主要な貢献、結果、および意義について詳細をまとめます。

1. 問題定義：dLLM における強化学習の課題

拡散大規模言語モデル（dLLMs）は、従来の自己回帰（AR）モデルと異なり、複数のトークンを並列にデコードできるため、推論遅延の削減やスケーラビリティにおいて有望です。しかし、dLLM を強化学習（RL）を用いて人間の嗜好やタスク報酬に整合させることは、以下の理由から極めて困難でした。

対数尤度の非計算可能性（Intractability）: dLLM の方策 $\pi_\theta(x|c)$ の対数尤度 $\log \pi_\theta(x|c)$ は、拡散過程の逆過程を解析的に計算できないため、直接算出できません。
既存手法の限界: 従来の RL 手法（GRPO や DPO など）は、この対数尤度の代わりに「証拠下限（ELBO）」や「1 ステップ推定」を代理として使用しています。
- ELBO の問題: ELBO は真の対数尤度の下限（Lower Bound）に過ぎません。正の報酬を持つサンプルに対しては ELBO を最大化することで真の尤度を上げられますが、負の報酬（悪い出力）を持つサンプルに対して ELBO を最小化しても、必ずしも真の対数尤度が減少するとは限りません。
- バイアスの発生: この近似は方策勾配に大きなバイアスを生じさせ、特に負のフィードバック（悪い出力へのペナルティ）を効果的に学習できず、最終的な方策の性能を制限します。

2. 手法：SPG（Sandwiched Policy Gradient）

著者らは、対数尤度の推定バイアスを解消し、安定した RL 学習を実現するために、**「挟み込み（Sandwiched）」**された方策勾配法を提案しました。

2.1 対数尤度の上下限を用いた最適化

SPG の核心は、生成されたシーケンスの対数尤度に対して、報酬の符号に応じて異なる境界（Bound）を最適化することです。

正のアドバンテージ（良い出力）の場合: 既存の手法と同様に、証拠下限（ELBO, $L_{ELBO}$ ）を最大化します。これにより、真の尤度を高める方向に学習を進めます。
負のアドバンテージ（悪い出力）の場合: 従来の ELBO 最小化ではなく、証拠上限（EUBO, $L_{EUBO}$ ）を最小化します。
- EUBO の導出: レニー変分境界（Rényi Variational Bound）に基づき、対数尤度の上限を導出しました。
- 理論的根拠: $L_{ELBO} \leq \log \pi_\theta \leq L_{EUBO}$ であるため、正のサンプルで下限を最大化し、負のサンプルで上限を最小化することで、元の目的関数（真の期待報酬の最大化）に対する有効な下限を構築できます。これにより、負の報酬に対するペナルティが理論的に正当化され、バイアスが低減されます。

2.2 ブロック単位マスキング戦略（Block-Wise Masking）

モンテカルロ推定における EUBO の計算は、サンプリングの不安定性や計算コストの問題を抱えていました。これを解決するため、以下の戦略を採用しました。

ブロック単位でのマスキング: 完全なランダムマスキングではなく、シーケンスをブロックに分割し、特定のブロックをランダムに選択してその中をマスキングする手法を採用します。
分布の整合性: 生成時（rollout）には、dLLM が一般的に採用する「ブロック単位半自己回帰（semi-autoregressive）」デコード戦略を使用します。学習時のモンテカルロ推定でも同様のブロック単位マスキングを行うことで、方策ロールアウトと最適化におけるデータ分布を一致させ、推定の安定性と効率を向上させました。

2.3 境界の混合（Mixture of Bounds）

EUBO のみを使用すると推定が不安定になる場合があるため、負のアドバンテージを持つサンプルに対しては、EUBO と ELBO を混合した目的関数を使用します。

混合係数 $\omega$ : $L_{Mix} = \omega \cdot L_{EUBO} + (1-\omega) \cdot L_{ELBO}$
理論的利点: 提案された理論（Proposition 1）により、最適な混合係数 $\omega^*$ を用いることで、単一の境界（ELBO または EUBO 単独）を使用する場合と比較して、勾配の分散を厳密に減少させることが示されています。これにより、学習の安定性と収束性が向上します。

3. 主要な貢献

新しい方策勾配アルゴリズム（SPG）の提案: 報酬に基づいて対数尤度の上下限を最適化する「挟み込み」アプローチにより、dLLM における RL のバイアスを大幅に低減しました。
ブロック単位マスキング技術: モンテカルロ推定の安定性を高め、方策最適化の効率を向上させる新しいサンプリング戦略を提案しました。
SOTA 性能の達成: 数学的・論理的推論タスクおよびコーディングタスクにおいて、既存の dLLM 向け RL アルゴリズムを凌駕する性能を達成しました。

4. 実験結果

著者らは、LLaDA-8B-Instruct ベースモデルを用いて、4 つの推論ベンチマーク（GSM8K, MATH500, Countdown, Sudoku）および 2 つのコーディングベンチマーク（HumanEval, MBPP）で実験を行いました。

推論タスクでの性能向上:
- GSM8K: 3.6% 向上（SOTA 対比）
- MATH500: 2.6% 向上
- Countdown: 18.4% 向上（大幅な改善）
- Sudoku: 27.0% 向上（劇的な改善）
コーディングタスク:
- HumanEval: 1.9% 向上
- MBPP: 4.7% 向上
ロバスト性: 異なる推論戦略（ブロックサイズやデコード順序の変更）に対しても、SPG は一貫して高い性能を発揮し、既存手法を大きく上回りました。
学習ダイナミクス: 報酬の増加が速く、安定して収束することが確認されました。また、勾配ノルムが小さく安定しており、学習の不安定性が解消されていることが示されました。

5. 意義と結論

本論文は、拡散言語モデルの強化学習における根本的な課題である「対数尤度の非計算可能性」に対する原理的な解決策を提示しました。

理論的貢献: 負の報酬サンプルに対して ELBO を最小化する従来のアプローチの限界を指摘し、EUBO を最小化することで真の尤度を減少させることを保証する理論的枠組みを提供しました。
実用的貢献: 並列デコードという dLLM の最大の利点を活かしつつ、RL による微調整を可能にしました。これにより、数学的推論や論理的思考、コード生成など、複雑なタスクにおいて dLLM の実用性が飛躍的に向上しました。
将来展望: SPG は、dLLM の RL 学習における標準的な手法となり得る可能性があり、より大規模なモデルや複雑なタスクへの適用が期待されます。

総じて、SPG は dLLM の能力を最大限に引き出すための重要なブレイクスルーであり、拡散モデルと強化学習の融合における新たな基準を確立したと言えます。