Diffusion language models learn when to correct errors for better reasoning
Learning When Diffusion Language Models Should Regenerate
Computation and Language
Summary
Some language models create text step-by-step but may make errors that are not easy to spot by just checking how confident they are in each word. The authors found that hidden information inside the model holds clues about mistakes. They built a tool called DiVR that predicts if a sentence will be correct or not, so it knows when to try rewriting a part of the text. This method improves the accuracy of solving problems like Sudoku and other reasoning tasks without changing the original model.
What this means in practice
- •For natural language processing engineers: Improve reasoning accuracy in language models by selectively regenerating uncertain text blocks using a learned verifier.
- •For ai-powered puzzle solvers: Enhance puzzle and logic problem solvers by using DiVR to decide when to fix parts of generated solutions for better correctness.
Authors
Minhae Oh, Jeonghye Kim, Jinwoo Koo, Yeongoon Kim, Jimin Seo, Minsik Seol, Jungwoo Lee
Abstract
Remasking allows diffusion language models to revisit generated text, but effective correction requires knowing when and what to revise. Our diagnosis on Sudoku reveals two gaps. Hidden states contain error information that token confidence does not reliably expose, and remasking an isolated error usually reproduces it when the surrounding context remains fixed. These findings motivate DiVR , a verifier that predicts whether a partially generated sequence will reach a correct answer. Trained on the final correctness of the model's own trajectories, DiVR decides when to regenerate a block and which alternative to retain, while the generator itself stays unchanged. On five reasoning benchmarks with LLaDA-8B-Base, DiVR raises the average accuracy of the fine-tuned generator from 35.3\% to 43.4\% at $1.9$-$3.5\times$ the cost of greedy decoding, above search baselines that use $3.8$-$7.4\times$. It also complements post-training such as supervised fine-tuning and reinforcement learning, and generalizes to other generators, including Dream.