PaperlayerView exact revision on arXiv

Source-linked AI summary

Let It Go or Learn to Self-Correct: Continuous Diffusion for Constrained Discrete Tasks

Mariia Drozdova, Stéphane Liem Nguyen, François Fleuret

arXiv:2609.09009v1cs.LGcs.AI

TL;DR

The paper asks why standard diffusion sampling can fail on globally constrained discrete tasks despite informative clean predictions. It compares a residual-free sampler with self-correction training and finds improved constraint satisfaction, including Sudoku validity rising from 31% to 95% without retraining. The authors conclude that inference and training should better accommodate states produced during inference.

  • Problem

    The paper studies whether diffusion denoisers learn global constraints when standard reverse updates can preserve locally plausible but globally invalid states.

  • Method

    The paper removes the direct current-state residual through Tweedie reprojection and trains denoisers on noisy versions of their own predictions.

  • Results

    31% to 95%: one-hot Sudoku validity improves without retraining, with consistent gains across other discrete tasks and improved standard samplers after self-correction training.

  • Takeaways & Limitations

    Continuous diffusion models can learn useful global constraint structure, but constrained discrete reasoning benefits from reducing early-state commitment or training recovery from inference-time errors.

  • Takeaways & Limitations

    The evaluation focuses on controlled constrained-discrete tasks, and Tweedie reprojection may discard meaningful continuous detail in perceptually rich domains.

Abstract

from arXiv · show

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.

1 Introduction

The paper examines why continuous diffusion can produce locally plausible but globally invalid solutions on constrained discrete tasks, and tests whether changing sampling or training improves constraint satisfaction.

  • The study evaluates diffusion on constrained discrete tasks including Sudoku, graph connectivity, Latin squares, and N-queens.
  • Standard DDPM sampling can preserve locally plausible states that violate global constraints, raising whether the denoiser or reverse update is responsible.
  • Removing the direct current-state residual from the reverse update, without retraining, improves one-hot Sudoku validity from 31% to 95%.The resulting sampler is called Tweedie reprojection.
  • The result indicates that clean denoiser proposals contain useful constraint structure, but standard trajectories may fail to exploit them after drifting toward invalid states.
  • The study treats Tweedie reprojection as a diagnostic for controlled constrained-discrete benchmarks rather than a general sampler for perceptually rich domains.
  • Self-correction training exposes the denoiser to noisy versions of its own predictions and improves standard samplers by reducing training-inference mismatch.

2 Preliminaries

The preliminaries define the forward diffusion path and explain how DDPM reverse steps use the model’s clean-sample estimate alongside the current noisy state.

  • Diffusion notation uses t = 0 for clean data and t = T for noise, with discrete or continuous time depending on the sampler.
  • The forward variance-preserving process maps a clean sample x0 to a noisy sample xt through α(t) and β(t).
  • At sampling time, DDPM replaces the unknown clean sample in its reverse posterior with the model’s current clean-sample estimate.
  • Each DDPM reverse step conditions on both the current noisy state xt and the clean prediction x̂0(xt, t).
  • The paper uses x-prediction, whose optimal prediction under the stated loss is E[x0 | xt], interpretable through Tweedie’s formula.
  • Continuous-time samplers can be represented by a probability-flow ODE or a reverse-time SDE, with σ(t) controlling stochasticity theoretically without changing target marginals.Learned denoisers and finite-step solvers can nevertheless make empirical performance depend on σ(t).

3 Method

The method addresses training-inference mismatch in constrained discrete diffusion from both sides: it reduces the reverse update’s dependence on the current state and trains the denoiser on its own imperfect predictions.

  • Motivation: Standard reverse updates can preserve wrong discrete commitments because inference visits states unlike the forward-noised valid samples used in training.The mismatch can produce locally plausible but globally invalid intermediate states.
  • Anchor-free update: Tweedie reprojection removes the direct x_t-dependent residual and re-noises the denoiser’s clean prediction with forward marginal variance.It uses x_t only through the denoiser prediction x̂_0(x_t,t).
  • Self-correction training: Self-correction trains the model to recover x_0 from noisy versions of its own intermediate predictions, which may violate constraints.The procedure samples t_1, predicts x̂_0, re-noises it at t_2, and applies a recovery loss.
  • Self-correction training: The combined objective retains ordinary denoising through L_simple while adding recovery from self-induced states through L_rec.This preserves single-step denoising on forward-noised data while training correction from sampler-induced states.
  • Local distributional interpretation: Under a fixed continuous proposal, self-correction matches the proposal-centered distribution locally while retaining the original clean sample as the target.The result is a local interpretation rather than a claim about the entire reverse trajectory.

4 Experiments

Experiments evaluate continuous samplers and self-correction across constrained discrete benchmarks, showing that removing the reverse residual gives broad gains and self-correction substantially improves standard samplers.

  • Setup: The evaluation covers Sudoku, Sudoku-Extreme, graph connectivity, Latin squares, and N-queens using continuous representations and task-specific models.The study also includes conditional generation and inpainting/completion settings.
  • Setup: Results compare DDPM, Euler, Euler-Maruyama, EM decay, and Tweedie reprojection, with sampler selection and noise tuning specified by the evaluation protocol.Table 1 reports mean ± standard deviation over three independently trained seeds and independent test seeds.
  • Baseline samplers: 31% to 95%: on 21-clue Sudoku, Tweedie reprojection improves success over DDPM without self-correction, with similar gains across other datasets.The comparison supports useful clean predictions under the standard denoising objective.
  • Self-correction: 31% to 87%: self-correction raises DDPM success on 21-clue Sudoku and narrows the gap between standard sampling and Tweedie reprojection.On Sudoku-Extreme, DDPM rises from about 19% to 64% pass@10.
  • Cross-task results: Self-corrected continuous configurations reach 84% pass@10 on Sudoku-Extreme and nearly perfect graph-connectivity accuracy, while EM decay reaches 90% pass@10 on Sudoku-Extreme.These comparisons are not compute-matched to the cited external methods.
  • Representation checks: 65.7% versus 0.8%: on MNIST-Sudoku Hard, Tweedie reprojection improves accuracy without retraining the released baseline checkpoint.It also exceeds the best reported SRM sampling-order strategy at 51.6%.

5 Analysis

The analysis attributes sampler differences to globally invalid states and shows that the DDPM residual often preserves such mistakes, while self-correction targets the missing inference-time states.

  • Failure mode: The same denoiser can produce locally plausible outputs whose decoded grids remain globally invalid, including near-one-hot states and recognizable MNIST digits.Exact validity therefore distinguishes local symbol quality from global constraint satisfaction.
  • Training-inference mismatch: Training uses forward-noised valid objects, whereas inference visits noisy versions of imperfect proposals that may have little probability under ordinary forward noising.This creates the relevant exposure-bias mismatch.
  • Residual mechanism: The DDPM reverse center combines the clean prediction with an x_t-dependent residual that can preserve existing constraint violations.The residual is useful when the current state lies on a reliable trajectory, but harmful when it contains discrete mistakes.
  • Residual mechanism: E+ is much larger than E−: along compared trajectories, the Tweedie center is more often valid when the DDPM center is invalid than the reverse.This identifies the residual’s dominant effect in the tested regime.
  • Training-side correction: Self-correction exposes the denoiser to noisy versions of its own intermediate predictions while retaining the valid solution as target.It improves reliability on sampler-induced states and narrows the DDPM-Tweedie gap.
  • Ablations: Self-correction outperforms generic input perturbation, self-conditioning, and random symbol corruptions, indicating that model-induced states are the most relevant training examples among these ablations.The comparison supports matching the inference sampling distribution rather than adding arbitrary noise.
  • Noise-scale analysis: Changing sampling variance alone does not close the DDPM-Tweedie gap, although EM decay can break some bad intermediate commitments.The sampler difference is therefore not explained solely by noise scale.

6 Related work

Related work spans diffusion samplers, explicit constrained generation, structured and discrete diffusion, exposure bias, and differentiable or recursive solvers. This paper instead diagnoses inference dynamics when constraints must be learned implicitly.

  • Diffusion models and samplers: Diffusion research includes DDPM, DDIM, ODE and SDE integrators, predictor-corrector methods, and Langevin-like noise injection.These methods provide alternative reverse-process designs beyond standard DDPM sampling.
  • Inference-time control: Constrained-generation methods use masking, inpainting, posterior guidance, projection, optimization, guidance, or search-based post-processing.They typically assume explicit constraints or a violation signal.
  • Structured and combinatorial tasks: Prior work applies diffusion to Sudoku, graphs, and combinatorial optimization, while discrete diffusion and structured variants often outperform continuous diffusion in these settings.This motivates examining the limitations of continuous diffusion on discrete structure.
  • Training-inference mismatch: Exposure-bias research links training on ground-truth inputs with inference on model predictions to error accumulation, including in diffusion trajectories.The present work studies this mismatch for constrained discrete sampling.
  • Non-diffusion solvers: Differentiable optimization layers and recursive refinement models offer alternative approaches to constrained reasoning and error compounding.The paper characterizes continuous diffusion inference rather than competing directly with these solvers.

7 Limitations

The study’s limitations center on scope and training coverage: Tweedie reprojection is evaluated for constrained discrete tasks, may discard continuous detail, and self-correction only partially covers inference-time states.

  • Scope: The experiments focus on constrained discrete tasks and do not establish that Tweedie reprojection suits perceptually rich domains.Such domains require modeling diverse texture, geometry, color, and fine details.
  • Sampler boundary: Tweedie reprojection discards the residual xt − α(t)ˆx0, potentially losing meaningful within-mode variation or fine details.State-preserving samplers may retain information that repeated reprojection removes.
  • Training coverage: Self-correction exposes the model to one-step model-induced states, but longer sampling trajectories may still reach poorly represented states.More complete training objectives or noise processes remain future work.

8 Conclusion

Across constrained discrete tasks, the same denoiser behaves differently under different reverse processes because sampling can produce states outside the forward-noised training distribution. Tweedie reprojection addresses this at inference, while self-correction addresses it during training, supporting better alignment between training and sampling states.

  • Conclusion: Across Sudoku, graph connectivity, Latin squares, and N-queens, the same trained denoiser can behave differently under different reverse processes.The conclusion concerns continuous diffusion models applied to constrained discrete tasks represented in continuous space.
  • Conclusion: Standard sampling can create locally plausible but globally invalid states that are not well covered by training.This is identified as a training-inference mismatch between forward-noised valid objects and sampler-induced states.
  • Conclusion: Tweedie reprojection removes the direct xt-dependent residual, while self-correction trains recovery from the model’s own intermediate predictions.The two interventions address the mismatch from inference and training sides, respectively.
  • Conclusion: The results suggest that continuous diffusion models can learn global constraints, but constrained discrete reasoning requires better alignment between training and sampling states.This is the paper’s overall supported conclusion.

D Additional tasks details

The appendix describes the benchmark suite, data representations, architectures, training settings, discrete-diffusion references, inference procedures, and evaluation resources used for constrained discrete generation and completion.

  • Datasets: The benchmark suite includes Sudoku-Extreme, MNIST Sudoku, graph connectivity, Latin squares, and N-queens.The study evaluates constrained data generation and inpainting or completion, including two introduced datasets: Latin squares and N-queens.
  • Datasets: Dataset conditions use partial revealed cells for Sudoku, Latin squares, and N-queens, while graph connectivity provides the full adjacency matrix.Table 2 also marks certain settings as out-of-distribution.
  • Representation and architecture: Discrete configurations are represented as one-hot tensors, with token count and vocabulary size varying across tasks.Most models receive continuous samples and diffusion time as inputs.
  • Representation and architecture: The main architecture is a Transformer with positional and diffusion-time embeddings, while graph connectivity uses a Neural Logic Machine and MNIST Sudoku uses a pretrained SRM model.Shared hyperparameters are used across most Transformer-based tasks.
  • Training and evaluation: Evaluation compares DDPM, Euler, Euler-Maruyama, EM decay, and Tweedie reprojection, with EM noise selected by validation cross-validation.Discrete-diffusion references use ancestral, confidence-based, and remasking inference procedures.
  • Resources: Training runs used one RTX 3090, while evaluation typically took 5–15 minutes per checkpoint and Sudoku Extreme pass@10 took 30–50 minutes per sampler.Graph-connectivity EM cross-validation generally required 12–25 minutes for five-fold evaluation.

F.3 Sensitivity of EM decay to sampler hyperparameters

EM decay validity depends sensitively on the noise scale and the time at which decay begins, whereas Tweedie reprojection avoids these two hyperparameters. The appendix also contrasts their update mechanisms: DDPM retains current-state residual information, while Tweedie recenters on clean predictions.

  • Sensitivity results: 0.177 validity falls to 0.010 when EM noise scale σ increases from 14 to 20 at τstart = 0.8 on Sudoku-Extreme.For the self-corrected model, validity instead reaches 0.531 at (σ, τstart) = (20, 0.7).
  • Sensitivity results: 0.531 validity at (σ, τstart) = (20, 0.7) falls to 0.138 when decay begins at τstart = 0.8 for the self-corrected model.These results show that delaying decay does not consistently improve validity.
  • Sampler mechanism: DDPM can remain trapped near the current state when the denoiser direction is decorrelated from or anti-aligned with the current state.The toy example illustrates this anchored-update behavior but is not evidence that the same mechanism causes failures in the discrete tasks.
  • Sampler mechanism: Tweedie reprojection directly recenters the next state on the denoiser prediction, while DDPM retains a residual anchor to the current noisy state.The residual can preserve part of an off-manifold state even when the denoiser predicts near the valid set.
  • Discrete validity: On Sudoku, cellwise accuracy declines gradually with noise, but full-grid validity collapses earlier because one flipped cell invalidates the puzzle.The trained denoiser extends the useful noise range and maintains larger argmax margins than raw noisy inputs.

K.5 Proof and nearby-proposal extension

The proof framework explains self-correction as matching training inputs to proposal-centered states encountered during inference. Its nearby-proposal extension relaxes the idealized assumption that the denoiser makes an unchanged proposal.

  • Proof interpretation: The proof’s fixed-proposal regime makes self-correction locally match the proposal-centered reverse distribution while standard training remains centered on the clean sample.This provides a local interpretation of why self-correction targets states the sampler may revisit.
  • Nearby-proposal extension: The nearby-proposal extension assumes each reverse-step proposal remains within radius ρk of an ideal proposal and bounds the resulting one-step change through the DDPM mean coefficient.Positive posterior variances allow KL chain-rule and data-processing arguments over a reverse window.
  • Training-inference alignment: Re-noising the model’s proposal produces training inputs similar to states encountered during the relevant part of sampling.The diagnostic compares forward-noised inputs, self-correction inputs, and actual DDPM states.
  • Training-inference alignment: The event-frequency gaps between self-correction inputs and DDPM states are 3.4–6.8× smaller than gaps for standard forward-noised inputs.The result supports, but does not establish, equality of the full distributions.
  • Self-correction construction: Self-correction trains on noisy versions of model-induced predictions while retaining the original valid grid as the regression target.This changes the supervised input distribution to include invalid or partially wrong intermediate predictions.

L Self-Correction Loss Ablations on Sudoku

Self-correction chiefly improves DDPM by exposing the model to its own inference-induced errors, while Tweedie is already strong across loss-weight settings. Variants that reproduce closed-loop sampling errors outperform generic perturbation or memory-based alternatives.

  • Simple loss weight: Self-correction improves DDPM from roughly 29% to approximately 85% validity on Sudoku puzzles with 21 clues.Tweedie reprojection already performs strongly across the tested loss weights, so the main effect is on DDPM.
  • Simple loss weight: λsimple ∈ {0, 0.01, 0.1, 0.2, 0.5} yields a stable high-performing region, with λsimple = 0.1 used in the main experiments.The method is not sensitive to the exact value of λ in this sweep.
  • Loss-formulation variants: Input perturbation and self-conditioning are weaker than self-correction, especially for DDPM, while DDPM-aware variants are strongest.This suggests recovery signals should reflect model-induced states rather than generic noise or an added memory channel.
  • Invalid-state exposure: Random-symbol corruption improves DDPM from 34.5% to 64.0%, whereas self-correction reaches 78.2% in the same comparison.Model-induced recovery states are therefore better matched to reverse-sampling errors than arbitrary symbol corruptions.
  • Sampling-step ablation: Increasing reverse steps drives Tweedie toward near-perfect validity, but baseline DDPM remains around 0.29–0.31.The DDPM gap is not fixed by adding more opportunities for reverse correction.
  • Representation robustness: The same qualitative pattern appears across representations and decoders: Tweedie improves over DDPM, while self-correction also benefits DDPM.The evaluation includes uint4 codes, fixed random embeddings, and mini MNIST-Sudoku with a learned CNN decoder.

O.1 Soft metrics.

The paper supplements strict validity with continuous proximity, constraint violations, uniqueness, and coverage metrics. These diagnostics distinguish discrete validity from closeness to one-hot states and show that validity gains can coexist with broader solution coverage.

  • Soft metrics: Distance to one-hot is per-puzzle RMSE from the final continuous output to its nearest one-hot encoding, with lower values indicating greater discrete proximity.Lower distance does not imply that the decoded configuration satisfies all constraints.
  • Soft metrics: Constraint-violation rate measures the fraction of predicted cells involved in violated task constraints and generally follows the same trends as exact validity.The constraints cover Sudoku, Latin squares, and N-queens-specific violations.
  • Diversity metrics: Valid uniqueness is the number of distinct valid decoded boards divided by the number of valid generated boards.The ratio is computed separately for each training seed and averaged across seeds.
  • Coverage: Self-correction increases total coverage from 20.0% to 44.8% with EM decay and from 24.6% to 33.8% with Tweedie.Coverage is similar within and outside the training set, indicating broader coverage rather than repeated generation of a small training-board subset.
  • Multiple completions: Both baseline and self-correction recover 139 of 297 valid completions, or 46.8%, for the fixed-clue Sudoku experiment.Neither configuration recovers all completions, and self-correction does not improve aggregate coverage in this experiment.

P Trajectories under different samplers

Sampling trajectories can stabilize at different times depending on final validity, and self-correction delays commitment to decoded predictions. Trajectory examples also show that locally recognizable symbols do not guarantee globally valid solutions, while Tweedie reprojection can reach validity where DDIM does not.

  • Commitment time: Final-valid predictions stabilize earlier than final-invalid predictions across the pooled nine-task results.The comparison averages 256 samples per validity group across nine tasks.
  • Commitment time: Self-correction increases mean commitment time for final-invalid trajectories from 0.92 to 0.98 under DDPM and from 0.89 to 0.99 under Tweedie.The later stabilization is consistent with longer opportunities to revise predictions, though revisions do not necessarily produce valid solutions.
  • Trajectory outcomes: Tweedie failures can occur even after a valid clean proposal appears during sampling.Among Sudoku-Extreme failures, 38% are denoiser-related, while 62% include at least one valid proposal before ending invalid.
  • Trajectory outcomes: In a visualized MNIST-Sudoku example, DDIM produces recognizable cells but an invalid grid, whereas Tweedie reprojection ends in a valid solution.The DDIM grid repeats digit 1 within the central 3 × 3 block, despite local cell correctness.
Loading 2609.09009v1…