Self-Aware Markov Models for Discrete Reasoning

2026-03-17

AI Summary

Existing masked discrete diffusion models lack self-correction capabilities during inference and rely on a fixed number of denoising steps, making them inflexible to varying problem complexities. This work proposes a self-aware Markovian inference framework that learns a transition kernel conditioned on its own outputs, incorporating a re-masking mechanism and an adaptive stopping criterion to dynamically adjust both the number of computation steps and the inference trajectory. The approach requires only two lightweight prediction heads and is compatible with existing pre-trained models, allowing for straightforward fine-tuning. On Sudoku-Extreme, it achieves a 95% solve rate, substantially outperforming flow-based methods; on Countdown-4, it solves 96% of instances within an average of 10 steps, with some problems resolved in as few as two steps.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Mixed Discrete/Continuous Search
• Reasoning under Uncertainty: Sequential Decision Making

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: LLM based quality controls for crowd work

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.

Problem

• discrete reasoning
• masked diffusion models
• self-correction
• adaptive computation
• Markov models

Innovation

• Self-Aware Markov Models
• remasking
• adaptive stopping criterion
• discrete reasoning
• lightweight fine-tuning