Draft refinement improves complex reasoning with parallel updates
Breaking Diffusion: Parallel Reasoning with Draft Refinement
Artificial Intelligence
Summary
Many models reason by creating an answer step-by-step, but this new method lets the model revise the entire draft answer all at once using a single tool called a denoiser. The authors train this tool to both fix errors and keep good parts by seeing versions of answers with random mistakes. Then they estimate how confident the model is in each part of the answer, helping it decide if the revised answer is better. This approach works better than older methods on math and logic puzzles, showing a new way for models to solve complex problems more efficiently.
What this means in practice
- •For machine learning engineers: Enhance reasoning capabilities of AI models by implementing parallel draft refinements for faster and more reliable complex problem solving.
Authors
Jaesang Yu, Seoyeon Byun, Yun Wang, Yan Lu, Juho Kim
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.