Title: Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking
Venue: NeurIPS 2025
Accepted: yes
Keywords: Discrete Diffusion Models

Masked Diffusion Models

Generative Models
Abstract: Masked diffusion models (MDM) are powerful generative models for discrete data that generate samples by progressively unmasking tokens in a sequence. Each token can take one of two states: masked or unmasked. We observe that token sequences often remain unchanged between consecutive sampling steps; consequently, the model repeatedly processes identical inputs, leading to redundant computation. To address this inefficiency, we propose the Partial masking scheme (Prime), which augments MDM by allowing tokens to take intermediate states interpolated between the masked and unmasked states. This design enables the model to make predictions based on partially observed token information, and facilitates a fine-grained denoising process. We derive a variational training objective and introduce a simple architectural design to accommodate intermediate-state inputs. Our method demonstrates superior performance across a diverse set of generative modeling tasks. On text data, it achieves a perplexity of 15.36 on OpenWebText, outperforming previous MDM (21.52), autoregressive models (17.54), and their hybrid variants (17.58), without relying on an autoregressive formulation. On image data, it attains competitive FID scores of 3.26 on CIFAR-10 and 6.98 on ImageNet-32, comparable to leading continuous generative models.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces MDM-Prime, a partial masking scheme designed to improve Masked Diffusion Models (MDMs). The core innovation is to allow tokens to exist in intermediate states between fully masked and unmasked, which the authors conceptualise as decomposing tokens into finer \"sub-token\" representations. This approach facilitates a more fine-grained denoising process by enabling the model to make predictions based on partially observed information. Empirically, Prime demonstrates superior performance across tasks, significantly outperforming previous MDM and autoregressive models on the OpenWebText benchmark and achieving competitive results on image generation datasets.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nOverall, the paper is well-written and presents very simple yet effective ideas. It is very promising that the proposed method outperforms auto-regressive models in the openwebtext dataset and achieves comparable performance to diffusion on image modelling.\n\n**Weaknesses**\n\nMy main concern lies in the absence of theoretical analysis explaining why partial-mask modelling yields better performance compared to MDM. While the use of sub-tokens enables smoother transitions, the parameterization of p(y_0|y_t) remains independent across different tokens. Does the proposed approach provide a tighter lower bound on the likelihood?\n\nQUESTIONS:\n- It would be helpful to include a sampling algorithm in the paper for clarity. Specifically, does the posterior $q(x_s^{i,l} \\mid x_t^{i,l}, x_0^{i,l})$ follow the same formulation as Eq3? Additionally, how is sampling from $p(x_s \\mid x_t)$ performed in practice? Given that $p(x_0^i \\mid x_t)$ is no longer independent across sub-tokens, we do not have a closed-form expression for $p(x_s^i \\mid x_t^i)$. In this case, is the sampling procedure carried out by first sampling $x_0^i \\sim p(x_0^i \\mid x_t)$, followed by sampling $x_s^{i,l} \\sim q(x_s^{i,l} \\mid x_t^{i,l}, x_0^{i,l})$?\n- It would be helpful to include the training curve in the paper. Does the training cost increase as $ l $ increases? Additionally, how does the training cost of MDM-Prime compare to that of MDM?\n- The paper does not report the model size. Could you clarify whether the model size is similar to that of the baselines?\n- While MDM includes more idle steps, these can be practically skipped as the model's output is not time-dependent, meaning logits do not need recomputation. In this regard, given the same sampling steps, MDM-Prime has a higher computational cost compared to MDM due to having fewer idle steps. This raises a critical question: given the same computational budget, like flops, at the test time, does MDM-Prime still maintain its performance advantage over MDM?\n- In Figures A3 and A4, there are cases where $ l < 1 $ for CIFAR-10. Could you clarify what this means? My understanding is that $ l = 1 $ corresponds to the original MLDM parameterisation without sub-tokens.\n- Everything has two sides. While MDM-Prime currently performs very well in terms of likelihood, what do you think are its downsides?\n\nLIMITATIONS:\nOverall, the paper is well-written and the results are impressive, particularly the achievement of outperforming autoregressive models in language modelling. However, some experimental and implementation details are missing, and the paper lacks theoretical analysis to support the proposed method. I would be inclined to raise my score if the authors can address these concerns.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-28"}

comments: SUMMARY:
This paper introduces a partial masking technique which let each token in the masked diffusion model goes beyond just masked or unmasked state. In particular, each token is decomposed to subtokens, where the subtokens can be masked individually, making the original token partially masked. There is a bijective mechanism that allows the translation between token and subtokens. Empirical results show improvement over traditional MDM.

STRENGTHS AND WEAKNESSES:
Strength:

1. The idea of partial masking via sub token is novel and elegant. It makes a lot of sense to have partial masking for discrete diffusion models, as the continuous diffusion model has many intermediate state, but it is non trivial to design a mechanism to do so. The design of sub tokens and bijection mapping is a clever way to tackle the problem.

2. The empirical result is strong. It surpasses MDM by a large margin on perplexity and FID.

3. The paper is well written and the idea is presented nicely. The figures are also very helpful.

Weaknesses:

1. Some discussion on the computational complexity is needed. Does it add any overhead to training and inference?

QUESTIONS:
N/A

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 2
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces MDM-Prime, an enhancement to Masked Diffusion Models (MDMs) designed to address computational inefficiency arising from \"idle steps\" where the model's input remains unchanged during the sampling process. The core idea is to replace the binary masked/unmasked states with a more granular system that represents each discrete token as a sequence of sub-tokens using a base-b encoding. Masking is then performed at the sub-token level. The method achieves a lower perplexity on the OpenWebText dataset than previous MDM and autoregressive models.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The paper is well written. The illustrations are nicely plotted.\n- The motivation for the work is clear and compelling. The concept of \"idle steps\" provides a strong intuition for the inefficiency of standard MDMs. The proposed solution of partial masking through sub-token decomposition is an elegant and simple way to create a more fine-grained and efficient denoising process.\n- MDM-Prime requires only minimal architectural changes to a standard MDM. \n- The method achieves impressive results, notably outperforming autoregressive models in text perplexity (15.36 PPL vs. 17.54 PPL) on the OpenWebText dataset.\n\nWeaknesses:\n- The connection between improved model utilization and superior likelihood estimation is not thoroughly explored. While the reduction in \"idle steps\" is the primary motivation, the paper does not provide a deep analysis of why this leads to such a dramatic improvement in modeling capabilities, to the point of outperforming autoregressive models. \n- The paper does not report the impact of MDM-Prime on training and inference time. Handling sub-tokens and a more complex parameterization should introduce overhead. A direct comparison of throughput against the baseline MDM is necessary for a complete picture of the method's practical efficiency.\n\nQUESTIONS:\nNone\n\nLIMITATIONS:\n- For the text generation task, performance is evaluated solely based on perplexity. Augmenting the evaluation with other automated metrics or a human study would substantially strengthen the claims of superiority.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a partial masking scheme (Prime) to address redundant and perhaps suboptimal computation in the masked diffusion model unmasking generative process, wherein some tokens are sampled early and then copied over and over while others are unmasked to be consistent with those already chosen. To challenge this setup, instead of the original masking/unmasking scheme, they introduce latent variables that are ``partially masked\" tokens which allows the model to gradually sample a consistent set of tokens rather than specialize/commit some dimensions too early and copying them over the rest of the diffusion process. They propose a variational training scheme for handling these latent partial masks (\"intermediate states\") and report perplexities and FIDs on text and image data.\n\nWhile the two goals (1. minimizing steps of doing nothing 2. introducing a richer latent denoising space) may seem orthogonal at first, I think the idea is roughly that time spent doing could be time spent throwing extra computation into a more expressive model, and a more fine-grained latent space is one way to do this.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n- Generalize masking and unmasking to partial masking , where a partial mask  is a base b ssequence of length $\\ell$ of \"sub-tokens\".\n- adapting MDMs to MDM-prime requires some embedding changes but not much architecture or diffusion formalism changes\n- Good PPL/FIDs per number of function evaluations\n\nWeaknesses\n\nSee questions.\n\nQUESTIONS:\n- [conceptual clarification] do I understand correctly that by explicitly defining the mapping between original tokens and subtokens, partially-masked subtoken sequences essentially represent a set of possible tokens that the current noisey state for this token can be mapped to? And this set of \"tokens to denoise to\" shrinks as we unmask more subtokens? Therefore, this is sort of a diffusion over the space of sets of tokens?\n\n- [detailed clarification of NFEs vs sampling steps vs number of modeling calls] If i understand correctly, it seems like one needs a joint model over subtokens given state $y_t$. Is it right that this joint model factors over tokens ($y_0^i$ independent of $y_0^j$ given $y_t$) but is correlated across subtokens $y_0^{i,i}$ dependent with $y_0^{i,j}$ given $y_t$? As a result, in the likelihood bound, there is a marginalization over subtokens consistent with a given observed partial masking, is that right? Is that what equations (6) and (7) get at? If so, any considerations computationally ? Does this marginalization govern that $\\ell$ should be kept somewhat low? \n\nSpecifically, If multiple model calls are made in the marginalization sum, then what exactly do the NFEs in the tables mean? Does 500 NFEs mean 500 denoising steps, where each denoising step might call the model many times to marginalize? In either case, **the relationship between NFEs, number of denoising steps, and number of model calls should be clarified. If multiple model calls are made per what the table refers to as \"1 NFE\", then this should be made more prominent in the text.**. Clarification will be appreciated.\n\n- [exploratory question] On an intuitive level, it seems that which original tokens to \"pair\" via shared subtoken subsequences is potentially an important detail that could affect the learning problem. e.g. if \"cat\" is 0001 and \"bird\" is 0000 and the partial mask is 000m then the model knows from [000m] that the word  is \"cat\" or \"bird\" and statistical strength can be shared in sentences like \"Lucy has a pet [000m]\". But less such strength could be shared if it were \"bird\" or \"chair\". Any particular thoughts on automated ways to assign the mapping? How was it done in your work?\n\n- [exploratory question] Any thoughts about data-driven ways to set $\\ell$ ?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-04"}