Why Deterministic PRM Guidance Underperforms in Discrete Diffusion Reasoning

AI summary

Overview

Research area: Test-time compute scaling for discrete diffusion language models (dLLMs), specifically process reward model (PRM) guidance versus outcome reward model (ORM) reranking for mathematical and code reasoning.

Technical level: Intermediate. The paper assumes familiarity with best-of-N sampling, reward models, and diffusion-style denoising, but its core argument is a compute-accounting comparison that does not require deep diffusion theory.

Scope: A controlled, matched-forward-pass comparison of deterministic PRM-guided dLLM denoising against independent sampling plus ORM reranking, with a decomposition of the resulting accuracy gap into candidate-pool damage and terminal-selection error.

What This Paper Is About

Discrete diffusion language models produce a partially denoised solution at every step, so it seems natural to train a process reward model that scores those intermediate "snapshots" and steers denoising toward better answers. This paper shows that when denoising passes, PRM scoring passes, and ORM scoring passes are all charged in the same unit, the natural deterministic form of PRM guidance loses to a much simpler baseline: sample several complete solutions independently and pick the one an outcome reward model scores highest. The authors then trace the loss to two separable causes and release the snapshot corpus and evaluation toolkit so others can reproduce the comparison.

Key Contributions

1. A matched forward-pass protocol for comparing reward-guided and unguided dLLM reasoning at fixed inference budgets, charging denoising, PRM scoring, and ORM scoring in the same unit (one forward pass through a dLLM-scale model for one candidate state).

2. A matched-compute ranking that holds across math and code on Dream-v0-Instruct-7B: deterministic top-1 PRM guidance trails ORM Rerank by 9.95 pp on GSM8K at K=8 and 12.69 pp at K=32, by 9.85 pp on MATH, and by 12.16 pp on MBPP.

3. A decomposition of the gap into candidate-pool damage and terminal selection: PRM ROC-AUC decays from 0.77 to 0.54 with mask ratio, top-1 pruning cuts the Oracle ceiling by 13.75 pp, an SMC sampler at matched budget restores most of that ceiling yet still selects at the top-1 level, and a PRM retrained on final states matches the ORM on identical candidates.

4. An architectural finding for PRM design: bidirectional PRMs beat causal PRMs on Dream-7B and LLaDA-8B-Base, and a readout fix recovers most of the causal gap — last-token pooling raises final-state ROC-AUC from 0.61 to 0.73.

5. Released artifacts: a corpus of denoising states with outcome labels and an evaluation toolkit for reproducible comparisons at matched compute (code at github.com/dLLM-PRM-Gap/, data and model collection at huggingface.co/collections/YanZhanPKU/dllm-prm-gap).

Main Findings

• ORM reranking wins at matched compute on GSM8K. With Dream-v0-Instruct-7B, ORM Rerank@8 reaches 75.13% while PRM Guided (K=8) reaches 65.18%; at the larger budget, ORM Rerank@32 reaches 82.71% while PRM Guided (K=32) reaches 70.02%. The gap is 9.95 pp at K=8 and 12.69 pp at K=32.

• The ordering survives stronger baselines and wall-clock matching. PRM Guided still beats Majority@N at both budgets, so the PRM carries real signal, but ORM Rerank beats Majority@N by 15.1 pp. Measured wall-clock ratios agree with forward-pass predictions within 9%, and ORM Rerank@6 reaches 72.40% in about 126 s per problem versus 65.18% in 156 s for PRM Guided K=8. ORM Rerank@8 beats every PRM Guided budget run, including K=32; the best PRM Guided result, 73.67% at K=24, still trails it.

• Headroom remains for better verifiers. Even ORM Rerank@32 sits 8.4 pp below the 91.13% of Oracle@32. Verifier-weighted voting stays within 0.4 pp of ORM Rerank at every budget.

• The failure has two separable parts. Module one is pool damage: guidance prunes on a weak signal. Module two is terminal selection: the PRM trained across mask ratios separates correct from incorrect final solutions with ROC-AUC 0.78, against 0.96 for the ORM trained on final states alone.

• PRM signal decays monotonically with mask ratio. Bidirectional PRM ROC-AUC falls from 0.77 on nearly decoded states to 0.54 on almost fully masked ones. The decay persists when 10,000 training states are relabeled with eight fresh rollouts each and retrained: pooled ROC-AUC rises by only 0.011 and GSM8K accuracy moves by 0.61 pp, small next to the 10 pp gap to ORM Rerank. The two label sources disagree on 21% of states in the most-masked bucket against 0.5% in the least-masked one.

• Deterministic pruning collapses diversity and the Oracle ceiling. At K=N=8, PRM Hybrid keeps only 1.75 unique answers per problem against 4.31 for independent samples, and answer entropy falls fourfold. Oracle@8 drops from 81.05% for independent samples to 67.30% for PRM Hybrid, a 13.75 pp ceiling loss. Even a perfect selector over the PRM Hybrid K=8 pool reaches only 67.30%, 7.83 pp below ORM Rerank@8.

• Pool damage happens early in denoising. An offline counterfactual over stored trajectories finds a top-1 cut removes every lineage that eventually reaches a correct answer in 46% of cases at the initial stored state and 20% at the final one. Keeping four candidates brings the risk to 10% at the middle state and 6% at the final one.

• Repairing the pool does not repair the final choice. An SMC sampler resampling K=8 particles by PRM score at the same 1,040 passes raises Oracle@8 from 67.30% to 77.89%, recovering 10.59 of the 13.75 points lost by top-1 pruning. Its PRM-selected accuracy stays at 65.48% with weighted answer voting and 66.34% with the top-scoring particle — level with top-1 guidance and well below ORM Rerank@8's 75.13%.

• Pooled discrimination does not imply within-problem ranking. The cross-mask PRM has pooled ROC-AUC 0.78 on final states yet at N=8 reranks no better than random selection. Its signal appears in larger pools, beating random selection by 8.6 pp at N=16 and 22.1 pp at N=32. On the 1,182 test problems containing both correct and incorrect candidates, its scores correlate with correctness with a median Kendall τ of 0.44 and are higher on average for correct candidates in 88% of them.

• A PRM trained on final states matches the ORM. On the same GSM8K candidate pool, the final-state PRM reaches 75.40% at N=8 and 82.79% at N=32, versus 75.13% and 82.71% for the ORM; the cross-mask PRM reaches only 42.84% and 65.35%.

• MBPP isolates the cost of guidance itself. The MBPP PRM reaches 65.47% when reranking finished programs, on par with the ORM, but only 50.88% when it guides denoising — a 14.6 pp loss. Task-specific verifiers carry the ordering over to other tasks: on MATH at N=K=8, ORM Rerank reaches 30.65% against 20.80% for both PRM Rerank and deterministic PRM Guided; on MBPP, ORM Rerank reaches 63.04% against 50.88% for PRM Guided.

• Causal PRMs fail largely because of readout, not attention mask. On final states, the mean-pooled causal PRM reaches ROC-AUC 0.61 against 0.78 for the bidirectional PRM; last-token pooling raises it to 0.73 and closes about 70% of the gap. Doubling training from 15K to 31K steps narrows the classification-accuracy gap in the least-masked bucket from 13.6 to 9.5 pp, with the bidirectional PRM ahead in all ten mask buckets; trained on final states only, the causal scorer still trails by 7.8 pp. At N=32 the mean-pooled causal reranker falls to 40.56%, below a single sample, while the last-token variant rises to 50.64%. Under the ORM protocol, the bidirectional scorer reaches 91.83% with last-token pooling and 91.78% with mean pooling.

• The failure is visible on individual problems. On a single GSM8K problem (id 526), the causal PRM ranks a wrong trajectory first and the correct one last of three, while the ORM selects the correct one. The same inversion occurs on 639 of 1,319 test problems, or 48%.

• A second backbone reproduces the bidirectional advantage. On LLaDA-8B-Base, averaged over all tested configurations on the full GSM8K test set, bidirectional PRM guidance reaches 31.64% against 22.25% for causal PRM guidance; over the 20.77% Vanilla baseline, bidirectional guidance adds 10.87 pp and causal guidance only 1.48 pp.

• Right context alone does not explain the bidirectional advantage. Zeroing the right half of each snapshot changes the bidirectional PRM's ROC-AUC by about 0.001, and averaging a causal PRM over forward and reversed inputs adds nothing and leaves it 0.084 below the bidirectional PRM, because the reversed pass scores at chance (ROC-AUC 0.500).

• Extraction choices matter for the baseline numbers. The strict regex extractor gives Vanilla accuracy 43.14% on Dream-7B, while the lm-eval-harness flexible-extract regex gives 54.28% — a +11.14 pp offset from extraction alone, which is why strict extraction is used throughout.

Methodology in Plain English

The authors set up a single accounting unit — one forward pass through a dLLM-scale model for one candidate state — and charge every method in it. Vanilla sampling costs 128 passes per trajectory (T=128 denoising steps). The segmental PRM Guided procedure branches to K candidates every b steps, denoises all K, scores all K with the PRM, and keeps only the highest-scoring one; with b=64 and T=128 this gives 2 segments and 130K passes (1,040 at K=8, 4,160 at K=32). The matched ORM Rerank budgets are 1,032 and 4,128, so the comparison is matched within about 0.8%, and the small asymmetry favors PRM Guided.

The PRM is a frozen dLLM backbone with trainable LoRA adapters (r=16, α=32 on q_proj and v_proj) plus a two-layer MLP reward head over pooled solution hidden states and a 256-dimensional sinusoidal step-index embedding. It is trained with BCE loss for 2,000 steps at batch size 32 with a cosine schedule from 2×10⁻⁵ to 0, on on-policy intermediate states from Dream-7B denoising trajectories on GSM8K training problems, labeled with binary final-correctness. Training and test trajectories share the sampler and snapshot schedule, so their mask-ratio distributions match by construction. All training, tuning, and early stopping use the GSM8K train split.

Evaluation is on GSM8K test (1,319 problems; the best-of-32 pool contains 42,208 complete trajectories) with a strict regex answer extractor, using a shared Dream sampling configuration (temperature=0.5, alg_temp=0.5, top_p=1.0, T=128). MATH500 serves as an out-of-distribution test for the GSM8K-trained scorers, while the task-specific MATH and MBPP controls train their own verifiers. Baselines are separated into sampling (Vanilla, Majority@N), scoring and guidance (PRM Guided), and oracle headroom (Oracle@N/Oracle@K).

Each diagnostic question gets a direct measurement: accuracy at matched compute, PRM ROC-AUC across mask-ratio buckets, answer diversity and Oracle@K of the guided pool, reranking of one shared candidate pool by different scorers plus PRM selection after SMC restores the pool, and readout ablations on causal PRMs. Controls include a top-M retention variant (each retained state spawns K/M children so total width stays K), an ESS-tempered SMC variant with K particles weighted by tempered PRM scores after 64 denoising steps and resampled when effective sample size falls below a threshold, and a relabeling experiment with fresh rollouts.

Why This Matters

Impact on research. The paper reframes dLLM reward guidance as a budget-allocation question rather than a modeling question. It shows that intermediate-state scoring carries usable signal but is spent in the wrong place — early, where the scorer is least reliable, and at the terminal step, where a final-state verifier would do better. The decomposition (pool damage versus selection error) is argued to apply to any search that prunes on intermediate scores and selects with the same scorer, including PRM beam search for autoregressive models.

Real-world applications:

• Test-time compute allocation for reasoning models. Systems that spend extra inference compute on math or code problems can redirect that budget from guided pruning to independent sampling plus final-answer verification.

• Code generation with verifiable outputs. The MBPP result shows a PRM that judges finished programs as well as the ORM can still lose 14.6 pp when it guides denoising, so pipeline designers should place the PRM at the end of the loop rather than inside it.

• Reward-model training pipelines. Practitioners training verifiers on intermediate states can compare against the released snapshot corpus and toolkit before committing to guidance-style decoding.

• Evaluation practice. Reporting Oracle@K next to accuracy exposes candidate-pool collapse independently of the final selector, and charging scorers in the same forward-pass unit removes hidden scorer cost as a confounder.

Industry relevance. Anyone deploying best-of-N with verifiers, self-consistency, or beam-style search on diffusion or masked-token generators has a direct interest in these results, since the recommended configuration — sparse branching, keep more than one candidate, verify at the end — is cheaper and simpler than the guided alternative it beats.

Future Directions

• Test whether bidirectional attention wins through joint modeling of unordered positions. The authors list this as the most actionable of three open explanations and propose a permutation-LM scorer trained over arbitrary orderings yet autoregressive at inference as a direct test. A shift in training distribution is a second candidate explanation.

• Benchmark other diversity-preserving stochastic decoders. The SMC control covers only one such sampler; the released toolkit is intended for benchmarking others.

• Run an LLaDA-specific ORM comparison. LLaDA-8B-Base reproduces the bidirectional advantage, but the ORM comparison has not been run on that backbone.

• Build guidance algorithms that keep correct partial solutions alive and defer the final choice. The authors point to energy-tilted sampling (blending the dLLM transition with a PRM-weighted distribution via an annealing schedule) and adaptive branching schedules that condition the denoising interval b on the current mask ratio, consulting the PRM only where the ROC-AUC curve is high.

Target Audience

Researchers and engineers working on test-time compute scaling, reward models and verifiers, and diffusion or masked-token language models. It is most useful to practitioners deciding where to place a reward model in a decoding pipeline, and to those evaluating guided decoding claims who need a matched-compute protocol and diagnostics (Oracle@K, ROC-AUC by mask ratio, answer diversity) to separate sampling gains from selection gains. Readers comparing PRM beam search in autoregressive settings will also find the pool-damage versus selection-error decomposition transferable.

Authors’ abstract

Discrete diffusion language models (dLLMs) expose a denoised solution at every step, which makes process reward model (PRM) guidance look like a way to spend compute at test time. We show that once denoising, PRM scoring, and outcome reward model (ORM) scoring are charged in the same budget of forward passes, its deterministic form loses to a much simpler baseline. Our PRMs score intermediate denoising states and are trained on the correctness of the final answer. On Dream-v0-Instruct-7B with 8 candidates per GSM8K problem, keeping the candidate with the highest PRM score at every scoring step reaches 65.18%, while independent sampling plus an ORM reranker trained for the task reaches 75.13%. The gap grows to 12.69 percentage points (pp) with 32 candidates, and is 9.85 pp on MATH and 12.16 pp on MBPP. We trace it to two separable failures. First, guidance prunes on a weak signal: on GSM8K, PRM ROC-AUC falls from 0.77 to 0.54 as the mask ratio rises, a decay that persists when states are relabeled with fresh rollouts, and pruning lowers the best accuracy reachable from the candidate pool from 81.05% for independent samples to 67.30%. Second, on GSM8K and MATH, the PRM is a poor final judge: a sequential Monte Carlo sampler at the same budget restores that ceiling to 77.89%, yet selecting with the PRM gives 65.48%, on par with deterministic guidance, while a PRM retrained on final states matches the ORM on identical candidates. MBPP separates the two: there the PRM reaches 65.47% when reranking finished programs, on par with the ORM, but 50.88% when it guides denoising. The results point to two targets for dLLM guidance: keep correct partial solutions alive through early denoising, and leave the final choice to a verifier trained on final states. We release the corpus of denoising states with outcome labels and evaluation toolkit for reproducible comparisons at matched compute.