🤖 machine learning

Rethinking Reasoning with MDLMs: Early Exits, Post-hoc Reasoning, and Beyond

本論文は、Masked Diffusion Language Models (MDLM) に対する「reasoning-as-infilling（穴埋めとしての推論）」を導入するものであり、これは回答領域を明示的に指定することで、早期終了や事後的な推論といった独自の能力を解禁する手法であり、最終的に、MDLMが人間が書いた思考プロセスに匹敵する推論性能を達成できること、および自己回帰型モデルよりも正確な中間ステップのスコアリングを提供できることを実証している。

原著者： Zachary Horvitz, Raghav Singhal, Hao Zou, Carles Domingo-Enrich, Zhou Yu, Rajesh Ranganath, Kathleen McKeown

公開日 2026-08-11

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

CC BY 4.0

原著者： Zachary Horvitz, Raghav Singhal, Hao Zou, Carles Domingo-Enrich, Zhou Yu, Rajesh Ranganath, Kathleen McKeown

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

ロボットに難しい謎解きを解く方法を教えようとしている場面を想像してみてください。長い間、最も優れた方法は、ロボットに物語を一度に一単語ずつ読ませ、これまでに読んだすべての内容に基づいて、まさに「次の単語」を予測させることでした。これは、先読みを一切許されず、一度書いたものは書き直すこともできずに、一行ずつ答えを書いていかなければならないテストを受けている生徒のようなものです。この方法で多くのことはうまくこなせますが、欠点があります。一度間違った単語を書いてしまうと、その間違いに縛られてしまい、答えにたどり着くまでのステップを考えている最中に、最終的な答えを「見る」ことが容易にできないのです。

最近、科学者たちは「マスク型拡散言語モデル（MDLM）」と呼ばれる、異なる種類のロボットの脳について実験を行っています。単語を一つずつ読む代わりに、想像してみてください。ロボットに、いくつかの単語が付箋で隠されたページを与えます。その仕事は、その付箋の下に何があるかを推測することです。しかし、ここからが魔法の部分です。ロボットはページ全体を一度に見ることができ、すべての欠落した単語を同時に、あるいは好きな順序で推測することができます。それは、角の部分、真ん中、あるいは端の部分を、同時に埋めていくパズルを解くようなものです。この論文は、この「穴埋め」スタイルの思考が、従来の「一単語ずつ」のアプローチよりも、数学の問題や論理的な謎解きを解く上で本当に優れているかどうかを探求しています。研究者たちは、この新しいスタイルを使うことで、ロボットをより賢く、速く、そして自分が答えを知っているかどうかについてより正直にさせることができるのではないかと考えました。

ザカリー・ホーウィッツとラガヴ・シンガル率いる研究チームは、「推論としての穴埋め（reasoning-as-infilling）」という巧妙な新しいトリックを提案しています。単にロボットに「考えてから答える」よう求めるのではなく、特定のボックスを持つテンプレートを与えます。具体的には、思考のステップのための大きなボックスと、最後に明確にマークされた小さな回答用のボックスです。そして、ロボットが考え始める前に、回答ボックスを先に埋めてしまいます。これは一見すると逆転しているように聞こえますが、これがスーパーパワーを解き放ちます。

第一に、ロボットは回答ボックスがどこにあるかを正確に把握しているため、思考ステップを書きながら、自分自身の確信度を覗き見ることができます。もしロボットが、そのボックス内の答えに対して非常に高い確信を持った場合、思考を即座に停止して結果を出力することができます。論文によれば、GSM8kと呼ばれる数学テストにおいて、この「早期終了（early exit）」戦略により、精度をほとんど損なうことなく、1.3倍速く処理を完了できることが示されました。他のスピードアップ技術と組み合わせれば、4倍速く動かすことさえ可能です。これは、長いエッセイを書いている途中で、すでに結論が分かったことに気づき、最後の文章だけを書いてテストを提出する学生のようなものです。

第二に、この手法は「事後的推論（post-hoc reasoning）」を可能にします。通常、ロボットが数学の問題を間違えた場合、まだ答えを知らないため、正しい考え方を教えることは困難です。しかし、この新しい手法では、最初に正しい答えを与えておけば、ロボットはそこから逆算して、どのようにその答えに辿り着いたのかという完璧なステップ・バイ・ステップの説明を生成することができます。研究者たちは、これらの「逆向きに生成された」説明を用いてロボットを訓練することで、数学のスコアが14.9%向上することを発見しました。これは、人間がロボットのために完璧な解説を書くように依頼した場合と同等の向上であり、非常に大きな飛躍です。これは、ロボットが自分自身の「未来の知識」を使って自習できることを示唆しています。

最後に、この論文は、この手法がロボットに進行中の作業を判断させるのに役立つことを示しています。ロボットは思考ステップを書きながら、「もしこの通りに思考を完了させた場合、目指している答えに到達する可能性はどのくらいか？」と常にチェックすることができます。研究者たちは、この自己チェック・スコアが、従来のロボットが生成するスコアよりも、最終的な答えが正しいかどうかを予測する上ではるかに優れていることを発見しました。従来のスタイルのロボットは、最後まで自分が間違っていることに気づかないことが多いですが、この新しいロボットは、早い段階でミスを察知できるのです。

また、論文では、うまくいかないケースについても指摘しています。彼らは、ロボットにパズルのピースを完全にランダムな順序（例えば、左上から右下へとパターンを持たずに飛び回るなど）で埋めさせるテストを行いました。その結果、数学の問題においては、このような混沌としたアプローチはロボットの性能を悪化させ、しばしば間違った答えを最初に推測させ、その後でそれを正当化するための偽のストーリーをでっち上げさせてしまうことが分かりました。ロボットが軌道から外れないためには、ある程度の構造、つまり「推論ボックス」や「回答ボックス」のようなテンプレートが必要なのです。

要約すると、この論文は、AIへの「考え方」の問いかけ方を変えること、つまり白紙のページに書かせるのではなく、埋めるべきテンプレートを与えることで、AIをより速く、学習能力を高め、自己チェック能力を持たせることができると示唆しています。これはあらゆる問題を解決する魔法の杖ではありませんが、より賢く、効率的な推論マシンを構築するための有望な新しい道筋を提示しています。

技術要約：マスクされた拡散言語モデルによる推論の再考

問題提起

現在の大規模言語モデル（LLM）における推論パラダイムは、自己回帰的な生成を通じて訓練された次トークン予測（NTP）モデルに大きく依存している。これらは効果的ではあるものの、推論タスクにおいて固有の限界に直面している：

困難な事後分布サンプリング： 正解（既知の答え）を条件とした高品質な推論トレースの生成（事後分布サンプリング、 $p(r|c,a)$ ）は、自己回帰モデルにとっては計算量的に困難である。なぜなら、自己回帰モデルは左から右へと生成を行うため、複雑なプロンプトエンジニアリングや補助モデルなしには、未来のトークン（答え）を容易に条件付けることができないためである。
中間スコアリングの欠如： 部分的な推論トレースの正当性を評価するには、通常、個別のプロセス報酬モデル（PRM）を訓練する必要があり、これは計算オーバーヘッドと複雑さを増大させる。
非効率な推論： 標準的な自己回帰デコーディングは逐次的であり、並列的なトークン生成を妨げ、モデルが確信を持った際の早期終了の機会を制限する。

マスクされた拡散言語モデル（MDLM）は、モデルがシーケンス内のランダムにマスクされた位置をインフィル（補完）することを学習する代替的な訓練目的関数を提供する。しかし、並列デコーディングを超えた、MDLMの推論能力に関する具体的な能力については、まだ十分に探索・構造化されていない。

メソドロジー：インフィリングとしての推論

本論文では、MDLMの独自の特性を活用して推論プロセスを再構築するプロンプティング手法である**「Reasoning-as-Infilling（インフィリングとしての推論）」を導入している。これは、出力を単なるプレフィックスとして扱うのではなく、推論ブロックと回答ブロック**を明示的に区切るユーザー指定のテンプレートで事前に入力シーケンスを埋める手法である。

コアとなるメカニズムは以下の通りである：

テンプレートの事前充填（Template Pre-filling）： 入力コンテキストには、質問と、マスクされた推論トークンおよびデリミタ（区切り文字）に続くマスクされた回答トークンを含むテンプレートが含まれる（例：[Reasoning Block] <Delimiter> [Answer Block]）。
明示的な条件付け： 回答ブロックの位置が明示的に指定され、かつマスクされているため、MDLMは推論生成のあらゆる中間ステップにおいて、回答トークンの条件付き分布（ $p(a_j | r_{unmasked}, c)$ ）を計算することができる。

このフレームワークにより、主に以下の3つの能力が可能になる：

1. 早期脱出（Early Exits）

モデルは推論トレースの生成を通じて回答トークンの分布を保持しているため、システムは回答領域の不確実性を監視することができる。

メカニズム： 著者らは、回答ブロックのエントロピー上限（ $H_{UB}$ ）を確信度の指標として提案している。 $H_{UB}$ は、マスクされていない推論トークンが与えられたときの回答トークンの周辺エントロピーの総和として計算される。
アクション： $H_{UB}$ が閾値 $\gamma$ を下回った場合、モデルは残りの推論ステップをスキップし、直接回答を出力することで、推論を大幅に加速させる。

2. 事後推論（Post-hoc Reasoning / 事後分布サンプリング）

このフレームワークは、正解を条件とした推論トレースのサンプリングを可能にする。これは、NTPモデルにとっては困難な分布である。

メカニズム： 正解の回答（ $a^*$ ）を回答ブロックにあらかじめ充填しておくことで、MDLMは事後分布 $p_\theta(r | c, a^*)$ からサンプリングを行う。
適用： これらの「事後トレース」は、高品質な推論データとしてファインチューニングに利用できる。論文では、これらのモデル生成されたトレースを用いてファインチューニングを行うことで、モデルが当初問題を解けなかった場合でも、性能が向上することが示されている。

3. PRMを用いない中間スコアリング

このフレームワークは、外部の報酬モデルを使用せずに、部分的な推論トレースをスコア化する組み込みのメカニズムを提供する。

メカニズム： 参照回答 $a^*$ と部分的な推論トレース $r_{unmasked}$ が与えられたとき、スコアは回答トークンの対数確率の総和として計算される： $\phi_{IPRM} = \sum \log p_\theta(a^*_j | c, r_{unmasked})$ 。
利点： このスコアは最終的な正解率と強く相関しており、自己回帰モデルとは異なり、生成の極めて早い段階から利用可能である（自己回帰モデルにおける回答尤度は、生成の終盤まで予測力が弱い）。

主な結果

推論の加速（早期脱出）

GSM8kおよびMATH500を用いたLLaDA-8BおよびDream-7Bによる評価：

速度と精度のトレードオフ： 早期脱出のみを用いた場合、GSM8kにおいて精度低下を1%未満に抑えつつ、1.3倍の高速化を達成した。
結合された加速： 並列デコーディングアルゴリズム（例：Block Parallel Confidence Decoding）と組み合わせることで、精度を最小限に抑えつつ、最大で4.0倍の加速（LLaDAで3.5倍、Dreamで2.5倍）を実現した。
可変長サポート： この手法は可変長拡散モデル（例：DreamOn）と統合可能であり、高い精度を維持しながら推論予算を動的に調整できる。

事後推論とファインチューニング

トレースの品質： LLaDA-8B-InstructモデルがGSM8kの問題を解けなかった場合でも、事後分布からサンプリング（正解を条件付け）を行うと、生成された推論トレースはGPT-4oおよび特化したPRMによって約40%のケースで正しいと判定された。ベースモデルであっても、これらの失敗事例に対して30%以上の正しいトレースを生成した。
性能向上： LLaDA-8B-Baseモデルをこれらの事後生成トレースでファインチューニングした結果、GSM8kのテスト精度が**+14.9%**向上した。この向上幅は、高価な人間によるアノテーション済み推論トレースによる向上幅（+13.4%）に匹敵する。
比較： 事後推論トレースは、STaR形式のプレフィックス・ヒント・プロンプティングによって生成されたトレースよりも優れた性能を示した。

中間スコアリング

相関関係： 中間ステップにおけるMDLMの回答尤度スコアは、特化した7Bのプロセス報酬モデル（PRM）よりも、最終的な正解率との強い相関を示した。
自己回帰ベースライン： 対照的に、自己回帰モデル（Llama-3.1-8B, Qwen2.5-7B）からの中間回答尤度スコアは、生成の終盤に至るまで正解率の予測力が著しく低かった。

意義と主張

本論文は、Reasoning-as-Infillingフレームワークが、NTPモデルでは容易に実現できない一連の統一された能力をMDLMにもたらすと主張している：

効率性： 精度を犠牲にすることなく、早期脱出や並列デコーディングを通じて大幅な推論加速を可能にする。
データ生成： 正解を条件とした高品質な推論トレースを生成するための実行可能な手法を提供し、ファインチューニングのための人間によるアノテーションに代わるスケーラブルな選択肢となる。
自己スコアリング： モデルが内部分布を用いて、中間ステップで自身の推論プロセスをスコア化することを可能にし、特定のポストトレーニングのシナリオにおいて、別途のプロセス報酬モデルの必要性を排除する。

著者らは、MDLMの訓練目的関数が構造化されたプロンプティングと組み合わされることで、推論タスク（特に検証可能な報酬、効率的な推論、または正解からの推論データの生成を必要とするシナリオ）において有望な利益をもたらすと結論付けている。本研究は、MDLMの「病理的」な挙動（例：推論の前に回答をデコードしてしまうこと）を、答えが既知の場合のポストトレーニング用データ生成のための「機能」へと再定義できることを示唆している。