Standard masked discrete diffusion models face limitations in reasoning tasks due to their inability to correct their own mistakes on the masking path. Since they rely on a fixed number of denoising steps, they are unable to adjust their computation to the complexity of a given problem. To address these limitations, we introduce a method based on learning a Markov transition kernel that is trained on its own outputs. This design enables tokens to be remasked, allowing the model to correct its previous mistakes. Furthermore, we do not need a fixed time schedule but use a trained stopping criterion. This allows for adaptation of the number of function evaluations to the difficulty of the reasoning problem. Our adaptation adds two lightweight prediction heads, enabling reuse and fine-tuning of existing pretrained models. On the Sudoku-Extreme dataset we clearly outperform other flow based methods with a validity of 95%. For the Countdown-4 we only need in average of 10 steps to solve almost 96% of them correctly, while many problems can be solved already in 2 steps.
Figures & tables
Figure 1: Inference trajectories. DFM Gat et al. (2024) : once the sampler leaves the masking path (red, incorrect token), it fails to recover and does not converge to the target solution ( X ). Ours : trained on model-induced off-path states, the model detects the mistake and corrects it, recovering the correct final sequence (green); color intensity indicates the model’s confidence.
Figure 2: Inference on a 3×3 Sudoku subgrid from a complete puzzle. We show the state after k=1,2,3 refinement steps; the predicted progress τθ is displayed above each panel. Cell shading encodes confidence (mixing score cθ ); clue cells are marked with ∗ . Correct predictions are blue and incorrect ones red.
Method
Sudoku Extreme
Kaggle Unfiltered
DFM (Euler)
31.6
44.5
DFM + Top-K
68.8
89.2
DFM + Top-K Margin
68.5
88.5
DFM + ReMDM
51.3
71.9
CTTT
–
98.9
Ours
95.2
99.9
Table 1: Sudoku solving accuracy (%). All methods except CTTT use the same backbone architecture.
Figure 3: Sudoku accuracy vs. ensemble size.
# Ensemble
Avg. steps
1
64.2
3
105.1
5
123.5
10
146
Table 2: Average sampling steps vs. ensemble size.
Method
Params
CD4 (%)
Autoregressive
GPT-2 Scratch
85M
45.8
GPT-2 Scratch
303M
41.3
Stream-of-Search
250M
54.2
LLaMA
7B
41.1
LLaMA
13B
51.1
Table 3: Countdown-4 (CD4). Baselines reproduced from Ye et al. Ye et al. (2025) .
Discrete diffusion models iteratively denoise a noisy sequence of categorical data such as text, enabling any-order generation with self-correction and in-filling capabilities. This makes them promising candidates for reasoning, planning, and constrained-generation tasks like coding and math. Yet masked and uniform diffusion erase a token's identity in a single corruption event, so their noisy states cannot explicitly represent partially resolved uncertainty in which several alternatives remain plausible. We close this gap with Binary Discrete Diffusion, a family of continuous-time discrete-diffusion processes whose noisy state is a set of candidate tokens per position, represented as a binary membership vector. Binary diffusion changes candidate sets through additions and removals, allowing information to be lost and recovered gradually while the diffusion state itself represents partial uncertainty. On Sudoku-Extreme, Boxoban, and graph coloring, the symmetric formulation of binary diffusion consistently outperforms strong, well-tuned masked and uniform baselines under matched architectures and controlled training, sampling, and tuning budgets, with its largest gains on the hardest instances, while remaining competitive on language modeling (Text8, TinyStories, OpenWebText). These results support candidate sets as an effective intermediate representation for iterative reasoning and show that the corruption kernel can materially change the computation learned by the reverse model.
Rodrigo Casado Noguerales, Dimitri von Rütte, Bernhard Schölkopf +2
ETH Zurich · Max Planck Institute for Intelligent Systems · ELLIS Institute Tübingen
Structured reasoning requires making and revising interdependent decisions to reach a globally consistent solution. Existing architectures struggle with this: autoregressive models commit sequentially and cannot revise earlier decisions, while masked diffusion models often require careful decoding schemes to coordinate interdependent predictions. We introduce Flow Reasoning Models (FRMs), a novel framework for structured reasoning that adapts continuous flows over discrete structured outputs with a simple recurrent refinement mechanism. By self-conditioning a flow model on its own past outputs, we turn one-shot denoising into iterative solution refinement. This lets FRMs make and revise decisions in parallel, efficiently coordinating interdependent choices across solutions. Yet conventional self-conditioning becomes unreliable at greater recurrent depth due to exposure bias between one-step training predictions and recursively generated inference states. We address this mismatch with Fixed-Point Forcing (FPF), which trains FRMs on states produced by their own inference dynamics while preserving the standard flow-matching objective. FRMs achieve solve rates of 99.5%, 100.0%, and 99.9% on Sudoku-Extreme, Zebra, and Maze-Unique, respectively. On Sudoku-Extreme, FRMs achieve higher peak accuracy than the evaluated masked-diffusion and specialized reasoning baselines while remaining highly compute-efficient, matching the next-best method's 98.7% peak solve rate with 44× fewer inference FLOPs.
Alec Helbling, Andrey Bryutkin, Mauro Martino +3
1Georgia Tech · 3MIT-IBM Computing Research Lab · 4IBM Research +1
Masked diffusion models (MDMs) generate sequences by progressively unmasking tokens, but errors introduced during decoding are typically irreversible. Recent work has shown that iterative correction can improve sample quality, yet existing theory fails to explain this given a learned, imperfect denoiser model. We re-formulate MDM inference as a Markov chain on the clean sequence space, defined by a remask-denoise kernel that randomly re-introduces masks and then denoises. This helps us analyze the actual denoising trajectory towards the data distribution. The analysis motivates Random Remask-Denoise (R2D) inference, a purely test-time, finetuning-free, reward-free procedure. To further improve performance, we propose Trajectory Correction Training (TCT), a denoiser-training procedure that exposes the model to its own errors along the generation trajectory. On logical tasks, our approach exhibits a striking near-log-linear test-time scaling, achieving over 97% accuracy on hard Sudoku puzzles and over 99% accuracy on 3-SAT. Our approach also applies to pretrained models on actual language tasks, beating iterative correction baselines and the official Q-mode accuracy of LLaDA2.1-mini on HumanEval+.