Source-linked AI summary
Learn from Your Mistakes: Self-Correcting Masked Diffusion Models
Yair Schiff, Omer Belhasin, Roy Uziel, Guanghan Wang, Marianne Arriola, Gilad Turok, Ran Zilberstein, Michael Elad, Volodymyr Kuleshov
TL;DR
MDMs cannot revise unmasked tokens, so decoding errors can accumulate and degrade generation. ProSeCo trains one model to unmask and correct its own sampled outputs, then interleaves correction with unmasking. Across conditional and unconditional tasks, it improves speed-quality trade-offs and supports further inference-time quality scaling, while requiring extra training computation.
Problem
MDMs leave unmasked tokens fixed, allowing parallel-decoding errors to accumulate and making efficient token correction difficult.
Method
ProSeCo adds a correction objective trained on model-generated outputs and interleaves corrective refinement with standard unmasking.
Results
Across conditional and unconditional generation, ProSeCo outperforms vanilla MDMs and alternative correctors in speed-quality trade-offs and supports inference-time compute scaling for improved generation.
Takeaways & Limitations
ProSeCo provides a minimally modified MDM framework that can revise already decoded tokens and improve generation quality through additional inference-time computation.
Takeaways & Limitations
The main drawback is the added computational cost of a second forward pass during training.
Abstract
from arXiv · showhide
Masked diffusion models (MDMs) have emerged as a promising alternative to autoregressive models, enabling parallel token generation while achieving competitive performance. Despite these advantages, MDMs face a fundamental limitation: once tokens are unmasked, they remain fixed, leading to error accumulation and ultimately degrading sample quality. We address this by proposing a framework that trains a model to perform both unmasking and correction. By reusing outputs from the MDM denoising network as inputs for corrector training, we train a model to recover from potential mistakes. During generation we apply additional corrective refinement steps between unmasking ones in order to change decoded tokens and improve outputs. We name our training and sampling method Progressive Self-Correction (ProSeCo) for its unique ability to iteratively refine an entire sequence, including already generated tokens. We conduct extensive experimental validation across multiple conditional and unconditional tasks, demonstrating that \method~yields better quality-efficiency trade-offs (up to ~4x faster sampling) and enables inference-time compute scaling to further increase sample quality beyond standard MDMs (up to ~1.2x improvement on benchmarks).
1 Introduction
MDMs generate tokens in parallel but cannot revise unmasked tokens, so decoding errors accumulate and degrade quality. ProSeCo jointly trains unmasking and self-correction, then interleaves corrective refinement during generation to improve speed-quality trade-offs and enable inference-time scaling.
- Motivation: MDMs offer parallel generation and competitive quality, but fixed unmasked tokens allow decoding errors to accumulate and degrade samples.The paper positions MDMs as an alternative to autoregressive models across multiple scales and domains.
- Approach: ProSeCo trains a model to both unmask tokens and correct mistakes by reusing model-generated outputs as corrupted inputs.The method adds a corrector objective to standard MDM training and refines already unmasked positions.
- Approach: During inference, ProSeCo interleaves corrective steps with unmasking steps to iteratively refine the full sequence.The method can update positions that were already decoded.
- Results: ∼4x faster generation is achieved without quality degradation, while benchmark accuracy increases by up to ∼1.2x on math and coding tasks.These results are reported for fine-tuning large MDMs.
- Results: Across experiments, ProSeCo improves quality-efficiency trade-offs and supports inference-time compute scaling for further quality gains.The introduction reports conditional and unconditional evaluations supporting this broader conclusion.
2 Background
Discrete diffusion trains a denoiser to reverse a corruption process over token sequences. In MDMs, absorbing masking supports parallel denoising, but the denoiser does not learn to modify unmasked inputs, allowing errors to propagate.
- Discrete Diffusion: Discrete diffusion trains a denoising process to undo a predefined corruption process that moves data toward noise.The process begins with data samples and produces progressively noisier latent variables.
- Discrete Diffusion: Discrete diffusion represents data as token sequences from a finite vocabulary and applies corruption over that sequence space.The paper denotes sequence tokens and their vocabulary explicitly.
- Masked Diffusion Models: MDMs use marginals that interpolate between data and a special mask token through a monotonically decreasing noise schedule.Their absorbing corruption process keeps tokens masked once they enter the masked state.
- Masked Diffusion Models: MDM training uses a variational objective that matches a learned reverse transition with the true posterior of the corruption process.The objective is optimized with stochastic gradient descent and Monte Carlo estimates.
- Limitation: Because the denoising network does not modify unmasked inputs, generation errors propagate and accumulate, moving samples away from the true data distribution.This is the central limitation motivating self-correction.
3 Self-Correcting Masked Diffusion Models
ProSeCo equips one MDM with separate operational modes for unmasking and correcting, training correction on the model’s own sampled outputs. Its objective adds an auxiliary correction loss, and its sampling procedure interleaves correction with denoising.
- Model Design: A single model un masks partially masked inputs and updates already generated positions when operating in corrector mode.The two modes share the same model while serving different input states.
- Self-Correcting Objective: ProSeCo treats denoiser outputs as potentially corrupted sequences and trains the model to recover clean data from them.This reframes model-generated errors as noise for corrector training.
- Sampling Design: ProSeCo interleaves error-correction predictions and unmasking steps to refine any sequence position, including already decoded tokens.The resulting method is named Progressive Self-Correction.
- Self-Correcting Objective: The method adds a cross-entropy correction loss to the standard MDM objective so the model can identify and correct denoising mistakes.The standard MDM term remains responsible for producing useful candidates for correction.
- Theoretical Motivation: The corrector can be derived through a learned predictor-corrector sampler intended to address mismatches between model and data marginals.The paper states that optimizing the corrector yields the auxiliary correction loss.
- Design Decisions: Tying corrector and denoiser weights produces a unified self-correcting model without the memory overhead of a separate corrector.Training requires a second forward pass and uses stop-gradient on denoiser outputs for stability.
4 Sampling with ProSeCo
ProSeCo sampling alternates standard unmasking with configurable correction loops that can replace decoded tokens and provide improved predictions for later unmasking. The procedure exposes frequency and per-loop step budgets to control compute.
- Sampling Procedure: ProSeCo interleaves unmasking and correction iterations so corrector steps can update decoded positions and improve subsequent unmasking predictions.Correction therefore serves both token remediation and proposal improvement.
- Sampling Procedure: The sampling algorithm controls correction cost through the correction frequency ω and the number of steps per loop S.These parameters determine how often correction runs and how many refinement steps each loop performs.
- Compute Control: ProSeCo sampling exposes a configurable trade-off between correction computation and generation quality through its loop budgets.The algorithm specifies the model, sequence length, unmasking schedule, correction budget, and correction frequency as inputs.
- Correction Loop: Corrector iterations convert denoiser outputs into a corrector input, repeatedly sample updated sequences, and replace unmasked positions with corrected outputs.The final corrector logits can also be used by the unmasking posterior.
- Unmasking: During unmasking, ProSeCo can use ancestral sampling or confidence-based greedy selection to decide which tokens become unmasked.This part of the loop otherwise matches standard MDM inference.
5 Experiments
ProSeCo is evaluated across conditional and unconditional generation tasks, where it improves quality-efficiency trade-offs and uses corrective refinement to recover from parallel-decoding errors. Experiments show strong benchmark performance, faster generation, improved guided-generation frontiers, and robustness to sampling choices.
- Math & Code Benchmarks: ProSeCo outperforms all diffusion baselines on every benchmark and beats a comparably sized instruction-tuned AR model on 3 of 4 tasks.Its ProSeCo-trained objective also outperforms standard-MDM training before applying ProSeCo sampling.
- Math & Code Benchmarks: 2–4x speed-ups relative to LLaDA decoding are achieved without sacrificing accuracy by combining increased parallelism with corrective loops.The Fast configuration decodes 4–8 tokens per unmasking step and applies correction every second decoding iteration.
- Guided Molecule Design: ProSeCo extends the parallel-decoding quality frontier by recovering from mistakes that significantly degrade standard-MDM sample quality as parallelism increases.Corrector loops provide the additional refinement needed to offset errors introduced by parallel generation.
- Sampling Ablations: ProSeCo remains robust across corrector-loop frequency and step-count choices, beating baseline accuracy at every tested token-parallelism level.An ablation varies the frequency and number of correction steps per loop.
- Guided Molecule Design: ProSeCo pushes the novelty-property Pareto frontier toward higher property values without sacrificing sample diversity and quality for both ring count and drug-likeness.The guided-generation setup varies unmasking budgets and guidance strength and evaluates novel molecular samples using mean property value.
6 Related Works
Prior work adapts diffusion, self-conditioning, and correction mechanisms to discrete generation, but ProSeCo is positioned as a trained approach that corrects already decoded tokens using predictions of clean data.
- Discrete Diffusion: Discrete diffusion methods adapt diffusion to finite-vocabulary data, with later work using categorical uniform noise to address locked-in decoded tokens.ProSeCo is aligned with continuous-time variational-inference extensions of discrete diffusion.
- Self-conditioning & Step Unrolling: Self-conditioning feeds clean-data predictions from earlier steps into future outputs, whereas unrolled-prediction methods train on less-noisy latent trajectories.ProSeCo instead trains on predictions of clean data rather than partially masked trajectories.
- Corrector Methods: Corrector methods include training-free remasking, auxiliary heads that identify positions to remask, and methods that predict corrections to decoded tokens.ProSeCo belongs to the trained-correction family while jointly training decoding and correction abilities.
7 Discussion & Conclusion
The framework jointly trains diffusion models to unmask and self-correct, using corrective sampling steps to refine generation. Experiments show consistent improvements in speed-quality trade-offs and inference-time compute scaling, with added training cost.
- Discussion & Conclusion: ProSeCo jointly trains a diffusion model to unmask tokens and correct mistakes through minimal modifications to standard training and sampling.The framework is designed to support both decoding and self-correction within the same model.
- Discussion & Conclusion: A key limitation is the added computational cost of a second forward pass during training, especially compared with inference-time-only correction schemes.The authors identify this train-time cost as the main drawback and leave more sophisticated joint corrector-unmasking schemes for future work.
- Discussion & Conclusion: Interleaving corrector MCMC steps with predictor unmasking steps enables recovery from deviations before generation continues.The corrector chain samples denoiser predictions, predicts clean data, and re-noises the corrected output according to the forward process.
- Discussion & Conclusion: The corrector objective is derived by requiring the induced latent distribution to align with the true marginal distributions.The derivation uses a sufficient detailed-balance condition and minimizes a divergence between the model-induced corrector distribution and the desired form.
- Discussion & Conclusion: The corrector remains compatible with parallel generation because its parameterization factorizes independently across sequence positions.This design preserves parallel sampling while allowing decoded tokens to be adapted during correction iterations.
C.1 Math & Code Benchmarks
The evaluation fine-tunes and tests models across math and coding benchmarks using standardized prompts, decoding settings, and comparison baselines.
- Data and training: The SFT dataset contains approximately 32 million samples spanning mathematical reasoning, coding, science, instruction following, chat, and safety.Math and coding dominate the distribution, at 66.84% and 30.62%, respectively.
- Evaluation: Models are evaluated on HumanEval, MBPP, GSM8K, and MATH using 0-shot evaluation.Evaluation uses batch size 1 and a maximum generation length of 1024 with early stopping on [EOS].
- Decoding: For LLaDA models, evaluation uses semi-autoregressive decoding with a default block size of 32.The implementation applies block autoregressive decoding across the sequence and does not use KV caching.
- Baselines: The benchmark comparisons include AR, UDLM, MDLM, ReMDM, ReMeDi, and GStar, with some prior results taken from earlier work or reimplemented.ReMeDi Instruct results are reported from Huang et al. because the authors could not reproduce or improve them using the released model.
D.2 Ablation: Robustness to Corrector Parameters
ProSeCo remains robust across corrector-loop frequencies and step counts, while faster parallel decoding benefits from more frequent correction.
- Robustness: ProSeCo beats the best Baseline SFT results for every tested decoding-parallelism level and every tested combination of corrector parameters.The sweep varies ω and S over {1, 2, 4, 8} while tokens per step ranges from 1 to 8.
- Fast sampling: In fast-sampling regimes, matching or exceeding baseline accuracy requires more frequent corrector loops with ω ∈{1, 2}.These configurations can substantially improve efficiency relative to the baseline, whose best accuracy occurs at tokens/step = 1.
- Trade-offs: For a fixed corrector budget, more frequent shorter correction loops are typically more effective, while scaling correction frequency and steps improves quality at additional NFE cost.The fast and balanced configurations are selected from parallel settings that outperform the best baseline accuracy.
D.4 Ablation: Self-correction Sampling with Standard MDM
Self-correction sampling depends on ProSeCo training: standard MDM training does not produce useful predictions for already unmasked tokens.
- Ablation: Applying self-correction sampling to a standard-MDM SFT model performs worse than applying it to a model trained with the ProSeCo loss.The comparison is reported using Pass@1 accuracy.
- Mechanism: Standard MDM training does not teach the model to produce meaningful predictions at already unmasked token positions.ProSeCo training unlocks this ability, allowing the self-correction sampler to benefit from modified decoded tokens.
- Ablation design: The ablation compares self-correction sampling under standard MDM loss against the ProSeCo loss and its sampling algorithm.The comparison isolates the contribution of training for self-correction rather than adding correction only at inference.
D.5 Ablation: Selecting λt
The time-dependent corrector-loss weighting improves performance consistently and remains robust to different scaling factors.
- Experimental setup: The ablation evaluates corrector-loss weights during discrete classifier guidance for the ring-count property.The results are reported in Figure 8, with sampling repeated across fixed and time-varying λ settings.
- Weighting choice: Including the α̇_t(1−α_t) weighting in the corrector loss consistently improves results.The ablation compares fixed λ values with time-varying weighted λ_t values for ring-count guidance.
- Scaling robustness: Model performance is robust to different scaling factors applied to the time-dependent weighting.The tested scaling factors are 0.1, 1, and 10.
E.2 ProSeCo Unconditional Generation Samples
The appendix documents an unconditional ProSeCo sample configuration on OWT and lists accompanying dataset and software license information. It also includes example generated outputs for HumanEval and GSM8K.
- OWT unconditional generation: The OWT sample uses a total budget of T = 256, with 64 unmasking steps, a corrector loop every ω = 1 step, and S = 3 corrector steps per loop.
- OWT unconditional generation: Table 7 reports mean ± standard error for OWT unconditional generation metrics for ProSeCo.
- Licenses: The appendix lists corresponding licenses for the datasets and software packages used in the work.
- Conditional generation examples: The appendix provides a HumanEval generated-response example using the max accuracy configuration referenced in Table 3.
- Conditional generation examples: The appendix provides a GSM8K generated-response example using the max accuracy configuration and shows an expected answer of 6.