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

Technical Summary: DPRM: A Plug-in Token-Ordering Module for Diffusion Language Models

1. Problem Statement

Diffusion Language Models (DLMs) offer a flexible alternative to autoregressive models because they do not require a fixed left-to-right generation order. However, this flexibility introduces a critical algorithmic challenge: token ordering.

Existing DLM systems primarily rely on two ordering strategies:

• Random Masking: Suffers from a significant train–test mismatch, as the model is trained on random patterns but tested on structured progressive unmasking.
• Confidence-Driven Ordering: While efficient, it is myopic. It prioritizes tokens the model is already certain about, which can suppress the exploration of globally high-reward trajectories that may initially appear low-confidence. This creates a "quality–exploration trade-off."

2. Methodology: The DPRM Framework

The authors introduce DPRM (Doob $h$-transform Process Reward Model), a plug-in module designed to optimize the token-ordering policy without altering the host model's architecture, denoising objective, or data pipeline.

Core Concept: Doob $h$-transform
DPRM treats token ordering as a Markov process. It draws inspiration from the Doob $h$-transform, a mathematical technique used to tilt a base transition law toward a desired terminal reward. In the context of DLMs, instead of just picking the most "confident" token, DPRM picks tokens that maximize the expected future reward.

Practical Implementation: Progressive Online Approximation
Since calculating the exact conditional expectation of future rewards is computationally prohibitive, DPRM uses a two-stage online approximation:

1. Warmup Stage (Confidence-Driven): The module starts with standard confidence-based progressive ordering to ensure stability and alignment.
2. Guidance Stage (Reward-Aware): As the model collects data, DPRM transitions to reward-guided ordering using a bucketized controller. It maintains statistics (accumulated exponentiated rewards and counts) in "decision cells" defined by the current phase (denoising stage) and confidence bucket.
3. Soft-BoN Reweighting: To maintain efficiency, DPRM uses a "shortlist" of candidate tokens and applies a Soft Best-of-N (Soft-BoN) reweighting to approximate the exact reward-tilted Gibbs law.

3. Key Contributions

• Theoretical Guarantees:
  • Proved that the stagewise Soft-BoN approximation converges to the exact DPRM target at a rate of $O(1/N)$ in terminal KL divergence.
  • Demonstrated that the online bucketized controller tracks the exact DPRM score at empirical-Bernstein rates.
  • Established a sample-complexity advantage: DPRM is exponentially faster than random masking in the early stage and exponentially faster than confidence-only ordering in the late stage of training.
• Plug-and-Play Versatility: The module is "host-agnostic," meaning it can be inserted into various existing diffusion frameworks (pretraining, post-training, or test-time scaling) with minimal changes.
• Cross-Domain Application: The authors demonstrated the module's effectiveness across natural language, protein design, single-cell genomics, molecular drug design, and DNA sequence design.

4. Experimental Results

The paper evaluates DPRM across seven distinct host settings:

• Natural Language (Reasoning & Scaling):
  • Pretraining (DPRM-PUMA): Improved GSM8K validation mean from 29.34 to 34.27 (+16.8%).
  • Post-training (DPRM-DMPO): Significant gains on hard reasoning tasks (e.g., MATH Hard improved from 44.3 to 47.9).
  • Test-time Scaling (DPRM-Prism): Improved GSM8K accuracy during hierarchical search.
• Scientific Domains:
  • Proteins (DPLM): Improved structural metrics like TM-score and reduced RMSD.
  • Single-cell (DCM): Drastically improved token recovery and zero-expression accuracy.
  • Molecular/DNA (GenMol/SDPO): Improved fragment-constrained validity and reward-specific metrics (e.g., HepG2 scores), though results showed that scientific domains often involve complex multi-objective trade-offs.

5. Significance

This work identifies token ordering as a fundamental and controllable axis in diffusion-based generative modeling. By moving beyond local confidence heuristics toward global reward-aware guidance, DPRM provides a principled way to bridge the gap between local denoising and global objective optimization. The ability to improve performance across vastly different modalities (from math reasoning to DNA design) using a single, unified ordering module marks a significant step toward more robust and efficient diffusion-based foundation models.