DPRM

Process reward models for masked diffusion

DPRM

A Plug-in Token-Ordering Module for Discrete Diffusion Models

Dake Bu1,2 · Wei Huang3,4 · Andi Han5 · Si Wu6 · Hau-San Wong1,* · Qingfu Zhang1 · Taiji Suzuki3,7 · Atsushi Nitanda2,8,*

1City University of Hong Kong 2CFAR and IHPC, A*STAR 3RIKEN AIP 4The Institute of Statistical Mathematics 5University of Sydney 6South China University of Technology 7The University of Tokyo 8Nanyang Technological University

*Corresponding authors

Confidence is an efficient token-order heuristic, but local certainty can be myopic. DPRM uses terminal utility to decide which position should be committed next.

Paper Code Artifacts
DPRM ranks candidate token positions using confidence and estimated terminal utility.
+13.2%Omni CLIP-L/14
+14.6%PUMA accuracy
+12.9%Countdown Hard
9 hostsone ordering interface

Omni-Diffusion

One token-order change alters the image trajectory

Both rows share the same canvas through step 95. At step 96, confidence commits its lowest-entropy position; DPRM selects a lower-confidence position with higher estimated terminal value.

Confidence and DPRM trajectories for three children playing on a beach.
Three children on a beach. DPRM commits an unfilled region that preserves object separation. The later canvas resolves the children and their toys while improving both CLIP checks.
Confidence and DPRM trajectories for an owl family in a forest.
Owl family. DPRM commits a peripheral body region and separates the family more clearly in the completed image.
Confidence and DPRM trajectories for a boy playing a flute to kittens.
Boy playing flute to kittens. DPRM first commits a lower-confidence sidewalk region that preserves the audience relation, after which the same continuation renders the boy, flute, and kittens more coherently.
CLIP-L/14
Confidence0.18661
DPRM0.21125
CLIP-B/32
Confidence0.23836
DPRM0.24854

512 untouched prompts. Paired 95% intervals for both improvements exclude zero.

Frozen confirmation split

Four-policy image gallery

Random order, Omni default, a deterministic uniform order, and DPRM use matched generation settings. These reader-selected cases emphasize visual quality and prompt fidelity; they are displayed only after the controller and aggregate evaluation were frozen.

Four-policy comparison for Downtown San Francisco at sunrise in detailed ink wash.
Downtown San Francisco at sunrise, detailed ink wash.
Four-policy comparison for a Harley-Davidson motorcycle with a flame decal.
A Harley-Davidson motorcycle with a flame decal.
Four-policy comparison for a dragon.
A dragon.
Four-policy comparison for a painting of the Mona Lisa on a white wall.
A painting of the Mona Lisa on a white wall.
Four-policy comparison for a sketch of a train.
A sketch of a train.
Four-policy comparison for a woman with long hair.
A woman with long hair.
Four-policy comparison for a pig in a field.
A pig in a field.
Four-policy comparison for a volcano exploding next to a marina.
A volcano exploding next to a marina.
Four-policy comparison for a yellow Corvette on a mountain road.
A yellow Corvette on a mountain road.
Four-policy comparison for a hot-air balloon.
A hot-air balloon.

Nine host settings

Ordering gains across modalities

HostMetricConfidenceDPRM
Omni-DiffusionCLIP-L/140.186610.21125
Omni-DiffusionCLIP-B/320.238360.24854
LLaDA-VRealWorldQA47.35%48.92%
LLaDA-VNumeric / count32.05%41.03%

Seven hosts use the same denoiser-call budget as confidence. Prism estimates verifier-conditioned values during test-time search, while Omni uses five continuations from the current canvas because values learned from earlier prompts did not transfer reliably across unrelated images.

Saved trajectories and controls

What changes when order changes?

Entropy-only and random-order controls do not reproduce the gains. The saved traces instead test whether terminal reward identifies useful uncertain positions rather than merely preferring uncertainty.

Two PUMA reasoning traces comparing confidence and DPRM reveal order.
Reasoning. DPRM reveals wider context, backfills earlier positions, and delays numeric commitment on two DPRM-only GSM8K wins.
LLaDA-V same-canvas intervention showing cone before number.
Visual question answering. Revealing the object word before the numeric slot preserves the context needed for the correct count.
Countdown confidence-bin and token-level diagnostics.
Hard order classes. Confidence selects easy states early; reward tilt restores selected low-confidence regions later.
Theory-facing diagnostics for Countdown, DNA, and VQA.
Mechanism controls. Hard/OOD, shuffled-value, entropy-only, and format-aware controls isolate terminal-utility effects.
Scientific-domain radar charts for protein, cells, molecules, and DNA.
Scientific preferences. Weighted and Tchebycheff rewards expose controllable Pareto responses across four scientific hosts.

Reproducibility

Code, checkpoints, raw records

GitHubNine host integrations Hugging FaceRetained checkpoints and records PaperMethod, theory, and protocols