E-MoE: Enhanced Mixture-of-Experts for Non-Factorized Diffusion Language Models

Arseny Ivanov; Alexander Kolesov; Alexander Korotin; Ivan Oseledets; Mikhail Goncharov

cs.CLSep 29, 2026

E-MoE: Enhanced Mixture-of-Experts for Non-Factorized Diffusion Language Models

Authors: Arseny Ivanov, Alexander Kolesov, Alexander Korotin, Ivan Oseledets, Mikhail Goncharov

Organizations: AXXX, Moscow, Russia · Applied AI Institute, Moscow, Russia

Abstract

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

\mathbf{z}=\{\mathbf{z}_{d}\}_{d=1}^{D}

, with

\mathbf{z}_{d}=(\mathbf{z}_{d}^{\ell}\}_{\ell=1}^{L}

, specifies the expert assignment for every position across layers.

Panel of Figure 3 — Figure 3: Training stage: we model prior $p_{\theta}(\textbf{z}_{t_{i}}|\textbf{x}_{t_{i}})$ and posterior $q_{\theta}(\textbf{z}_{t_{i}}|\textbf{x}_{t_{i}},\textbf{x}_{0})$ by the same router, unlike VADD. The router $\rho_{\theta}(\cdot,t_{i})$ takes as input a noisy sample $\textbf{x}_{t_{i}}$ to approximate $p_{\theta}$ in the noisy pass, and a concatenation of $\textbf{x}_{t_{i}}$ and $\textbf{x}_{0}$ for $q_{\theta}$ in the clean pass. Then, the experts return predictions for each masked $\textbf{x}_{t_{i}}^{\ell}$ using appropriate $\textbf{z}_{t_{i}}$ . 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 )
	$\displaystyle\prod_{\ell}p_{\theta}(\textbf{x}^{\ell}\mid\textbf{x}_{t})$	$\displaystyle\int p(\textbf{z})\prod_{\ell}p_{\theta}(\textbf{x}^{\ell}\mid\textbf{x}_{t},\textbf{z})\,d\textbf{z}$	$\displaystyle\sum_{\textbf{z}}p_{\theta}(\textbf{z}\mid\textbf{x}_{t})\prod_{\ell}p_{\theta}(\textbf{x}^{\ell}\mid\textbf{x}_{t},\textbf{z})$
Latent space	—	$\mathbb{R}^{d}$	$\{1,\dots,\text{E}\}^{L\times D}$
Latent is chosen	—	once per sequence	per token and layer
Learned, data-dependent prior	$\times$	$\times$	✓
Posterior $q$ comes from	—	$q_{\phi}$ from VAE (additional net)	$q_{\theta}$ 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 $\in\{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 (

\uparrow

) of generated samples across sampling steps (NFE). Values are averaged over 3 seeds. Best per column in bold .

Model	BPD $\downarrow$	Params
MDLM	0.077	2.07M
VADD	0.064	2.50M
E-MoE (Ours)	0.062	2.49M

Table 3: Test Bits-per-dimension (

\downarrow

) and total parameter count on Binarized-MNIST.

	MDLM		SEDD		VADD		BD3-LM		E-MoE (ours)
NFE	Gen-PPL $\downarrow$	$H\uparrow$	Gen-PPL $\downarrow$	$H\uparrow$	Gen-PPL $\downarrow$	$H\uparrow$	Gen-PPL $\downarrow$	$H\uparrow$	Gen-PPL $\downarrow$	$H\uparrow$
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 (

\downarrow

). Right: MAUVE (

\uparrow

)

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\times 32$	Optimizer	AdamW
Vocabulary size K	$2$	Learning rate	$2\times 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.

Shared
	MDLM	VADD	E-MoE (Ours)
Backbone	DiT-small: $\text{D}=12$ blocks, hidden size $768$ , $12$ heads, dropout $0.1$
Data	$\text{K}=30{,}522$ ( bert-base-uncased ), $L=128$ , sentence packing
Optimizer	AdamW, $\beta=(0.9,0.999)$ , weight decay $0$ , gradient clipping $1.0$
Learning rate	$3\times 10^{-4}$ , $2{,}500$ -step linear warmup, then constant
Batch size	$512$ ( $4$ GPUs $\times$ $128$ ), EMA $0.9999$ , bf16

Appendix

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 $\in\{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 $\in\{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 $\leq 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_{\pm 0.04}$	$1.06_{\pm 0.03}$	$\mathbf{4.67}_{\pm 0.31}$
2	$2.95_{\pm 0.29}$	$3.76_{\pm 0.13}$	$\mathbf{24.58}_{\pm 0.62}$
4	$17.09_{\pm 0.94}$	$25.16_{\pm 2.56}$	$\mathbf{58.25}_{\pm 1.69}$
8	$57.04_{\pm 1.72}$	$63.93_{\pm 2.19}$	$\mathbf{78.68}_{\pm 0.55}$
16	$83.75_{\pm 1.91}$	$83.73_{\pm 0.94}$	$\mathbf{88.37}_{\pm 0.41}$
32	$90.53_{\pm 0.86}$	$86.35_{\pm 0.96}$	$\mathbf{90.59}_{\pm 0.68}$

Appendix

Table 8: MAUVE (

\uparrow

, %) on LM1B. Subscripts are the standard deviation over

3

k

-means seeds. Best diffusion model per row in bold

	NFE $=1$		NFE $=2$
	PPL $\downarrow$	$H\uparrow$	PPL $\downarrow$	$H\uparrow$
E-MoE, mixture on ( $\tau_{s}=1$ )	626.1	4.35	389.0	4.36
E-MoE, greedy routes ( $\tau_{s}\to 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 (

\tau_{s}=1

) and greedy (

\tau_{s}\to 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.

Explore similar work

Multi-Mask Diffusion Language Models for Few-Step Generation
Jul 22, 2026Sijin Chen, Yinuo Ren, Heyang Zhao +3Masked Diffusion Language Models Diffusion Language Models
ByteDance Seed · Princeton University · Stanford University +2
Infinite Mask Diffusion for Few-Step Distillation
May 11, 2026Jaehoon Yoo, Wonjung Kim, Chanhyuk Lee +1Masked Diffusion Language Models Diffusion Models
Korea Advanced Institute of Science and Technology (KAIST)
DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling
May 22, 2026Jean-Marie Lemercier, Tomas Geffner, Karsten Kreis +3Diffusion Language Models Dataset Distillation
NVIDIA

Jul 22, 2026, cs.CLJ/K move · Enter open · S save

Multi-Mask Diffusion Language Models for Few-Step Generation

Sijin Chen, Yinuo Ren, Heyang Zhao, Ziheng Cheng+2

ByteDance Seed · Princeton University · Stanford University+2

Masked Diffusion Language Models Diffusion Language Models Few-Step Generation

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

Multi-Mask Diffusion Language Models for Few-Step Generation

Infinite Mask Diffusion for Few-Step Distillation

DiLaDiff: Distilled Latent-Augmented Diffusion for Language Modeling