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

Arseny Ivanov1,2, Alexander Kolesov2,1, Alexander Korotin2,1, Ivan Oseledets1,2, Mikhail Goncharov1
1AXXX   2Applied AI Institute
Accepted at NeurIPS 2026 DiffuLM Workshop
Paper arXiv Code 🤗Models 🤗HF Paper

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.

The Factorization Barrier

A masked diffusion model conditions on the whole sequence, but at every step it predicts an independent marginal for each masked position. Sampling several positions from these marginals treats them as conditionally independent and discards the dependencies between them — an error that grows with the number of tokens unmasked at once, precisely where parallel unmasking would pay off.

E-MoE overcomes this with a discrete shared latent: the expert-routing decisions an MoE backbone already computes. Each expert is still factorized, but tokens unmasked in the same step are drawn from the same mixture component, and the mixture over routes recovers the joint distribution.

Factorization error and mixture over expert routes
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=(z_d^{\ell})_{\ell=1}^{L}$, specifies the expert assignment for every position across layers.

Method

E-MoE builds the reverse process as a mixture of factorized distributions over a discrete latent $\mathbf{z}$: $\;p_\theta(\mathbf{x}_s\mid\mathbf{x}_t)=\sum_{\mathbf{z}} p_\theta(\mathbf{z}\mid\mathbf{x}_t)\prod_{\ell} p_\theta(\mathbf{x}^\ell_s\mid\mathbf{x}_t,\mathbf{z})$. The latent is the per-token, per-layer routing decision of an MoE backbone, so no extra network is needed. The same router evaluated on the noisy sequence is the learned, data-dependent prior $p_\theta(\mathbf{z}\mid\mathbf{x}_t)$, and evaluated with the clean sequence it is the training-time posterior $q_\theta(\mathbf{z}\mid\mathbf{x}_t,\mathbf{x}_0)$.

E-MoE training: noisy pass gives the prior, clean pass gives the posterior
Training stage: we model prior $p_\theta(\mathbf{z}_{t_i}\mid\mathbf{x}_{t_i})$ and posterior $q_\theta(\mathbf{z}_{t_i}\mid\mathbf{x}_{t_i},\mathbf{x}_0)$ by the same router, unlike VADD. The router takes as input a noisy sample $\mathbf{x}_{t_i}$ to approximate $p_\theta$ in the noisy pass, and a concatenation of $\mathbf{x}_{t_i}$ and $\mathbf{x}_0$ for $q_\theta$ in the clean pass. Then, the experts return predictions for each masked token using the corresponding route. Inference stage: we use only the noisy pass.

Training minimizes reconstruction of the masked tokens under routes drawn from the posterior, plus a router-matching term: a sum of categorical KL divergences between the clean and noisy routers over all layers and positions. Discrete routes are sampled with the Gumbel-softmax trick and a straight-through estimator. At sampling time the route is drawn from the prior in the same forward pass, so one NFE is one forward pass.

$\displaystyle \mathcal{L}_{\text{E-MoE}} = \mathbb{E}_{p_{\text{data}}(\mathbf{x}_0),\,t_i}\,\mathbb{E}_{q(\mathbf{x}_{t_i}\mid\mathbf{x}_0)}\,\mathbb{E}_{q_\theta(\mathbf{z}\mid\mathbf{x}_{t_i},\mathbf{x}_0)}\Big[-\frac{1}{t_i}\sum_{\ell\in\mathcal{M}(\mathbf{x}_{t_i})}\log p_\theta(\mathbf{x}^\ell_0\mid\mathbf{x}_{t_i},\mathbf{z}) \;+\; T\cdot\mathrm{KL}_{\text{route}}\Big]$
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.
MDLMVADDE-MoE (ours)
Reverse process$\prod_\ell p_\theta(\mathbf{x}^\ell\mid\mathbf{x}_t)$$\int p(\mathbf{z})\prod_\ell p_\theta(\mathbf{x}^\ell\mid\mathbf{x}_t,\mathbf{z})\,d\mathbf{z}$$\sum_{\mathbf{z}} p_\theta(\mathbf{z}\mid\mathbf{x}_t)\prod_\ell p_\theta(\mathbf{x}^\ell\mid\mathbf{x}_t,\mathbf{z})$
Latent space—$\mathbb{R}^d$$\{1,\dots,E\}^{L\times D}$
Latent is chosen—once per sequenceper token and layer
Learned, data-dependent prior✗✗✓
Posterior $q$ comes from—VAE (additional net)the same net
Forward passes per training step122
Forward passes per sampling step111

Results

Text generation on LM1B

All models share the same small DiT backbone and are trained on LM1B for 1M steps at global batch 512; E-MoE uses as many active parameters per token as MDLM. The gain is largest where the factorization barrier binds: at NFE = 1 and 2, E-MoE lowers generative perplexity by 2× or more relative to every baseline, including the capacity-matched MDLM-MoE, at the same sample entropy, and at NFE = 2 it reaches a MAUVE of 24.6% against 3.0% for MDLM and 3.8% for VADD.

Generative perplexity over NFE on LM1B MAUVE over NFE on LM1B
Two views of the same LM1B sweep. Left: Generative Perplexity ($\downarrow$). Right: MAUVE ($\uparrow$).
Few-step generation on LM1B. Generative perplexity (GPT-2 Large, $\downarrow$) and unigram sample entropy $H$, mean $\pm$ std over 5 disjoint groups of 1000 samples. MDLM-MoE uses the same MoE backbone as E-MoE with deterministic routing and no latent. BD3-LM numbers are from 2000 samples; superscripts give its block size. Autoregressive baseline: 68.3 $\pm$ 0.9 ($H$ = 4.32) at NFE = 128.
NFEMDLMSEDDVADDBD3-LMMDLM-MoEE-MoE (ours)
11440 ± 13 / 4.371604 ± 14 / 4.381279 ± 20 / 4.35—1293 ± 13 / 4.33644 ± 6 / 4.35
2993 ± 19 / 4.371065 ± 16 / 4.38769 ± 11 / 4.33—816 ± 14 / 4.33379 ± 5 / 4.35
4469 ± 4 / 4.36461 ± 2 / 4.35372 ± 2 / 4.33—383 ± 2 / 4.32235 ± 3 / 4.34
8260.0 ± 0.8 / 4.35242.5 ± 5.4 / 4.33223.6 ± 5.5 / 4.331130.616 / 4.34214.7 ± 1.2 / 4.32174.8 ± 2.1 / 4.34
16179.4 ± 3.1 / 4.35167.4 ± 3.5 / 4.33165.7 ± 1.5 / 4.331004.18 / 4.32149.0 ± 3.1 / 4.32148.4 ± 2.0 / 4.34
32149.9 ± 1.9 / 4.35138.0 ± 1.2 / 4.33141.1 ± 2.5 / 4.32790.74 / 4.29125.1 ± 1.3 / 4.33135.3 ± 1.8 / 4.34
64133.7 ± 2.0 / 4.35124.2 ± 1.5 / 4.33126.6 ± 0.6 / 4.32—112.4 ± 1.1 / 4.33130.1 ± 0.8 / 4.34
128127.6 ± 1.3 / 4.35118.9 ± 1.2 / 4.33117.8 ± 0.6 / 4.32—107.0 ± 1.4 / 4.32124.9 ± 1.6 / 4.34

2-D toy distributions

Generation results on 2-D toy examples
Generation results on 2-D toy examples. Ground truth and generations of MDLM, VADD and E-MoE at NFE = 1 on 8-modes and Swiss-roll.

BibTeX

@article{ivanov2026emoe,
  title   = {E-MoE: Enhanced Mixture-of-Experts for Non-Factorized Diffusion Language Models},
  author  = {Ivanov, Arseny and Kolesov, Alexander and Korotin, Alexander and
             Oseledets, Ivan and Goncharov, Mikhail},
  journal = {arXiv preprint arXiv:2609.37533},
  year    = {2026}
}