Diffusion language models learn when to rewrite: DiVR lifts a fine-tuned generator's five-benchmark average from 35.3% to 43.4%
Synopsis
Masked diffusion language models can remask generated text, but deciding where to revise by confidence or self-judgment often leaves the error in place. DiVR instead learns an outcome verifier on the generator's own decodes and, at each block boundary, judges whether the partial solution can still reach a correct answer, resampling the whole block when the score is low. On the supervised fine-tuned LLaDA-8B-Base generator, the five-benchmark math and code average rises from 35.3% to 43.4% at the cost of greedy decoding.
Interpretation
Revision in diffusion language models no longer rests on token confidence; an outcome verifier trained on the generator's own decodes decides whether a whole block should be written again. Standard masked diffusion decoding never revisits a committed token, so a wrong answer committed early with high confidence tends to remain; existing remasking and search methods choose what to revise by a schedule, confidence, likelihood or prompted self-verification, none of which is trained to predict whether the final answer is correct. On the supervised fine-tuned LLaDA-8B-Base generator, the five-benchmark math and code average rises from 35.3% to 43.4% at the cost of greedy decoding, against 39.1% for SR-SMC and 37.3% for Prism at higher computation.
The verifier shares the generator's backbone and trains only a low-rank adapter and a small head, reading the partially masked canvas and scoring whether that trajectory will end correct. Labels come only from the final-answer correctness of the generator's own sampled trajectories, with no token-level labels or reference reasoning steps, so the verifier learns the error modes of the generator it guides. On GSM8K, holding the regeneration procedure and budget fixed and replacing only the trigger signal, commitment confidence reaches 72.2%, prompted Yes/No 67.9%, a head-only verifier 74.2%, and full DiVR 78.8%.
Repairs concentrate at the start of a solution, because the first block fixes how the problem is interpreted and later tokens tend to confirm it. Re-predicting a single wrong cell with the rest of the canvas kept leaves the original wrong digit at a mean probability of 0.869; on 77 failed Sudoku puzzles, a gold oracle that reopens every wrong answer cell repairs only 6, and reopening preceding reasoning as well raises this to 13. On GSM8K, 74% of the 154 repaired problems depart from the greedy trajectory in the first block and 92% within the first two; on Countdown, 58 of 59 repairs begin in the first block.
Perspective
The demonstrated scope is inference-time revision for three masked diffusion generators (LLaDA-8B-Base, LLaDA-8B-Instruct and Dream-v0-Instruct-7B) on five math and code benchmarks: GSM8K, MATH-500, Countdown, HumanEval and MBPP, with one verifier per domain trained on GSM8K for math and MBPP for code. A reader should treat it as evidence about an inference-time method, not a conclusion about open-ended text generation or training-stage improvements.
The threshold was chosen on the GSM8K development split; on the GSM8K test set accuracy stays within 77.5% to 79.2% for every threshold from 0.1 to 0.9 while cost rises. The paired intervention analysis covers one regeneration event per state, not the sequential three-event policy. Changing only the verifier-training seed gives 77.5% and changing only the re-decoding seed gives 77.6%, which are separate sensitivity checks rather than independent end-to-end replications.