Simplified and Generalized Masked Diffusion for Discrete Data

By Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, Michalis K. Titsias. Neural Information Processing Systems, 2024

576 citations recorded by the source corpus.

Masked (or absorbing) diffusion is actively explored as an alternative to autoregressive models for generative modeling of discrete data. However, existing work in this area has been hindered by unnecessarily complex model formulations and unclear relationships between different perspectives, leading to suboptimal parameterization, training objectives, and ad hoc adjustments to counteract these issues. In this work, we aim to provide a simple and general framework that unlocks the full potential of masked diffusion models. We show that the continuous-time variational objective of masked diffusion models is a simple weighted integral of cross-entropy losses. Our framework also enables training generalized masked diffusion models with state-dependent masking schedules. When evaluated by perplexity, our models trained on OpenWebText surpass prior diffusion language models at GPT-2 scale and demonstrate superior performance on 4 out of 5 zero-shot language modeling tasks. Furthermore, our models vastly outperform previous discrete diffusion models on pixel-level image modeling, achieving 2.75 (CIFAR-10) and 3.40 (ImageNet 64x64) bits per dimension that are better than autoregressive models of similar sizes. Our code is available at https://github.com/google-deepmind/md4.

Subfields: Computer Vision and Pattern Recognition

View Paper: https://www.semanticscholar.org/paper/1bed8c7541381b1f79027c240b64c9276573fc3c

Citations in the corpus

• Simple and Effective Masked Diffusion Language Models (2024, unknown): “…et al. [1] show that the posterior corresponding to (13) is given as follows: which we simplify to the following: q ( z s | z t , x ) A.1.3 p θ ( z s | z t ) Austin et al. [1] approximate the reverse process in the following manner: where x θ ( z t ,t ): V× [0 , 1] → ∆ K is an approximation for x .”
• Large Language Diffusion Models (2025, unknown): “Lou et al. [17] showed that masked diffusion, as a special case of discrete diffusion, achieves perplexity comparable to or surpassing ARMs at GPT-2 scale.”
• Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data (2024, unknown): “Sahoo et al. (2024) derive the same cross-entropy losses with Shi et al. (2024).”
• BLOCK DIFFUSION: INTERPOLATING BETWEEN AU-TOREGRESSIVE AND DIFFUSION LANGUAGE MODELS (2025, unknown): “Further simplifications (Sahoo et al., 2024a; Shi et al., 2024; Ou et al., 2025) tighten the ELBO and enhance performance.”
• Dream 7B: Diffusion Large Language Models (2025, unknown): “Building on this foundation, several studies have improved the training objectives for discrete diffusion (Zheng et al., 2024b; Lou et al., 2024; Shi et al., 2024; Ou et al., 2025; Zhao et al., 2024; Sahoo et al., 2024).”
• Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decoding (2025, unknown): “More recently, work by MDLM [28, 27, 41] and RADD [22] has shown that for MDMs, different parameterizations are equivalent.”