Continuous diffusion models learn to fix errors in discrete puzzles
Let It Go or Learn to Self-Correct: Continuous Diffusion for Constrained Discrete Tasks
Machine LearningArtificial Intelligence
Summary
The paper looks at how a type of AI model called continuous diffusion works on tricky puzzles like Sudoku and other tasks with strict rules. These models usually clean up noisy data step-by-step, but early mistakes in puzzles are hard to fix later. The authors found that changing how the model picks answers and training it to fix its own errors makes it much better at solving these puzzles. Their work shows that to handle puzzles with strict rules, models need to either avoid early mistakes or learn to correct them.
What this means in practice
- •For ai developers: Improve AI models that solve discrete puzzles and tasks by training them to self-correct during sampling, increasing solution validity.
- •For software engineers in automated reasoning: Design inference processes for constraint-based systems that reduce early commitment to errors, enhancing overall correctness in output.
Authors
Mariia Drozdova, Stéphane Liem Nguyen, François Fleuret
Abstract
Denoising Diffusion Probabilistic Models (DDPMs) generate samples by starting from noise and repeatedly denoising while keeping each update close to the current noisy state. This behavior is effective in many continuous domains, but its role is less clear for globally constrained discrete tasks, such as Sudoku, graph connectivity, Latin squares, and N-queens. In such settings, early discrete errors can be difficult to undo. As a result, standard diffusion sampling may preserve early mistakes, even when the model's clean predictions are informative. We compare standard samplers to sampling directly from the model's clean prediction. Without retraining, this single change improves Sudoku validity from 31% to 95%, with consistent gains across the other discrete tasks. We hypothesize that staying close to the current noisy state is harmful because the reverse trajectory can drift off the forward noising distribution the model was trained on. To reduce this train-test mismatch, we further introduce self-correction training, which exposes the model to its own predictions, improving robustness to errors that arise during inference. This substantially improves the performance of standard samplers. Our results suggest that continuous diffusion models can learn nontrivial global constraints, but discrete reasoning tasks require better alignment between training and inference: either through samplers that reduce commitment to early decisions, or through training that teaches the model to correct its own inference-time errors.