PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2602.01842v3cs.LG

TL;DR

Test-time scaling is underexplored for discrete diffusion language models because most methods assume autoregressive decoding, whereas dLLMs generate through parallel denoising. PRISM uses hierarchical trajectory search, partial remasking, and self-verified feedback to allocate compute efficiently. Across math and code benchmarks, it matches or approaches Best-of-N performance with substantially lower denoising compute, while self-verification remains imperfect on out-of-distribution inputs.

  • Problem

    Most test-time scaling methods target autoregressive decoding, leaving effective and efficient compute allocation for discrete diffusion language models underexplored.

  • Method

    PRISM combines hierarchical trajectory search, local branching with partial remasking, and Self-Verified Feedback from Yes/No self-evaluation prompts on intermediate completions.

  • Results

    Across four math and code benchmarks on three dLLMs, PRISM matches or approaches Best-of-N baselines with markedly reduced denoising compute.

  • Takeaways & Limitations

    PRISM demonstrates that dLLMs can achieve competitive reasoning and code-generation performance with significantly lower inference cost than vanilla width-scaling methods.

  • Takeaways & Limitations

    Self-verification may be imperfect or overconfident on out-of-distribution inputs, and PRISM inherits biases and limitations from its pretrained models.

Abstract

from arXiv · show

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.

1. Introduction

Test-time scaling is effective for reasoning but is predominantly designed for autoregressive decoding, creating an efficiency challenge for parallel discrete diffusion models. PRISM addresses this with adaptive trajectory search, partial remasking, and self-verification, achieving gains across math and code tasks with reduced denoising compute.

  • Autoregressive test-time scaling expands left-to-right search trees, where early mistakes are difficult to correct without backtracking.
  • Discrete diffusion language models iteratively denoise masked sequences with global bidirectional context, making them distinct candidates for planning and self-correction.
  • Increasing denoising steps is often impractical for dLLMs, while width scaling requires completing all candidate trajectories and can be computationally expensive.
  • PRISM combines Hierarchical Trajectory Search, local branching with partial remasking, and Self-Verified Feedback to adaptively explore and select trajectories without external reward models.
  • PRISM yields consistent gains over N=1 and matches or approaches Best-of-N baselines across four math and code benchmarks on three dLLMs with markedly reduced NFE.

2. Related Work

Related work establishes that dLLMs require inference-time compute allocation suited to partially denoised trajectories rather than autoregressive prefixes. PRISM therefore combines diffusion-aligned verification with hierarchical, budgeted search instead of exhaustive completion or probabilistic resampling.

  • Discrete Diffusion Language Models: Discrete diffusion models replace left-to-right decoding with Markovian denoising over token sequences and support non-autoregressive refinement with global bidirectional conditioning.
  • Discrete Diffusion Language Models: Prior dLLM scaling methods often provide marginal gains or require substantial computational overhead.
  • Test-Time Scaling and Verification: PRISM studies how to allocate test-time compute under multi-step denoising dynamics without relying on external verifiers.
  • Test-Time Scaling and Verification: Existing test-time scaling allocates compute through parallel candidate generation, refinement, and selection, with verification supplying control signals for pruning decisions.
  • Test-Time Scaling and Verification: Standard process verifiers can be brittle on dLLM intermediate states because those states are partially masked rather than left-to-right textual prefixes.
  • Test-Time Scaling and Verification: PRISM uses Yes/No self-verification on partially denoised states and couples it with hierarchical trajectory search to concentrate computation during structure formation.
  • Relation to PG-DLM and SMC-style inference: Unlike PG-DLM's reward-weighted probabilistic inference, PRISM uses heuristic SVF ranking, sparse top-S pruning, and local remasking under a fixed inference budget.

3. Method

PRISM combines Hierarchical Trajectory Search, local branching through partial remasking, and Self-Verified Feedback to allocate dLLM inference compute efficiently while retaining diverse candidate solutions.

  • Discrete Diffusion Language Models: PRISM uses masked diffusion and block-wise denoising, progressively unmasking the current block while keeping finalized preceding blocks fixed.The model predicts all positions, but the sampling schedule commits only the active block.
  • Hierarchical Trajectory Search: During thinning, HTS geometrically reduces the active trajectory pool and uses SVF scores to select top-S survivor seeds for branching.SVF-guided pruning is applied at intervals, with denoising continuing between pruning steps.
  • Self-Verified Feedback: SVF reuses the same dLLM as a binary verifier, scoring completed intermediate hypotheses through prompts that request a Yes-or-No correctness judgment.SVF is evaluated sparsely after a warm-up period because each evaluated hypothesis requires an additional forward pass.
  • Hierarchical Trajectory Search: HTS divides inference into stochastic exploration, progressive thinning, and final refinement, reallocating compute as trajectories become more structured.The trajectory pool remains wide during high-noise exploration, contracts during thinning, and undergoes final refinement on a small survivor set.
  • Local Branching: Local branching preserves high-confidence tokens as a logic skeleton while selectively remasking low-confidence positions to explore alternative realizations.The operator estimates token-wise uncertainty, samples a remasking subset from low-confidence positions, and continues denoising from the partially specified state.
  • Efficiency: HTS reduces denoising cost toward near-linear scaling because the trajectory pool contracts from N to K < N before final refinement, outperforming O(NT) linear search.Its cost combines exploration over N trajectories, hierarchical thinning, and refinement over K trajectories.

4. Experiments

Experiments evaluate PRISM across four math and code benchmarks, three dLLMs, baseline and verifier comparisons, efficiency trade-offs, and component analyses. PRISM consistently improves accuracy while reducing denoising compute relative to exhaustive Best-of-N search.

  • Setup: Experiments cover GSM8K, MATH500, HumanEval, and MBPP on LLaDA-8B-Instruct, Dream-7B-Instruct, and LLaDA-2.0-mini.Baselines include single-trajectory decoding and Best-of-N with N ∈ 4, 8, 16; math uses accuracy and code uses pass@1.
  • Main Results: At least 26% improvement over single-trajectory decoding is reported across all benchmarks and foundation models, at a cost comparable to Best-of-4.The evaluation uses PRISM with K=8 for this aggregate comparison.
  • Main Results: On LLaDA-8B, PRISM increases GSM8K accuracy from 67.58% to 85.30% and MATH500 accuracy from 26.40% to 42.80%.HumanEval and MBPP also improve by 24.39 and 16.40 points, respectively.
  • Efficiency-Accuracy Trade-off: PRISM reaches 85.30% on GSM8K with 1,048 NFE, while Best-of-16 reaches 87.50% with 4,096 NFE.This is an over 4× reduction in denoising cost with marginal accuracy degradation; on MATH500 and MBPP, PRISM often matches or surpasses Best-of-16 under one-third of the budget.
  • Target Width: Increasing target width K improves performance, with K=4 and K=8 offering the best balance between accuracy and efficiency.K=2 already provides noticeable gains with minimal overhead.
  • Verification and Ablation: SVF evaluations are fewer than 10% of total NFE in most settings, providing a verification signal without external reward models.SVF-guided pruning outperforms random pruning by 3.56, 4.32, and 3.71 points for K=2, 4, and 8, while removing partial remasking causes a smaller 1.13–1.89-point drop.

5. Conclusion

PRISM enables efficient test-time scaling for discrete diffusion language models by concentrating computation in the critical early-to-mid denoising window. It achieves competitive reasoning and code-generation performance at lower inference cost than vanilla width scaling.

  • PRISM concentrates computation on the critical early-to-mid denoising window through hierarchical search.
  • PRISM achieves competitive mathematical reasoning and code-generation performance with significantly lower inference cost than vanilla width-scaling methods.

Impact Statement

Prism reallocates computation through hierarchical search and uses lightweight self-verification to improve generation quality under constrained inference budgets. The approach can reduce additional memory overhead and inherits risks from general-purpose language models and pretrained models.

  • Prism targets improved reasoning and generation quality under a constrained inference budget by reallocating computation and using lightweight self-verification.The framework is designed for discrete diffusion language models.
  • Hierarchical search reallocates computation, while self-verification replaces external verifiers and can reduce additional memory overhead.
  • Self-verification may be imperfect or overconfident on out-of-distribution inputs, and stronger inference-time reasoning could generate harmful or misleading content.
  • Prism introduces no new data collection or user profiling but inherits biases and limitations from its underlying pretrained models.

A. Entropy Analysis

The entropy analysis tracks token-averaged predictive uncertainty across denoising trajectories and benchmarks. Entropy generally falls sharply early, then more gradually, while code-generation tasks retain greater trajectory variability than GSM8K.

  • The analysis samples eight independent stochastic DREAM-7B-INSTRUCT trajectories for each of GSM8K, HumanEval, Math-500, and MBPP.The trajectories use different random seeds under identical sampling hyperparameters.
  • At each timestep, token-level predictive entropies are averaged across the L positions to summarize uncertainty about token identities.Lower entropy indicates a sharper, more confident predictive distribution.
  • Across all four benchmarks, entropy drops sharply at early timesteps and then decays more gradually, with occasional non-monotonic bumps.The bumps reflect stochastic exploration and local ambiguity.
  • Code-generation benchmarks show greater trajectory-to-trajectory entropy variance than GSM8K, indicating multiple plausible partial programs before convergence.The observed difference is between HumanEval and MBPP versus GSM8K.

B. Hyperparameter Analysis

The hyperparameter analysis presents PRISM trade-offs between inference compute, measured by NFE, and benchmark performance. Separate figures cover HumanEval, GSM8K, Math500, and MBPP.

  • Figure 5 compares PRISM HumanEval Pass@1 with inference compute measured by NFE.
  • Figure 6 compares PRISM GSM8K Accuracy with inference compute measured by NFE.
  • Figure 7 compares PRISM Math500 Accuracy with inference compute measured by NFE.
  • Figure 8 compares PRISM MBPP Pass@1 with inference compute measured by NFE.

B.1. Analyses on HumanEval

HumanEval ablations identify moderate settings across Prism’s pruning and refinement controls as the strongest configurations. The best reported setting uses a pruning window of W = 0.1–0.6 and final target width K = 8, achieving 79.27% Pass@1 at 3.3× speedup.

  • Pruning window: 79.27% Pass@1 is achieved with pruning window W = 0.1–0.6, the best result among Prism configurations in the analysis.Pruning too early or too late degrades performance.
  • Pruning interval: Pruning every i = 3 inference steps performs best; more frequent pruning can discard promising trajectories, while sparser pruning weakens compute reallocation.
  • Decay: Decay factor d = 1.8 yields the strongest results, whereas slower or more aggressive thinning reduces Pass@1.
  • Survivors: The survivor width S must balance diversity and focus: too few survivors harm diversity, while too many dilute branching.
  • Final target: K = 8 provides a strong efficiency–accuracy trade-off, reaching 79.27% at 3.3× speedup while larger K increases NFE.

B.2. Analyses on GSM8K

The GSM8K analysis examines Prism’s pruning-window, interval, decay, survivor, and final-target controls, alongside task-specific self-verification prompts. The supplied results emphasize that moderate pruning and refinement settings balance accuracy, diversity, and computation.

  • Ablation design: GSM8K ablations vary the pruning window, interval, decay factor, survivor width, and final target width under fixed settings for the remaining controls.
  • Self-verification: Math and code tasks use separate judge prompts that include the problem statement and truncated completion, with a single-word Yes/No verifier output.
  • Comparisons: The qualitative comparisons include the LLaDA 8B Instruct baseline across four benchmarks.

D.1.2. DREAM 7B INSTRUCT

The Dream 7B examples contain both correct and incorrect arithmetic reasoning across candle-profit and clothing-cost tasks. The outputs illustrate substantial variation in intermediate calculations and final answers.

  • Candle-profit example: Another candle calculation uses $1.00 supply cost per candle and concludes a $20 total net profit for 20 candles.
  • Candle-profit example: A candle-profit example reports $60 after calculating $40 earnings and subtracting $100 in supplies, despite the stated per-10-candle supply cost.
  • Clothing-cost example: For Mishka’s clothing purchase, one Dream output states $84.50 for shoes and concludes $84.5 total expenditure.
  • Clothing-cost example: The supplied itemized calculation instead gives $49.50 for shorts, $67.50 for pants, $126.00 for shoes, and $243 total.

D.2. HumanEval

The HumanEval examples show Prism-related code outputs spanning string XOR, mean absolute deviation, greatest common divisor, and duplicate detection, alongside mathematical reasoning examples. The passages include both coherent implementations and visibly inconsistent or incorrect completions.

  • Code examples: The string-XOR examples define binary-string inputs and show implementations iterating over paired characters, while another completion contains malformed typing and indexing code.
  • Code examples: Mean absolute deviation is described as average absolute distance from the mean, with a function signature shown in one completion.
  • Code examples: The greatest-common-divisor examples use the recursive Euclidean pattern returning gcd(b, a % b), with examples gcd(3, 5) = 1 and gcd(25, 15) = 5.
  • Math examples: The polar-coordinate example identifies the point (0, 3) and reports polar coordinates beginning with (3, π.
  • Math examples: For the digit-restricted multiple-of-30 problem, the supplied model outputs disagree, giving 2020 and 2220 as purported answers.
  • Code examples: Duplicate detection is implemented by tracking seen elements and returning True when an element repeats, otherwise returning False.
Loading 2602.01842v3…