Beyond Masked and Unmasked:
Discrete Diffusion Models via Partial Masking

The work proposes MDM-Prime, an enhanced masked diffusion model (MDM) incorporating a Partial masking scheme (Prime) that introduces intermediate token states between masked and unmasked representations. By enabling a fine-grained denoising process and improving model utilization, MDM-Prime achieves superior performance on text and image generation.

The study of discrete data generation has been largely shaped by autoregressive models (ARM), which capture the distribution of sequence data by factorizing it according to a prespecified left-to-right order. Recently, promising results from [1,2] have demonstrated that order-agnostic models, such as masked diffusion models (MDM) (e.g., [1-4]), can be effectively extended to large-scale generation tasks, opening a new venue for discrete generative modeling.

Figure 1. An illustrative example of (a) MDM and (b) MDM-Prime. Each token and its corresponding sub-token sequence (constructed via base-$b$ encoding) can take one of three states: unmasked, masked, or intermediate. This example contains $C = 4$ possible token classes, labeled as `bird’, `cat’, `dog’, and `frog’. $\ell = 2$ indicates that each token is represented using two sub-tokens, and $b = \sqrt[\ell]{C} = 2$ denotes the number of classes per sub-token. The symbol $\texttt{m}$ represents a masked token or a masked sub-token. The bottom-right sections of (a) and (b) illustrate the state transition trees. MDM-Prime supports transitions through intermediate states while retaining the ability to directly reach unmasked states. The bottom-left portions depict the sampling process for a token sequence of length $L = 4$.

interpolation between the masked and unmasked states, in the diffusion process. We refer to this method as the Partial masking scheme (Prime), and denote MDM augmented with Prime as MDM-Prime. Prime represents each token as a sequence of sub-tokens using a base-$b$ encoding, with masking performed at the sub-token level. Since sub-tokens can be masked independently during the forward diffusion process, this method introduces intermediate states that partially reveal token information. An illustrative example is shown in the top of Fig. 1, where unmasked states `0-3' are first encoded as `00-11', and the intermediate states are obtained by masking one of the sub-tokens in the sub-token sequence. The intermediate states enable MDM-Prime to perform a fine-grained denoising process. For example, as illustrated in the `State Transition Tree' of Fig. 1 (b), a four-choice prediction can be decomposed into two binary decisions during the sampling process (e.g., $\texttt{mm} \to \texttt{m}1 \to 11$). Compared to standard MDM transitions (i.e., Fig. 1 (a)), MDM-Prime is capable of making predictions based on partially observed token information while deferring the final token revelation until later sampling steps. With its ability to transition through intermediate states, MDM-Prime demonstrates improved model utilization during sampling, as reflected in the reduced number of idle steps in Fig. 2 (i.e., the purple curves). This in turn leads to enhanced performance, as later presented in the experiment section. The contributions of this work are as follows:

This section provides background on the forward and reverse processes of MDM.

Figure 3. An illustrative example of the forward and reverse masked diffusion processes.

▶ Notations:

Foward Diffusion Process

The forward diffusion process is performed through an element-wise conditional sampler $q(\vx_{t}|\vx_0)=\prod_{i=1}^L q(x_{t}^i|x^i_0)$ constructed by interpolating between $\delta_{\texttt{m}}(\cdot)$ and $\delta_{x_0^i}(\cdot)$, defined as follows: [3,4]

$$ \begin{equation} \label{eq:kernel_t0} q(x^i_{t}|x^i_0)= (1-\alpha_t)\delta_{\texttt{m}} (x^i_{t})+\alpha_t \delta_{x^i_0}(x^i_{t}), \end{equation} $$

where $\alpha_t \in [0,1]$ is a strictly decreasing scheduling function with boundary conditions $\alpha_0 \approx 1$ and $\alpha_1 \approx 0$. Intuitively, each perturbed token $x_t^i$ retains the original value $x_0^i$ with probability $\alpha_t$, and is replaced by $\texttt{m}$ with probability $1-\alpha_t$. At time $t=1$, the latent variable $\vx_1=[\texttt{m},\cdots,\texttt{m}]$ consists entirely of masked tokens. An illustrative example of this process is shown in Fig. 3.

Reverse Diffusion Process

Let $s$ and $t$ be two time variables that satisfy $0 \leq s < t \leq 1$. The reverse diffusion process is performed by iterating through $p(\vx_{s}|\vx_t)$, starting from $\vx_1$. The distribution $p(\vx_{s}|\vx_t)$ can be derived using $q(\vx_{t}|\vx_{s})$ and $q(\vx_{s}|\vx_{t},\vx_{0})$. In particular, $q(\vx_{t}|\vx_{s})=\prod_{i=1}^L q(x_{t}^i|x^i_s)$ is defined to be absorbing [5] on the masked state, and is derived from Eq. (\ref{eq:kernel_t0}) as follows:

$$ \begin{equation} \label{eq:kernel_ts} \begin{aligned} q(x^i_t|x^i_{s})=\begin{cases} \frac{\alpha_{s}-\alpha_t}{\alpha_{s}} \delta_{\texttt{m}}(x^i_{t})+\frac{\alpha_t}{\alpha_{s}}\delta_{x^i_{s}}(x^i_{t})& \text{if }x^i_s \in \XX,\\ \delta_{x^i_s}(x^i_{t}) & \text{if } x^i_s = \texttt{m}. \end{cases} \end{aligned} \end{equation} $$

Based on Eqs. (\ref{eq:kernel_t0}) and (\ref{eq:kernel_ts}), the posterior distribution $q(\vx_{s}|\vx_{t},\vx_{0})=\prod_{i=1}^L q(x^i_{s}|x^i_{t},x^i_{0})$ can be derived using Bayes’ rule, and is expressed as [3,4]:

$$ \begin{equation} \label{eq:kernel_ts0} \begin{aligned} q(x^i_{s}|x^i_{t},x^i_{0})= \begin{cases} \delta_{x^i_t}(x^i_{s}) & \text{if } x^i_t \in \XX,\\ \frac{1-\alpha_{s}}{1-\alpha_t}\delta_{\texttt{m}}(x^i_{s})+\frac{\alpha_{s}-\alpha_{t}}{1-\alpha_t} \delta_{x^i_0}(x^i_{s})& \text{if }x^i_t = \texttt{m}. \end{cases} \end{aligned} \end{equation} $$

Since $p(\vx_{s}|\vx_t)=\EE_{p(\vx_0|\vx_t)}[q(\vx_{s}|\vx_{t},\vx_{0})]$ [5], many recent works [2-4] choose to model $p(\vx_0|\vx_t)$ using $p_\theta(\vx_0|\vx_t)$, where $\theta$ denotes the model parameters. To facilitate computational efficiency, this distribution is typically factorized as $p_\theta(\vx_0|\vx_t)=\prod_{i=1}^L p_\theta (x_0^i|\vx_t)$ [1-7]. The parameter $\theta$ can be optimized by minimizing a weighted cross-entropy loss, written as [3,4]:

$$ \begin{equation} \label{eq:diffusion_elbo} \begin{aligned} \mathcal{L}(\vx_0;\theta)=\int_0^1 \frac{\alpha'_t}{1-\alpha_t} \mathbb{E}_{q(\vx_t|\vx_0)}\left[\sum_{i=1}^L \log p_\theta (x_0^i|\vx_t) \right] dt, \end{aligned} \end{equation} $$

where $\alpha_t'=\frac{d}{dt} \alpha_t$. Since unmasked elements in $\vx_t$ retain their values over time according to Eq. (\ref{eq:kernel_ts0}), the carry-over parameterization [3,4] can be applied by explicitly setting the corresponding entries in $\vx_0$ to the unmasked values in $\vx_t$. Formally, this is expressed as $p_\theta(x_0^i | \vx_t) \triangleq \delta_{x_t^i}(x_0^i)$ for position $i$ where $x_t^i \in \XX$.

We introduce a new MDM framework that incorporates Prime, referred to as MDM-Prime. We begin by introducing an invertible function for constructing sub-token sequences and define a novel masked diffusion process over these sub-tokens. We then describe a novel parameterization and a simple model architecture design that operate over sub-token inputs.

Figure 4. An illustrative example of the training process of MDM-Prime. In this example, $L=4$ and $\ell=2$.

▶ Notations:

Discrete Diffusion via Partial Masking

The core idea of Prime is to represent each token $x^i_0$ with a sub-token sequence $\vy^i_0=[y^{i,1}_0,\cdots,y^{i,\ell}_0]$, allowing the creation of intermediate states during the element-wise masking of the forward diffusion process (see Fig. 4). Given a target length $\ell>1$, this can be achieved through an invertible function $f$, which maps each token $x^i_0 \in \XX$ to its base-$b$ encoding $\vy^i_0 \in \YY^\ell$.

To learn $p_\theta(\vy_0)$ using MDM, each data point $\vx_0$ is first transformed into $\vy_0 = f(\vx_0)$, and its corresponding latent variable $\vy_t$ is sampled from $q(\vy_{t}|\vy_0)=\prod_{i=1}^L\prod_{j=1}^{\ell} q(y_{t}^{i,j}|y^{i,j}_0)$ according to Eq. (\ref{eq:kernel_t0}) by substituting $x$'s with $y$'s. The reverse diffusion process is parameterized by $p_\theta(\vy_0 | \vy_t)=\prod_{i=1}^L p_\theta(\vy_0^i| \vy_t)$, and the model is trained to minimize the loss $\mathcal{L}(\vy_0; \theta)$ as in Eq. (\ref{eq:diffusion_elbo}):

$$ \begin{equation} \label{eq:prime_elbo} \begin{aligned} \mathcal{L}(\vy_0;\theta)=\int_0^1 \frac{\alpha'_t}{1-\alpha_t} \mathbb{E}_{q(\vy_t|\vy_0)}\left[\sum_{i=1}^{L} \log p_\theta(\vy^{i}_0|\vy_{t}) \right] dt. \end{aligned} \end{equation} $$

Our paper shows that the number of idle steps decreases as $\ell$ increases, which guarantees an improved model utilization of MDM-Prime during the reverse diffusion process.

Parameterization

This section discusses our implementation for $p_\theta(\vy^{i}_0|\vy_{t})$, which comprises a decoder for modeling the distribution of $\vy^{i}_0$ and an encoder for processing the input $\vy_{t}$.

$$ \begin{equation} \label{eq:parameterization} \begin{aligned} p_\theta(\vy^{i}_0|\vy_{t})= \begin{cases} \frac{\exp (E_\theta (\vy_0^i|\vy_t))}{ \sum_{\vy^i\in \VV(\vy_t^i)} \exp (E_\theta (\vy^i|\vy_t)) },& \text{if $\vy^{i}_0\in \VV(\vy_t^i)$},\\ 0,& \text{if $\vy^{i}_0\notin \VV(\vy_t^i)$}, \end{cases} \end{aligned} \end{equation} $$

where $\VV(\vy_t^i)\triangleq \{ \vy^i=[y^{i,1}, \cdots, y^{i,\ell}] \in f(\XX) \,\,\,s.t.\,\,\, (y^{i,j} = y^{i,j}_t) \vee (y^{i,j}_t = \texttt{m})\}$ denotes a set of outputs that is consistent with $y_t^{i,j}$, and $E_\theta:\YY^{\ell}\times \tilde{\YY}^{\ell\times L}\to \RR$ is a scalar logit. Proposition A.4 in our paper guarantees that Eq. (\ref{eq:parameterization}) satisfies Eq. (\ref{eq:marginalization}). An illustrative example is provided in Fig. 5. As the reverse diffusion process progresses, the number of unmasked sub-tokens in $\vy_t^i$ increases, leading to a substantial reduction in $|\VV(\vy_t^i)|$. This results in a decreasing number of candidate classes of $\vy^i_0$ over time, and thus explicitly reducing uncertainty in the prediction task.

In summary, adapting a standard MDM to MDM-Prime requires only minimal architectural modifications on the embedding layer. This simple strategy preserves the overall architectural design of the standard MDM, enabling a fair comparison with our baseline in the following experiments.

This section presents empirical evaluations to examine the effectiveness of the proposed method.

Text Generation

Models are trained on OpenWebText (OWT), tokenized with the GPT-2 tokenizer. We adopt MDLM [3] as our baseline, apply Prime with varying $\ell$ to obtain MDLM-Prime.

Figure 7. Visualization of the sampling processes of ARM, MDLM, and MDLM-Prime. The masked ratio is measured on a per-token basis, with higher values indicated by darker shades of blue. The samples are generated with prefix and suffix texts sourced from an online article describing Rabindranath Tagore’s poems.

Image Generation

Models are trained and evaluated on CIFAR-10 and ImageNet-32 ($L=32\times 32 \times 3$, $C = 256$ pixel intensities), using the ablated diffusion model (ADM) architecture [9], as in [7]. Sample quality is evaluated using FID and IS.

Figure 8. CIFAR-10 samples generated by MDM-Prime with NFE=512.

▶ Improvements to Sample Quality: Table 1 compares MDM-Prime ($\ell = 2$) against two baselines, MDM ($\ell = 1$) and MDM-Mixture (a mixture-distribution extension via an auxiliary variable, similar to [10]), as well as several existing generative models. MDM-Prime outperforms both MDM baselines, matches StyleGAN+ADA [11] on CIFAR-10 with fewer NFE, and improves over ScoreFlow (VP) [12] on ImageNet-32 by an FID of 1.36.

MDM-Prime-v2 introduces principled subtokenizers, index shuffling and binary encoding, implemented as zero-FLOP lookup tables during data preprocessing. Trained at 1.1B parameters on 540B SlimPajama tokens, it achieves 49.42% average accuracy on eight commonsense reasoning benchmarks, improving over the MDM baseline (SMDM) by 4.87 points.

The MDM-Prime paper reports that the model exhibits more favorable NLL than MDM and other diffusion models. The training objective used in this comparison, however, is not an evidence lower bound (ELBO). We discuss this issue and its implications in a dedicated errata note. A corrected version has been implemented in our Github repo chen-hao-chao/mdm-prime.