PaperlayerView exact revision on arXiv

Source-linked AI summary

Breaking Diffusion: Parallel Reasoning with Draft Refinement

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

arXiv:2610.26542v1cs.AIcs.CLcs.LG

TL;DR

Complete-answer reasoning needs a way to revise parallel token updates and determine whether the resulting draft is better supported. DRLM uses one denoiser for direct parallel proposals and self-token-corrected candidate scoring; across three reasoning benchmarks it outperforms evaluated autoregressive and discrete diffusion baselines.

  • Problem

    Parallel token revisions can be individually plausible without forming a coherent answer, and evaluating their contextual support in a complete draft is a separate challenge.

    Source: §1 · p. 1

  • Method

    DRLM trains one denoiser on uniformly corrupted drafts to propose parallel revisions and uses self-token-corrected cavity estimates to score completed candidates in a propose–read–retain loop.

    Source: §1 · pp. 1–2, §3.3 · p. 5

  • Results

    Across Countdown, Sudoku-Extreme, and Maze-Hard, DRLM outperforms evaluated autoregressive and discrete diffusion baselines, reaching 44.4% on Countdown-5 with 6M parameters.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    The results support complete-draft parallel refinement with whole-draft evaluation as an effective alternative to reverse diffusion for complex reasoning.

    Source: §7 · p. 9

    At finite noise, cavity energy is a learned, noise-smoothed pseudo-likelihood score, not a correctness certificate.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

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.

1 INTRODUCTION

DRLMs treat a complete provisional answer as context for reconsidering earlier decisions, while addressing the challenge of evaluating whether parallel revisions improve the draft. One denoiser proposes revisions and evaluates candidates through a self-token-corrected contextual score.

1 INTRODUCTION

  • A complete provisional answer makes the whole response available as context, so earlier decisions can be reconsidered in light of later ones.The paper frames a draft as something to inspect, compare, and revise, not merely a record of thought.Source: §1 · p. 1
  • Autoregressive models commit earlier tokens before a complete answer exists, while diffusion reverse transitions constrain updates and plausible token revisions may not yield a coherent answer.The stated challenge is assessing revisions as well as proposing them.Source: §1 · p. 1
  • DRLM trains one denoiser on uniformly corrupted drafts to reconstruct every position, then directly proposes parallel revisions without a reverse diffusion transition.Because incorrect tokens are unmarked, training teaches both preservation and repair.Source: §1 · p. 1
  • Dividing out each observed token’s known likelihood estimates contextual support, and a fresh read scores each revised candidate before the loop retains it only if cavity energy falls.The same denoiser supplies the one-pass score; the loop returns the lowest-energy draft within its fixed budget.Source: §1 · p. 2
  • With 6M parameters, DRLM outperforms evaluated autoregressive and discrete diffusion baselines across Countdown, Sudoku-Extreme, and Maze-Hard, reaching 44.4% on Countdown-5.It exceeds the strongest diffusion baseline by +14.3, +34.1, +28.4, and +33.7 percentage points on Countdown-4, Countdown-5, Sudoku-Extreme, and Maze-Hard, respectively.Source: §1 · p. 2

2 BACKGROUND

Discrete diffusion corrupts sequences toward a prior and trains denoisers through reverse transitions. Its standard decoding schedules mediate token updates, motivating direct complete-draft revision while leaving candidate evaluation as a separate challenge.

2 BACKGROUND

  • Interpolating discrete diffusion independently corrupts each sequence position toward a prior, with keep probability decreasing over time.A mask prior hides tokens, whereas a uniform vocabulary prior replaces them with ordinary tokens that may coincide with the originals.Source: §2 · p. 2
  • Standard diffusion training minimizes an evidence lower bound whose diffusion term compares exact and learned reverse transitions, thereby training the denoiser through its parameterized transitions.The denoiser uses the condition and full noisy sequence to predict clean-token probabilities.Source: §2 · p. 2
  • Standard ancestral decoding samples successive reverse states; in masked diffusion, the schedule controls when tokens are revealed, and revealed tokens remain unchanged as context changes.Remasking samplers relax the unchanged-token restriction by allowing visible tokens to return to the mask state.Source: §2 · p. 3
  • Even uniform-noise diffusion constrains token changes through a schedule-dependent reverse-kernel bound, so clean-token predictions do not directly specify the next sequence.This motivates direct complete-draft refinement, but deciding whether to keep a candidate remains separate from proposing it.Source: §2 · p. 3

3 DRAFT REFINEMENT LANGUAGE MODELING

DRLM uses one denoiser to propose and evaluate complete drafts, correcting for the evidence contributed by each observed token. Its method combines all-position reconstruction training, one-pass cavity scoring, and iterative retention of lower-energy candidates.

3.1 LEARNING COMPLETE-DRAFT PREDICTION

  • Uniformly corrupted drafts provide complete training inputs with unmarked replacements, and all-position reconstruction trains the denoiser to preserve or repair every token.The training channel defines the inputs, not a generation trajectory that must be reversed.Source: §3.1 · p. 3
  • Parallel prediction reconsiders all positions synchronously from the same draft and specifies the next complete draft directly, without a reverse diffusion transition.Figure 1 summarizes the all-position training and propose-and-read process.Source: §3.1 · p. 4, Figure 1 · p. 4

3.2 FROM DENOISER PREDICTIONS TO CAVITY ENERGY

  • Dividing out the observed token’s known likelihood converts full-posterior predictions into cavity estimates of support from the condition and remaining draft, all from one forward pass.The resulting cavity energy is a noise-smoothed pseudo-likelihood score, not a correctness certificate; the correction is for evaluation, while proposals use full-posterior predictions.Source: §3.2 · p. 5

3.3 REFINEMENT BY PROPOSING, READING, AND RETAINING

  • Each round re-noises the incumbent, proposes a complete candidate, and reads it in its revised context; only a strictly lower-energy candidate replaces the incumbent.The incumbent energy is non-increasing, and the Maze-Hard figure illustrates refinement into a valid shortest path.Source: §3.3 · pp. 5–6, Figure 2 · p. 6

4 EXPERIMENTS

Experiments evaluate complete-draft refinement on Countdown, Sudoku-Extreme, and Maze-Hard, where constraint-coupled decisions must remain consistent. DRLM outperforms the reported autoregressive and discrete diffusion baselines across these tasks.

4 EXPERIMENTS

  • The three tasks test joint consistency and revision: one incorrect position invalidates a solution, making them tests of global planning rather than local pattern completion.Countdown requires a valid equation chain; Sudoku-Extreme evaluates all 81 cells, and Maze-Hard requires an optimal path preserving walls and endpoints.Source: §4 · p. 6, §4.1 · p. 6, §4.2 · p. 7

4.1 COUNTDOWN

  • Countdown comparisons include GPT-2 and fine-tuned LLaMA autoregressive models alongside four parameter-matched discrete diffusion methods.The diffusion baselines are MGDM, LDDM, SUNDAE, and UDLM.Source: §4.1 · p. 6
  • 44.4% on CD-5 exceeds LLaMA-13B (7.4%) and LDDM (10.3%); DRLM also achieves 87.6% on CD-4 and 99.5% on CD-3.The CD-4 result exceeds LLaMA-13B SFT by +36.5 points and UDLM by +14.3 points; the CD-3 result exceeds UDLM by +0.8 points.Source: Table 1 · p. 7, §4.1 · p. 7

4.2 SUDOKU AND MAZE-HARD

  • For Sudoku-Extreme and Maze-Hard, DRLM is compared with fine-tuned LLaMA-13B and parameter-matched UDLM, MGDM, and LDDM.The models use the same training splits and evaluation criteria as DRLM.Source: §4.2 · p. 7
  • 93.9% Sudoku-Extreme exact-solution accuracy and 73.0% Maze-Hard optimal-path accuracy show DRLM outperforming the reported baselines on both tasks.The reported gains include +28.4 points over UDLM on Sudoku-Extreme and +33.7 points over LDDM on Maze-Hard.Source: Table 2 · p. 7, §4.2 · p. 7

5 ANALYSIS

Analysis shows that DRLM produces accurate drafts early, while both its training objective and direct refinement contribute gains. Energy-controlled retention also matters: candidates are evaluated in their revised context before replacing the incumbent.

Model-bounded versus schedule-bounded generation

  • DRLM agrees with its final output at 98.4% of positions after the first round, while every position remains eligible for revision.MGDM and UDLM instead track schedule-dependent agreement profiles.Source: §5 · p. 8

Accuracy–Compute curve

  • 44.4% Pass@1 on CD-5 after 32 rounds, versus final UDLM 6.7% and MGDM 5.8%; DRLM reaches 14.0% on its first call.Refinement improves accuracy through all 32 rounds.Source: §5 · p. 8

Objective and decoder contribute separately

  • At 128 model evaluations on CD-5, replacing only the decoder raises accuracy from 6.9% to 36.6%, while replacing only the objective gives 22.5%; combining both reaches 44.4%.Both decoder and objective gains appear with either counterpart, and all four paired effects exclude zero.Source: Table 3 · p. 8, §5 · p. 8

Energy-controlled retention drives refinement

  • Accepting every candidate lowers Pass@1 from 44.4% to 24.0% on CD-5 and from 93.9% to 68.6% on Sudoku-Extreme.Restarting each round from a uniform canvas recovers only part of the gain across tasks; on Maze-Hard, refining the retained draft remains essential.Source: §5 · pp. 8–9, Table 4 · p. 9

6 RELATED WORK

Prior work established both diffusion-based parallel prediction and complete-sequence iterative refinement. DRLM differs by proposing complete-draft revisions and retaining them only when its model assigns lower energy.

Discrete Diffusion

Non-Autoregressive Iterative Refinement

  • Complete-sequence refinement predates diffusion, including draft denoising and parallel updates; unlike methods that keep every revision, DRLM accepts only lower-energy drafts.DRLM also differs from HRM, TRM, and PTRM by using the complete discrete draft itself as recurrent state [Jolicoeur-Martineau, 2025; Sghaier et al., 2026].Source: §6 · p. 9

7 CONCLUSION

The conclusion presents DRLM as parallel revision with whole-draft evaluation, while Appendix C details how standard ancestral diffusion kernels constrain token updates. Its masked-kernel analysis shows that revelation is schedule-controlled and visible tokens are carried over unchanged.

7 CONCLUSION

  • Across Countdown, Sudoku-Extreme, and Maze-Hard, DRLM outperforms the evaluated autoregressive and discrete diffusion baselines.The conclusion attributes this result to combining parallel revision with whole-draft evaluation for complex parallel reasoning.Source: §7 · p. 9

C DIFFUSION KERNELS AND TOKEN-REVISION CONSTRAINTS

  • Appendix C formalizes standard ancestral decoding through plug-in reverse kernels, schedule-controlled masked revelation, and a revision ceiling for uniform-noise diffusion.Its diffusion times are distinct from the read level used in Section 3.2.Source: C DIFFUSION KERNELS AND TOKEN-REVISION… · p. 13

C.1 PLUG-IN REVERSE KERNELS

  • The plug-in kernel reweights a denoiser prediction by the likelihood of the current token, forming a factorized learned transition when applied independently at each position.The substitution includes a normalizer rather than averaging separately normalized posteriors.Source: §C.1 · p. 14

C.2 MASKED REVELATION AND CARRY-OVER

  • In the standard masked kernel, the schedule sets the revelation probability and the model sets the revealed value; once visible, a token is never changed by later reverse steps.The stated scope excludes remasking samplers and confidence-ordered decoders.Source: §C.2 · p. 14

C.3 UNIFORM-NOISE REVISION CEILING

The appendix characterizes how uniform-noise reverse updates revise tokens and establishes the population target and synchronous prediction behavior of the all-position objective.

C.3 UNIFORM-NOISE REVISION CEILING

  • For interior steps, the uniform plug-in kernel’s token-change probability is strictly below one, while its revision ceiling is vacuous at boundary settings.The redraw form mixes staying with a redraw from ν; for b < 1, every token has positive redraw probability.Source: §C.3 · pp. 14–15

D FULL-POSTERIOR LEARNING AND DIRECT PREDICTION

  • All-position cross-entropy has the full clean-token posterior as its population optimum, with minimum loss at most log N.Equality in the bound holds only when predictions match the posterior at every position almost surely.Source: §D.1 · p. 15

D.1 POPULATION OPTIMUM OF THE ALL-POSITION CROSS-ENTROPY

  • The objective characterizes positionwise marginals, assumes no independence among clean tokens, and does not identify a joint sequence distribution.The proposition concerns the population optimum; a trained network approximates it within its capacity, data, and optimization budget.Source: §D.1 · p. 15

D.2 PRESERVATION, REPAIR, AND SYNCHRONOUS PREDICTION

  • Uniform replacement trains matching positions for preservation and mismatching positions for repair, without imposing a copy rule.The mode map updates all positions synchronously from one common draft and is neither an autoregressive sweep nor a joint-sequence maximizer; one application need not lower cavity energy.Source: §D.2 · p. 15

D.3 TRAINING CONVENTIONS

The paper specifies shared corruption conventions, derives a cavity estimate by removing the observed token’s likelihood contribution, and uses a fresh model read to score each completed candidate.

D.3 TRAINING CONVENTIONS

  • Each training sequence shares one keep probability across positions, while condition tokens remain uncorrupted and the loss covers all canvas positions.The loss includes padding and end-of-sequence positions in Countdown; corruption and scoring use the canvas alphabet.Source: §D.3 · p. 16

E.1 PROOF OF THE CAVITY DECOMPOSITION

  • The cavity conditional excludes the observed token while retaining the condition and all other noisy draft tokens as context.The clean token at that position remains the variable being predicted.Source: §E.1 · p. 16

E.2 LIKELIHOOD INVERSION AND EXACT RECOVERY

  • With an exact denoiser, likelihood inversion recovers a cavity estimate independent of the observed token; with an approximate predictor, the correction removes only the channel’s known self-likelihood.The correction does not remove arbitrary model error, so the estimate may still depend on the observed token.Source: §E.2 · pp. 16–17

E.3 SCORING COMPLETED CANDIDATES

  • A fresh forward pass scores the completed candidate at a common read level, computing all positionwise cavity estimates without updating any token.The proposal pass cannot serve as this read because it was conditioned on the preceding draft, not the candidate.Source: §E.3 · p. 17

E.4 CLEAN-DATA PSEUDO-LIKELIHOOD LIMIT

The appendix relates cavity scores to clean-data conditionals, specifies their implementation and decoding schedule, and clarifies what the resulting energy and best-so-far rule do—and do not—guarantee.

E.4 CLEAN-DATA PSEUDO-LIKELIHOOD LIMIT

  • As α approaches 1, the exact cavity conditional converges to the clean conditional given the other draft tokens, and the summed score approaches pseudo-likelihood when those contexts have positive probability.At α = 0, the cavity conditional reduces to the condition-only marginal; the limit is undefined for zero-probability contexts and can be slow for rare ones.Source: §E.4 · p. 17

E.5 STABLE LOGIT-SPACE IMPLEMENTATION

  • The scoring implementation computes all positionwise cavity estimates from one forward pass, using corrected logits and a log-sum-exp energy calculation.The implementation restricts logits to the canvas alphabet and assigns infinite energy to candidates containing a non-canvas token.Source: §E.5 · p. 17

E.6 BOUNDARY CASES AND INTERPRETATION

  • The energy is not a normalized sequence likelihood, and a lower value does not certify that a draft is correct.The method uses the interior read level αs = 7/8; αs = 1 cannot be inverted, while scores near that boundary amplify prediction errors.Source: §E.6 · p. 18

F.1 DECODING SPECIFICATION

  • Each round re-noises the incumbent, applies proposal passes along a keep-probability schedule, then reads at αs without a token update.The first round starts from a uniform canvas independent of the placeholder; later rounds use fresh independent noise.Source: §F.1 · p. 18

F.2 BEST-SO-FAR INVARIANT

  • The incumbent energy is non-increasing because only lower-energy completed candidates are retained, but this does not establish a global minimum or correctness.Intermediate proposal drafts are not scored, and decoding runs for its fixed budget rather than stopping when energy ceases to decrease.Source: §F.2 · p. 19

F.3 MODEL EVALUATIONS AND CACHED STATE

  • R rounds use RJ model evaluations per condition, and every completed round leaves a complete incumbent output.The incumbent is not re-read; conditions in a batch accept or reject candidates independently.Source: §F.3 · p. 19

F.4 DECODING HYPERPARAMETERS

  • The default decoding configuration uses a budget of RJ = 128 model evaluations per problem.The supplied table caption identifies its scope as decoding configuration.Source: Table 5 · p. 19

G.1 TASKS AND DATA

The experiments define evaluation criteria for Countdown, Sudoku-Extreme, and Maze-Hard, and describe the baseline and diagnostic settings used for comparisons.

G.1 TASKS AND DATA

  • Countdown test targets are out of distribution, with nine target values absent from training; a solution must use valid arithmetic and reach the target.Inputs contain 3, 4, or 5 source numbers and a target; outputs contain up to k − 1 equations.Source: §G.1 · p. 19

G.2 MODEL AND TRAINING

  • The DRLM training configuration is shared across tasks unless noted.The supplied table caption identifies the configuration’s scope but gives no specific settings.Source: Table 7 · p. 20

G.3 BASELINES

  • Diffusion baselines use differing objectives and decoders, while the autoregressive and recurrent baseline results come from specified prior reports or fine-tuning setups.UDLM uses the same backbone, data, optimizer, and budget as DRLM but an ancestral sampler; MGDM uses its 6M recipe and top-k decoding. Recurrent Sudoku and Maze metrics differ in training set size and path-matching criteria.Source: §G.3 · pp. 20–21

H ADDITIONAL EXPERIMENTS

  • Additional diagnostics change one component at a time and report Pass@1 at 128 model evaluations, averaged over three decode seeds.The reported ± values are standard deviations over seeds.Source: H ADDITIONAL EXPERIMENTS · p. 21

H.1 EFFECT OF THE SELF-LIKELIHOOD CORRECTION

The experiments compare cavity scoring with raw self-token scoring, test fixed-budget decoder settings, examine candidate coverage, and cross the training objective with the decoder.

H.1 EFFECT OF THE SELF-LIKELIHOOD CORRECTION

  • At αs = 7/8, cavity scoring gains +5.2 points over raw energy on Sudoku-Extreme; the gap reaches +5.8 at αs = 0.98 and 0.995.Cavity accuracy rises from 93.9% to 94.3% across the tested read levels, while raw accuracy slips from 88.8% to 88.5%.Source: Table 8 · p. 21, §H.1 · p. 21

H.2 SENSITIVITY AT A FIXED BUDGET

  • J = 2 is the worst tested proposal-pass setting on every task, while the defaults are best tested; restarting from noise is the worst re-noising setting on every task.On CD-5, keeping more of the incumbent (ρe = 0.3) is also clearly worse than the default; Sudoku configuration was selected on 2,000 puzzles before evaluation on the remaining 420,786.Source: §H.2 · pp. 21–22, Table 9 · p. 22, Table 10 · p. 22

H.3 BUDGET, ACCEPTANCE, AND COVERAGE

  • Pass@1 stays within one point of candidate coverage on every task, and most remaining failures occur when no correct candidate is proposed within budget.Coverage is the fraction of problems with at least one correct completed candidate, upper-bounding what any selection rule among those candidates could achieve.Source: §H.3 · p. 22, Table 11 · p. 22

H.4 OBJECTIVE AND DECODER

  • Both the training objective and decoder improve accuracy under either choice of the other, but their effects are not additive.Refinement helps the UDLM model more than the DRLM model, while the objective matters more under ancestral sampling.Source: Table 12 · p. 23, §H.4 · p. 23

I QUALITATIVE EXAMPLES

The examples show how retained parallel revisions change drafts across Maze-Hard, Countdown, and Sudoku-Extreme, and illustrate cases where lower energy can still favor an invalid answer.

Repair scale

  • On successful repairs, revisions change a median of 2.5 of 900 Maze-Hard cells, 13 of 26 CD-4 characters, 14 of 33 CD-5 characters, and 28 of 81 Sudoku cells.Maze-Hard uses ρe = 0.7; Countdown and Sudoku-Extreme use ρe = 0.1.Source: I QUALITATIVE EXAMPLES · p. 23

Countdown

  • Only 17% of revisions fixing invalid CD-4 drafts change at most three characters, compared with 4% on CD-5.The examples therefore show that single-token repairs are the exception on Countdown.Source: I QUALITATIVE EXAMPLES · p. 23

Sudoku-Extreme

  • For all 527 subsample Sudoku puzzles solved within 128 evaluations, wrong-cell counts never increase across proposal passes in the solving round.A median of only 20% of wrong cells conflict with another cell in their row, column, or box; 70% of error-containing units still hold all nine digits.Source: I QUALITATIVE EXAMPLES · p. 23, Figure 5 · p. 25
  • Before Sudoku’s solving round, retained drafts replace one another with decreasing energy while staying about as far from the solution.The displayed rounds omit passes after the first solved one.Source: I QUALITATIVE EXAMPLES · pp. 24–25, Figure 5 · p. 25

Failure cases

  • A lower-energy candidate can still be invalid: one Maze-Hard round adds a cell to a valid shortest path and replaces it.The authors also report a Countdown failure where the learned score prefers an invalid draft because it estimates contextual support, not arithmetic validity.Source: I QUALITATIVE EXAMPLES · p. 23, Figure 6 · p. 25

Related papers

Loading 2610.26542v1…