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

Jinbin Bai, Yixuan Li, Yuchen Zhu, Yi Xin, Qingyu Shi, Aosong Feng, Xiaohong Liu, Molei Tao, Jianru Xue, Xiangtai Li, Ming-Hsuan Yang
Proceedings of the 43rd International Conference on Machine Learning, PMLR 306:5348-5388, 2026.

Abstract

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 LLaDA-S, 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) replaces external verifiers with Self-Verified Feedback (SVF) obtained via self-evaluation prompts on intermediate completions, and (iii) introduces Local branching with partial remasking to explore diverse implementations while preserving a high-confidence tokens. Across four mathematical reasoning and code generation benchmarks on three dLLMs, including LLaDA 8B Instruct, Dream 7B Instruct, and LLaDA 2.0-mini, our LLaDA-S achieves a favorable performance-efficiency trade-off, matching best-of-$N$ performance with substantially fewer function evaluations (NFE). The code will be released.

Cite this Paper

BibTeX
@InProceedings{pmlr-v306-bai26f, title = {{PRISM}: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models}, author = {Bai, Jinbin and Li, Yixuan and Zhu, Yuchen and Xin, Yi and Shi, Qingyu and Feng, Aosong and Liu, Xiaohong and Tao, Molei and Xue, Jianru and Li, Xiangtai and Yang, Ming-Hsuan}, booktitle = {Proceedings of the 43rd International Conference on Machine Learning}, pages = {5348--5388}, year = {2026}, editor = {Zhang, Tong and Dudik, Miroslav and Jaggi, Martin and Agarwal, Alekh and Li, Sharon and Schuurmans, Dale and Zhu, Jerry and Berkenkamp, Felix and Dong, Hanze and Bietti, Alberto}, volume = {306}, series = {Proceedings of Machine Learning Research}, month = {06--11 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v306/main/assets/bai26f/bai26f.pdf}, url = {https://proceedings.mlr.press/v306/bai26f.html}, abstract = {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 LLaDA-S, 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) replaces external verifiers with Self-Verified Feedback (SVF) obtained via self-evaluation prompts on intermediate completions, and (iii) introduces Local branching with partial remasking to explore diverse implementations while preserving a high-confidence tokens. Across four mathematical reasoning and code generation benchmarks on three dLLMs, including LLaDA 8B Instruct, Dream 7B Instruct, and LLaDA 2.0-mini, our LLaDA-S achieves a favorable performance-efficiency trade-off, matching best-of-$N$ performance with substantially fewer function evaluations (NFE). The code will be released.} }
Endnote
%0 Conference Paper %T PRISM: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models %A Jinbin Bai %A Yixuan Li %A Yuchen Zhu %A Yi Xin %A Qingyu Shi %A Aosong Feng %A Xiaohong Liu %A Molei Tao %A Jianru Xue %A Xiangtai Li %A Ming-Hsuan Yang %B Proceedings of the 43rd International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2026 %E Tong Zhang %E Miroslav Dudik %E Martin Jaggi %E Alekh Agarwal %E Sharon Li %E Dale Schuurmans %E Jerry Zhu %E Felix Berkenkamp %E Hanze Dong %E Alberto Bietti %F pmlr-v306-bai26f %I PMLR %P 5348--5388 %U https://proceedings.mlr.press/v306/bai26f.html %V 306 %X 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 LLaDA-S, 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) replaces external verifiers with Self-Verified Feedback (SVF) obtained via self-evaluation prompts on intermediate completions, and (iii) introduces Local branching with partial remasking to explore diverse implementations while preserving a high-confidence tokens. Across four mathematical reasoning and code generation benchmarks on three dLLMs, including LLaDA 8B Instruct, Dream 7B Instruct, and LLaDA 2.0-mini, our LLaDA-S achieves a favorable performance-efficiency trade-off, matching best-of-$N$ performance with substantially fewer function evaluations (NFE). The code will be released.
APA
Bai, J., Li, Y., Zhu, Y., Xin, Y., Shi, Q., Feng, A., Liu, X., Tao, M., Xue, J., Li, X. & Yang, M.. (2026). PRISM: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models. Proceedings of the 43rd International Conference on Machine Learning, in Proceedings of Machine Learning Research 306:5348-5388 Available from https://proceedings.mlr.press/v306/bai26f.html.

Related Material