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

원문 초록

Discrete diffusion models admit many token orders, yet most systems rely on confidence-based decoding. Confidence is a strong and efficient heuristic, but it can be myopic because local certainty does not measure a position's effect on terminal quality. We introduce DPRM (Doob h-transform Process Reward Model), a plug-in token-ordering module induced by a terminal-reward-tilted trajectory distribution. The Doob transform converts terminal reward into a process reward for each candidate position. A compact estimator indexed by generation progress, confidence, and optional positional or task state makes this correction reusable across partial states. DPRM changes only token order; the host architecture, denoising objective, supervision, and token sampler remain fixed. We derive the exact reward-tilted law and bound its shortlist and online approximations. Across nine recent open-source hosts spanning language, multimodal generation, and scientific sequences, DPRM improves reasoning, numeric VQA, visual-codebook ordering, and preference-conditioned generation. Relative to confidence, it improves PUMA GSM8K by 14.6%, Omni-Diffusion CLIP-L/14 by 13.2%, Prism voted accuracy from 82.41% to 83.85%, and RealWorldQA numeric/count accuracy by 8.97 points. Matched scientific gains include 8.3% in DCM nonzero recovery, 15.0% in GenMol QED, and 53.3% in SDPO-DNA total utility. Entropy and random-order controls show that the gains do not come from favoring uncertain tokens indiscriminately. Token-order traces and shared-state interventions instead show how terminal reward can correct a locally confident but globally poor choice. Code is available at: https://github.com/DakeBU/DPRM-DLLM.