Enhanced mixture-of-experts improve few-step diffusion language models
E-MoE: Enhanced Mixture-of-Experts for Non-Factorized Diffusion Language Models
Computation and Language
Summary
Generating sequences of words quickly and accurately is challenging for some language models, especially when they try to predict multiple words at once. The authors found that previous methods that treated each word independently limited quality when generating sequences in just a few steps. They introduced a new approach called Enhanced Mixture-of-Experts (E-MoE), which combines multiple expert models to better capture the relationships between words during generation. This approach improves the quality of generated text without needing more computational power than before.
What this means in practice
- •For natural language engineers: Improve few-step sequence generation quality for efficient language model deployment where speed and accuracy are critical.
- •For image captioning developers: Use improved cross-token correlation modeling to enhance multimodal systems combining text and image generation.
Authors
Arseny Ivanov, Alexander Kolesov, Alexander Korotin, Ivan Oseledets, Mikhail Goncharov
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.