Masked diffusion models (MDMs) generate sequences by progressively unmasking several tokens per denoising step, but their reverse process is typically factorized over positions, limiting sample quality in the few-step regime where diffusion's speed advantage over autoregressive decoding matters most. A recent line of work introduces a continuous Gaussian latent, trained as a variational autoencoder, to capture correlations across positions, but such approaches are prone to posterior collapse, where the latent is silently ignored. We propose Enhanced Mixture-of-Experts (E-MoE), which builds the reverse process as a mixture of factorized distributions over a discrete shared latent given by the expert-routing decisions of a Mixture-of-Experts (MoE) backbone, without increasing active parameters over the factorized baseline. Across synthetic multi-modal benchmarks, binarized MNIST, and LM1B, E-MoE improves few-step generation over factorized baselines.
Figures & tables
Figure 1: Few-step generative perplexity on LM1B (lower is better, H = sample entropy). E-MoE gives markedly lower generative perplexity than both factorized MDLM ( Sahoo et al., 2024 ) and continuous-latent VADD ( Xie et al., 2026 ) for all NFEs 1–16, at matched sample entropy.
Figure 2: Factorization error and mixture over expert routes. Top: a factorized reverse process can produce spurious pairings when multiple tokens are sampled at once. Bottom: the routing code z={zd}d=1D , with zd=(zdℓ}ℓ=1L , specifies the expert assignment for every position across layers.
Figure 3: Training stage: we model prior pθ(zti∣xti) and posterior qθ(zti∣xti,x0) by the same router, unlike VADD. The router ρθ(⋅,ti) takes as input a noisy sample xti to approximate pθ in the noisy pass, and a concatenation of xti and x0 for qθ in the clean pass. Then, the experts return predictions for each masked xtiℓ using appropriate zti . Inference stage: we use only the noisy pass (Section 3.2 ).
Figure 4
Factorized
Mixture of factorized distributions
MDLM ( Sahoo et al., 2024 )
VADD ( Xie et al., 2026 )
E-MoE (ours) (§ 3 )
ℓ∏pθ(xℓ∣xt)
∫p(z)ℓ∏pθ(xℓ∣xt,z)dz
z∑pθ(z∣xt)ℓ∏pθ(xℓ∣xt,z)
Latent space
—
Rd
{1,…,E}L×D
Latent is chosen
—
once per sequence
per token and layer
Learned, data-dependent prior
×
×
✓
Posterior q comes from
—
qϕ from VAE (additional net)
qθ from the same net
Table 1: A summary of the design choices behind E-MoE and its two baselines. VADD needs a continuous latent with a fixed Gaussian prior and a separate recognition network, while E-MoE reuses the routing decisions the backbone already makes, so it trains as simply as the factorized baseline.
Figure 4: Generation results on 2-D toy examples. Ground truth and generations of MDLM ( Sahoo et al., 2024 ) , VADD ( Xie et al., 2026 ) , and E-MoE at NFE =1 on 8-modes and swiss-roll. Full sweeps over NFE ∈{1,2,8,32} are in Appendix C.1 .
8-modes
Swiss-roll
Model
NFE =1
2
4
8
16
32
NFE =1
2
4
8
16
32
MDLM
37.7
68.1
84.2
90.9
94.9
97.0
49.3
73.1
86.5
91.9
95.2
96.9
VADD
91.8
95.3
97.3
97.7
98.5
98.5
93.8
95.9
97.1
97.6
98.2
98.4
E-MoE (ours)
94.4
96.6
97.7
98.3
98.6
98.7
90.3
92.6
95.7
96.4
98.6
97.9
Table 2: Validity ( ↑ ) of generated samples across sampling steps (NFE). Values are averaged over 3 seeds. Best per column in bold .
Model
BPD ↓
Params
MDLM
0.077
2.07M
VADD
0.064
2.50M
E-MoE (Ours)
0.062
2.49M
Table 3: Test Bits-per-dimension ( ↓ ) and total parameter count on Binarized-MNIST.
MDLM
SEDD
VADD
BD3-LM
E-MoE (ours)
NFE
Gen-PPL ↓
H↑
Gen-PPL ↓
H↑
Gen-PPL ↓
H↑
Gen-PPL ↓
H↑
Gen-PPL ↓
H↑
1
1433.8
4.37
1578.8
4.37
1270.8
4.34
—
—
643.8
4.35
2
997.5
4.37
1059.3
4.37
763.2
4.33
—
—
383.7
4.35
4
477.8
4.36
456.0
4.34
370.5
4.33
—
—
235.4
4.34
8
260.5
4.35
243.4
4.33
220.7
4.33
1130.6 16
4.34
174.9
4.34
16
179.4
4.35
166.7
4.33
168.4
4.33
1004.1 8
4.32
146.4
4.34
Table 4: Few-step generation on LM1B. Generative perplexity (Gen-PPL) and sample entropy H (data: 4.32 ), averaged over 2000 samples with categorical sampling in fp64 . Superscripts mark BD3-LM models trained with different block sizes. Best per row in bold .
Figure 5: Two views of the same LM1B sweep. Left: Generative Perplexity ( ↓ ). Right: MAUVE ( ↑ )
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Tokens per coordinate K
50
Sequence length L
2
Training points
20,000
Backbone blocks
2
Attention heads
4
Model dim.
128
MLP dim.
512
Appendix
Table 5: Toy hyperparameters, shared by all models.
Resolution
32×32
Optimizer
AdamW
Vocabulary size K
2
Learning rate
2×10−4
Sequence length L
1024
Weight decay
0.01
UNet base channels
72
Warmup steps
5000
Batch size
64
LR schedule
cosine
Training steps
200,000
Gradient clipping
1.0
EMA decay
0.9999
Appendix
Table 6: Binarized-MNIST hyperparameters, shared by all three models.
Table 7: LM1B configuration. The three models share data, backbone and optimization, and differ only in the latent. “Active” counts the parameters applied to a token at inference.
Figure 6: Full NFE sweep generations on 8-modes , NFE ∈{1,2,8,32} . MDLM smears mass into a blurred grid even at NFE =32 . VADD and E-MoE both recover the eight clusters from NFE =1 .
Figure 7: Full NFE sweep generations on Swiss-roll , NFE ∈{1,2,8,32} . MDLM spreads its mass over the whole disk at low NFE and concentrates on the spiral only by NFE =32 . VADD and E-MoE both recover the spiral from NFE =1 .
Figure 8: Binarized-MNIST samples across NFE , the first 16 draws from a fixed seed per cell. MDLM gives unstructured speckle at NFE ≤2 and resolves digits only by NFE =10 – 20 , while E-MoE produces recognizable digits from a single step.
NFE
MDLM
VADD
E-MoE (Ours)
1
0.93±0.04
1.06±0.03
4.67±0.31
2
2.95±0.29
3.76±0.13
24.58±0.62
4
17.09±0.94
25.16±2.56
58.25±1.69
8
57.04±1.72
63.93±2.19
78.68±0.55
16
83.75±1.91
83.73±0.94
88.37±0.41
32
90.53±0.86
86.35±0.96
90.59±0.68
Appendix
Table 8: MAUVE ( ↑ , %) on LM1B. Subscripts are the standard deviation over 3k -means seeds. Best diffusion model per row in bold
NFE =1
NFE =2
PPL ↓
H↑
PPL ↓
H↑
E-MoE, mixture on ( τs=1 )
626.1
4.35
389.0
4.36
E-MoE, greedy routes ( τs→0 )
1079.5
4.34
834.2
4.32
MDLM
1433.8
4.37
997.5
4.37
Appendix
Table 9: Switching the mixture off with the routing temperature. The same E-MoE checkpoint with sampled ( τs=1 ) and greedy ( τs→0 ) routes, and MDLM for reference.
Figure 9: Training terms on LM1B. Left: reconstruction term, right: KL term (routing KL for E-MoE, Gaussian KL for VADD), in nats per token and smoothed, over the first 700 k steps.
Masked diffusion models (MDMs) are a promising family of language generators, but achieving high-quality few-step generation remains challenging. In MDMs, all forward trajectories collapse to a single fully masked state, leaving no terminal entropy for consistency-style few-step generation. While recent few-step alternatives based on uniform-state diffusion avoid this degeneracy, it becomes harder to distinguish clean tokens from noise than MDMs, which usually harms modeling quality and training efficiency. In this work, we propose a multi-mask diffusion model (MultiMDM) that preserves the masking structure towards few-step generation. In the forward process, each clean token is first pushed towards a designated mask and then gradually mixes over the mask set. As a result, the backward process has a drafting capability by predicting a designated mask before refining to a clean token. We derive a closed-form ELBO training objective for MultiMDM that supports continual training from pretrained MDMs. In addition, we formulate a purely discrete-state consistency distillation scheme, with a shared-Gumbel coupling to reduce pathwise entropy. Experiments on pretraining and distillation show that MultiMDM provides an effective foundation for principled few-step generation.
Sijin Chen, Yinuo Ren, Heyang Zhao +3
ByteDance Seed · Princeton University · Stanford University +2
Masked Diffusion Models (MDMs) have emerged as a promising alternative to autoregressive models in language modeling, offering the advantages of parallel decoding and bidirectional context processing within a simple yet effective framework. Specifically, their explicit distinction between masked tokens and data underlies their simple framework and effective conditional generation. However, MDMs typically require many sampling iterations due to factorization errors stemming from simultaneous token updates. We observe that a theoretical lower bound of the factorization error exists, which standard MDMs cannot reduce due to their use of a deterministic single-state mask. In this paper, we propose the Infinite Mask Diffusion Model (IMDM), which introduces a stochastic infinite-state mask to mitigate the theoretical bound while directly inheriting the benefits of MDMs, including the compatibility with pre-trained weights. We empirically demonstrate that MDM fails to perform few-step generation even in a simple synthetic task due to the factorization error bound, whereas IMDM can find an efficient solution for the same task. Finally, when equipped with appropriate distillation methods, IMDM surpasses existing few-step distillation methods at small step counts on LM1B and OpenWebText. Code is available at https://Ugness.github.io/official_imdm.
Jaehoon Yoo, Wonjung Kim, Chanhyuk Lee +1
Korea Advanced Institute of Science and Technology (KAIST)
Diffusion language models intrinsically fail to capture correlations between decoded tokens, which leads to a harsh trade-off between sampling quality and throughput. To solve this issue, we propose DiLaDiff, a variant of masked diffusion language models with three components: (1) a continuous latent space with semantic capabilities, learned by an auto-encoder fine-tuned from an existing masked diffusion language model; (2) a latent diffusion model learning the prior over the encoder distribution; (3) a consistency model distilling the learned prior into a few-step latent generative model. We show that, even without distillation, our latent-guided diffusion model outperforms the masked diffusion baseline while significantly accelerating inference. Consistency distillation further lowers the computational overhead of continuous diffusion, such that the latent is generated in negligible time compared to discrete decoding.
Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3