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.
Model
Params
CD-3
CD-4
CD-5
Autoregressive
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.
Model
Sudoku- Extreme
Maze-Hard
Autoregressive
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 α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 × passes per round, R×J
32×4
32×4
16×8
Read level αs
7/8
7/8
7/8
Re-noising keep probability ρ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 ρ=0 in round 1 and ρ=ρ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 ( ×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.
Table 7: DRLM training configuration, shared across tasks unless noted.
αs
κα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 α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
αs=3/4
87.7
44.9
71.5
αs=7/8
87.6
44.4
73.0
Appendix
Table 9: Pass@1 at RJ=128 when one hyperparameter is changed.
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) × refinement decoder
44.4±1.0
UDLM (NELBO) × refinement decoder
36.6±0.6
DRLM (CE) × ancestral sampler
22.5±0.3
UDLM (NELBO) × ancestral sampler
6.9±0.5
UDLM × refinement decoder, without conversion
37.3±0.8
DRLM × ancestral sampler, without conversion
21.2±0.3
Appendix
Table 12: Objective × decoder on CD-5 at 128 model evaluations: Pass@1 per cell ( ± : standard deviation over decode seeds) and paired effects (points) with 95% bootstrap intervals over problems.
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.
Round
Candidate
Valid
Energy
Kept
Numbers 9,7,13,21 ; target 89 : the plan is replaced
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 α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.
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
Diffusion language models (DLMs) can revise tokens bidirectionally, but standard decoding procedures often adapt them to left-to-right generation by producing text block by block. We study a simple plug-and-play inference pattern: first generate a complete draft, then refine the full response using bidirectional diffusion. Using LLaDA2.1-Flash and LLaDA2.1-Mini, we evaluate two configurations. In Flash-Flash, the same Flash model serves as both drafter and refiner, testing whether an existing model can improve its own block-autoregressive output through global refinement. In Mini-Flash, inspired by speculative decoding, we introduce speculative correction: Mini drafts a full response, and Flash revises it as an editable initialization. Flash-Flash improves GSM8K-384 accuracy from 0.848 to 0.899 while running 1.20 times faster than the selected Flash block-autoregressive baseline, and improves MBPP-384 from 0.545 to 0.693. Latency-window-matched Flash-only controls indicate that these gains persist after targeted tuning of block-autoregressive decoding. Causal ablations indicate that completed drafts provide useful initializations: refinement from a fully masked span performs poorly, full global refinement provides a clear additional gain on GSM8K, and local refinement captures much of the gain on MBPP and MATH. Mini-Flash provides useful quality-latency trade-offs, including MATH-384 performance of 0.294 versus 0.300 for Flash while running 2.17 times faster. These results support a Pareto-frontier interpretation rather than the claim that the heterogeneous cascade uniformly matches Flash quality. Overall, same-model draft-and-refine provides evidence that bidirectional refinement is a useful decoding primitive for DLMs, while speculative correction demonstrates a training-free route to fast DLM generation.
Brian K Chen, Chong Wu, Kenji Kawaguchi
National University of Singapore · City University of Hong Kong
Correcting flawed reasoning traces remains a significant challenge for Large Language Models (LLMs), whose autoregressive generation can propagate early mistakes into subsequent reasoning. We introduce DCGC, a Masked Diffusion Model (MDM) framework for global correction that uses an imperfect solution draft from an upstream solver as auxiliary context. DCGC combines task-specific Supervised Fine-Tuning (SFT) with a novel inference-time mechanism called Dynamic Dual-CFG. This mechanism separates problem-only and joint problem-draft branches and scales the draft-conditioned residual using a relative confidence gap. Across math, code, and knowledge reasoning benchmarks, DCGC outperforms standard sampling and simpler CFG variants, with additional results suggesting transfer to different diffusion backbones. In test-time setting where ground-truth failure labels are unavailable, DCGC improves full test set accuracy by correcting low-consensus upstream outputs, highlighting its utility as a verifier-free global correction module for difficult reasoning instances.
Minhae Oh, Nakyung Lee, Jungwoo Lee
Department of Electrical and Computer Engineering, Seoul National University