Beyond Masked and Unmasked: Discrete Diffusion Models via Partial Masking

We sincerely appreciate the reviewer’s valuable feedback and questions. The responses to the reviewer’s questions are as follows.

Questions

Q1. My 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?

Response. We thank the reviewer for the comment, and are pleased to share some preliminary thoughts about the potential relationship between the likelihood bounds (Eqs. (4) and (5)).

Both objectives aim to approximate the same target, $-\log p\_{\theta}(\boldsymbol{x}\_0)$, using a discrete diffusion process with the same scheduling function $\alpha\_t$. The key difference lies in the expectations:

$\mathbb{E}\_{q(\boldsymbol{x}\_t \| \boldsymbol{x}\_0)} \left[\sum\_{i=1}^L \log p^{\text{MDM}}\_{\theta}(x\_0^i \| \boldsymbol{x}\_t) \right] \\, \\, \\, \\, \text{ and }\\, \\, \\, \\, \mathbb{E}\_{q(\boldsymbol{y}\_t \| \boldsymbol{y}\_0)}\left[\sum\_{i=1}^L \log p^{\text{Prime}}\_{\theta}(\boldsymbol{y}\_0^i \| \boldsymbol{y}\_t)\right].$

A possible direction for theoretical exploration is to analyze the per-step difference between these expectations:

$\mathbb{E}\_{q(\boldsymbol{y}\_t \| \boldsymbol{y}\_0)} \left[ \sum\_{i=1}^L \log p^{\text{Prime}}\_{\theta}(\boldsymbol{y}\_0^i \| \boldsymbol{y}\_t)\right] - \mathbb{E}\_{q(\boldsymbol{x}\_t|\boldsymbol{x}\_0)} \left[\sum\_{i=1}^L \log p^{\text{MDM}}\_{\theta}(x\_0^i \| \boldsymbol{x}\_t) \right] \triangleq d.$

If this difference ($d$) is consistently positive, it could suggest that MDM-Prime provides a tighter bound. However, since $p^{\text{MDM}}\_{\theta}$ and $p^{\text{Prime}}\_{\theta}$ are parameterized by distinct neural networks operating on different input spaces, establishing this relationship analytically is nontrivial and requires specific assumptions about the model families.

To investigate this empirically, we estimated the expectation difference using trained models at a number of discretized timesteps. The results show positive values as follows:

We leave a rigorous theoretical analysis that fully explains this behavior for future work.

Q2. (a) It would be helpful to include a sampling algorithm in the paper for clarity. (b) Does the posterior $q(x\_s^{i,l} \| x\_t^{i,l}, x\_0^{i,l})$ follow the same formulation as Eq. (3)? (c) Additionally, how is sampling from $p(x\_s \| x\_t)$ performed in practice? Given that $p(x\_0^i \| x\_t)$ is no longer independent across sub-tokens, we do not have a closed-form expression for $p(x\_s^i | x\_t^i)$. In this case, is the sampling procedure carried out by first sampling $x\_0^i \sim p(x\_0^i \| x\_t)$, followed by sampling $x\_s^{i,l} \sim p(x\_s^{i,l} \| x\_t^{i,l}, x\_0^{i,l})$?

Response. (a) We thank the reviewer for the suggestion and will enhance our description regarding the sampling process in Section 3.1 in an updated manuscript.

(b) The posterior $q(y\_s^{i,j} | y\_t^{i,j}, y\_0^{i,j})$ follows the same formulation as Eq. (3) and is defined as follows:

We sincerely appreciate the reviewer’s valuable feedback and questions. The responses to the reviewer’s questions are as follows.

Comments

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

Response. A runtime analysis of MDLM-Prime with varying $\ell$ is presented in Appendix A.5.2 and Table A3. Larger $\ell$ generally incur higher training cost due to more complex filtering in carry-over parameterization. The table below presents the runtime of MDLM ($\ell = 1$) and MDLM-Prime with varying $\ell$, along with the corresponding perplexity (PPL). Based on these results, we recommend $\ell=4$ as an ideal balance between efficiency and performance:

For inference, we perform additional experiments to measure the generative perplexity (Gen PPL) under approximately aligned sampling time. The corresponding number of non-cacheable function evaluations (N-NFE) is also presented in the table. The results indicate that MDLM-Prime exhibits lower Gen PPL when N-NFE is higher.

• MDLM | Runtime (sec.) | 2.89 | 6.02 | 8.87 | | - | - | - | - | | Discretized Timesteps | 262 | 644 | 1,685 | | N-NFE | 256 | 512 | 768 | | Gen PPL ($\downarrow$) | 112.0 | 99.0 | 87.0 |

• MDLM-Prime ($\ell=4$) | Runtime (sec.) | 2.83 | 5.99 | 8.72 | | - | - | - | - | | Discretized Timesteps | 230 | 500 | 772 | | N-NFE | 230 | 500 | 768 | | Gen PPL ($\downarrow$) | 123.0 | 90.5 | 77.0 |

Thank you again for your thoughtful review.

We sincerely appreciate the reviewer’s valuable feedback and questions. The responses to the reviewer’s questions are as follows.

Comments

C1. 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.

Response. We thank the reviewer for the comment. In Section 1, we provide the core intuition behind our method, and in Appendix A.2.4, we analyze how the target length $\ell$ (which controls the number of idle steps) affects the model’s likelihood estimation as measured by perplexity. While our current focus centers on providing empirical evidence and intuitive motivation, we acknowledge that a deeper analysis would constitute a valuable contribution and consider this an important direction for future work.

C2. 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.

Response. A runtime comparison between the MDLM baseline and MDLM-Prime with varying $\ell$ is presented in Appendix A.5.2 and Table A3 (the throughput can be derived by taking the inverse of the reported runtime values). Larger $\ell$ generally incur higher training cost due to more complex filtering in carry-over parameterization. The table below presents the runtime of MDLM ($\ell = 1$) and MDLM-Prime with varying $\ell$, along with the corresponding perplexity (PPL). Based on these results, we recommend $\ell=4$ as an ideal balance between efficiency and performance:

For inference, we perform additional experiments to measure the generative perplexity (Gen PPL) under approximately aligned sampling time. The corresponding number of non-cacheable function evaluations (N-NFE) is also presented in the table. The results indicate that MDLM-Prime exhibits lower Gen PPL when N-NFE is higher.

• MDLM | Runtime (sec.) | 2.89 | 6.02 | 8.87 | | - | - | - | - | | Discretized Timesteps | 262 | 644 | 1,685 | | N-NFE | 256 | 512 | 768 | | Gen PPL ($\downarrow$) | 112.0 | 99.0 | 87.0 |

• MDLM-Prime ($\ell=4$) | Runtime (sec.) | 2.83 | 5.99 | 8.72 | | - | - | - | - | | Discretized Timesteps | 230 | 500 | 772 | | N-NFE | 230 | 500 | 768 | | Gen PPL ($\downarrow$) | 123.0 | 90.5 | 77.0 |

C3. 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.

Response. We appreciate this valuable suggestion. We selected perplexity as the primary metric in the main manuscript due to its widespread adoption and its suitability for the evaluation of likelihood-based models such as MDM-Prime. In addition to standard perplexity, we also report generative perplexity (Gen PPL) in Appendix A.5.3, which evaluates the quality of model-generated samples through pretrained large language models. We agree with the reviewer and consider the reviewer's suggestion regarding additional automated metrics or human evaluation an interesting and promising future research direction.

We sincerely appreciate the reviewer’s valuable feedback and questions. The responses to the reviewer’s questions are as follows.

Questions

Q1. [conceptual clarification] (a) 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? (b) Therefore, this is sort of a diffusion over the space of sets of tokens?

Response. (a) Yes, this is a valid and correct interpretation. When some sub-tokens are masked, the noisy state corresponds to a subset of possible original tokens consistent with the unmasked parts. As more sub-tokens are revealed, this subset shrinks, narrowing down the candidate tokens. (b) Yes, this process can be intuitively viewed as diffusion over sets of tokens. Each intermediate state (i.e., partial sub-token sequence) defines a set of possible tokens, and the diffusion moves from larger sets (more uncertainty) to singleton sets (fully denoised tokens).

Q2. [detailed clarification of NFEs vs sampling steps vs number of modeling calls] (a) 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 dependent with given $y_t$?

(b) As a result, in the likelihood bound, there is a marginalization over subtokens $y_0^{i,i}$ consistent with a $y_0^{i,j}$ given $y_t$ 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?

(c) Specifically, 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.

Response. (a) Yes, this is correct. Eq. (7) defines a joint distribution $p_\theta (\boldsymbol{y}^i_0| \boldsymbol{y}_t)$ over sub-tokens at position $i$ given $\boldsymbol{y}_t$.

(b) No, the likelihood bound does not require explicit marginalization over sub-token combinations. Instead, the distribution $p_\theta (\boldsymbol{y}_0^i | \boldsymbol{y}_t)$ in Eq. (7) is defined as a softmax over the set $\mathcal{V}(\boldsymbol{y}_t^i)$, and this formulation is guaranteed to satisfy the marginalization condition in Eq. (6), as shown in Proposition A.4.

To construct the softmax over $\mathcal{V}(\boldsymbol{y}_t^i)$, we apply a filtering mechanism that masks out the logits entries corresponding to inconsistent $\boldsymbol{y}_0^i$ based on the observed $\boldsymbol{y}_t^i$. This procedure, as detailed in Appendix A.4 and illustrated in Fig. A2 (a), performs logical AND operations across $\ell$ binary filters to the logits, which does not require additional model evaluations.

Regarding the computational cost, larger $\ell$ generally incur higher training cost due to more complex filtering in carry-over parameterization. The table below presents the runtime of MDLM ($\ell = 1$) and MDLM-Prime with varying $\ell$, along with the corresponding perplexity (PPL). Based on these results, we recommend $\ell=4$ as an ideal balance between efficiency and performance:

(c) Each NFE corresponds to a single denoising step, which includes one forward pass of the model. The marginalization condition is handled via filtering and does not require additional model calls as described above.

Q3. [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?

Response. We appreciate the reviewer’s insight regarding the impact of sub-token assignments. In our current implementation, we directly apply the base-b encoding to the original token indices (i.e., [0, ..., 255] for images and [0, ..., 50,256] for text). An interesting extension for future work is to reorder token indices based on semantic similarity, which could be automatically determined by a word embedding model, prior to base-b encoding. This would allow related tokens (like "cat" and "bird") to share more sub-token prefixes, enhancing the model’s ability to support fine-grained imputation in textual settings, analogous to the sub-token completion task demonstrated for images in Fig. 8.

Q4. [exploratory question] Any thoughts about data-driven ways to set $\ell$?

Response. An effective approach to setting $\ell$ is to monitor the Idle Step Ratio (ISR), which can be calculated using Eq. (A5). As shown in our experiments in Appendix A.2.4, the ‘elbow point’ in the ISR curve often aligns with optimal performance. Therefore, $\ell$ can be selected by gradually increasing its value and choosing the smallest point at which ISR plateaus. This selection method can be performed efficiently prior to training a model.