Breaking Diffusion: Parallel Reasoning with Draft Refinement

Jaesang Yu; Seoyeon Byun; Yun Wang; Yan Lu; Juho Kim

cs.AIOct 19, 2026

Breaking Diffusion: Parallel Reasoning with Draft Refinement

Authors: Jaesang Yu, Seoyeon Byun, Yun Wang, Yan Lu, Juho Kim

Organizations: KAIST · Microsoft Research Asia · SkillBench

Abstract

Drafting an answer allows a reasoning model to reconsider each decision in the context of the others. Yet repeated diffusion denoising does not ensure that individually plausible token updates form a better answer together. We introduce Draft Refinement Language Models (DRLMs), which use a single denoiser to propose parallel revisions and evaluate a candidate draft. We train the denoiser to reconstruct every token from uniformly corrupted drafts, teaching both preservation and repair. Its predictions specify revisions directly, without a diffusion reverse transition. Evaluation requires a different use of these predictions: confidence in an existing token reflects both contextual support and evidence from observing the token itself. Using the known corruption process, we correct for this self-token evidence to estimate how well each token is supported by the prompt and the remaining draft. Aggregating these estimates yields a cavity energy computed in one forward pass, without an additional scoring model. Each round proposes a revision, reads the completed candidate in its revised context, and retains it only if its energy is lower. Across Countdown, Sudoku-Extreme, and Maze-Hard, our method outperforms autoregressive and discrete diffusion baselines. With 6M parameters, DRLM reaches 44.4% on Countdown-5 and 93.9% on Sudoku-Extreme. The results establish complete-draft parallel refinement as an effective alternative to reverse diffusion for complex reasoning.

Figures & tables

Figure 1: Overview of Draft Refinement Language Models (DRLMs) . (a) All-position supervision on uniformly corrupted drafts teaches preservation and repair. (b) Starting from uniform initialization, the decoder proposes parallel revisions and reads each completed candidate using the same denoiser, at a fixed read level and without changing its tokens. Dividing out the self-token likelihood yields cavity estimates for contextual scoring. A candidate replaces the best-so-far draft only when its cavity energy decreases; otherwise, the incumbent draft is retained. Subsequent rounds re-noise the retained draft, and decoding returns the best-energy draft after a fixed round budget.

Figure 2: Generative process on a Maze-Hard instance. Given a maze with a start (circle) and a goal (star), DRLM generates a draft in the first round, then refines it into a valid shortest path. Each later round re-noises the draft, updates all cells in parallel, and keeps the result if its energy is lower. Panels after the input show the boxed region.

Autoregressive
Model	Params	CD-3	CD-4	CD-5
GPT-2 Scratch	6M	94.1	31.9	4.3
	85M	95.9	45.8	5.1
	303M	96.4	41.3	4.5
LLaMA SFT	7B	95.7	41.1	6.7
	13B	96.5	51.1	7.4

Table 1: Results on Countdown (CD). CD-

k

uses

k

input numbers. Best results are in bold.

Autoregressive
Model	Sudoku- Extreme	Maze-Hard
LLaMA SFT (13B)	1.6	1.5
Discrete diffusion
UDLM (6M)	65.5	0.0
MGDM (6M)	12.3	10.4
LDDM (6M)	24.1	39.3

Table 2: Results on Sudoku-Extreme and Maze-Hard. Best results are in bold.

Figure 3: Analysis. (a) Agreement between each intermediate state and the decoder’s own final output versus the current keep probability

\alpha_{s}

. DRLM reaches

98.4\%

agreement after one round, while diffusion trajectories closely track their noise schedules. (b) Pass@1 of the current output after each sequential pass on Countdown-4 and Countdown-5, up to each method’s default budget. DRLM exceeds both baselines’ final accuracy after two passes on CD-4 and after one on CD-5.

	Reverse-diffusion	Refinement
Training objective	sampler (128)	decoder (128)
NELBO (UDLM)	6.9	36.6
All-position CE (DRLM)	22.5	44.4

Table 3: Objective–Decoder ablation (Pass@1).

Decoder	CD-4	CD-5	Sudoku-Extreme	Maze-Hard
Accept every candidate	82.9	24.0	68.6	70.3
Restart from noise each round	85.3	30.6	77.0	66.2
DRLM (energy-controlled retention)	87.6	44.4	93.9	73.0

Table 4: Retention ablation (Pass@1).

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

	Countdown (CD-3/4/5)	Sudoku-Extreme	Maze-Hard
Canvas alphabet $N$ / canvas length $L$	$18$ / $64$	$9$ / $81$	$5$ / $900$
Rounds $\times$ passes per round, $R\times J$	$32\times 4$	$32\times 4$	$16\times 8$
Read level $\alpha_{s}$	$7/8$	$7/8$	$7/8$
Re-noising keep probability $\rho_{e}$	$0.1$	$0.1$	$0.7$
Positions changed by re-noising equation 23	$85.0\%$	$80.0\%$	$24.0\%$
Inner schedule	equation 22 , with $\rho=0$ in round $1$ and $\rho=\rho_{e}$ afterwards

Appendix

Table 5: Decoding configuration at the default budget of

RJ=128

model evaluations per problem.

	Countdown (CD-3/4/5)	Sudoku-Extreme	Maze-Hard
Training examples	$500$ k per setting	$3{,}831{,}994$	$1{,}000$ ( $\times 8$ symmetries)
Test examples	$1{,}000$ per setting	$422{,}786$	$1{,}000$
Condition / canvas length	$20$ / $64$	$81$ / $81$	$900$ / $900$
Model vocabulary / $N$	$31$ / $18$	$11$ / $9$	$6$ / $5$

Appendix

Table 6: Tasks. Lengths are in tokens;

N

is the canvas alphabet used for corruption and scoring.

Architecture	bidirectional pre-LN Transformer: width $192$ , $14$ layers, $4$ heads, MLP width $768$ , GELU
	no biases in attention and MLP projections; untied output head
Embeddings	shared token embedding; learned position embeddings for condition and canvas;
	level features (Appendix D.3 ) added to canvas embeddings
Parameters	$6.27$ M (Countdown), $6.28$ M (Sudoku-Extreme), $6.59$ M (Maze-Hard)
Objective	all-position cross-entropy equation 3 , one $\alpha\sim U(0,1)$ per sequence
Optimizer	AdamW, learning rate $10^{-3}$ , default betas, weight decay $0.01$ , gradient-norm clipping $1.0$

Appendix

Table 7: DRLM training configuration, shared across tasks unless noted.

$\alpha_{s}$	$\kappa_{\alpha_{s}}$	Cavity	Raw	Difference (Cavity $-$ Raw)
$3/4$	$28$	$94.0$	$88.8$	$+5.2$
$7/8$	$64$	$93.9$	$88.7$	$+5.2$
$15/16$	$136$	$94.0$	$88.6$	$+5.4$
$0.98$	$442$	$94.3$	$88.5$	$+5.8$
$0.995$	$1{,}792$	$94.3$	$88.5$	$+5.8$

Appendix

Table 8: Self-likelihood correction on Sudoku-Extreme. Pass@1 at

128

evaluations when the decoder accepts candidates by the cavity or by the raw energy, both read at level

\alpha_{s}

Setting	CD-4	CD-5	Maze-Hard
$J=2$	$82.7$	$26.3$	$67.5$
$J=4$	$87.6$	$44.4$	$72.3$
$J=8$	$86.9$	$41.7$	$73.0$
$J=16$	$86.2$	$36.7$	$72.4$
$\alpha_{s}=3/4$	$87.7$	$44.9$	$71.5$
$\alpha_{s}=7/8$	$87.6$	$44.4$	$73.0$

Appendix

Table 9: Pass@1 at

RJ=128

when one hyperparameter is changed.

Configuration	Pass@1
Subsample: $J=4$ , $\rho_{e}=0.1$	$93.9\pm 0.2$
Subsample: $J=4$ , $\rho_{e}=0.7$	$74.5\pm 0.1$
Subsample: $J=8$ , $\rho_{e}=0.1$	$91.3\pm 0.4$
Subsample: $J=8$ , $\rho_{e}=0.7$	$75.1\pm 0.3$
Selected, held-out complement ( $420{,}786$ puzzles, 3 seeds)	$93.9\pm 0.0$
Selected, full test set ( $422{,}786$ puzzles)	$93.9\pm 0.0$

Appendix

Table 10: Sudoku-Extreme: configuration selection on the

2{,}000

-puzzle subsample, and the selected configuration on the held-out complement and the full test set.

	Pass@1 after this many model evaluations
Task	$4$	$8$	$16$	$32$	$64$	$128$	Coverage	Accepted, last round (%)
CD-4	$82.8$	$83.8$	$84.8$	$85.9$	$86.6$	$87.6$	$88.6$	$0.10$
CD-5	$24.6$	$26.5$	$29.3$	$33.5$	$38.2$	$44.4$	$44.8$	$1.93$
Maze-Hard	–	$62.9$	$69.3$	$70.7$	$72.2$	$73.0$	$73.9$	$0.47$

Appendix

Table 11: Pass@1 by model evaluations, coverage at

128

evaluations, and the fraction of problems whose incumbent is replaced in the last round.

Cell or effect	Pass@1 or difference
DRLM (CE) $\times$ refinement decoder	$44.4\pm 1.0$
UDLM (NELBO) $\times$ refinement decoder	$36.6\pm 0.6$
DRLM (CE) $\times$ ancestral sampler	$22.5\pm 0.3$
UDLM (NELBO) $\times$ ancestral sampler	$6.9\pm 0.5$
UDLM $\times$ refinement decoder, without conversion	$37.3\pm 0.8$
DRLM $\times$ ancestral sampler, without conversion	$21.2\pm 0.3$

Appendix

Table 12: Objective

\times

decoder on CD-5 at

128

model evaluations: Pass@1 per cell (

\pm

: standard deviation over decode seeds) and paired effects (points) with

95\%

bootstrap intervals over problems.

Panel of Figure 4 — Figure 4: Maze-Hard repairs by a single retained proposal. Panels after the input show the boxed region. Red fill marks path cells not in the final path, and red outlines mark final-path cells not yet drawn.

Numbers $9,7,13,21$ ; target $89$ : the plan is replaced
Round	Candidate	Valid	Energy	Kept
1	21-7=12,9*13=102,102-13=89	no	$15.73$	yes
2	9*13=117,7+21=28,117-28=89	yes	$0.06$	yes
Numbers $81,4,1,64$ ; target $96$ : one token is repaired, and a worse candidate is rejected
1	81-1=80,64/4=14,80+16=96	no	$18.72$	yes
2	81-1=80,64/4=16,80+16=96	yes	$0.00$	yes

Appendix

Table 13: Countdown-4 rounds. Each candidate is read at

\alpha_{s}=7/8

and kept only if its energy is lower than the incumbent’s. Rounds whose candidate repeats the incumbent are omitted.

Figure 5: Sudoku-Extreme refinement on two puzzles of the

2{,}000

-puzzle subsample. Bold digits are givens, blue digits match the unique solution, and red cells differ from it. Re-noising acts on all

81

cells, givens included. Each round has three proposal passes before its read; passes after the first solved one are omitted.

Figure 6: A Maze-Hard failure case. Round $2$ retains a valid shortest path while round $3$ proposes the same path with one extra cell, whose lower energy leads the decoder to replace the valid draft.

Diffusion Language Models

Diffusion language models generate and refine text through iterative denoising, enabling bidirectional context and parallel token updates.

See all 443 papers

Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching
Feb 26, 2026Roy Miles, Aysim Toker, Andreea-Maria Oncescu +2Structured Reasoning in Language Models Test-Time Scaling
Huawei London Research Center · MVP Lab
Speculative Correction: Draft-then-Refine Decoding for Diffusion Language Models
Jul 21, 2026Brian K Chen, Chong Wu, Kenji KawaguchiLanguage Model Decoding Speculative Decoding for Diffusion Language Models
National University of Singapore · City University of Hong Kong
DCGC: Draft-Conditioned Global Correction for Complex Reasoning with Masked Diffusion Models
Aug 26, 2026Minhae Oh, Nakyung Lee, Jungwoo LeeLLM Self-Correction Diffusion Language Model Inference
Department of Electrical and Computer Engineering, Seoul National University

Feb 26, 2026J/K move · Enter open · S save

Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching

Roy Miles, Aysim Toker, Andreea-Maria Oncescu, Jiankang Deng+1

Huawei London Research Center · MVP Lab

Structured Reasoning in Language Models Test-Time Scaling Diffusion Language Models

Open paper AskPDF arXiv

Reasoning with large language models often benefits from generating multiple chains-of-thought, but existing aggregation strategies are typically trajectory-level (e.g., selecting the best trace or voting on the final answer), discarding useful intermediate work from partial or "nearly correct" attempts. We propose Stitching Noisy Diffusion Thoughts, a self-consistency framework that turns cheap diffusion-sampled reasoning into a reusable pool of step-level candidates. Given a problem, we (i) sample many diverse, low-cost reasoning trajectories using a masked diffusion language model, (ii) score every intermediate step with an off-the-shelf process reward model (PRM), and (iii) stitch these highest-quality steps across trajectories into a composite rationale. This rationale is then used to recompute only the final answer. This modular pipeline separates exploration (diffusion) from evaluation and solution synthesis, avoiding monolithic unified hybrids while preserving broad search. Across math reasoning benchmarks, we find that step-level recombination is most beneficial on harder problems, and ablations highlight the importance of the final solver in converting stitched but imperfect rationales into accurate answers. Using low-confidence diffusion sampling with parallel, independent rollouts, our training-free framework improves average accuracy by up to 23.8% across six math and coding tasks. At the same time, it achieves up to a 1.8x latency reduction relative to both traditional diffusion models (e.g., Dream, LLaDA) and unified architectures (e.g., TiDAR). The code is available https://github.com/roymiles/diffusion-stitching

Abstract

Diffusion Language Models

Explore similar work

Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching

Speculative Correction: Draft-then-Refine Decoding for Diffusion Language Models

DCGC: Draft-Conditioned Global Correction for Complex Reasoning with Masked Diffusion Models