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

E-MoE. Enhanced Mixture-of-Experts for Non-Factorized Diffusion Language Models. By Arseny Ivanov and four others, from the Applied AI Institute. Published on the twenty-ninth of September. No code repository is linked to it.

Masked diffusion models generate a sequence by unmasking several tokens at each denoising step. Their reverse process, though, is typically factorized over positions: each position is predicted without regard to the others in the same step. That limits sample quality in the few-step regime, which is exactly where diffusion's speed advantage over autoregressive decoding matters most.

One recent line of work adds a continuous Gaussian latent, trained as a variational autoencoder, to capture correlations across positions. Those approaches are prone to posterior collapse, where the model silently ignores the latent altogether.

Enhanced Mixture-of-Experts, or E-MoE, finds its latent somewhere already present. A Mixture-of-Experts backbone makes expert-routing decisions, and those decisions are discrete. E-MoE treats them as a shared latent, and builds the reverse process as a mixture of factorized distributions over it.

Each component of the mixture is still factorized, but because the components are selected through a shared routing latent, the overall reverse process is no longer factorized over positions. And it comes at no extra cost in active parameters compared with the factorized baseline.

Results are reported on synthetic multi-modal benchmarks, on binarized MNIST, and on LM1B. Across all three, E-MoE improves few-step generation over factorized baselines.

Rather than bolting on a continuous latent that may go unused, E-MoE reads the correlation structure out of a choice the backbone was making anyway: which experts to route to.