Masked diffusion models have shown promising performance in generating high-quality samples in a wide range of domains, but accelerating their sampling process remains relatively underexplored. To investigate efficient samplers for masked diffusion, this paper theoretically analyzes the MaskGIT sampler for image modeling, revealing its implicit temperature sampling mechanism. Through this analysis, we show that MaskGIT is asymptotically equivalent to a choose-then-sample (CTS) formulation, instantiated as the "moment sampler," which explicitly separates index selection from token sampling. This CTS reformulation is essential: it yields unbiased token sampling and exposes an algorithmic design space for index selection, both of which are inaccessible in MaskGIT's original formulation. Regarding token sampling, we reveal that MaskGIT implicitly adopts a low-temperature sampler, which explains why MaskGIT often degrades with more sampling steps. The CTS reformulation of MaskGIT allows us to fix the temperature sampling to ensure unbiasedness. We also improve the index selection in CTS through two key innovations: a partial caching technique for transformers that approximates longer sampling trajectories without proportional computational cost, and a hybrid approach formalizing the exploration-exploitation trade-off in adaptive unmasking. Experiments in image and text domains demonstrate our theory as well as the efficiency of our proposed methods, advancing both theoretical understanding and practical implementation of masked diffusion samplers.
Figures & tables
Figure 1: Overview of our contributions. We show that the MaskGIT sampler is asymptotically equivalent to a choose-then-sample (CTS) formulation (Theorem 2 ), instantiated as the moment sampler, which decouples index selection and token sampling. This CTS reformulation yields an unbiased sampler when temperature is removed (Proposition 3 ) and provides a principled explanation for MaskGIT’s degradation at large step counts (Section 3 ). Moreover, it opens an algorithmic design space for index selection that enables partial caching and hybrid bi-criteria methods (Section 4 ).
Figure 2: Illustration of partial caching approximation applied to an L -layer transformer, where σ=softmax(⋅/dk) , with dk being the dimension of key and query vectors.
Figure 3: Fréchet Inception Distance (FID, ↓ ) and Inception Score ( ↑ ) against the number of steps for various samplers with MAGE. Both metrics were computed by 50,000 generated images. We can see that Moment closely approximates MaskGIT with the same temperature in both metrics; their quantitative approximation precision is given in Table 1 .
\Block 2-1Ref. MaskGIT
temperature
\Block 2-1 Random
MaskGIT
Moment
Temp
α=3.0
6.0
12.0
3.0
6.0
12.0
3.0
6.0
12.0
α=3.0
67.1
-
37.8
50.4
4.1
38.5
47.3
5.7
36.4
47.5
α=6.0
121.9
77.0
-
25.1
83.7
4.4
20.0
74.3
3.1
21.6
α=12.0
113.9
106.8
24.5
-
113.4
27.6
4.3
98.8
24.0
3.0
Table 1: Mean relative difference (%) of ImageNet FIDs against reference MaskGIT samplers across 8, 16, 32, 64 steps in unconditional experiment with MAGE. Given reference sequence (xi)i , mean relative difference of sequence (yi)i was computed as average of ∣1−yi/xi∣ . Difference under 10% is bolded. Moment and Temp closely approximate MaskGIT with the same temperature.
Appendix figures & tables16 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: Additional experimental results. ( Left ) Generative Perplexity of various samplers with temperature sampling. ( Right ) Generative Perplexity of our proposed samplers against sampling time per batch on H100 GPU.
\Block 2-1Sampler
\Block 2-1Precision
8 steps
32 steps
128 steps
Gen. PPL
Entropy
Gen. PPL
Entropy
Gen. PPL
Entropy
Vanilla
32-bit
125.62
5.40
41.91
5.31
27.63
5.17
64-bit
137.95
5.42
46.57
5.35
33.80
5.28
Fixed
32-bit
131.01
5.41
45.10
5.32
33.15
5.26
(= Random )
64-bit
130.76
5.41
46.66
5.35
34.29
5.29
Appendix
Table 2: Comparison of different numerical precision in Vanilla and Fixed samplers. Vanilla is a standard sampler for discrete diffusion, where it independently determines whether or not unmasking a certain position. Fixed is a sampler that pre-determines the number of unmasked positions at each step, and it determines which positions to unmask uniformly at random. Both follows the uniform schedule (Appendix D.1 ) in expectation.
\Block 2-1Steps
\Block 2-1 Random
MaskGIT
Moment
Temp
U-Moment
α=3.0
6.0
12.0
3.0
6.0
12.0
3.0
6.0
12.0
3.0
6.0
12.0
8
49.86
18.95
19.21
27.18
19.47
17.42
24.26
20.83
18.74
25.26
48.42
48.75
48.92
16
31.85
19.36
11.29
13.32
20.71
11.52
12.73
19.86
12.09
13.25
29.41
30.21
30.99
32
22.55
24.61
11.45
9.84
25.76
11.77
9.73
23.83
11.73
10.00
20.71
21.53
22.24
64
18.40
27.23
12.36
9.03
27.80
12.77
9.13
25.32
12.33
9.27
16.79
17.52
18.29
Appendix
Table 3: ImageNet FIDs for various methods in unconditional experiment with MAGE.
\Block 2-1Steps
\Block 2-1 Random
MaskGIT
Moment
Temp
U-Moment
α=3.0
6.0
12.0
3.0
6.0
12.0
3.0
6.0
12.0
3.0
6.0
12.0
4
30.11
14.34
19.17
23.83
12.85
17.92
23.01
17.65
22.82
25.97
31.96
31.89
31.65
8
12.92
6.96
6.60
7.67
7.26
6.41
7.21
7.02
8.20
9.89
13.53
13.54
13.40
16
7.85
8.65
6.79
6.08
9.43
7.39
6.12
7.24
6.18
6.41
7.90
7.84
7.87
32
6.37
10.28
8.10
6.72
11.15
8.94
7.09
8.43
6.45
5.91
6.32
6.26
6.26
Appendix
Table 4: FIDs for various methods and steps in class-conditional experiment with MaskGIT-PyTorch.
\Block 2-1Ref. MaskGIT
temperature
\Block 2-1 Random
MaskGIT
Moment
Temp
α=3.0
6.0
12.0
3.0
6.0
12.0
3.0
6.0
12.0
α=3.0
60.7
-
20.4
35.2
8.1
15.1
31.1
14.5
35.7
47.9
α=6.0
47.5
21.3
-
17.0
29.9
7.1
12.9
6.3
18.2
29.5
α=12.0
32.4
36.1
16.4
-
43.1
23.9
3.9
19.7
4.2
13.9
Appendix
Table 5: Mean relative difference (%) of ImgeNet FIDs against reference MaskGIT samplers across 4, 8, 16, 32 steps in class-conditional experiment with MaskGIT-PyTorch. Mean relative difference was computed in the same way as Table 1 . Bolded under 10% . Moment consistently approximates MaskGIT with the same temperature.
Figure 9: FID against number of steps for various samplers with MaskGIT-PyTorch. FID was computed with 50,000 (50 per class) class-conditional generation results.
Figure 12: Efficiency comparison of SDTT-large and SDTT with Random and Hybrid samplers. ‘Large+’ means SDTT-large; otherwise it shows the results of SDTT. Generative Perplexity was measured by Llama3-8B for both models.
Metric
Method
Steps
8
16
32
64
128
256
Gen. PPL
Random
126.36
64.31
45.45
37.67
33.78
31.67
U-Moment
107.04
55.36
38.95
31.96
28.40
26.62
Hybrid
73.02
46.73
37.55
34.09
31.87
29.64
s-Hybrid
73.46
48.76
40.25
35.30
31.97
30.84
Halton
55.09
42.53
37.41
35.76
34.32
33.13
Appendix
Table 6: Generative Perplexity and Entropy of s-Hybrid and other samplers applied to SDTT.
Figure 15: Comparison of various merging schedulers applied to Hybrid sampler of the SDTT model.
Figure 16: Comparison of generated samples by MAGE.
Figure 17: Generated by SDTT with MaskGIT , α=3.0 (Gen. PPL: 13.88, Entropy: 3.320).
Figure 18: Generated by SDTT with Hybrid , α=3.0 (Gen. PPL: 67.98, Entropy: 5.360). [U+2620] is actually a single Unicode character of crossbones.
Figure 19: Generated by SDTT with MaskGIT , α=6.0 (Gen. PPL: 31.89, Entropy: 4.465).
Figure 20: Generated by SDTT with Hybrid , α=6.0 (Gen. PPL: 73.02, Entropy: 5.399).
Figure 21: Generated by LLaDA-8B with MaskGIT , α=1.0 , with 32-step and 256-step sampling.
Figure 22: Generated by LLaDA-8B with Hybrid , α=1.0 , with 32-step and 256-step sampling. In the result of 32-step generation, we used the verbatim package of L a T e X to properly display the spacings. [3 characters] are the Unicode characters [U+6389][U+5802][U+5802].
Masked diffusion language models can reduce inference steps by revealing multiple tokens per denoising iteration, but this parallelism is fragile: positions that are individually confident may be unsafe to commit together when their predictions are coupled. Existing training-free samplers such as Top-k, Fast-dLLM, and EB-Sampler mainly control how many tokens to reveal, while often ranking candidates by token-wise scores that ignore interactions within the selected set. We propose ADAS, a training-free reranking rule that leaves the base sampler's stopping rule unchanged and greedily discounts each token-wise confidence score according to its attention to already selected positions, weighted by their prediction uncertainty. Across LLaDA-8B-Base and Dream-7B-Base on the reasoning benchmarks GSM8K and MATH500 and the code benchmarks HumanEval and MBPP, plugging ADAS into all three samplers improves low-NFE performance at matched denoiser evaluations by 9.11 and 10.46 percentage points on average, respectively, with 3.1% per-forward runtime overhead. Code is available at https://github.com/yusufsahin99/ADAS.
Yusuf Sahin, Ahmed Rockey Saikia, Volkan Cevher +1
University of Bern, Bern, Switzerland · EPFL, Lausanne, Switzerland
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.
Inference-time scaling is a promising paradigm to improve generative models, especially when outputs must satisfy structural constraints or optimize downstream rewards. We consider Masked Diffusion Model (MDM) and introduce MDM-VGB, a discrete diffusion sampler that augments unmasking generation with theoretically principled reward-guided remasking. Inspired by the recent success of the classical Jerrum-Sinclair backtracking Markov chain in reward-tilted generation, MDM-VGB extends the backtracking random walk from a fixed prefix tree to a masked-state graph, allowing tokens to be unmasked and remasked at arbitrary positions. The resulting sampler favors unmasking and remasking moves that lead to higher-value partial configurations, enabling both effective high-reward generation and efficient repair of low-reward samples. We prove that MDM-VGB is robust to process-verifier noise and achieves quadratic complexity, while popular test-time heuristics such as best-of-N can incur exponential complexity due to error accumulation. Our theoretical findings are corroborated by strong empirical performance, particularly on popular constraint-satisfaction and scientific benchmarks such as Sudoku and QM9.