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

2026-04-27

AI Summary

Diffusion language models suffer from the absence of a fixed generation order, rendering existing token scheduling strategies vulnerable to train-test mismatch and myopic decision-making. This work proposes DPRM, a novel module that, without altering the backbone architecture or training objective, introduces the Doob h-transform into diffusion language models for the first time to construct a reward-tilted Gibbs revelation policy that dynamically optimizes token revelation order. Integrating process reward modeling, staged Soft-BoN approximation, bucketed online control, and empirical Bernstein convergence analysis, DPRM enables efficient online inference. Experiments demonstrate that DPRM consistently outperforms baselines across pretraining, post-training, test-time scaling, and single-cell masked diffusion tasks, with particularly pronounced gains in complex reasoning and biological sequence generation scenarios.

Technology Category

• Planning, Routing, and Scheduling: Planning with Language Models
• Natural Language Processing: (Large) Language Models
• Computer Vision: Diffusion Models for Vision

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

Diffusion language models 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 language models. 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 language models and establish DPRM as a general-purpose module for improving it. Code is available at https://github.com/DakeBU/DPRM-DLLM.

Problem

• diffusion language models
• token ordering
• Doob h-transform
• process reward model
• generation strategy

Innovation

• Doob h-transform
• token ordering
• diffusion language models
• Process Reward Model
• online policy adaptation