プリズム：離散拡散言語モデルのための階層的探索と自己検証による効率的なテスト時スケーリング

Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models

February 2, 2026

著者: Jinbin Bai, Yixuan Li, Yuchen Zhu, Yi Xin, Qingyu Shi, Aosong Feng, Xiaohong Liu, Molei Tao, Jianru Xue, Xiangtai Li, Ming-Hsuan Yang

cs.AI

要旨

推論時計算は、大規模言語モデル(LLM)の推論能力を向上させる実用的な手法として再び注目を集めている。ほとんどのテスト時スケーリング(TTS)アルゴリズムは自己回帰的なデコーディングに依存しているが、この手法は系列全体を並列にデコードする離散拡散言語モデル(dLLM)には不向きである。その結果、dLLMの生成能力を最大限に引き出す効果的かつ効率的なTTS手法の開発は、未開拓の課題として残されている。この問題に対処するため、我々はdLLM向けの効率的なTTSフレームワークであるPrism(Pruning, Remasking, and Integrated Self-verification Method)を提案する。Prismは、(i) ノイズ除去プロセスの初期から中期にかけて計算リソースを動的に刈り込み再配分する階層的軌道探索(HTS)を実行し、(ii) 高信頼度トークンを保持しつつ多様な実装を探索するための部分リマスキングを伴う局所分岐を導入し、(iii) 外部検証器を、中間生成結果に対する自己評価プロンプトを通じて得られる自己検証フィードバック(SVF)で置き換える。LLaDA 8B Instruct、Dream 7B Instruct、LLaDA 2.0-miniを含む3つのdLLMと4つの数学推論・コード生成ベンチマークにおいて、我々のPrismは良好な性能と効率性のトレードオフを達成し、Best-of-N性能を大幅に少ない関数評価回数(NFE)で実現した。コードはhttps://github.com/viiika/Prismで公開されている。

English

Inference-time compute has re-emerged as a practical way to improve LLM reasoning. Most test-time scaling (TTS) algorithms rely on autoregressive decoding, which is ill-suited to discrete diffusion language models (dLLMs) due to their parallel decoding over the entire sequence. As a result, developing effective and efficient TTS methods to unlock dLLMs' full generative potential remains an underexplored challenge. To address this, we propose Prism (Pruning, Remasking, and Integrated Self-verification Method), an efficient TTS framework for dLLMs that (i) performs Hierarchical Trajectory Search (HTS) which dynamically prunes and reallocates compute in an early-to-mid denoising window, (ii) introduces Local branching with partial remasking to explore diverse implementations while preserving high-confidence tokens, and (iii) replaces external verifiers with Self-Verified Feedback (SVF) obtained via self-evaluation prompts on intermediate completions. Across four mathematical reasoning and code generation benchmarks on three dLLMs, including LLaDA 8B Instruct, Dream 7B Instruct, and LLaDA 2.0-mini, our Prism achieves a favorable performance-efficiency trade-off, matching best-of-N performance with substantially fewer function evaluations (NFE). The code is released at https://github.com/viiika/Prism.