🤖 AI

Masked Diffusion Language Models are Strong and Steerable Text-Based World Models for Agentic RL

本論文は、大規模なデータセットと新しいGRPOトレーニングフレームワークに支えられ、自己回帰的なバイアスを克服することで、より大規模なLLMと比較して著しく高いロールアウトの多様性とゼロショット転移性能を実現する、エージェンティック強化学習のための優れた制御可能なテキストベースの世界モデルとして、Masked Diffusion Language Models (MDLM) を導入するものである。

原著者： Darshan Deshpande

公開日 2026-07-21

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

CC BY 4.0

原著者： Darshan Deshpande

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

あなたは、ロボットに複雑なビデオゲーム（例えば、刻々と変化する巨大な都市のナビゲーションや、忙しいコーヒーショップの運営など）を教えようとしていると想像してください。学習するためには、ロボットは何百万回もの練習を行う必要があります。現実の世界では、練習のたびに物理的なコーヒーショップを建設しなければなりませんが、それは不可能です。そこで科学者たちは、「シミュレーター」を使用します。これは、ロボットが何かを壊すことなく、試行錯誤し、学習できるデジタル世界です。これが、強化学習（RL）の核心です。つまり、サンドボックスの中で試行錯誤しながら学ぶAIの学習プロセスです。

しかし、ここに落とし穴があります。これらのサンドボックスを構築するのは困難です。通常、人間がすべてのルールやシナリオを手作業で作成しなければならず、それは時間がかかる上に、ロボットが見ることができる状況の幅を制限してしまいます。もしロボットがある特定のシナリオに習熟しすぎると、思考する方法を実際に学ぶのではなく、正確な手順をそのまま暗記することで「ズル」をし始めます。これは、練習テストの答えを丸暗記したものの、問題が少し変わっただけで本番の試験に失敗してしまう学生のようなものです。これを解決するために、研究者たちは「ワールドモデル（世界モデル）」、つまり、自律的にデジタル世界を想像し、シ Simulate（シミュレート）できるAIシステムを構築しようとしています。これにより、無限に多様な練習シナリオを生み出すことができます。大きな疑問は、これらのAIシミュレーターは、混乱することなく、複雑でルールが多い環境において次に何が起こるかを予測できるほど賢くなれるのか、という点です。

本論文は、Masked Diffusion Language Model (MDLM) と呼ばれる新しい種類のAIシミュレーターを紹介し、それが通常使用される標準的なAIモデルよりも優れた「ゲームマスター」であることを主張しています。

問題点：一方通行の道 vs 二方向の道

なぜこれが重要なのかを理解するために、標準的なAI（**自己回帰（Autoregressive/AR）モデルと呼ばれます）がどのように考えるかを想像してみてください。それは、左から右へと一文字ずつ物語を書く人のようです。一度言葉を書いたら、戻って書き直すことはできません。「銀行はカードを拒否しました」と書いたとしたら、その文章に縛られてしまいます。もし後になって「銀行はカードを承認しました」と書くべきだったと気づいても、もう手遅れです。物語はすでに壊れてしまっています。複雑なシミュレーションにおいて、これはグローバルな不整合（global incoherence）**と呼ばれる問題を引き起こします。AIは完璧な始まりを書くかもしれませんが、最後まで到達する頃には、細部が世界のルールと一致しなくなってしまうのです。それは、犯人を特定したものの容疑者の名前を忘れてしまった探偵や、料理を作ったもののコンロの火を消し忘れたシェフのようなものです。

著者らは、この標準的な「左から右へ」という考え方がボトルネックになっていると示唆しています。彼らは異なるアプローチ、すなわち**Masked Diffusion（マスク付き拡散）**を提案しています。物語を最初から最後まで書く代わりに、テキストのページがあり、いくつかの単語が黒い箱（マスク）で隠されている状況を想像してください。あなたの仕事はその隠された単語を推測することですが、ここでのコツは、箱の「前」にある単語と「後」にある単語の両方を見て推測できるという点です。全体像をガイドとして使いながら、中間、最後、あるいは最初を同時に埋めることができるのです。これにより、AIは始まりが終わりと一致しているかを確認でき、物語全体が整合性を持つようにすることができます。

実験：より優れたサンドボックスの構築

研究者たちは単に議論しただけでなく、実際に構築しました。彼らは、コーディングツールからカスタマーサービスに至るまで、9つの異なる環境における239,403もの異なる「練習の軌跡（トライアクトリー）」を含む膨大なデータセットを作成しました。そして、この新しいMasked Diffusionモデルをこのデータで訓練し、利用可能な最高の標準的な「左から右へ」のモデルと比較検証しました。

結果として、新しいモデルは驚くほど強力であることが分かりました。Masked Diffusionモデルは、比較対象となった標準的なモデル（300億以上のパラメータを持つものも含む）よりもはるかに小規模（わずか80億パラメータのものもある）であったにもかかわらず、より一貫性があり、現実的なシミュレーションを生成しました。

これを例えるなら、標準的な巨大モデルは、一文字ずつタイピングする非常に速いタイピストのようなものです。もし早い段階で打ち間違いをすると、文章全体が台無しになります。新しいMasked Diffusionモデルは、粘土の塊全体を一度に見つめ、最終的な彫刻がどの角度からも正しく見えるように、同時にさまざまな箇所を削っていく彫刻家のようなものです。

結果：より賢い練習、より優れたプレイヤー

本当のテストは、これらのシミュレーターを使用して実際のAIエージェント（タスクを学ぼうとしている「ロボット」）を訓練した時に行われました。彼らはMasked Diffusionモデルによって作成された世界でエージェントを訓練し、その後、全く未知の環境（ScienceWorld、ALFWorld、AppWorldなど）でテストを行いました。

結果は驚くべきものでした。新しいMasked Diffusionシミュレーターで訓練されたエージェントは、従来の標準的なモデルで訓練されたエージェントと比較して、成功率が最大**47%**向上しました。これは、研究者が新しい環境について具体的な教示を与えなくても達成されました。エージェントは、練習環境がより多様で論理的であったため、より優れた汎用スキルを自然に学習したのです。

例えば、通常は複雑なタスクの成功率が極めて低い（わずか5.7%）12億パラメータの小さなAIエージェントが、新しいシミュレーターでの訓練後、成功率が**53.6%**へと跳ね上がりました。これは劇的な飛躍であり、シミュレーターがより豊かで多様な挑戦を提供したことで、エージェントが適応する方法をより深く学べたことを示唆しています。

なぜ重要なのか

本論文は、AIシミュレーターの構築方法を変える必要があると示唆しています。従来の「左から右へ書く」手法は限界に達しており、ループに陥ったり、シミュレーションを壊す論理的エラーを引き起こしたりしています。未来を「一方向の道」から「二方向の道（双方向のデノイジング）」へと変えることで、よりスマートで、かつ効率的なシミュレーターを構築できます。

著者らはまた、「人間によるチェック」も実施し、専門家にシミュレーションの評価を依頼しました。Masked Diffusionモデルは、リアリズム（5点満点中4.75）と正確性（4.25）において非常に高いスコアを獲得し、人間がシミュレートされた世界を信憑性があり有用であると感じたことを示しました。

しかし、論文ではこれがまだ完璧で完成された製品ではないことも注意深く述べています。新しいモデルは、正しいAPIキーの生成や、非常に長いアイテムリストの処理などで混乱するなど、特定の詳細において依然として課題を抱えています。しかし、核心となる発見は明白です。AIが未来をどう「考える」か（一方通行の道から二方向の道へ）を変えることで、より優れたトレーニングの場を構築できるのです。これは、ソフトウェアのバグ修正からカスタマーサービスの管理まで、複雑な仕事をこなせる、より速く、より賢い次世代のAIエージェントの実現につながります。なぜなら、彼らはより論理的な世界で練習してきたからです。

テクニカル・サマリー：ステアラブルなテキストベース・ワールドモデルとしてのマスク型拡散言語モデル

問題提起

エージェント・システムの強化学習（RL）には、多様で適応的なトレーニング環境が必要である。しかし、現在の環境は固定されたタスクや報酬の難易度を持つ手作業によるもの（hand-curated）であることが多く、モデルの向上に伴い、効果的な学習信号が得られなくなるという課題がある。エージェントがロールアウト中に直面する状況の多様性が減少すると、特定のワークフローやツールの構造に過学習するモード崩壊（mode collapse）が発生する。ワールドモデル（将来の状態を予測する生成モデル）は、トレーニングデータを多様化するためのスケーラブルな解決策を提供するが、ワールドモデルとして使用される既存の自己回帰（AR）言語モデルは、根本的な左から右へのバイアスを抱えている。この因果的な制約により、グローバルに相互依存する状態アンカー（例：ツールのスキーマ、過去のターン、期待される結果）に対して同時に条件付けることができず、その結果、接頭辞（prefix）としては一貫しているものの、全体としては不整合な出力（例：返金アクションの成功とエラーステータスの発生が同時に出力されるなど）を生成してしまう。

メソドロジー

1. テキストベース・ワールドモデリングの定式化

著者らは、ワールドモデリングをステアラブルな遷移ダイナミクス問題として定式化した。彼らは環境状態 $s_t$ を、以下の5つのグラウンデッドな構成要素に分解している：

$e_0$ : 初期環境状態（ファイル、データベース、UI）。
$c_t$ : タスク・コンテキスト（目標、履歴）。
$\mathcal{T}$ : ツール・スキーマ（シグネチャおよび制約）。
$\mathcal{R}$ : ドメイン・ルール（有効な遷移）。
$D$ : ステアリング指令（安全性制約、高レベルな指示）。

目的は、 $\mathcal{T}, \mathcal{R}, D$ が軌道を通じて固定されている条件下で、分布 $p_\theta(c_{t+1} | e_0, c_{\le t}, a_{\le t}; \mathcal{T}, \mathcal{R}, D)$ を学習することである。

2. マスク型拡散言語モデル (MDLMs)

本論文では、AR生成の代わりにマスク型拡散言語モデル (MDLMs) を採用している。

メカニズム: MDLMsは、破損したトークンシーケンスを反復的にデノイジング（ノイズ除去）することを学習する。学習中、クリーンな次状態のシーケンスは確率 $\tau$ でマスクされ、モデルは履歴および（将来のアンカーを含む）未マスクされたコンテキストに基づいて、マスクされたトークンを予測する。
ARに対する利点: ARモデルが左から右への順序にコミットしなければならないのに対し、MDLMは双方向のデノイジングを行う。これにより、予測される状態を、状態内の先行するコンテキストと後続のコンテキストの両方に条件付けることができる（例：既知の status: "error" アンカーを使用して、error_code や resolution の生成を制御する）。これにより、トークンの並列生成が可能になり、接頭辞に起因するエラーの蓄積を回避できる。
多様性: デノイジング・プロセスは、トークンの選択と生成順序の両方に確率性を導入するため、ARデコーダーで一般的な「接頭辞モード崩壊（prefix mode collapse）」を防ぐことができる。

3. データセットおよびトレーニング・フレームワーク

データセット: 9つのオープンソース環境（SWE-SMITH, TAUBENCH, TOOLATHLONなど）と12のフロンティア・モデルファミリー（Qwen, GPT, Claude, Geminiを含む）にわたる、精選された239,403個のグラウンデッドな状態・行動軌跡。
トレーニング: コサイン・スケジュールを用いたAdamWを使用してファインチューニングを行った。著者らは、MDLMによって生成されたロールアウト上でRLエージェントを訓練するために、決定論的な状態チェックを備えたプラグアンドプレイ型のGRPO (Group Relative Policy Optimization) トレーニング・フレームワークを導入している。

主な貢献

定式化とデータセット: テキストベースのワールドモデリングを、5つの特定の構成要素を持つステアラブルな遷移問題として体系的に定式化した最初の研究であり、23.9万の軌跡を含む大規模なデータセットを伴う。
比較研究: AR言語モデルとMDLMの包括的な比較を行い、MDLMが双方向のアンカー認識デノイジングを通じて、優れた一貫性とグラウンデッドネスを実現することを実証した。
性能向上: MDLMs（特にSDAR-8Bのような8Bパラメータモデル）が、ワールドモデリングの忠実度において、4倍のパラメータ数を持つARモデル（例：Qwen-3.5-35B）を凌駕することを実証した。
ゼロショット転移: 環境固有のファインチューニングなしに、OOD（分布外）環境（ScienceWorld, ALFWorld, AppWorld）でRLエージェントを訓練するための、MDLM生成ロールアウトの適用に成功した。
人間による評価: リアリズム、正確性、有用性を検証するための厳格な人間による評価を行い、MDLMのロールアウトの信頼性を確認した。

結果

生成の忠実度 (RQ1)

指標: BLEU-1、ROUGE-L、およびMAUVE（分布の整合性）を用いて評価。
知見: ゼロショット設定において、MDLMは分布の整合性（MAUVE）においてARのカウンターパートを上回った。ファインチューニング後、MDLMはすべての分割において優位に立った。
- SDAR-8B（8Bパラメータ）は、インドメインMAUVE（0.982 vs 0.932）およびOCCUBENCH MAUVE（0.979 vs 0.960）において、Qwen-3.5-35B-A3B（35Bパラメータ）を上回った。
- これは、LLMによるワールドモデリングにおける主要なボトルネックが、モデルの容量ではなく、因果的な帰納バイアスであることを示唆している。

下流のRL性能 (RQ2)

セットアップ: Qwenベース（AR）またはSDARベース（MDLM）のワールドモデルからのロールアウトを用いて、GRPOによるエージェント（LFM2.5-1.2B, Qwen3-4B, Mistral-7B）を訓練した。
知見: SDAR-WMのロールアウトを用いた訓練は、ベースラインおよびQwen-WMのロールアウトと比較して、大幅な絶対的改善を示した：
- ALFWorld: LFM2.5-1.2Bは、SDAR-WMを使用することで、5.7%（ベース）から53.6%（+47.9% 絶対増加）へと向上した。
- ScienceWorld: Mistral-7Bは、3.3%から48.4%（+45.1% 絶対増加）へと向上した。
- これらの改善は、テスト環境に関するゼロショット設定において発生した。

人間による評価 (RQ3)

指標: リアリズム、結果の正確性、および訓練の有用性（1–5のリッカート尺度）。
結果: SDARは高いスコアを記録した：4.75（リアリズム）、4.25（正確性）、4.50（有用性）。
一致度: 全ての次元においてクリッペンドルフの $\alpha$ が0.89を超え、強力なアノテーター間の一致を示した。
失敗モード: アノテーターは、数値キーの型強制（coercion）の問題、APIキーの破損、および初期失敗後に繰り返しエラーに陥る傾向を指摘した。

意義と主張

本論文は、マスク型拡散言語モデルが、オートレグレッシブ（自己回帰）モデルと比較して、テキストベースのワールドモデリングのためのより優れたアーキテクチャであることを主張している。双方向の条件付けを活用することで、MDLMはツールのスキーマやドメインルールとのグローバルな一貫性を維持し、多様でグラウンデッドな環境シミュレーションを生成することができる。

著者らは、このアプローチがエージェントRLのためのスケーラブルで多様かつステアラブルな学習信号を可能にすると断言している。極めて重要な点は、高精度なワールドモデリングには膨大なパラメータ数が必要ではなく、むしろアーキテクチャの選択（拡散 vs 自己回帰）が決定的な要因であるということを実証したことである。本研究は、「想像上の世界でのRL」への道筋を提供し、エージェントが手作業で作られた静的なトレーニング場を必要とせず、シミュレートされた多様な環境で訓練できるようにするものである。

認められた限界: 論文では、MDLMは有望ではあるものの、APIキーのような構造化されたフィールドの扱いに依然として苦慮しており、複雑なナビゲーションタスクにおいて「ページネーションのドリフト（pagination drift）」を示すことが指摘されている。これらの問題は、ベースモデルの安全性アライメントや、拡散におけるブロックサイズの制約に起因しており、将来的な最適化の領域として示唆されている。