💬 NLP

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models

本論文では、マスクされた拡散言語モデルにおける直接的なトークン修正を可能にするパラメータフリーのサンプラーであるD3IMと、それによって生じる保存バイアスを軽減するためのポストトレーニング手法であるSCOPEを紹介し、これらを総じて推論およびコーディングのベンチマークにおいて大幅な性能向上を実現する。

原著者： Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu, Yue Dong, Greg Ver Steeg

公開日 2026-06-02

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

CC BY 4.0

原著者： Longxuan Yu, Shaorong Zhang, Yu Fu, Hui Liu, Yue Dong, Greg Ver Steeg

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

大きな問題：「凍結された間違い」

想像してみてください。あなたは、非常に賢いけれど少し神経質なロボットの助手と一緒に物語を書いているとします。ロボットが単語を推測するたびに、それは紙の上にその単語を書き込みます。

これらのロボットの標準的な仕組み（Masked Diffusion Language Modelsと呼ばれます）では、一度ロボットが単語を書き込み、そこに「自信」を持った時点で、その単語をその場に**凍結（フリーズ）**させてしまいます。ロボットはその単語の上にテープを貼ってしまうのです。たとえ後になって、「待てよ、この文章の中でこの単語は意味が通じないぞ」と気づいたとしても、テープが貼ってあるため、その単語を変えることはできません。ロボットは、まだ書き込まれていない「空白」の部分を埋めることしかできないのです。

これは、数学やコーディングのような難しいタスクにおいて大きな問題となります。もしロボットが早い段階で間違った数字を書いてしまった場合（例えば、「2」と書くべきところを「5」と書いてしまった場合）、ロボットは行き詰まってしまいます。ロボットはその間違った数字を前提として残りの答えを組み立てようとするため、結果として完全に間違った答えに辿り着いてしまうのです。

解決策：修正のための2つのステップ

著者であるLongxuan Yu氏とその仲間たちは、ロボットには実際には「考えを変える能力」があるものの、ゲームのルールがそれを妨げていることに気づきました。彼らは、この問題を解決するために、新しいゲームの遊び方（D3IM）と、ロボットのための特別なトレーニング演習（SCOPE）という、2つの解決策を導入しました。

パート1：新しいゲームのルール (D3IM)

比喩： 全員が文章の次の単語を予想する「ホットポテト（熱いジャガイモ）」ゲームを想像してください。

古いルール： 一度単語を叫んだら、最後までそれを持ち続けなければなりません。もし間違った単語を叫んでしまっても、もう取り返しがつきません。
D3IMのルール（「白紙の状態」）： ゲームのあらゆるステップにおいて、ロボットはすでに叫んだ単語も含め、文章内のすべての単語を見渡します。そして、「自分はこの単語が今でもベストだと思うか？」と自問自答します。
- もしロボットがまだ自信を持っているなら、その単語はそのまま残ります。
- もしロボットが「いや、これは違う」と思ったなら、その単語を消去（再び空白に戻す）するか、即座により良い単語に差し替えます。
- 特別なツールを使ったり許可を求めたりする必要はありません。現在の自信に基づいて、すべてを再評価するだけです。

落とし穴： ロボットはこのゲームが得意ですが、それは「自分が間違っているとき」を知っている場合に限られます。しかし、元のトレーニングでは、ロボットは自分の間違いを認める方法を教わっていませんでした。むしろ、たとえその推測が悪くても、最初の推測を信じるように訓練されていたのです。そのため、新しい「白紙の状態」のゲームをさせても、ロボットはそれが正しいと思い込んで、間違った単語を凍結させ続けてしまうのです。

パート2：特別なトレーニング (SCOPE)

比喩： これは「間違いを含んだリハーサル」のようなものです。
ロボットの「自分のエラーを信じすぎてしまう」という悪い癖を直すために、著者らはSCOPE（Self-Conditioned On Prediction Errors）と呼ばれるトレーニング手法を作成しました。

仕組み： トレーニング中、ロボットは文章を書くよう求められますが、その後、自分の書いたものを見て、自分が間違いを犯したふりをすることを強制されます。
ロボットは文章を書きます。次に、トレーナーが言います。「よし、君の単語をいくつか隠すよ。さあ、それらをもう一度予想してみて」。
ひねり： トレーナーは、ロボットが間違えたにもかかわらず、非常に高い自信を持っていた単語を特別に選び出します。ロボットは、自分の間違った答えを見て、それが間違いであると認識し、それを正しいものに置き換えなければなりません。
これによって、ロボブルットは重要な教訓を学びます。「自信満々に言ったからといって、それが正しいとは限らない。私は考えを変える準備をしておく必要があるのだ」。

結果：共に機能するチーム

新しいゲームのルール（D3IM）と特別なトレーニング（SCOPE）を組み合わせると、ロボットは自己修正を行う天才になります。

トレーニングなしの場合： 特別なトレーニングなしで新しいルールだけを与えると、ロボットの性能は逆に悪化します。言葉を変えようとはしますが、自分のミスを察知する能力を信頼できていないため、同じ間違いを繰り返してしまいます。
トレーニングありの場合： ロボットは自分の「凍結された間違い」を見つける術を学びます。これにより、数学の問題の中で間違った数字を見つけたときに、「いや、これは違う」と言い、文章を構築している最中に、それを正しいものへと差し替えることができるようになるのです。

実世界の成果

論文では、このモデルをLLaDA-8B（賢い言語モデル）を用いて、困難なタスクでテストしました。

数学 (GSM8K & MATH): スコアが大幅に上昇しました。例えば、難易度の高い数学テストでは、正解率が**55.3%から68.3%**へと跳ね上がりました。
コーディング (HumanEval & MBPP): 動作するコードを書く能力が劇的に向上し、あるテストでは**14.0%から29.3%**へと改善しました。

まとめ

主な発見は、**「考えを変える能力を持つことだけでは不十分であり、いつ考えを変えるべきかを判断するように訓練されなければならない」**ということです。

著者らは、モデルに自身の自信に満ちたエラーを認識させること（SCOPE）と、それらのエラーを即座に優れたものへと差し替えるメカニズムを与えること（D3IM）によって、モデルが以前よりもはるかに複雑な推論問題を解決できることを示しました。これは、ロボットをより賢くすることではなく、初期の間違いに対して「頑固にならない」ように教えることなのです。

技術要約：修正するか、凍結するか：自己修正型マスク拡散言語モデルのためのサンプラー一致型学習

1. 問題提起

LLaDAに代表されるマスク拡散言語モデル（MDLM）は、マスクされたシーケンスを反復的にデノイジング（除去）し、各ステップで全位置を再予測することでテキストを生成する。このアーキテクチャは本質的に「修正能力」を備えている。つまり、あるトークンが可視化された後でも、デノイザーは後のステップでその位置に対して異なる予測を割り当てることが理論上可能である。

しかし、標準的なMDLMのデコーディング戦略は、一度確定したトークンを「凍結」してしまう**単調なアンマスキング（monotonic unmasking）**を利用している。一度トークンが露出すると、それはクリーンなトークンとして扱われ、後続のステップではマスクされた位置のみが更新される。これは以下の2つの決定的な問題を引き起こす：

永続的な早期エラー： 生成プロセスの初期段階で行われた間違い（例：数学における誤った数字や、コードにおける誤った演算子）は、そのまま永続的なものとなり、推論重視のタスクにおいて性能を著しく低下させる。
修正のギャップ： 既存の修正アプローチは、ヒューリスティックなメカニズム（信頼度閾値や検証パス）に依存するか、あるいは再予測の前に確定したトークンを再び[MASK]へとリマスクする必要がある（例：ReMDM）。しかし、補助モジュールやリマスクの手順を介さずに、**直接的な「可視から可視への修正（visible-to-visible revision）」**を可能にする原理的なサンプラーは、未だ探索されていない。

さらに、著者らは**保存バイアス（preservation bias）**と呼ばれるモデル側の障害を特定している。たとえ学習目的がトークンの保存を強制していなくても、標準的なマスク目的関数で学習されたモデルは、自己生成されたコンテキストに直面した際、自らの誤った確定トークンを修正するのではなく、それを再現してしまう傾向がある。

2. 手法

本論文は、連携して動作するように設計された2つの解決策、D3IM（新しいサンプラー）とSCOPE（学習手順）を提案している。

2.1 D3IM：クリーン・スレート・リマスキング（Clean-Slate Remasking）

D3IMは、**コレクター型の逆更新（corrector-style reverse update）**として導出された、パラメータフリーのサンプラーである。確定したトークンを保持する標準的なサンプラー（保存重み $c_t > 0$ ）とは異なり、D3IMは $c_t = 0$ と設定する。

メカニズム： 各デノイジングステップにおいて、モデルはすべての活性化された位置を予測する。D3IMはこれらの予測を信頼度によってランク付けし、上位 $K$ 個の位置のみを保持する。
修正チャネル：
- トークンからトークンへ（Token-to-Token）： 以前に確定したトークンが、[MASK]状態を経ることなく、新しい高信頼度の予測によって直接上書きされる。
- トークンからマスクへ（Token-to-Mask）： 低信頼度の確定トークンが[MASK]へと格下げされ、後続のステップで再予測される。
制約： 次の状態は、以前の可視トークンをコピーすることではなく、現在の予測のみによって決定される。これには、自己生成されたコンテキスト下において、モデルの信頼度ランキングが信頼できるものであることが求められる。

2.2 SCOPE：予測エラーに基づく自己条件付け

D3IMをベースのMDLMに直接適用すると、モデルが自身の高信頼度のエラーを修正する能力を欠いているため、性能が低下する。SCOPEは、モデルの信頼度挙動をD3IMのサンプリング・ダイナミクスに適合させるために設計された、軽量なポストトレーニング手順である。

プロセス：
1. モデルは、部分的にマスクされた入力に対してフォワードパスを行う（自己条件付け）。
2. マスクされた位置に対して予測をサンプリングする。
3. これらの予測のうち、信頼度による上位 $\rho$ の割合が可視トークンとして「確定」され、残りはマスクされたままとなる。決定的なのは、このプロセスがモデルを自身の高信頼度のエラーにさらす点である。
4. その後、モデルは確定した（潜在的に誤った）トークンと、残りのマスクされた位置の両方について、正解（グラウンドトゥルース）のシーケンスを復元するように学習される。
目的： 可視となった自己生成トークンを固定されたコンテキストとしてではなく、一つの「仮説」として扱うようモデルを訓練することである。これにより、保存バイアスを軽減し、D3IMのトップ $K$ 選択における信頼度ランキングを信頼できるものにする。

3. 主な貢献

D3IMサンプラー： 補助モジュールやリマスクなしで、可視から可視への直接的なトークン修正を可能にする、パラメータフリーの直接修正サンプラーであり、学習済みのMDLMに対する有効なコレクター更新として導出された。
保存バイアスの特定： 標準的なMDLMの学習では、学習目的が修正を許容している場合であっても、モデルが自身の誤った可視予測を再現してしまう行動的傾向があることを示した。
SCOPE学習： D3IMのサンプリング・ダイナミクスをシミュレートする、軽量なポストトレーニング手順である。これは保存バイアスを効果的に軽減し、D3IMが意図通りに機能することを可能にする。
サンプラーとモデルの整合性： 修正能力は汎用的なものではなく、高度な推論戦略（例：D3IMのグローバルなトップ $K$ ルール）を機能させるためには、学習目的がそのサンプラーの決定規則に特化して一致している必要があることを確立した。

4. 実験結果

実験は、4つのベンチマーク（GSM8K（数学）、MATH-500（数学）、HumanEval（コード）、MBPP（コード））を用いてLLaDA-8B-Instructに対して行われた。

性能向上： 64ステップのデノイジングにおいて、SCOPE + D3IMの組み合わせは、標準的なアンマスキングを用いる元のLLaDA-8Bを大幅に上回った。
- GSM8K: +13.0% (68.3% vs. 55.3%)
- MATH-500: +4.8% (23.6% vs. 18.8%)
- HumanEval: +15.3% (29.3% vs. 14.0%)
- MBPP: +10.4% (30.8% vs. 20.4%)
スケーリング： 数学およびコードのタスクにおいて、デノイジングステップ数が増えるにつれて性能向上が増大する（例：GSM8Kは256ステップで73.8%に達する）。
整合性の必要性：
- SCOPEなしのD3IM： 性能が大幅に低下し（例：GSM8Kは48.0%に低下）、学習されたモデルがどのトークンを修正すべきかを信頼性を持って判断できないことを裏付けた。
- D3IMなしのSCOPE： SCOPEで学習されたモデルを標準的な単調アンマスキングで使用した場合、結果は混在しており、学習の恩恵はD3IMサンプラーの修正メカニズムに特有であることを示している。
メカニズム診断：
- 保存バイアスの解消： モデルが誤ったトークンを確定させたストレス・テストにおいて、元のモデルはエラーを70.1%の確率で保持する。SCOPEはこの値を17.0%に減少させ、正解を50.2%の確率で復元する。
- 軌跡分析： SCOPEがない場合、D3IMの修正は「コンテンツから空白への振動（content-to-whitespace oscillation）」（トークンと改行の間での切り替え）を引き起こす。SCOPEがある場合、修正は特定の意味的なエラー（例：単位変換、数字の修正）を対象とし、迅速に収束する。

5. 意義と主張

本論文は、MDLMにおける自己修正推論のボトルネックは、トークンを再訪する「能力」ではなく、自己生成されたコンテキスト下における信頼度ランキングの信頼性にあると主張している。

原理的な修正： D3IMは、ヒューリスティックな閾値や間接的なリマスクを超えた、直接的なトークン修正のための理論的に裏付けられた手法を提供する。
学習と推論の整合性： 本研究は、高度な推論戦略（例：グローバルなトップ $K$ 再ランキング）を実行するには、特定の学習目的（SCOPE）が必要であることを強調している。標準的な継続事前学習や一般的な自己修正メソッドは、これらの新しいサンプリング・ダイナミクスに自動的には適応しない。
効率性： 提案されたソリューションは計算効率が高く、軽量なLoRAファインチューニング（1台のA100で約20分、500ステップ）のみを必要とし、推論時のアーキテクチャ変更や補助モジュールを必要としない。

著者らは、D3IM+SCOPEは推論とコード生成を大幅に改善するものの、トークンレベルの修正に限定されており、初期の軌跡が根本的に誤っている場合に推論経路全体を再構築することはできないと結論づけている。今後の課題として、豊かな修正の軌跡をプロセス報酬モデリングや強化学習に活用する可能性が挙げられている。