MoE routing as a discrete latent cuts few-step diffusion LM perplexity by 2.6x at matched entropy

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

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

cs.CL

2026-09-29

E-MoE reuses MoE routing decisions as a discrete shared latent, cutting few-step LM1B generative perplexity 2-2.6x over MDLM and VADD at matched entropy and active parameters.

What problem this solves

Masked diffusion models (MDMs) sell parallelism over autoregressive decoding: each denoising step reveals several tokens in one forward pass. The catch is that the reverse process factorizes over positions. Every masked position gets its own marginal distribution, and tokens revealed in the same step are sampled as if conditionally independent. Each token can look fine on its own while the combination never occurs in data. The fewer the steps, the more tokens per step, the worse this factorization error gets, and few-step generation is exactly where diffusion's speed argument lives.

The closest prior fix, VADD (Xie et al., 2026), adds a continuous Gaussian latent trained VAE-style with a separate recognition network. That means a fixed parametric prior, an extra network to train, and a known failure mode: posterior collapse, where the decoder quietly ignores the latent.

Method

E-MoE starts from an observation: an MoE backbone already computes routing decisions for every token at every layer, and those decisions are a ready-made discrete latent.

• The reverse process becomes a mixture over routes. Given a route z (which expert handles each token at each layer), the conditional stays factorized; marginalizing over z, which at sampling means drawing z once from the prior, couples everything revealed in one step.
• The same router runs twice. On the noisy sequence it yields the prior p(z|xt); on the noisy sequence concatenated with the clean one it yields the posterior q(z|xt, x0). The training objective is a derived ELBO bound: a reconstruction term plus a routing KL that decomposes into per-layer, per-token categorical KLs computed from the router itself. No VAE, no hand-designed prior.
• Discrete routing is sampled with Gumbel-Softmax and a straight-through estimator: top-2 selection, one expert active in the forward pass, gradients flowing through both renormalized weights.
• Training costs two forward passes per step (same as VADD); sampling costs one forward pass per NFE, with the route and the token logits computed in the same pass.

Sharing parameters between prior and posterior reduces the whole training problem to one demand: the route chosen from the corrupted sequence must align with the route chosen when the clean sequence is visible. On LM1B, active parameters are 139.4M against MDLM's 139.3M.

Results

LM1B, DiT-small backbone, L=128, 1M steps, 65B tokens:

Sample entropy is matched at 4.34-4.35 (data: 4.32), so the perplexity gain is not bought with diversity. MAUVE separates the models more sharply: 24.6% for E-MoE at NFE=2 against 2.95% for MDLM and 3.76% for VADD. The edge narrows with more steps; from NFE=32 on, all three sit at 86-92% MAUVE, with E-MoE trailing the best by up to 3 points. The autoregressive reference (128 steps) reaches 67.97 perplexity and 96.5% MAUVE, still far ahead.

Elsewhere: 0.062 BPD on binarized MNIST versus 0.064 (VADD) and 0.077 (MDLM); 94.4% validity at NFE=1 on the 8-modes toy against 37.7% for MDLM.

Two ablations pin down the mechanism. Switching the same checkpoint from sampled to greedy routing (temperature to zero, the mixture collapses to a single expert) raises perplexity 1.72x at NFE=1 and 2.14x at NFE=2, moving back toward MDLM with the weights untouched. During training, VADD's latent KL falls to about 0.14 nats per token, near collapse despite annealing, while E-MoE's routing KL holds around 0.65 with no warmup and no auxiliary losses.

Why it matters

Few-step quality is the main front if diffusion LMs want to compete with AR on decoding latency. E-MoE addresses it inside the architecture: no recognition network, no prior to design, active parameters unchanged, training pipeline as simple as the dense baseline. For MoE diffusion models such as LLaDA-MoE the routing decisions already exist, so the change is mostly the training objective. It is a mechanism-level change rather than post-hoc distillation, which needs a pretrained teacher.

Limitations

The authors state two: the benefit fades as NFE grows, since each step reveals fewer tokens; and the method bets on router alignment, degenerating to a single factorized component if the noisy router cannot recover clean routing or the routing distribution collapses. What the paper does not cover: validation stops at LM1B, unconditional generation, 139M active parameters, nearly two orders of magnitude below LLaDA scale; total parameters are 536M against 139M, so memory and deployment are not free; MNIST needs a load-balancing loss while LM1B does not, and KL weights and temperature schedules differ across tasks, with no discussion of hyperparameter sensitivity; and MAUVE trails competitors from NFE=64 on, a fact the main text mentions only briefly.

Terms

• MDM: masked diffusion model, generation by progressively unmasking tokens, several per forward pass.
• NFE: number of function evaluations, forward passes of the denoiser, i.e. sampling steps.
• Factorization error: the reverse process predicts independent per-position marginals, discarding correlations between tokens revealed in the same step.
• Posterior collapse: VAE failure mode where the posterior degenerates to the prior and the decoder ignores the latent.
• Gumbel-Softmax: relaxation that makes discrete categorical sampling differentiable, paired here with a straight-through estimator.
• Gen-PPL: generative perplexity, quality score of sampled text under an external judge (GPT-2-large here).

Source

• arXiv 2609.37533v1
• PDF
• Hugging Face

Related papers

• MMGR: video models score 0% on Sudoku, 3.67% at best on mazes
• Unconverged Cola2 tops Llama-3-8B on two tasks after 346B tokens
• Advance-named cruxes overlap real disagreements at 13.4%, below a 39.2% shuffle
• Sampling reorders three diffusion LMs by bias; scores move over 40%
• One 4B router unifies single-, multi-round, and agentic LLM routing; 61.4% avg, +11% over best
• EvoForesight leads FutureX three weeks running, with rating gains up to 195%

All paper explainers