Self-Aware Markov Models for Discrete Reasoning

Gregor Kornhardt; Jannis Chemseddine; Christian Wald; Gabriele Steidl

cs.LGMar 17, 2026

Self-Aware Markov Models for Discrete Reasoning

Authors: Gregor Kornhardt, Jannis Chemseddine, Christian Wald, Gabriele Steidl

Organizations: Technische Universität Berlin, Straße des 17. Juni 136, 10623 Berlin, Germany

Abstract

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{\times}3$ Sudoku subgrid from a complete puzzle. We show the state after $k\!=\!1,2,3$ refinement steps; the predicted progress $\tau_{\theta}$ is displayed above each panel. Cell shading encodes confidence (mixing score $c_{\theta}$ ); clue cells are marked with $\ast$ . 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.

Autoregressive
Method	Params	CD4 (%)
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) .

Masked Diffusion Models

Masked diffusion models generate sequences by repeatedly predicting and revising masked tokens, using bidirectional context for parallel sampling.

See all 184 papers

Binary Discrete Diffusion: Candidate-Set Diffusion for Logic and Planning Tasks
Oct 20, 2026Rodrigo Casado Noguerales, Dimitri von Rütte, Bernhard Schölkopf +2Diffusion Models Diffusion-Based Planning
ETH Zurich · Max Planck Institute for Intelligent Systems · ELLIS Institute Tübingen
Flow Reasoning Models: Turning Flows Into Efficient Recurrent Reasoners
Jun 28, 2026Alec Helbling, Andrey Bryutkin, Mauro Martino +3Self-Conditioning Structured Reasoning
1Georgia Tech · 3MIT-IBM Computing Research Lab · 4IBM Research +1
Random Remasking Scales Masked Diffusion Inference
Oct 22, 2026Zhaoyi Zhou, Frank Wu, Andrea ZanetteMasked Diffusion Models Diffusion Language Model Inference
Carnegie Mellon University

Oct 20, 2026J/K move · Enter open · S save

Binary Discrete Diffusion: Candidate-Set Diffusion for Logic and Planning Tasks

Rodrigo Casado Noguerales, Dimitri von Rütte, Bernhard Schölkopf, Thomas Hofmann+1

ETH Zurich · Max Planck Institute for Intelligent Systems · ELLIS Institute Tübingen

Diffusion Models Diffusion-Based Planning Masked Diffusion Models

Open paper AskPDF arXiv

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.

Abstract

Masked Diffusion Models

Explore similar work

Binary Discrete Diffusion: Candidate-Set Diffusion for Logic and Planning Tasks

Flow Reasoning Models: Turning Flows Into Efficient Recurrent Reasoners

Random Remasking Scales Masked Diffusion Inference