近年、大規模言語モデルの性能向上の焦点は、「学習時のスケーリング」から「推論時のスケーリング」へと移り変わりつつあります。Best-of-N、Self-Consistencyから、より複雑な探索と検証のフレームワークに至るまで、Test-Time Scalingは大規模モデルの複雑な推論能力を引き出す重要なパラダイムとなっています。
しかし、長らく見過ごされてきた問題があります。それは、これらの手法のほとんどが、モデルが自己回帰的に生成されることを前提としている点です。
離散拡散言語モデル（Discrete Diffusion Language Models, dLLMs）にとっては、状況が全く異なります。dLLMは左から右へ一語ずつ生成するのではなく、マスクされた系列から出発し、複数ステップのノイズ除去を通じて徐々に完全な回答を復元します。この並列的で非自己回帰的な生成方式は、本質的にグローバルな双方向コンテキストを備えており、計画や自己修正にも適しています。しかしその一方で、従来の自己回帰モデル向けに設計された木探索、プロセス報酬モデル、Best-of-N推論は、そのまま効率的に適用できません。
この問題に対処するため、論文はPRISM: Pruning, Remasking, and Integrated Self-verification Methodを提案します。これは、離散拡散言語モデルのために特別に設計された、効率的なTest-Time Scalingフレームワークです。その中核的な目標は明確です。単にモデルを「何度も走らせる」のではなく、ノイズ除去プロセスにおいて、より有望な軌跡を識別し、動的に枝刈りし、局所的に分岐させ、モデル自身を用いて軽量な検証を行うことで、より少ない推論予算でBest-of-Nの効果に迫る、あるいはそれを超えることです。
論文タイトル: Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models
arXiv URL: https://arxiv.org/abs/2602.01842
コードURL: https://github.com/viiika/Prism
従来のBest-of-Nは高コスト、PRISMがdLLMの推論探索を再設計
dLLMにとって、単純なBest-of-Nの代償は極めて直接的なものです。仮にN本の軌跡をサンプリングし、各軌跡にTステップのノイズ除去が必要な場合、総関数呼び出し回数はO(NT)となります。これは、たとえ途中で多くの軌跡の品質が明らかに劣っていても、全ての候補が最後まで完全に実行され、予算を消費することを意味します。
PRISMの重要なアイデアは、推論プロセスを初期のランダム探索、中期の段階的枝刈り、後期の精緻化という3つの段階に分解することです。
ノイズが大きい段階では、モデルの出力はまだ不安定であるため、PRISMは多様性を維持するために広範な候補集合を保持します。ノイズ除去の早中期の段階で、回答の「論理的な骨格」が形成され始めると、PRISMは自己検証シグナルを用いて低品質な軌跡を枝刈りし、計算リソースをより有望な候補に再配分します。最終的には、少数の軌跡だけを残して精緻化を完了させます。論文では、このプロセスを階層的軌跡探索（Hierarchical Trajectory Search, HTS）と呼んでいます。
この設計により、PRISMの実際の計算複雑性はO(N + KT)に近づきます。ここでKは最終的に保持される小さな候補幅です。従来のBest-of-NのO(NT)と比較すると、これは「すべての経路を最後まで走らせる」方法を、「まず広く探索し、その後で集中砲火する」という方法に変えたことに相当します。
ゼロから始めるのではなく、低信頼度トークンで局所分岐
PRISMの第二の重要な構成要素は、部分的な再マスクによる局所分岐（Local Branching via Partial Remasking）です。直感的に言うと、モデルは中期のノイズ除去段階で、既に一部の高信頼度トークンを形成しています。これらのトークンは、多くの場合、回答の安定した構造や論理的骨格に対応します。一方で、低信頼度のトークンは、不確かな推論の詳細、実装方法、または局所的な表現に対応する可能性があります。
PRISMは、軌跡全体を粗く破棄して再サンプリングするのではなく、高信頼度の部分を保持します。そして、低信頼度の位置のみを再マスクし、これらの局所的な変化から新しい分岐を生成します。この方法の利点は、既存の高品質な構造を保持しつつ、異なる詳細な実装方法を探索し続け、単一のパスへの早期収束を避けられることです。論文の図2には、このプロセスが直感的に示されています。段階的枝刈り段階で、PRISMは高スコアの軌跡の周辺で局所分岐を行い、部分的な再マスクを通じて新しい候補を生成します。
外部検証器に依存しない：モデル自身が自己採点
多くのTest-Time Scaling手法は、候補回答の品質を判断するために、追加の報酬モデルや検証器に依存します。しかし、これにはデプロイ時に追加モデルをロードする必要があり、VRAM、レイテンシ、エンジニアリングの複雑さが増大するという、顕著なシステムオーバーヘッドが伴います。
PRISMは、自己検証フィードバック（Self-Verified Feedback, SVF）を提案します。これは、同じdLLMを二値分類検証器として直接再利用するものです。具体的には、モデルはまず中間ノイズ除去状態に基づいて完全な候補回答を生成します。次に、Yes/Noの検証プロンプトを構築し、その回答が正しい可能性があるかどうかをモデルに判断させます。PRISMは、YesとNoのロジットを二値の正規化スコアに変換し、これを軌跡のランク付け、枝刈り、最終選択に使用します。
この設計の価値は、検証を「追加モデル」から「同一モデルによる一度の軽量な自己チェック」へと変えたことにあります。論文ではさらに、SVFの呼び出し回数はノイズ除去のNFEと比較して非常に少なく、実験では通常、総NFEの10%未満であることが示されています。そのため、低い追加オーバーヘッドで効果的な探索シグナルを提供できます。
実験：数学推論とコード生成で顕著なコストパフォーマンス向上を達成
論文では、4つのベンチマークでPRISMを評価しています。数学推論タスクのGSM8K、MATH-500、そしてコード生成タスクのHumanEval、MBPPです。実験は、LLaDA-8B-Instruct、Dream-7B-Instruct、LLaDA-2.0-miniという3つの離散拡散言語モデルを対象に行われました。
LLaDA-8B-Instructにおいて、PRISM（K=8）はGSM8Kを67.58%から85.30%へ、MATH-500を26.40%から42.80%へと向上させました。コードタスクでは、HumanEvalが24.39ポイント、MBPPが16.40ポイント向上しました。さらに重要なのは、これらの向上がBest-of-Nの計算量を線形に増加させることで得られたものではない点です。例えばGSM8Kでは、PRISMは1048 NFEで85.30%を達成しましたが、Best-of-16は87.50%に達するために4096 NFEを必要とし、4倍以上のノイズ除去計算の節約を実現しました。
論文の図1では、PRISMはBest-of-Nと比較して、複数のタスクでより優れた性能対計算量の曲線を示しています。同等の精度において、GSM8K、MATH500、HumanEval、MBPPでそれぞれ2.9倍、6.5倍、1.8倍、1.7倍の速度優位性を示しています。
論文では、他の推論時スケーリング手法との比較も行われています。TruthfulQAにおいて、PRISMのROUGE-1/2/Lは31.8/35.5/31.9、推論時間は1048.0秒に達しました。これに対し、LLaDA-ReMDMは29.5/31.8/29.5、推論時間は1354.8秒でした。これは、PRISMがタスク性能を向上させるだけでなく、より優れた推論効率を維持できることを示しています。
外部検証器との比較では、SVFはGSM8Kで85.30%を達成し、元の8Bモデルのロードのみで済みました。Qwen3-8B検証器は87.35%に達する可能性がありますが、追加モデルのロードが必要で、総パラメータ数は16Bになります。論文では、SVFの利点はすべての外部検証器を絶対的に代替することにあるのではなく、より軽量でデプロイが容易なdLLM推論拡張パスを提供することにあるとしています。
意義：非自己回帰型言語モデルに推論時スケーリングの道を開く
PRISMの核心的な貢献は、単に新しい探索ヒューリスティックを提案したことではなく、dLLMにおけるTest-Time Scalingがどのように起こるべきかを再定義したことにあります。
自己回帰モデルでは、推論探索は通常「プレフィックス」を中心に展開されます。しかし、離散拡散モデルでは、中間状態は部分的にマスクされたグローバルな系列であり、従来のプレフィックスに基づくプロセス報酬や木探索は、そのままでは自然に適合しません。PRISMは、探索、枝刈り、局所分岐、自己検証のすべてを、dLLMのノイズ除去のダイナミクスに再配置しました。構造形成段階で集中的に予算を配分し、低信頼領域で代替表現を探索し、追加モデルなしで検証を完了します。
これは、dLLMがもはや「並列生成により高速」な代替パラダイムであるだけでなく、推論、計画、自己修正に適した新しいタイプの言語モデルアーキテクチャになり得ることを意味します。LLaDA、Dream、Mercury、Gemini Diffusionなどのモデルが、離散拡散言語モデルをより大規模なものへと押し進める中で、PRISMは重要な方向性を示しています。それは、非自己回帰モデルもまた、現在の主流LLMのように、推論時の計算を通じて継続的に能力を向上させられるという道です。
この観点から見ると、PRISMは単なる計算量節約型のBest-of-N代替案ではなく、離散拡散言語モデルが高効率な推論システムへと進化するための、重要なパズルの一片なのです。
著者紹介
本論文は、Jinbin Bai氏をはじめとする研究者チームによって執筆されました。著者チームは、離散拡散やマスク生成モデリングといった新世代の生成パラダイムに長期的に焦点を当てており、研究の方向性は、高解像度テキスト画像生成、統一マルチモーダル生成、離散拡散モデルの嗜好アライメントと推論最適化、そしてインタラクティブな世界モデルなど多岐にわたります。
これまでにチームは、マスク生成トランスフォーマーの高解像度テキスト画像生成における可能性を探求したMeissonic [1]を提案し、その後、離散拡散モデリングを画像生成からより統一されたマルチモーダル生成フレームワークへと推し進めるMuddit [2]を発表しました。このたびICML 2026に採択されたPRISMは、この研究の系譜を推論段階にまでさらに拡張し、階層的探索、自己検証フィードバック、局所的な再マスクを通じて、離散拡散モデルが追加の検証器なしで、どのように効率的なTest-Time Scalingを実現するかに焦点を当てています。
[1] Meissonic: Revitalizing Masked Generative Transformers for Efficient High-Resolution Text-to-Image Synthesis, ICLR 2025, https://arxiv.org/abs/2410.08261
[2] Muddit: Liberating Generation Beyond Text-to-Image with a Unified Discrete Diffusion Model, ICLR 2026, https://arxiv.org/abs/2505.23606
© THE END
転載は本公式アカウントからの許可を得てください。
投稿または取材のご依頼: liyazhou@jiqizhixin.com