Unifying Masked Diffusion Models with Various Generation Orders and Beyond

TL;DR

This work tackles the dependence of generation quality on unmasking order in masked diffusion models (MDMs) by introducing order-expressive masked diffusion models (OeMDM) and a learnable-order variant (LoMDM). OeMDM provides a generalized NELBO that treats the generation order as a learnable scheduler, enabling unmasking strategies that unify autoregressive, block-diffusion, and diffusion-based decoding within a single framework. LoMDM extends this by jointly learning a context-aware forward scheduler and a diffusion backbone via a single objective, yielding order-aware training and inference with improved sample quality. Empirically, LoMDM achieves lower perplexities than BD3LM, GenMD4, and MDLM across LM1B and OpenWebText, shows strong zero-shot generalization, and attains MDLM-level performance after roughly 18% of the training steps, highlighting substantial efficiency gains and practical impact for discrete diffusion-based language modeling.

Abstract

Masked diffusion models (MDMs) are a potential alternative to autoregressive models (ARMs) for language generation, but generation quality depends critically on the generation order. Prior work either hard-codes an ordering (e.g., blockwise left-to-right) or learns an ordering policy for a pretrained MDM, which incurs extra cost and can yield suboptimal solutions due to the two-stage optimization. Motivated by this, we propose order-expressive masked diffusion model (OeMDM) for a broad class of diffusion generative processes with various generation orders, enabling the interpretation of MDM, ARM, and block diffusion in a single framework. Furthermore, building on OeMDM, we introduce learnable-order masked diffusion model (LoMDM), which jointly learns the generation ordering and diffusion backbone through a single objective from scratch, enabling the diffusion model to generate text in context-dependent ordering. Empirically, we confirm that LoMDM outperforms various discrete diffusion models across multiple language modeling benchmarks.

Figures (7)

• Figure 1: Conceptual illustration of learnable-order masked diffusion model (LoMDM) and other language models. Black text denotes already generated tokens, while the colored tokens indicate the generation candidates, with a lower color represents low generation order priority. In training time, LoMDM jointly learns what to generate and where to generate next, and in inference-time, LoMDM selects where to unmask next and predict a token.
• Figure 2: Illustration of $\alpha_{\mathrm{arm},\epsilon}(t)$ that makes OeMDM to generate in L2R order. The explicit function formulation is in Appendix \ref{['sec:appendix_proof_arm']}
• Figure 3: Model structure of LoMDM. We view backbone of diffusion model $\theta$ as a feature extractor of $\mathbf{z}_t$ or $\mathbf{x}$, and train $\theta$,$\alpha_\phi$, and $\hat{\alpha}_\psi$ jointly. Depending on the input type, final layers are switched off or on. For example in the above figure, the input is $\mathbf{z}_t$ so the final diffusion MLP layer and $\hat{\alpha}_\psi$ is activated. Meanwhile, if input was $\mathbf{x}$, only $\alpha_\phi$ would be activated. We detach the gradient of $\alpha_\phi$ and $\hat{\alpha}_\psi$ from flowing to the diffusion backbone (N-Layer transformer blocks in figure.)
• Figure 4: Pearson correlation per training step for LoMDM trained on OWT. We report correlations among $A_\phi^{(i)}(\mathbf{x},t)$, $\hat{A}_\psi^{(i)}(\mathbf{z}_t,t)$, and $\langle \mathbf{x}_\theta^{(i)}(\mathbf{z}_t,t), \mathbf{x}^{(i)}\rangle$. When measuring correlation with $\langle \mathbf{x}_\theta^{(i)}(\mathbf{z}_t,t), \mathbf{x}^{(i)}\rangle$, we compute it only over masked positions in $\mathbf{z}_t$, since $\mathbf{x}_\theta^{(i)}(\mathbf{z}_t,t)$ is zero at unmasked positions.
• Figure 5: Test PPL per wall-clock-time during training on OWT. We truncate the curves at the point where LoMDM matches the 1M-step MDLM performance (PPL = 23.0). At this cutoff, MDLM had reached PPL = 24.9 with $\sim$0.30M steps, while our method had reached PPL = 23.0 with $\sim$0.18M steps.

Theorems & Definitions (22)

• Definition 3.1: free-form scheduler function class
• Proposition 3.1: NELBO of in continuous time
• Proposition 3.1: Autoregressive models as a special case of
• proof : proof sketch.
• Definition 2.1: Masked sequence set
• Definition 2.2: Absorbing mask trajectory
• Lemma 3.1: Conditional absolute continuity of OeMDM
• proof
• Proposition \ref{NELBO}: NELBO of OeMDM in continuous time
• Proposition \ref{NELBO}: Finiteness of the OeMDM NELBO