💬 NLP

Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space

本論文は、生成順序とトークン空間を共同で探索することで、数学的推論およびコーディングのベンチマークにおいて既存のベースラインを凌駕する、拡散言語モデルのための新しいデコーディング手法であるOrder-Token Searchを提案する。

原著者： Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

公開日 2026-02-06

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

原著者： Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

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

複雑なパズル、例えば数学の問題やコーディングの課題を解こうとしている場面を想像してみてください。そこには、人間が左から右へと一文字ずつタイピングしていくのとは異なる、魔法の助手（拡散言語モデル：Diffusion Language Model）がいます。この助手は、まず「MASK」というプレースホルダーで埋め尽くされた白紙の状態からスタートし、そこにどの単語が入るべきかを一斉に推測して埋めていきます。

問題は、この助手が少し混沌としていることです。助けは、好きな順番で空白を埋めることができます。例えば、最初の空白には正しい単語を当てはめられたものの、2番目の空白で間違った単語を推測してしまい、行き詰まってしまうことがあります。また、単語自体は合っているのに順番を間違えてしまい、行き止まりに突き当たってしまうこともあります。

旧来の手法：推測と枝刈り

以前は、人々はこの問題を解決するために主に2つの方法を試みましたが、どちらにも欠陥がありました。

「自信満々」アプローチ（The "Confident" Approach）: 助手が、自分が最も確信を持っている空白から先に埋めていく方法です。これは、最も足場がしっかりしているように見える道だけを進むハイカーのようなものです。
- 利点: 通常、素早く良い答えに到達できます。
- 欠点: もしその「しっかりした道」が崖へと続いていた場合（間違った答えだった場合）、ハイカーはそこで立ち往生してしまいます。より良いかもしれない他の道を探索することはありません。
「ランダム」アプローチ（The "Random" Approach）: 助手が完全にランダムに空白を選んで埋めていく方法です。
- 利点: 非常に多様な経路を探索するため、最終的に正しい解にたどり着く可能性は非常に高いです。
- 欠点: あまりに散漫であるため、最初の一手で最善の経路を選ぶことは滅多にありません。それは、円を描いて彷徨うハイカーのようなものです。宝を見つけることはできるかもしれませんが、その過程で多くの空の穴を掘ることにもなるでしょう。

新しい解決策：オーダー・トークン探索（Order-Token Search）

この論文の著者たちは、**オーダー・トークン探索（Order-Token Search）と呼ばれる新しい手法を導入しました。これは、協力して働く「探検家チーム」**のようなものです。

一人のハイカーや混沌とした群衆を送り出すのではなく、この手法は小さなチーム（「ビーム」）の探検家を送り出します。彼らの仕組みは以下の通りです。

分岐する経路（探索）: 定期的な間隔で、チームは分かれ道を作ります。各探検家は異なる戦略を試みます。
- 探検家Aは、最初の欠落した単語を埋めることに決めます。
- 探検家Bは、最後の欠落した単語を埋めることに決めます。
- 探検家Cは、真ん中のスペースに別の単語を試します。
- 比喩: 彼らは、次にどこに書くか（順序）と、何を書くか（トークン）の両方を同時に探索しています。
スコアカード（尤度推定器）: これが魔法の部分です。チームには特別な審判（尤度推定器）がついています。この審判は、最終的な答えだけを見るのではありません。審判は、探検家が行ったすべてのステップを観察します。
- その動きは論理的だったか？
- その部分的な文章は、それまでに書かれた内容と照らし合わせて意味が通るか？
- 比喩: リレーのコーチがレースを見守っている場面を想像してください。もしランナーが早い段階で転んだら、コーチは最後まで走り終えるのを待たずに、即座に介入します。なぜなら、その「ステップ」自体が間違っていたからです。
行き止まりの切り捨て（枝刈り）: 審判は、あらゆる探検家の進捗をスコア化します。もしある探検家が、成功する可能性が低いと思われる経路を進んでいる場合、チームはその経路を切り捨て、より優れた軌道に乗っている探検家にリソースを集中させます。

なぜこれが重要なのか

論文では、この手法を難しい数学の問題（GSM8KやMATH500データセット）やコーディングタスク（HumanEval）でテストしました。

結果: 「探検家チーム」（オーダー・トークン探索）は、「自信満々」や「ランダム」といった従来の方法よりも、一貫して正解を見つけ出すことができました。
比較: この手法は非常に優れたパフォーマンスを示し、AIに数ヶ月に及ぶ高価な再学習を強いる手法（diffu-GRPOなど）に匹らい、あるいはそれを上回る結果を出しました。これは、AIの「脳」自体を再学習させる必要はなく、テスト時の「考え方」を変えるだけで、より賢いAIを実現できることを意味しています。

数独に関する特記事項

論文では、この手法を**数独（Sudoku）**のパズルにも適用してみました。興味深いことに、うまくいきませんでした。著者は、数独には厳格なグローバルルール（例：「列の中に同じ数字を繰り返さない」など）が必要であり、AIの内部にある「スコアカード」では、それを理解できないことを説明しています。それは、崖の位置が示されていない地図を渡されたハイカーのようなものです。どれほど上手く探索したとしても、地図が間違っていれば、崖から落ちるのを避けることはできません。このことは、一部のタスクにおいては、単に探索手法を変えるだけでなく、AIそのものを異なる形で訓練する必要があることを示唆しています。

まとめ

この論文は、AIに**「書き方の異なる複数の方法（順序）」と「複数の異なる単語（トークン）」**を同時に探索させ、さらにスマートなスコアリングシステムを用いて悪いアイデアを早期に切り捨てることで、拡散言語モデルを再学習させることなく、より優れた結果を得られることを示しています。それは、混沌とした推測ゲームを、真実への構造化された効率的な探索へと変えるのです。

技術要約：生成順序とトークン空間の結合探索による拡散言語モデルデコーディングの改善

問題提起

拡散言語モデル（DLM）、特にマスク拡散モデル（MDM）は、自己回帰（AR）モデルとは異なる明確な利点を持っています。それは**順序に依存しない生成（order-agnostic generation）**です。左から右へと厳密にトークンを生成するARモデルとは異なり、MDMは複数の位置を並列にデノイズできるため、適応的な生成順序が可能です。しかし、現在のデコーディング手法はこの柔軟性を十分に活用できていません。

既存のデコーディング戦略は、通常、ヒューリスティックな再マスク規則に従って、実行ごとに**単一の軌跡（trajectory）**に固執します。

低信頼度再マスク（Low-Confidence Remasking）: 高信頼度のトークンを早期に確定させる手法。これは単一サンプルの精度（pass@1）においては強力な結果をもたらすことが多いものの、探索を狭い軌跡の集合に制限してしまい、マルチサンプル・カバレッジ（pass@k）の成長を鈍化させ、誤った局所最適解に陥る可能性があります。
ランダム再マスク（Random Remasking）: 軌跡の多様性を最大化するために、再マスクのパターンを一様にサンプリングする手法。これはより広い解空間を探索することでpass@kを向上させますが、競合する候補の中から正しい軌跡を選択するための原理的なメカニズムを欠いているため、pass@1が弱くなる傾向があります。

これにより、探索（exploration）と活用（exploitation）のトレードオフが生じています。現在の手法は、局所的な信頼度を過度に活用して（正しい解を見逃し）か、あるいは盲目的に探索して（最良の解を選択できず）いるかのどちらかに偏っています。本論文は、標準的なデコーディングがこの結合空間における明示的な探索を行わないため、DLMが持つ多様な生成順序とトークン値を探索する潜在能力が実現されていないと主張しています。

手法：順序・トークン探索（Order-Token Search）

これに対処するため、著者らは生成順序とトークン値の結合空間においてビームサーチを行うデコーディングアルゴリズムである「Order-Token Search（OTS）」を提案しています。このアルゴリズムは、候補となる軌跡を拡張することと、有望でないものを枝刈りすることを交互に行います。

1. 探索プロセス（拡張）

Order-Token Searchは、 $K$ 個の部分的なデノイズ済み候補からなるビームを保持します。固定された決定シーケンスに従う標準的なサンプリングとは異なり、このアルゴリズムはユーザーが指定した間隔で、2つの不確実性の源に対して明示的に分岐します。

生成順序（Generation Order）: 次にどの位置（トークン）をアンマスクするか。
トークン値（Token Values）: それらの位置にどのようなトークンを割り当てるか。

計算量を管理するため、探索は**ブロック拡散（block diffusion）**を利用します。すべてのデノイズステップで分岐する（これは膨大なオーバーヘッドを引き起こします）のではなく、アルゴリズムは連続するトークンブロックの境界で候補を拡張します。各ブロック内では独立したデノイズが行われます。これにより、決定ポイントを全シーケンス長からブロック数へと削減し、推論コストを管理可能な範囲に抑えつつ、順序とトークン割り当ての両方において大幅な多様性を実現します。

2. 枝刈りプロセス（尤度推定）

核心となる革新は、部分的な軌跡をスコアリングし枝刈りするために使用される尤度推定器にあります。標準的なMDMの目的関数（ノイズの多い状態からクリーンな完全シーケンスを再構成する尤度）は、多くのマスクが残っている中間状態に対しては不安定です。なぜなら、それはモデルが訓練時に最適化されていない困難なインフィリング問題を解くことを強いるからです。

著者らは、増分的なアクションレベルのスコアリング関数を導入しました。

全体の状態をスコア化する代わりに、アルゴリズムは、より腐食した状態 $x_t$ から、より少ない腐食状態 $x_s$ への遷移において行われた特定のデノイズ・アクションをスコア化します。
スコア $s(x_t; x_s)$ は、モデルのフルシーケンス予測 $x_0$ （コンテキストとして機能）を条件として、新たに明らかになったトークン（ $x_s$ 内のもの）をモデルがいかに自信を持って再予測できるかを測定します。
このアプローチは、ある軌跡の過去の決定（順序とトークンの選択）がモデルによって支持されているかどうかを評価するため、枝刈りに対してより低分散で識別力の高い信号を提供します。

候補の累積スコアは、これらの探索間隔における増分スコアの合計です。アルゴリズムは各拡張の後にトップ $K$ の候補を保持し、デコーディングを継続する前に、誤った、あるいは最適ではないパスを効果的にフィルタリングします。

主な貢献

トレードオフの特定: 本論文は、標準的なMDMデコーディング手法が、探索と活用のスペクトルの両極端に位置していることを経験的に示しています。低信頼度再マスクは過剰に活用し（高いpass@1、低いpass@kの成長）、ランダム再マスクは盲目的に探索します（高いpass@k、低いpass@1）。
Order-Token Search アルゴリズム: 生成順序とトークン値の結合空間を探索する新しいデコーディングアルゴリズムの提案。安定した増分尤度推定器を用いて部分的な軌跡を枝刈りすることで、軌跡の多様性を向上した単一サンプルの精度へと変換します。
経験的検証: 本手法は、数学的推論およびコーディングのベンチマークにおいてベースラインを一貫して上回る性能を示しており、diffu-GRPOのような高コストなポストトレーニング手法によって得られる利得に匹敵またはそれを凌駕しています。

実験結果

著者らは、LLaDA-8B-InstructおよびLLaDA-1.5（RLポストトレーニング版）を用い、GSM8K、MATH500、Countdown、およびHumanEvalにおいてOrder-Token Search（OTS）を評価しました。

性能向上: OTSは、ベンチマーク全体で一貫して最高の平均精度を達成しました。
- LLaDA-8Bにおいて、OTSは低信頼度再マスクのベースラインと比較して、GSM8Kで**+3.1%、MATH500で+3.8%、Countdownで+7.9%、HumanEvalで+6.8%**のpass@1向上を実現しました。
- LLaDA-1.5においても同様の相対的改善を達成し、最良のベースライン（Low-conf + MV）の37.2%に対し、**40.4%**の平均精度に達しました。
ポストトレーニングとの比較: OTSは、reinforcement learning（強化学習）によるポストトレーニング（diffu-GRPO）を受けたモデルと同等またはそれを上回る性能を示しました。これは、推論時の探索が、推論能力を向上させるためにモデルの再学習と同等に効果的であることを示唆しています。
アブレーション研究:
- 尤度推定器: 特定の増分スコアリング関数（式2）が極めて重要でした。全開示ブロックの結合尤度やARスタイルの尤度を用いたバリアントは大幅に性能が低下しており、MDMにおける安定した枝刈りにはカスタマイズされた推定器が必要であることが確認されました。
- 探索空間: OTSは、「Order Search」（順序のみを探索）および「Token Search」（トークンのみを探索、ARスタイルの尤度を使用）を上回りました。これは、結合探索の必要性と、ARのヒューリスティックがMDMと適合していないことを浮き彫りにしています。
効率性: OTSは、多数決（majority voting）よりもテスト時の計算量（NFE）に対してより効果的にスケールします。計算予算が等しい場合、OTSは構造化された探索を活用することで、非構造的なサンプリングよりも高い精度（例：Countdownにおいて29.3% vs 19.9%）を達成しました。

意義と主張

本論文は、結合的な順序・トークン探索がDLMデコーディングを進展させるための鍵であると主張しています。生成順序とトークン値の空間を明示的に探索し、安定した尤度推定器に基づいて枝刈りを行うことで、この手法は拡散モデルに備わっているマルチ・トラジェトリーの潜在能力を、信頼できる単一サンプルのパフォーマンスへと変換します。

著者らは、本研究を現在のMDMデコーディングにおける「探索–活用」のボトルネックに対する解決策として位置付けています。彼らのアプローチは、追加のモデル訓練を必要とすることなく、貪欲な手法を捉える局所最適解から脱出し、ランダムサンプリングが見逃してしまう正しい解を選択することを可能にすると断言しています。本研究は、拡散モデルの並列的かつ順序に依存しない性質に特化した推論時探索戦略が、推論およびコーディングタスクにおいて系統的な改善をもたらすことを確立しています。