One-sentence Summary

Researchers from the Applied AI Institute propose Enhanced Mixture-of-Experts (E-MoE), which builds the reverse process as a mixture of factorized distributions over a discrete shared latent derived from Mixture-of-Experts (MoE) expert-routing decisions without increasing active parameters over the factorized baseline, improving few-step generation across synthetic multi-modal benchmarks, binarized MNIST, and LM1B.

Key Contributions

• Enhanced Mixture-of-Experts (E-MoE) is introduced as a masked diffusion model whose reverse process is a mixture of factorized token distributions over a discrete shared latent given by the MoE backbone's expert-routing decisions.
• E-MoE uses a learned discrete routing latent in place of the continuous Gaussian latent and recognition network of prior work, avoiding posterior collapse without increasing active parameters over the factorized baseline.
• Experiments on synthetic multi-modal benchmarks, binarized MNIST, and LM1B show that E-MoE improves few-step generation quality over factorized baselines.

Introduction

Autoregressive language models provide tractable likelihood but require inherently sequential decoding. Masked diffusion models offer faster parallel unmasking, yet their usual factorized reverse process predicts independent token marginals, so tokens generated in the same step are treated as conditionally independent and correlations are lost. Prior approaches address this with a continuous shared latent, but that requires a Gaussian prior and an extra recognition network, increasing model complexity. The authors propose Enhanced Mixture-of-Experts (E-MoE), which reuses the discrete routing decisions of a mixture-of-experts backbone as the shared latent, enabling a correlated reverse process without added parameters or inference cost, and derive a tractable ELBO for training.

Method

The authors propose a masked discrete diffusion model that builds the reverse process over a discrete shared latent space. To overcome the factorization barrier without training additional auxiliary models like a VAE, they leverage the routing decisions of a Mixture-of-Experts (MoE) backbone. This discrete latent inherently dictates which sparse mixture of experts should explain the unmasking step at each layer.

The MoE architecture comprises $E$E experts across $D$D layers. For a sequence $\mathbf{x}_{t_i}$xti​​, the latent code $\mathbf{z}$z is modeled as a set of layer-wise and token-wise routing assignments. The authors define both the posterior and prior distributions over these routing decisions using per-layer and per-token factorization. The posterior distribution conditions on both the noisy input $\mathbf{x}_{t_i}$xti​​ and the clean target $\mathbf{x}_0$x0​, while the prior conditions solely on the noisy input $\mathbf{x}_{t_i}$xti​​.

During training, the model requires two forward passes to align the routing distributions. As shown in the figure below:

In the first pass, referred to as the clean pass, the router processes the concatenation of the noisy sample and the clean sequence to compute the posterior routing distribution. In the second pass, known as the noisy pass, the same router processes only the noisy sample to approximate the prior distribution. The core training objective is to align the route selected when the clean sample is available with the route selected from the corrupted sequence alone.

The training objective consists of two main components. The first term calculates the negative log-likelihood of the clean tokens predicted by the selected experts. The second term is a router-matching loss, formulated as the sum of categorical KL divergences between the posterior and prior routing distributions across all layers and positions. Since the shared latent $\mathbf{z}$z is discrete, the authors employ the Gumbel-Softmax trick with a straight-through estimator to enable differentiable sampling from the routing distributions. This involves adding i.i.d. Gumbel noise to the router logits and applying a temperature parameter to control the relaxation towards a one-hot expert assignment.

During the inference stage, the model operates efficiently by utilizing only the noisy pass. It samples the latent routing decisions directly from the prior distribution and computes the token predictions in a single forward pass, ensuring that one network function evaluation corresponds to one forward pass through the MoE backbone.

Experiment

The experiments evaluate E-MoE against factorized MDLM and VADD across three settings with matched backbones, optimizers, and training budgets: two-dimensional toy densities, binarized MNIST generation, and LM1B text generation. Toy results show that factorized sampling produces spurious coordinate combinations, while E-MoE's shared discrete latent preserves true clusters and manifolds. In image and text generation, E-MoE yields more coherent low-step samples, achieving the best bits per dimension on MNIST and substantially better generative perplexity and MAUVE on text at one and two sampling steps. Overall, the shared discrete latent helps break the factorization barrier and improves few-step generation without increasing active parameters.

E-MoE keeps the mixture-of-factorized-distributions view of VADD but replaces its continuous Gaussian latent with a discrete routing latent that the backbone already computes. This removes the need for a separate recognition network and extra active parameters while learning a data-dependent prior chosen per token and layer. In experiments, the shared discrete latent captures cross-token correlation and improves few-step generation over both factorized and continuous-latent baselines. E-MoE reuses routing decisions already computed by the backbone as a discrete latent, so it needs no separate recognition model or additional active parameters. Unlike VADD's fixed continuous prior sampled once per sequence, E-MoE learns a data-dependent discrete prior selected per token and layer. E-MoE produces more coherent few-step samples than factorized MDLM and improves likelihood and generative perplexity while matching VADD's parameter count.

On two synthetic two-dimensional densities, E-MoE and VADD both substantially improve validity over factorized MDLM in the few-step regime. E-MoE consistently outperforms VADD on 8-modes and remains competitive on Swiss-roll, with the advantage shrinking as sampling steps increase. At one sampling step, factorized MDLM yields substantially lower validity than VADD and E-MoE on both synthetic densities. E-MoE achieves the best validity on 8-modes at every step budget and is competitive on Swiss-roll, where VADD leads at most step counts but differences narrow as steps increase.

On binarized MNIST, E-MoE reports the lowest bits-per-dimension among the compared methods, followed closely by VADD and then MDLM. E-MoE achieves this with a parameter count essentially matching VADD, while MDLM is somewhat smaller but less effective on likelihood. E-MoE achieves the best test bits-per-dimension, slightly lower than VADD and clearly lower than factorized MDLM. E-MoE and VADD have nearly identical parameter counts, while MDLM uses fewer parameters but yields a higher bits-per-dimension.

Across few-step unconditional text generation on LM1B, E-MoE achieves the best generative perplexity at every reported number of function evaluations. Its largest improvements over factorized and latent baselines occur at one and two steps, where it substantially reduces generative perplexity while maintaining sample entropy close to the data level. The advantage narrows as more sampling steps are used. E-MoE achieves the lowest generative perplexity across all reported NFE settings on LM1B. At one and two NFE, E-MoE reduces generative perplexity by about 2 to 2.6 times relative to MDLM, SEDD, and VADD, with sample entropy similar to those baselines and near the data entropy. As NFE increases, E-MoE's advantage over the strongest baselines shrinks, consistent with fewer tokens being unmasked per step.

E-MoE is evaluated on synthetic two-dimensional densities, binarized MNIST likelihood modeling, and few-step unconditional text generation on LM1B. Its discrete routing latent removes the need for a separate recognition network and captures cross-token correlation, improving validity and likelihood over factorized MDLM while matching VADD's parameter count. The largest gains appear in few-step text generation, where E-MoE achieves the best generative perplexity, especially at one and two sampling steps, with the advantage narrowing as more steps are used.