SPMDM: Enhancing Masked Diffusion Models through Simplifing Sampling Path

Q1. Other Partitioning Schemes

Thank you for suggesting alternative segmentation strategies. This is indeed a promising direction. Indeed, we have explored similar approaches, including using punctuation marks as delimiters for text segmentation, which is similar to your proposed syntactic-structure-based partitioning. However, the main challenge is that such syntactic structures are not available during inference with fully masked sequences.

Q2. Formulation

Thank you for your thoughtful consideration. We'd like to provide a comprehensive response to the concerns you raised.

About the subsequence-dependent noise scheduler. For each subsequence $\mathbf{x}^k$, we independently sample noise scale parameters $t_k \sim \mathcal{U}[0, 1]$.

The fundamental difference from standard masked diffusion models lies in the forward process formulation. In standard MDMs (D3PM, MDLM), the forward process is defined as:

$$q_{t|0}(\mathbf{x}\_t^\ell \mid \mathbf{x}\_0^\ell) = \operatorname{Cat}(\alpha_t \mathbf{x}\_0^\ell + (1 - \alpha_t) \mathbf{m}).$$

In contrast, SPMDM's forward process operates at the subsequence level:

$$q_{\mathbf{t}\_k|0}(\mathbf{x}\_{t_k}^k \mid \mathbf{x}\_0^k) = \prod_{\ell}^L q_{t_k|0}(\mathbf{x}\_{t_k}^{k,\ell} | \mathbf{x}\_{0}^{k,\ell}), \ q_{t_k|0}(\mathbf{x}\_{t_k}^{k,\ell} \mid \mathbf{x}\_0^{k,\ell}) = \operatorname{Cat}(\alpha_{t_k} \mathbf{x}\_0^{k,\ell} + (1 - \alpha_{t_k}) \mathbf{m}).$$

This formulation introduces heterogeneous noise levels across subsequences within a single training sample.Meanwhile, the changes in the forward process also lead to modifications in the reverse process and the training objective (NELBO). We provide the detailed derivation and proof in Section 3.2 and Appendix B for your reference.

About the detailed derivation . For Equation (4), the original process is $q_{t|0}(\mathbf{x}\_t^\ell \mid \mathbf{x}\_0^\ell) = \operatorname{Cat}(\alpha_t \mathbf{x}\_0^\ell + (1 - \alpha_t) \mathbf{m}),$ and our modification is introduces heterogeneous noise levels across subsequences, which leads to $q_{t_k|0}(\mathbf{x}\_{t_k}^{k,\ell} \mid \mathbf{x}\_0^{k,\ell}) = \operatorname{Cat}(\alpha_{t_k} \mathbf{x}\_0^{k,\ell} + (1 - \alpha_{t_k}) \mathbf{m}).$

Following previous work, we derive the reverse process Equation (5) corresponding to the forward process by applying Bayes’ theorem.

The modification above results in a different training objective (NELBO), as shown in Equation (6). A detailed derivation is provided in Appendix B. Since the derivation is somewhat complex, we may not have covered all the details. Please let us know if any parts are unclear.

Here we provide a brief explanation of how our model captures inter-subsequence relationships. As reflected in our training loss:

$$\mathcal{L}\_{\text{NELBO}} = \sum_{k=1}^{K} \mathbb{E}\_{t_k \sim [0,1]} \mathbb{E}\_{q} \left[ \frac{\alpha_{t_k}'}{1 - \alpha_{t_k}} \log p_\theta(\mathbf{x}^k \mid \mathbf{x}\_{t_k}^k, \mathbf{x}\_\mathbf{t}^{-k}) \right],$$

We shift the modeling unit from individual tokens to subsequences, and assign different noise scales to each subsequence. This design encourages the model to utilize information from other subsequences with lower noise levels, thereby facilitating the learning of inter-subsequence dependencies.

Q3. Computational Overhead

Thank you for raising the question regarding computational overhead. As discussed in Section 3.3, the potential computational overhead of our method compared to MDLM arises from two sources: (1) managing multiple subsequence-specific noise schedules (i.e. step 3 in Algo.1, step 5 in Algo.2), and (2) reintroducing the timestep embedding module (i.e. step 6 in Algo.2). Please kindly refer to the updated version of Algo.2 in our response to reviewer kRCe for clarification.

Regarding the first point, managing separate noise schedules requires only basic matrix operations in PyTorch, which execute in constant time. Consequently, this overhead is negligible in practice.

For the second point, we found that the official MDLM implementation does not actually remove the timestep embedder—it simply feeds it an all-zero input. Therefore, our reintroduction of timestep embedding incurs no additional computational cost compared to MDLM.

To empirically validate SPMDM's efficiency, we conducted speed benchmarks on the Countdown dataset using 12K examples (batch size 32) on a single RTX 4090:

For sampling (i.e., Algo.2), we measured the generation time for a single sample with 32 sampling steps. The results are reported in the table below:

These results confirm that our method maintains computational efficiency comparable to baseline approaches while achieving superior performance.

Q4. Open-ended Language Modeling

As suggested by the reviewer, we now add an additional experiment on open-ended language modeling using the Text8 dataset (a small, character-level language modeling task).

We follow [1] for network hyperparameters and dataset splits and compare with methods that employ a similar model size. Given the time constraints of the rebuttal period, we were limited to evaluating on this smaller dataset. The results are presented below (baseline performances are from [1]):

The results show that SPMDM achieves performance improvements on plain text generation, though the gains are relatively small. We attribute this to the nature of the Text8 dataset. Unlike reasoning tasks or structured datasets, Text8 lacks explicit inter-subsequence relationships that SPMDM is designed to exploit.

Furthermore, uniform segmentation may inadvertently disrupt natural linguistic dependencies in Text8. This observation reinforces the importance of semantic or syntactic-based segmentation strategies mentioned in your earlier comment (Q1).

Q5. Effectiveness of Adaptive Sampling

Thank you for your thoughtful consideration. As discussed in Section 3.1, the subsequence-specific noise schedule in SPMDM is designed to guide the sampling trajectory to better align with the simple path characteristics. Similarly, adaptive sampling (including both top-K probability sampling and top-K probability margin sampling) aims to unmask the easier and/or more confident tokens first, and thus also encourages the model to follow simpler paths as demonstrated in Figure 1. In other words, both the subsequence-specific noise schedules and the adaptive sampling strategy are motivated by the same underlying principle. Therefore, combining the two can enhance the model's overall performance, as shown in figure 3(b).

References

[1] Simple Guidance Mechanisms for Discrete Diffusion Models

W1, W2, Q1. Optimal Subsequence Length

Thank you for your insightful comments.

For datasets with clear structural patterns (e.g., Sudoku), we leveraged the inherent structure for segmentation. Specifically, since Sudoku inputs are 9×9 grids, we used the row length (i.e., 9) as the basis for subsequence segmentation. Here, the question length is 81 (9×9), and we segment the concatenated sequence of question and solution using a subsequence length of 81. As noted in Section 4.4, when the subsequence length is set to 162 (i.e., 81 + 81), the model effectively reduces to MDLM.

For datasets without obvious structure (e.g., Countdown, GSM8K), we used a fixed subsequence length of 8. Tables 2 and 4 demonstrate that this simple choice yields strong performance across all datasets.

We acknowledge that optimal sequence length selection is important. We are currently investigating adaptive methods and will include results in the final version if they improve performance.

Q2. Variable-length Generation Scenarios

Thank you for this insightful question. You are correct that our proposed SPMDM method cannot be directly applied to variable-length generation scenarios like Block Diffusion [1].

The key difference is that Block Diffusion handles variable-length generation through autoregressive modeling of inter-block dependencies, which SPMDM currently lacks. However, as shown in Tables 1-3, our method achieves superior performance compared to Block Diffusion due to its full bidirectional modeling capability, while Block Diffusion is limited by its semi-unidirectional approach and incurs higher computational costs during training.

We acknowledge that variable-length generation is an important problem. We are actively exploring approaches to extend our method to generate high-quality variable-length sequences while maintaining the performance advantages of SPMDM.

References

[1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models

W1. Choice of Baselines

Thank you for your valuable feedback. We apologize for the misunderstanding and recognize that our sampling procedure (Algorithm 2) may not have been clearly written in the manuscript. The following shows the revised Algorithm 2. In particular, we change $p_\theta(\cdot \mid \mathbf{x}\_{t_k}^{k}, \mathbf{x}\_{\mathbf{t}}^{-k})$ to $p_\theta(\cdot \mid \mathbf{x}\_{t_k}^{k}, \mathbf{x}\_{\mathbf{t}}^{-k},t_k)$ so as to explicitly show its dependence on the time step $t_k$.

---

$$\begin{aligned} 1.\ & **Input:** \text{Network}\ p_\theta, \text{subsequence length}\ L, \text{time}\ [0, 1], \text{sampling steps}\ T, \text{ordering oracle function}\ \mathcal{F} \\\\ 2.\ & **Initialize**\ \mathbf{x}\_\mathbf{1} \sim \{\mathbf{m}\}^N,\ \mathbf{t} \gets \mathbf{1},\ \Delta t \gets 1 / T \\\\ 3.\ & **for ** n=0 \text{ to } T ** do** \\\\ 4.\ & \quad \forall k: \text{Count unmasked tokens } n_k \\\\ 5.\ & \quad \forall k: t_k \gets 1 - n_k /L \\\\ 6.\ & \quad \forall k: \hat{\mathbf{x}}\_{0}^{k} \sim p_\theta(\cdot \mid \mathbf{x}\_{t_k}^{k}, \mathbf{x}\_{\mathbf{t}}^{-k}, t_k) \\\\ 7.\ & \quad **if** \text{ using adaptive sampling strategy } **then** \\\\ 8.\ & \quad \quad \text{Sample a set of masked token indices } \mathcal{S} = \mathcal{F}(\theta, \mathbf{x}\_\mathbf{t}) \\\\ 9.\ & \quad \quad \forall (i, \ell) \in \mathcal{S}: \mathbf{x}\_\mathbf{t}^{i, \ell} = \hat{\mathbf{x}}\_\mathbf{0}^{i, \ell} \\\\ 10.\ & \quad **else** \\\\ 11.\ & \quad \quad \forall k: s_k \gets \max(t_k - \Delta t, 0) \\\\ 12.\ & \quad \quad \forall k: \text{For all masked tokens, with probability } \frac{s_k}{t_k}, \hat{\mathbf{x}}\_\mathbf{0}^{k, \ell} \gets \mathbf{m} \\\\ 13.\ & \quad \quad \text{Update } \mathbf{x}\_\mathbf{t} \gets \hat{\mathbf{x}}\_\mathbf{0} \\\\ 14.\ & \quad **end if** \\\\ 15.\ & **end for** \\\\ 16.\ & **Return ** \mathbf{x}\_\mathbf{t} \end{aligned}$$

---

Moreover, in the original submission, experimental results on the proposed method (denoted "Ours" in section 4.4) use steps 11-13 (note that these steps are directly inherited from the MDLM [1]), while the experimental results on the method denoted "Ours + Adaptive" in Figure 3(b) use steps 7-9 with top-K probability sampling. In other words, the commonly-used adaptive sampling strategy mentioned by the reviewer has already been used as a baseline.

As suggested by the reviewer, we now conduct additional experiment with top-K probability margin sampling, using the setup as in Section 4.4 (Figure 3(b)). For easier comparison, we convert Figure 3(b) to the table below, and add in the results on top-K probability margin sampling (note that 'Ours + Adaptive' in figure 3(b) now corresponds to the entry 'Ours + Top-K Probability Sampling').

As can be seen, while 'Ours + Top-K probability margin sampling' achieves the best performance, the improvement over 'Ours + Top-K probability sampling' is small. We will add this new comparison and further clarifications in the final version.

Q1. Training-free Setting Performance

Thank you for your question. As suggested, we also add an experiment on the training-free setting with a pretrained MDLM. However, as mentioned in the response above, $p_\theta(\cdot \mid \mathbf{x}\_{t_k}^{k}, \mathbf{x}\_{\mathbf{t}}^{-k},t_k)$ depends on the time step $t_k$. However, in the standard MDLM, time step $t_k$ is NOT used in the training. Hence, in order to be used on a pretrained MDLM, we do not feed the time step into $p_\theta$ (of the pretrained MDLM). The algorithm (modified from the algorithm above) is shown in the following:

---

$$\begin{aligned} 1.\ & **Input: ** \text{Network}\ p_\theta, \text{subsequence length}\ L, \text{time}\ [0, 1], \text{sampling steps}\ T \\\\ 2.\ & **Initialize**\ \mathbf{x}\_\mathbf{1} \sim \{\mathbf{m}\}^N,\ \mathbf{t} \gets \mathbf{1},\ \Delta t \gets 1 / T \\\\ 3.\ & **for ** n=0 \text{ to } T ** do** \\\\ 4.\ & \quad \hat{\mathbf{x}}\_{\mathbf{0}} \sim p_\theta(\cdot \mid \mathbf{x}\_{\mathbf{t}}) \\\\ 5.\ & \quad \forall k: \text{Count unmasked tokens } n_k \\\\ 6.\ & \quad \forall k: t_k \gets 1 - n_k /L \\\\ 7.\ & \quad \forall k: s_k \gets \max(t_k - \Delta t, 0) \\\\ 8.\ & \quad \forall k: \text{For all masked tokens, with probability } \frac{s_k}{t_k}, \hat{\mathbf{x}}\_\mathbf{0}^{k, \ell} \gets \mathbf{m} \\\\ 9.\ & \quad \text{Update } \mathbf{x}\_\mathbf{t} \gets \hat{\mathbf{x}}\_\mathbf{0} \\\\ 10.\ & **end for** \\\\ 11.\ & **Return ** \mathbf{x}\_\mathbf{t} \end{aligned}$$

---

The following table shows the testing on the Countdown dataset. For comparison, we also include the proposed SPMDM below.

As can be seen, while the proposed subsequence-based sampling method can be applied on a pretrained MDLM and yields performance improvements, the gains are significantly smaller than those obtained by the proposed SPMDM (which modifies both both the training and sampling phases). The results and discussion here will be included in the final version.

References:

[1] Simple and Effective Masked Diffusion Language Models

Q1. How does SPMDM perform in comparison with block diffusion methods?

Thank you for your thoughtful comments. The key difference is: Block Diffusion (BDLM) [1] methods use autoregressive modeling, while SPMDM employs different noise scheduling across subsequences. BDLM is semi-unidirectional rather than fully bidirectional. While this may benefit general text modeling, it negatively impacts reasoning tasks that require bidirectional context understanding.

Additionally, BDLM's hybrid diffusion-AR approach incurs higher training costs, whereas SPMDM adds minimal training overhead while achieving superior performance on reasoning tasks.

As for the empirical performance, the following is a summary of the observations from the real-world tasks. On Countdown tasks (Table 2), SPMDM achieves 0.9%, 7.2%, and 7.9% higher accuracy than BDLM on CD-3, CD-4, and CD-5 respectively. On Sudoku (Table 3), SPMDM surpasses BDLM by 8.1% accuracy.

We will add the above discussion and more detailed comparison with BDLM in the final version.

Q2. For planning tasks, the priority is to learn earlier reasoning tokens. Does this learning order also appear in SPMDM's inter-subsequence relationships?

Thank you for your question. Yes, this learning order do appear in SPMDM's inter-subsequence relationships. To demonstrate how SPMDM captures inter-subsequence relationships through its denoising order, we now add the following experiment on the Countdown task.

First, we pick an example from the Countdown-3 test sets and show the step-by-step denoising process of SPMDM as follows:

1,82,8 ### 73 ### [MASK]2-[MASK][MASK][MASK][MASK][MASK][MASK]1[MASK][MASK]=[MASK][MASK]
1,82,8 ### 73 ### [MASK]2-1[MASK][MASK][MASK][MASK][MASK]1[MASK][MASK]=[MASK][MASK]
1,82,8 ### 73 ### 82-1[MASK]8[MASK][MASK][MASK]1[MASK]8=[MASK][MASK]
1,82,8 ### 73 ### 82-1[MASK]8[MASK][MASK][MASK]1[MASK]8=7[MASK]
1,82,8 ### 73 ### 82-1[MASK]81[MASK][MASK]1[MASK]8=7[MASK]
1,82,8 ### 73 ### 82-1=81[MASK][MASK]1[MASK]8=7[MASK]
1,82,8 ### 73 ### 82-1=81,[MASK]1-8=7[MASK]
1,82,8 ### 73 ### 82-1=81,[MASK]1-8=73
1,82,8 ### 73 ### 82-1=81,81-8=73

As can be seen, SPMDM first decodes the first question, followed by the second equation.

In addition, we conduct a quantitative experiment on the Countdown test sets. Specifically, following the experimental setup in Section 4.2, we set the sequence length to 48 and the subsequence length to 8, using the CD-3 problem as the prompt input. We calculate the average denoising timestep for tokens within each subsequence throughout the generation process.

We set the number of sampling steps to 12. The experimental results are presented in the table below, where Index indicates the subsequence index, and Avg Unmask Steps refers to the average denoising timestep of tokens within each subsequence:

Note that, based on the current partition with a subsequence length of 8, the first two subsequences are prompts. As a result, for sequence index 0 and 1, the average unmask steps are always zero.

We observe that SPMDM prioritizes denoising the first equation (i.e, sequence index 2). This behavior is consistent with construction process of Countdown ground-truth solutions. The process begins by solving the first equation, which helps determine the valid factors that can be used in subsequent equations. The last two subsequences mainly consist of padding tokens, and thus do not exhibit any significant differences in denoising order.

Therefore, we believe that SPMDM successfully captures the underlying dependencies between subsequences.

Q3. How does the preferred order differ for different tasks?

To illustrate this, we perform the experiment in Q2 above on the Sudoku task. Recall that each row in the Sudoku question is treated as a subsequence. We set the number of sampling steps to 18, and the results are shown in the table below.

As can be seen, the average denoising timestep of the subsequences exhibits a uniform distribution, which is different from that on the Countdown tasks we saw above. This is because the rows in the Sudoku question vary in difficulty, thus prioritizing the denoising of simpler rows (subsequences) leads to a simpler overall sampling path. Given a sufficiently large amount of data, the difficulty distribution across rows is likely to be approximately uniform.

Overall, combining with the observations from Q2 above, we find that SPMDM prefers prioritizing the denoising of subsequences that are easier to complete.

References

[1] Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Models