Guided Star-Shaped Masked Diffusion

Viacheslav Meshchaninov; Egor Shibaev; Artem Makoian; Ivan Klimov; Nikita Balagansky; Daniil Gavrilov; Aibek Alanov; Dmitry Vetrov

cs.LGOct 9, 2025

Guided Star-Shaped Masked Diffusion

Authors: Viacheslav Meshchaninov, Egor Shibaev, Artem Makoian, Ivan Klimov, Nikita Balagansky, Daniil Gavrilov, Aibek Alanov, Dmitry Vetrov

Organizations: Constructor University · Higher School of Economics · T-Tech · FusionBrain Lab

Abstract

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

Panel of Figure 1 — Figure 1: Pipeline comparison of MDLM, Star, and G-Star sampling. Each panel illustrates how a sampler moves from a noisy state $\mathbf{x}_{T}$ toward a clean sample $\mathbf{x}_{0}$ . (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 $\hat{\mathbf{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_{\phi}$ that selectively remasks likely incorrect tokens. This targeted remasking focuses refinement on probable errors while preserving the stability of already plausible context.

Panel of Figure 2 — Algorithm 1 Training the Error Predictor $g_{\phi}$

Panel of Figure 3 — 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 (

t_{on}

) on generation quality. We plot the final MAUVE score for a hybrid sampler that switches from MDLM to Star at time

t_{on}

Figure 4: Performance comparison in few-step generation regimes. Guided sampler (G-Star+) consistently outperforms the unguided Star+.

Panel of Figure 5 — 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 $\downarrow$
Algorithm	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+ ${}_{t_{on}=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

\eta

. 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

\eta

, 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 (

T_{\text{remask}}

) . 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\approx 4-32

Evaluation dataset	Domain	PPL val	Accuracy	AUC-ROC
OWT	General	$\leqslant 22.89$	$0.88$	$0.94$
TinyStories	Stories	$\leqslant 12.72$	$0.88$	$0.98$
OpenWebMath	Math	$\leqslant 33.72$	$0.85$	$0.92$
CNN/DailyMail	News	$\leqslant 25.69$	$0.88$	$0.94$
The Stack	Python code	$\leqslant 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\pm 0.03$	$512$
AR (no KV cache)	512	$2.51\pm 0.01$	$512$
MDLM / ReMDM / P2 / Star	128	$2.26\pm 0.09$	$128$
DDPD	128	$4.63\pm 0.11$	$256$
G-Star-loop 1B,F	128	$2.60\pm 0.10$	$133$
G-Star-loop 12B,H	128	$3.43\pm 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

\pm

standard deviation in seconds.

Benchmark	Sequence Length ( $L$ )	$\alpha_{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

\alpha_{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.

Diffusion Sampling

Diffusion sampling studies how to turn a trained diffusion model into high-quality samples efficiently and reliably.

See all 278 papers

Learn from Your Mistakes: Self-Correcting Masked Diffusion Models
Feb 12, 2026Yair Schiff, Omer Belhasin, Roy Uziel +6Masked Diffusion Models Diffusion Language Model Decoding
Cornell · NVIDIA
Improving Sampling for Masked Diffusion Models via Information Gain
Feb 20, 2026Kaisen Yang, Jayden Teoh, Kaicheng Yang +2Diffusion Sampling Diffusion Language Model Decoding
Department of Computer Science and Technology, Tsinghua University · Massachusetts Institute of Technology · Shanghai Jiao Tong University +2
Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models
May 31, 2026Longxuan Yu, Shaorong Zhang, Yu Fu +3Masked Diffusion Models Diffusion Model Fine-Tuning
University of California, Riverside · Microsoft

Feb 12, 2026J/K move · Enter open · S save

Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

Yair Schiff, Omer Belhasin, Roy Uziel, Guanghan Wang+5

Cornell · NVIDIA

Masked Diffusion Models Diffusion Language Model Decoding Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

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).

Abstract

Diffusion Sampling

Explore similar work

Learn from Your Mistakes: Self-Correcting Masked Diffusion Models

Improving Sampling for Masked Diffusion Models via Information Gain

Revise, Don't Freeze: Sampler-Matched Training for Self-Correcting Masked Diffusion Language Models