Breaking Diffusion: Parallel Reasoning with Draft Refinement

Breaking Diffusion. Parallel Reasoning with Draft Refinement. By Jaesang Yu and four others, from KAIST, Microsoft Research Asia, and one other institution. Published on the nineteenth of October. No code repository is linked to it.

Writing out a whole draft of an answer has an appeal for a reasoning model. With every piece on the page at once, each decision can be reconsidered in light of all the others. Diffusion models seem built for this, since they denoise a full sequence again and again. The catch the authors point to is that repeated denoising carries no guarantee of improvement. Each token update can look plausible on its own, and yet the updates taken together need not form a better answer. Many locally reasonable edits can still add up to a worse draft.

Draft Refinement Language Models, or DRLMs, replace that loop with one that proposes and then checks. A single denoiser does two jobs. It proposes parallel revisions to a draft, and it evaluates a candidate draft. The training is simple to state. The denoiser learns to reconstruct every token from drafts that have been uniformly corrupted, which teaches it both to preserve what is right and to repair what is wrong. Its predictions then specify revisions directly, without a diffusion reverse transition.

Evaluation needs the same predictions read a different way. When the model is confident in a token that is already in the draft, that confidence mixes two sources: support from the surrounding context, and evidence from simply seeing the token itself. Because the corruption process is known, the method corrects for that self-token evidence, leaving an estimate of how well each token is supported by the prompt and the rest of the draft.

Those estimates are aggregated into what the authors call a cavity energy, and it comes from one forward pass, with no extra scoring model. A round of refinement then runs as follows. The model proposes a revision, reads the completed candidate in its revised context, and keeps it only if its energy is lower than before. A revision that does not improve the draft as a whole is thrown away.

The tests are three puzzle-like reasoning tasks: Countdown, Sudoku-Extreme, and Maze-Hard. On all three the method outperforms both autoregressive and discrete diffusion baselines. The model is small, at six million parameters, and it reaches forty-four point four percent on Countdown-5 and ninety-three point nine percent on Sudoku-Extreme.

Diffusion's promise for reasoning was the whole draft in view. This work keeps that and drops the reverse process. One network proposes, the same network judges the complete result, and an edit survives only when the judgment says the draft got better. The authors read their results as establishing complete-draft parallel refinement as an effective alternative to reverse diffusion for complex reasoning.