Unifying Masked Diffusion Models with Various Generation Orders and Beyond

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.

Conventional language models generate text from left to right, one token at a time. Masked diffusion language models offer a different approach: they start from a sentence full of blanks and gradually fill in the missing words. Recent studies have shown that generation order plays an important role in these models, but existing methods usually rely on random order, block-wise left-to-right order with random order inside each block, or learn the order only after the model has already been trained. To establish a foundation for understanding generation order, we introduce a unified framework that shows how different generation orders can be fundamentally integrated into masked diffusion models. This framework also provides a deeper understanding of how training should change depending on the chosen generation order. Based on this framework, we develop a new model that learns both what words to generate and where to generate them, all within a single training process. This allows the model to choose a context-dependent generation order rather than relying on a hand-designed rule. Our results show that this learned ordering improves language modeling performance across multiple benchmarks, bringing diffusion-based language models closer to becoming a practical alternative to standard left-to-right language models.