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

TL;DR

Prism tackles the challenge of test-time scaling for discrete diffusion language models by introducing Hierarchical Trajectory Search to adaptively allocate compute, Local Branching via partial remasking to preserve diverse yet coherent candidates, and Self-Verified Feedback to prune without external verifiers. The method delivers near-linear compute scaling, $C_{ ext{HTS}} \approx O(N + KT)$, and demonstrates substantial accuracy gains over single-trajectory decoding while matching or approaching Best-of-$N$ performance with far fewer denoising evaluations across math and code benchmarks on multiple dLLMs. Through extensive experiments, Prism shows robust improvements in GSM8K, MATH500, HumanEval, and MBPP with a favorable efficiency-accuracy trade-off, and SVF typically incurs minimal overhead compared to external verifiers. The work highlights how diffusion-aligned TTS can unlock strong reasoning capabilities in non-autoregressive LMs while reducing memory and compute demands, with code released for reproducibility.

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 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.

Figures (13)

• Figure 1: Comparison between Best-of-N and Prism on LLaDA-8B-Instruct. The red curve illustrates Best-of-N scaling, while the blue curve depicts Prism scaling, with a dashed line indicating the difference in inference compute (NFE) with comparable accuracy.
• Figure 2: Overview of Prism. (a) Given a prompt, multiple diffusion trajectories are generated in parallel, and intermediate completions are evaluated by Self-Verified Feedback (SVF) using the same dLLM. (b) Hierarchical Trajectory Search (HTS) allocates inference compute dynamically across different stages: stochastic exploration, progressive thinning with SVF-guided pruning and branching, and final refinement on a small survivor set. During thinning, local branching via partial remasking selectively re-masks low-confidence tokens to explore diverse realizations while preserving a high-confidence logic skeleton.
• Figure 3: Prism strategy trade-off between HumanEval Pass@1 and inference compute (NFE). An enlarged version is in App.\ref{['app:ablation']}.
• Figure 4: Token-averaged predictive entropy trajectories of Dream-7B-Instruct on GSM8K. Each curve corresponds to one independently sampled decoding trajectory under identical inference settings, and the y-axis reports $\mathcal{H}(t)$ from Eq. \ref{['eq:token_avg_entropy']} (entropy averaged over all token positions at each timestep). The light gray curves indicate trajectories that are pruned during thinning (shown only up to the timestep where they are discarded). Entropy decreases rapidly at the beginning, followed by a smoother decay with mild mid-trajectory fluctuations, and all runs converge to a low-entropy regime near the end of decoding, indicating increasing confidence in token identities as denoising progresses.
• Figure 5: Token-averaged predictive entropy trajectories of Dream-7B-Instruct on HumanEval. Each curve corresponds to one independently sampled decoding trajectory under identical inference settings, and the y-axis reports $\mathcal{H}(t)$ from Eq. \ref{['eq:token_avg_entropy']} (entropy averaged over all token positions at each timestep). The light gray curves indicate trajectories that are pruned during thinning (shown only up to the timestep where they are discarded). Compared with GSM8K, the curves exhibit a more pronounced high-entropy plateau and larger inter-trajectory variance in the early-to-mid timesteps, consistent with multiple competing program structures remaining plausible for longer. Despite such variability, all trajectories eventually enter a low-entropy phase and converge toward completion.