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