DPRM: A Plug-in Doob h transform-induced Token-Ordering Module for Diffusion Language Models

Abstract: Diffusion LLMs generate without a fixed left-to-right order, making token ordering a central algorithmic choice: which tokens should be revealed, retained, revised or verified at each step? Existing systems mainly use random masking or confidence-driven ordering. Random masking creates train--test mismatch, while confidence-only rules are efficient but can be myopic and suppress useful exploration. We introduce DPRM (Doob h-transform Process Reward Model), a plug-in token-ordering module for diffusion LLMs. DPRM keeps the host architecture, denoising objective and supervision unchanged, and changes only the ordering policy. It starts from confidence-driven progressive ordering and gradually shifts to Doob h transform Process Reward guided ordering through online estimates. We characterize the exact DPRM policy as a reward-tilted Gibbs reveal law, prove O(1/N) convergence of the stagewise Soft-BoN approximation, and show that the online bucketized controller tracks the exact DPRM score at empirical-Bernstein rates. Under tractable optimization assumptions, DPRM also yields a sample-complexity advantage over random and confidence-only ordering. DPRM improves over confidence-based baselines in pretraining, post-training, test-time scaling, and single-cell masked diffusion, with particularly strong gains on harder reasoning subsets. In protein, molecular generation and DNA design, the effect is more multi-objective: ordering-aware variants significantly improve selected structural or fragment-constrained metrics while not uniformly dominating the host baseline on every quality metric. These results identify token ordering as a fundamental control axis in diffusion LLMs and establish DPRM as a general-purpose module for improving it. Code is available at https://github.com/DakeBU/DPRM-DLLM.

• The paper introduces DPRM as a token-ordering module that leverages Doob’s h-transform to guide reveal decisions in diffusion language models.
• It employs a progressive shift from confidence-driven to reward-guided unmasking, improving sample efficiency and convergence rates on challenging tasks.
• Empirical results demonstrate significant accuracy gains across language and scientific domains, supported by rigorous theoretical guarantees and diagnostic evaluations.

DPRM: Doob $h$-Transform Token-Ordering for Diffusion LMs

Motivation and Overview

Diffusion LLMs (DLMs) eschew the left-to-right generative order of their autoregressive counterparts, exposing token ordering as a central control variable. Existing heuristics—random masking and confidence-based unmasking—are suboptimal, either generating a pronounced train–test mismatch or producing myopic, locally greedy exploration. The authors introduce the Doob $h$-transform Process Reward Model (DPRM), a token-ordering module applicable to any host discrete diffusion model, which systematically transitions from confidence-driven unmasking to a reward-tilted policy using online reward estimation. Unlike previous approaches, DPRM leaves the host model, architecture, loss, and data pipeline unchanged, modifying solely the reveal order.

Figure 1: DPRM as a plug-in token-ordering module, implementing a progressive transition from confidence-based to reward-tilted ordering via Soft-BoN reweighting and online process reward estimation.

DPRM is motivated by the insight that the token-reveal process is a Markov chain on partially observed states, and that optimal reveal sequences should be chosen not just by the model’s local confidence, but by their anticipated contributions to future (or terminal) reward, made tractable through Doob’s $h$-transform. In practice, computing the exact future reward for all reveal actions is infeasible, so the authors propose an empirical, bucketized, and shortlist-based approximation.

Progressive Online DPRM: Algorithmic Details

DPRM defines the optimal transition kernel at each reveal stage as a Gibbs law tilted by the future expected reward. For each candidate token action, the score is adjusted by an online estimate of the log-moment generating function of terminal reward conditioned on that action:

$$g_i = \log \psi_i(s) + \beta \widehat R_{\phi, b_i}$$

where $\psi_i(s)$ is the local confidence proposal in state $s$, $\beta$ is a guidance strength factor, and $\widehat R_{\phi, b_i}$ is the mean log-moment reward for the corresponding phase $\phi$ and confidence bucket $b_i$. During early training (warmup), the controller prioritizes confidence and gradually interpolates to reward-guided ordering as bucket statistics become reliable.

The algorithm draws a shortlist of candidate actions, estimates their process-reward-corrected scores, and re-ranks accordingly before selecting tokens to reveal. This structure preserves strict alignment between train–test masking states, supports theoretical tractability, and introduces only computational overhead at the ordering stage.

Theoretical Guarantees

The paper provides rigorous theoretical support for DPRM. Key results include:

• Train–test alignment: The teacher-forced progressive masking preserves the Bayes-optimal denoiser and improves sample efficiency exponentially over random masking.
• Online bucket estimation: The gap between practical and optimal process-reward-corrected scores is controlled via empirical-Bernstein rates, with estimation error decaying as $h$0 where $h$1 is the number of bucketed samples.
• Soft-BoN convergence: Stagewise shortlist-based Soft-BoN approximations converge in KL divergence to the exact reward-tilted law as $h$2, where $h$3 is the shortlist size.
• Sample complexity superiority: Under plausible optimization assumptions, DPRM achieves an exponential speedup over both random and purely confidence-based ordering, particularly in late-stage convergence on hard or out-of-distribution subsets.

Empirical Results: Language and Scientific Domains

The evaluation of DPRM spans seven host frameworks, keeping the model, objective, and data pipeline fixed and varying only the ordering controller. The module consistently improves over confidence-based and random reveal policies.

• Language Reasoning (GSM8K, MATH, Countdown): DPRM-PUMA improves GSM8K validation from 29.34 to 34.27 (+16.8%); DMPO-DPRM elevates MATH Hard from 44.3 to 47.9 (+8.1%) and Countdown Hard from 29.6 to 33.4 (+12.8%). The largest gains are on hard and out-of-distribution reasoning subsets.

  Figure 2: GSM8K results for PUMA vs. DPRM-PUMA, with DPRM boosting accuracy across both unmasking protocols.

  

  Figure 3: GSM8K pass@K curves by difficulty level indicating DPRM’s enhancement is concentrated on harder levels and at larger $h$4.

  

  Figure 4: MATH pass@K curves by difficulty showing DPRM consistently outperforms confidence-only ordering, especially on hard and OOD problems.

  

  Figure 5: Countdown task pass@K curves; DPRM rescues the collapse of vanilla DMPO and dominates at all operand levels.

• Test-time Scaling (Prism): DPRM-Prism yields monotonic improvements in per-rank GSM8K accuracy (Figure 6) and shifts the NFE–accuracy trade-off curve (Figure 7), at the cost of additional computation.

  Figure 6: DPRM-Prism ensures improvement at every rank of output candidate on GSM8K.

• Protein, Gene-Expression, Molecule, DNA Sequence Generation: DPRM, when applied to protein inverse folding, gene-expression recovery, de novo molecule generation, and DNA regulatory sequence design, delivers domain-specific improvements: reduction in RMSD, increase in TM-score, higher token recovery and zero-expression accuracy (Figure 8), improved molecule validity and fragment quality, and preservation of critical reward metrics not attainable with confidence-based ordering alone (Figures 9-14). The module enables explicit trade-off management across quality axes.

  Figure 9: Forward-folding protein structure metrics; all ordering-aware variants outperform baseline.

  

  Figure 8: Gene-expression imputation with DCM—DPRM(random) achieves maximal recovery and zero-expression accuracy.

  

  Figure 10: Balancing between HepG2, ATAC, k-mer and log-likelihood metrics in DNA sequence design—DPRM maintains higher regulatory quality than confidence-only ordering.

Diagnostics and Optimization Insights

To validate the sample-complexity and alignment theory, the paper provides detailed diagnostics using testbeds like the Countdown task. These show that:

• Confidence order is a strong early-stage proxy for true CE loss (Figure 16a), but myopic at late stages.
• DPRM increases exploration in low-confidence regions (Figure 12).
• DPRM’s advantage is magnified at higher candidate budget $h$5 and on difficult/OOD data (Figure 13).

  Figure 11: Early- and late-stage diagnostics—confidence alignment is effective initially, but DPRM carries more mass in low-confidence bins as $h$6 increases.

  

  Figure 13: DPRM’s pass@K advantage is concentrated among hard and OOD examples, especially at larger $h$7.

Implications, Limitations, and Future Trajectories

This work establishes token ordering as a reusable, high-leverage axis for controlling learning and generation in masked diffusion models across modalities. From a theoretical perspective, DPRM frames token selection within the language of Markov control and optimal variance sampling, importing the Doob $h$8-transform as a general framework for integrating downstream reward awareness into discrete generation processes.

Pragmatically, the separation of ordering from architecture or loss makes DPRM highly modular and suitable for plug-in upgrades. The empirical results highlight an important contradiction: confidence-only policies are neither information-theoretically nor computationally optimal for diffusion-style models. The reward-tilted Gibbs law targeted by DPRM yields faster and more reliable convergence, especially in the low-entropy or difficult-subset regime, where myopic confidence plateaus.

Nevertheless, current reliance on bucketized, trajectory-averaged estimates limits theoretical guarantees to the regime where the process reward statistics are well-sampled. There may also be settings where the reward used in the Doob transform is challenging to design or compute. Adapting the method for fine-grained and multi-objective settings, and for non-linguistic modalities (visual, multimodal, structured graphs), is a natural direction.

Potential extensions include tighter finite-sample bounds, adaptive switching schedules between confidence and reward, more expressive action-state abstraction in reward estimation, and application to new domains such as masked visual diffusion or predictive representation learning (e.g., I-JEPA, V-JEPA).

Conclusion

DPRM introduces a general, theoretically-justified, and empirically robust module for token ordering in diffusion LLMs, leveraging the Doob $h$9-transform to inject future reward information into reveal decisions. The approach outperforms random and confidence-based masking in both language and scientific domains, optimizes exploration–exploitation trade-offs, and defines new standards for sample efficiency and controllable generative modeling. Its modular design and broad applicability position token-ordering control as a foundational component for future discrete diffusion frameworks (2604.24357).