The performance of pre-trained masked diffusion models is often constrained by their sampling procedure, which makes decisions irreversible and struggles in low-step generation regimes. We introduce a novel sampling algorithm that works with pre-trained models and, after a lightweight fine-tuning of a single layer, significantly improves sample quality and efficiency. Our method reformulates the generation process using a star-shaped paradigm, which inherently allows for error correction. To make this process effective, we augment it with a learnable re-masking scheduler that intelligently identifies and revises likely errors. This approach yields a substantial quality boost, particularly when using a small number of sampling steps. We extensively ablate key components of our approach and show its usability in different scenarios. In comprehensive experiments on text, and code generation, our sampling algorithm outperforms or matches existing methods.
Figures & tables
Figure 1: Pipeline comparison of MDLM, Star, and G-Star sampling. Each panel illustrates how a sampler moves from a noisy state xT toward a clean sample x0 . (a) MDLM follows an irreversible trajectory: once a token is unmasked, it is copied forward and cannot be revised. (b) Star routes each transition through a fully denoised draft x^0 and then randomly remasks tokens, enabling revision but wasting refinement steps on already correct tokens. (c) G-Star uses the same star-shaped transition, but replaces random remasking with an error predictor gϕ that selectively remasks likely incorrect tokens. This targeted remasking focuses refinement on probable errors while preserving the stability of already plausible context.
Algorithm 1 Training the Error Predictor gϕ
Algorithm 2 Guided Sampling Step
Figure 2: Analysis of the star-shaped (Star) sampler’s dynamics. (Left) Perplexity and (Right) step-to-step similarity over the generation trajectory for three configurations: MDLM, Star, and our hybrid approach (Star+), which switches from MDLM to Star at step 90 (dotted line).
Figure 3: The impact of the star-shaped sampler’s activation time ( ton ) on generation quality. We plot the final MAUVE score for a hybrid sampler that switches from MDLM to Star at time ton .
Figure 4: Performance comparison in few-step generation regimes. Guided sampler (G-Star+) consistently outperforms the unguided Star+.
Figure 5: Pareto fronts for different methods obtained by varying the denoiser temperature. The left plot compares MDLM, Star/ReMDM-loop, P2-planner, RDM, DDPD and G-Star-loop, while the right plot compares different G-Star configurations (12B–F, 12B–H, 1B–F) against Star/ReMDM-loop and MDLM for 128 denoising steps.
Dream-Ins. (Paper)
Dream-Ins. (Reproduced)
+ G-Star (Ours)
MMLU
67.0
69.9
71.2
MMLU-PRO
43.3
46.9
47.9
GSM8K
81.0
81.5
81.6
GPQA
33.0
31.0
32.8
HumanEval
55.5
53.7
54.9
MBPP
58.8
58.0
59.4
Table 1: Downstream benchmark results for Dream-Instruct 7B. The best result is marked in bold .
Algorithm
Qwen2.5B-Coder ppl ↓
32 steps
64 steps
128 steps
MDLM
29.8
25.5
26.7
ReMDM-loop η=0.02
30.1
25.0
20.4
ReMDM-cap η=0.04
27.3
22.5
19.1
G-Star-loop
22.5
17.8
17.8
G-Star+ ton=0.3
20.4
18.9
16.4
Table 2: Conditional code generation results on the Conala benchmark for different samplers and step counts. All methods use the same conditional MDLM generator; Qwen2.5B-Coder is used only as a frozen external evaluator to compute conditional perplexity. Best and second-best results are highlighted.
Appendix figures & tables12 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 6: Performance of ReMDM as a function of the hyperparameter η . Results are shown for three different remasking schedules. The dashed line indicates the performance of the baseline MDLM sampler. The plots reveal a high sensitivity to the choice of η , with suboptimal values often performing worse than the baseline.
Figure 7: Performance as a function of the refinement loop size. Increasing the refinement budget generally improves quality (lower PPL, higher MAUVE) but reduces diversity. Our guided G-Star-loop demonstrates a much steeper rate of improvement, achieving higher quality with fewer steps. The MAUVE score eventually peaks and declines as the loss of diversity outweighs quality gains.
Figure 8: Performance as a function of the error predictor temperature ( Tremask ) . The plots reveal a clear trade-off: lower temperatures improve quality (PPL) at the cost of diversity, while higher temperatures increase diversity at the cost of quality. The MAUVE score, which balances both, peaks at an optimal temperature of T≈4−32 .
Evaluation dataset
Domain
PPL val
Accuracy
AUC-ROC
OWT
General
⩽22.89
0.88
0.94
TinyStories
Stories
⩽12.72
0.88
0.98
OpenWebMath
Math
⩽33.72
0.85
0.92
CNN/DailyMail
News
⩽25.69
0.88
0.94
The Stack
Python code
⩽31.96
0.84
0.90
Appendix
Table 3: Results for an error predictor trained on OWT and evaluated on out-of-domain datasets. For each evaluation dataset, we report the validation perplexity (PPL val) of the OWT-pretrained diffusion model and the binary classification metrics of the error predictor.
Method
Steps T
Time [s]
NFEs
AR (KV cache)
512
2.28±0.03
512
AR (no KV cache)
512
2.51±0.01
512
MDLM / ReMDM / P2 / Star
128
2.26±0.09
128
DDPD
128
4.63±0.11
256
G-Star-loop 1B,F
128
2.60±0.10
133
G-Star-loop 12B,H
128
3.43±0.10
192
Appendix
Table 4: Wall-clock generation time and number of Transformer forward passes (NFEs) per 512-token sample on OpenWebText (batch size 1 ). Times are reported as mean ± standard deviation in seconds.
Benchmark
Sequence Length ( L )
αon
MMLU
128
0.88
MMLU-PRO
128
0.88
GSM8K
256
0.95
GPQA
128
0.88
HumanEval
768
0.98
MBPP
1024
0.98
Appendix
Table 5: Benchmark-specific sequence lengths and the noise level αon for the refinement loop.
Figure 9: The plot illustrates the trajectory of the number of generated (unmasked) tokens across 512 generation steps. The G-Star+ schedule (solid line) follows a linear progression, dedicating the majority of steps to draft generation and applying remasking only towards the end of the trajectory. In contrast, the G-Star-loop schedule (dashed line) adopts an aggressive drafting strategy, rapidly producing a full candidate sequence (steep red slope) to maximize the computational budget allocated to the remasking phase (green plateau), enabling deep iterative refinement.
Figure 10: Remasking pattern of the guided G-Star-loop sampler. The plot visualizes the masked token positions (orange dots) at each generation step. In contrast to the unguided sampler, our guided approach exhibits a highly structured remasking pattern. The error predictor directs the sampler to focus on specific, clustered regions of the text. This often results in the selection of contiguous segments for revision, as highlighted in the figure. This ability to perform coordinated, phrase-level corrections is a key advantage of our targeted approach.
Figure 11: Remasking pattern of the Unguided Star-loop sampler. The plot visualizes the masked token positions (orange dots) at each generation step. During the remasking phase (steps 114-240), the pattern of selected tokens is scattered and visually resembles random noise. This illustrates the indiscriminate nature of the unguided approach, where every token position has a roughly equal probability of being revised at each step, without regard to the semantic or syntactic structure of the text.
Figure 12: Remasking pattern of the DDPD sampler. The plot shows the token positions (orange dots) selected by the planner at each generation step. The selected-token pattern is largely scattered and visually resembles random noise, which is a consequence of initializing from a uniformly sampled sequence that introduces many errors. Despite this, some positions are still remasked more than once (indicated by gray segments).
Figure 13: Remasking pattern of the P2 self-planning sampler. The plot visualizes the masked token positions (orange dots) at each generation step. During the remasking phase (steps 228-480), remasking is applied only in the initial steps; after this early period, no further remasking occurs and the token configuration remains fixed.
Figure 14: Qualitative refinement trajectories. We show steps 90–95 from a 128-step generation of a 512-token sequence. Both unguided Star+ (bottom) and guided G-Star+ (top) start this phase from the same text draft produced by standard MDLM sampling. The panels display the beginning of the sequence; tokens remasked at each step are highlighted in red. Starting from the same draft, the two samplers immediately diverge. Star+ applies scattered token-level remasking, including to tokens that already form plausible context. G-Star+ instead produces a more structured refinement trajectory, concentrating remasking on likely erroneous regions while preserving more of the coherent draft.
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 ProSeCo yields better quality-efficiency trade-offs (up to ~4x faster sampling while still outperforming the best accuracy baseline method) and enables further gains through inference-time compute scaling (up to ~1.2x improvement on benchmarks when allowing for slower, high-NFE decoding).
Masked Diffusion Models (MDMs) enable flexible decoding orders, yet existing samplers remain largely greedy, selecting locally certain tokens without accounting for their downstream effects. We show that this myopia can increase cumulative uncertainty and lead to suboptimal generation. To address this, we propose the Info-Gain Sampler, a training-free decoding method that uses the bidirectional structure of MDMs to balance immediate uncertainty with the information gained over remaining masked positions. Across reasoning, coding, creative writing, and image generation tasks, Info-Gain Sampler consistently outperforms existing MDM samplers, improving average reasoning accuracy by 2.9--11.6 percentage points and achieving a 62.8% average win rate in creative writing. The code is available at https://github.com/yks23/Information-Gain-Sampler.
Kaisen Yang, Jayden Teoh, Kaicheng Yang +2
Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University +2
Masked diffusion language models (MDLMs) re-predict every position at each denoising step, but standard samplers commit tokens once revealed, leaving this revision capability unused. Existing approaches either add heuristic or learned mechanisms to revise committed tokens, or remask them back to [MASK] before re-predicting; a principled sampler that directly revises visible tokens without auxiliary modules remains underexplored. We introduce D3IM, a parameter-free sampler derived as a corrector-style reverse update that permits direct visible-to-visible revision without additional modules or auxiliary passes. D3IM also reveals a model-side obstacle we term preservation bias: the model tends to reproduce its own wrong committed tokens rather than correct them. We address this with SCOPE (Self-Conditioned On Prediction Errors), a lightweight post-training procedure that simulates D3IM's sampling process. On LLaDA-8B at 64 denoising steps, SCOPE+D3IM improves over the original LLaDA-8B with standard unmasking by +13.0 on GSM8K (68.3%), +4.8 on MATH-500 (23.6%), +15.3 on HumanEval (29.3%), and +10.4 on MBPP (30.8%), with gains that increase as more denoising steps are used on math and HumanEval.