Thank you for the positive feedback! We are glad that you find our contribution notable, our methodology thorough and our experimental results strong and robust. We address each comment below

Detailed pseudocode and specific implementation challenges

We thank the reviewer for the suggestion. Please refer to our response to questions shared by reviewers for pseudocode of training and sampling algorithms. We will add them in the final version. We have discussed implementation challenges in Appendix A of the paper. Please let us know if you have any additional questions about the implementation.

The range of tasks and datasets could be expanded, such as VQ-Diffusion [1] for token-based text-to-image.

We thank the reviewer for the suggestion and will pursue this in future work. It's important to note that our pixel-level image modeling task is more challenging than modeling the latent space of vector quantized image encodings. While some discrete diffusion models have tackled high-resolution generation using the latter approach, to our knowledge, we're the first to go beyond 32x32 resolution for pixel-level modeling using discrete diffusion.

Sample metrics for image generation

In our rebuttal pdf we added an FID evaluation using 50K randomly generated samples from our MD4 models on ImageNet 64x64. We compared different alpha schedules (Linear v.s. Cosine) and show that the cosine schedule leads to significant better sample quality. We further trained a class-conditional model on ImageNet 64x64 that boosts the FID to around 7. Although these are not state-of-the-art FIDs on ImageNet 64x64, we emphasize our models are optimized for likelihood using ELBO instead of sample quality, unlike the reweighted objectives used in continuous diffusion.

Thank you for the time and the feedback! We clarify a few points.

D3PM results & differences between MD4 and D3PM

• The reviewer is concerned if the D3PM results for zero-shot transfer tasks are comparable to MD4. We clarify that these D3PM results are from the SEDD paper (Lou et al., ICML 2024) which used the same setup as SEDD, including an absorbing state diffusion. This paper states: "... baselines are ... D3PM (with Absorbing Transition) … We retrain both models … reuse our model architecture and match hyperparameters (i.e. model size, training specifications)." To ensure direct comparability with MD4, we reimplemented SEDD and reproducing their results (see "SEDD Absorb (reimpl.)" in Table 1). This ensures a fair comparison between the models.

• The reviewer states that MD4's simplified ELBO form shouldn't provide benefits compared to D3PM as both are ELBOs. However, we emphasize that this simplification is crucial for a well-engineered, numerically stable ELBO implementation, leading to significantly better performance. To illustrate, we analyze D3PM's original implementation of the $KL(q(x_s|x_t, x_0)||p(x_s|x_t))$ terms in the ELBO (s = t - 1 for D3PM) without simplification:

  1. Sample $x_t \sim q(x_t|x_0)$
  2. Calculate $q_{s|0}(\\cdot|x_0)$ where $\cdot$ takes all values in the vocabulary.
  3. Calculate $q_{t|s}(x_t|\\cdot)$.
  4. Calculate the logits of $q(x_s|x_t, x_0)$ via $qlogits = \\log (q_{s|0}(\\cdot|x_0) + \\epsilon) + \\log (q_{t|s}(x_t|\\cdot) + \\epsilon)$.
  5. Repeat 2-4, replacing $x_0$ with NN output probabilities to get $p(x_s|x_t)$ logits ($plogits$).
  6. Compute $KL(q(x_s|x_t, x_0)\\|p(x_s|x_t))$ using log_q = log_softmax(plogits) and log_p = log_softmax(qlogits).

For masked diffusion, as shown by our simplification, many elements of $q(x_s|x_t, x_0)$ and $p(x_s|x_t)$ are zeros (see eq. (40) in App.). These terms can be removed but are instead included by D3PM. Thus, D3PM has to introduce an $\\epsilon$ to avoid inf - inf which gives NaNs in the KL computation. We perform an experiment in CIFAR10 to show the numerical issues: we replaced the MD4 objective in our code with D3PM’s ELBO and tried different $\\epsilon$s (1e-20, 1e-12, 1e-8, 1e-6). Only after increasing it to 1e-6 we avoided NaNs at the start of training. Still, we got NaNs after 300k iterations. The test BPD before NaNs was 3.03, while MD4 at this training iteration reports 2.89.

• We agree with the reviewer that, similar to VDM, the cont-time v.s. discrete-time difference will not make a huge difference in likelihood. However, our cont-time formulation remains a key contribution for these reasons:
  • It specifies the diffusion model using a simple monotonic function $\\alpha_t$ representing unmasking ratio at t. This enables exploration of schedules not used in discrete-time models like D3PM where the model is specified with less intuitive jump probabilities $\\beta_i$ rather than masking ratios. One of our key findings is that the cosine alpha_t schedule produces the best sample quality. In our rebuttal pdf, we include a comparison between the widely used linear schedule and the cosine schedule, showing the latter leads to significantly better FIDs in ImageNet 64x64.
  • It enables easy adoption of training techniques from Gaussian diffusions; e.g. the antithetic sampling used by VDM for estimating the time integral in the ELBO which reduces training variance—on CIFAR-10, this translates to ~0.02 improvement on BPD.

CIFAR10 comparison with AR

The reviewer states that “The AR baselines use old transformer models hence the comparison isn't fair.” We have included the best AR results we found for CIFAR10. We're open to comparing with newer models if the reviewer can provide references. Also note our reported MD4 result of 2.78 is not fully optimized, e.g., increasing the batch size to 256 improves the result to 2.75 using the same 20M parameter model.

CIFAR10 comparison with continuous diffusion

The reviewer states that “Current SOTA diffusion models ... achieve a NLL of 2.55 [2] which is far better than absorbing diffusion” We believe the 2.55 from [2] (which we assume refers to CIFAR-10, not ImageNet 32) isn't directly comparable to our results in Table 3 for these reasons:

• The 2.55 in [2] estimates exact log likelihood with probability flow ODE, while our results are upper bounds.
• The best variational bound result in [2] is 2.65, which relies on learnable schedules using NNs, while MD4 uses a fixed schedule.
• All continuous diffusion results in [2] Table 2 with BPD below 2.8 use extra image Fourier features (FFs) inputs from VDM. Comparing these with discrete diffusion is unfair since the latter assumes order-agnostic image data, i.e., the model is unaware of the proximity between pixels.

Also the impact of FFs becomes negligible at larger scale such as ImageNet 64x64. We have found that MD4 on ImageNet 64x64, by reducing the dropout rate to 0.0, gives 3.40 BPD which is the same as VDM relying on FFs.

We will revise the imprecise statement in the conclusion as "On text data, our masked diffusions outperform existing discrete and continuous diffusion models. For pixel-level image modeling, we significantly improve discrete diffusion results, outperforming similar-sized AR models and achieving comparable likelihoods to continuous diffusion models without Fourier features."

GenMD4 on zero-shot

We clarify that we did not state GenMD4 performs poorly on zero-shot tasks. In fact, the results there are mixed. We have included both validation PPL numbers on OWT and zero-shot PPL numbers for GenMD4 in our response to shared questions by reviewers.

Dropout rate for OWT

Following the reviewer’s suggestion, we conducted dropout rate tuning and found that lower rate is better: dropout rate 0.05 has eval PPL 24.63, dropout rate=0.02, eval PPL is 22.13, and dropout rate = 0.0, eval PPL is 21.86. We’ll update the results in the final version.

Thank you for the time you’ve taken to review our work and for the positive and constructive feedback! We are glad that you find our work "offers a novel theoretical formulation" and "achieve state-of-the-art performance" with " comprehensive experimental validation". We address each individual comment below.

Practical challenges in implementation of state-dependent masking schedules

• Regarding computational cost, the state-dependent masking schedules and corresponding GenMD4 objective require only twice the computation of MD4 with the same network size. While slightly more expensive, this increase remains within affordable limits.
• Regarding sampling, we clarify that x_T is a full mask state even for the state-dependent schedule case. This is because the learned schedule $\\alpha_t = 1 - t^w$ still satisfies $\\alpha_1 = 0$ regardless of the value of $w>0$. We also clarify that the sampling process is stochastic in each step, please refer to the sampling algorithm described in our response to questions shared by reviewers.

Effectiveness of state-dependent models in zero-shot transfer tasks

We have included the results of GenMD4 on zero-shot transfer tasks in our response to shared questions. The results are mixed: on some datasets (Lambada, PTB) GenMD4 results in significantly better zero-shot PPL numbers than MD4 while on other datasets (WikiText) GenMD4 is slightly worse. We believe this indicates some degree of overfitting to the dataset statistics of OpenWebText.

Comparative analysis regarding computational efficiency and training times

We follow the reviewer’s suggestion to compare the training times of MD4/GenMD4 on OpenWebText using 128 TPUv3. We can see that MD4 is as fast as our reimplementation of SEDD while achieving better results. GenMD4 is slightly slower than MD4 due to the 2x computational cost.

• Gaussian Diffusion: 3.5 steps / s
• SEDD: 4.2 steps / s
• MD4: 4.2 steps / s
• GenMD4: 3.5 steps / s

How the model ensures forward/reverse consistency and impact on training stability v.s. SEDD

For simplicity let’s consider the single-dimension case. As we have shown in Proposition 1, the true score function for a mask state $x_t=m$ has the following relationship with the conditional mean of $x_0$:

$$s(m, t)\_j = \\frac{\\alpha_t}{1 - \\alpha_t} E[x_{0,j}|x_t=m]$$

Note that here $x_0$ is a one-hot representation of the data therefore $\\sum_j x_{0,j} = 1$ and thus also $\\sum_j E[x_{0,j}|x_t=m] = 1$. This implies that

$$\\sum_j s(m, t)_j = \\frac{\\alpha_t}{1 - \\alpha_t}.$$

Since the transition rate matrix of the true reverse process depends on the score, the above equation implies a constraint that the reverse process has to satisfy in order to be consistent with the $\\alpha_t$-determined forward process.

As we have explained in Sec. 4 and App. F3, we can interpret both our method and SEDD as learning a model $s_{\\theta}$ for the true score function $s$. The differences are

• In MD4, $s_{\\theta}$ is parameterized as

$$s_{\\theta}(m, t)\_j = \\frac{\\alpha_t}{1 - \\alpha_t} \\mu_{\\theta}(m)\_j$$

where $\\mu_{\\theta}$ has a softmax output that produces a probability vector. Therefore, the score model also satisfies $\\sum_j s_{\\theta}(m, t)\_j = \\frac{\\alpha_t}{1 - \\alpha_t}\\sum_j \\mu_{\\theta}(m)\_j = \\frac{\\alpha_t}{1 - \\alpha_t}$.

• In SEDD, $s_{\\theta}$ is a free-form neural network model that outputs a real-valued vector, which has no guarantee it satisfies the constraint. This inconsistency between forward and reverse process leads to large variational gaps in the ELBO which we have empirically shown to lead to unstable training (see Figure 3).

Detailed description of the training and sampling algorithms

Please refer to our response above to questions shared by reviewers.

Sensitivity to hyperparameter choices and multiple runs

We test the sensitivity to multiple runs by re-training MD4-S with the same hyperparameter as the original one reported in the paper, where the original PPL on OpenWebText eval split is 22.126 and the re-trained MD4-S gets 22.134. In the pdf we uploaded for rebuttal, we include an analysis of the impact of masking schedule choice on sample quality, showing the benefits of a cosine schedule.

Thank you for the time you’ve taken to review our work and for the positive and constructive feedback! We are glad that you found our paper "offers a valuable approach to optimize the forward process”, “organizes prior work into a cohesive structure”, and "may serve as a source of inspiration" for future work. We respond to each of your comments below.

The handling of 𝑝(𝑥0|𝑥𝑡(1)) could be enhanced

Thank you for the suggestion. In our implementation t(1) is tiny ($\\alpha_{t(1)} = 10^{-4}$) and we expect p(x0|xt(1)) to be very close to q(xt(1)|x0). We have also tried treating p(x0|xt(1)) the same as other p(xs|xt) and did not observe significant differences.

Multidimensional data and backward factorization

Thank you for the suggestion. We agree with the reviewer that the true backward distribution becomes factorized only as $s \\to t$, which is why we adopt the continuous-time formulation throughout the work. We will follow the reviewer’s suggestion to make this more clear in the final version.

Non-zero $\\alpha_1$

We originally introduced the non-zero $\\alpha_1$ to resolve numerical issues in GenMD4, where the REINFORCE LOO gradient requires calculating $\\log \\alpha_t$. We have since then used the same schedule for MD4 as well. We agree with the reviewer that the non-zero $\\alpha_1$ is unnecessary for MD4. Although we have not found any “medium brightness” problems, we will still rerun the MD4 experiment with $\\alpha_1 = 0$ and update the results.

Appendix G2: Variance of Campbell et al. (2022)

We thank the reviewer for noting this. In the single dimension case, we agree that there is no additional variance via further simplification of the Campbell et al. bound. The additional variance comes from multidimensional data where the sum over $j$ becomes the sum over all neighbors in the forward process, i.e., the states that mask out a single dimension of $x_t$ that has not been masked yet. In this case, the Campbell et al. (2022) bound simplifies to

\\beta(t) \\sum_{i: x_{t, i} \\neq m} \\log (x_{t, i}^\\top \\mu(x_t(i\\to m))_d

where $x_t(i\\to m)$ denotes the state we get by masking out the i-th dimension of $x_t$. This equation requires $d$ function evaluations of the prediction model $\\mu$, thereby incurring additional variance if we estimate the sum with Monte Carlo. We will fix the discussion of Campbell et al. (2022) in the final version.

𝜏LDR with masked noise

We followed the reviewer’s suggestion to perform an additional experiment using 𝜏LDR with masked noise. We made sure the neural networks and training hyperparameters used by 𝜏LDR (masked) is the same with our MD4 experiment. The results are as follows (for comparison we also included the MD4 results)

The results show that 𝜏LDR with masked noise suffers from the high-variance objective even with the same networks and training hyperparameters from MD4.

Details of the sampling method

Please refer to our response to questions shared by reviewers.

Results of GenMD4 on image datasets

We ran GenMD4 on both image datasets with the same architecture and hyperparameters as in MD4. On CIFAR-10, we get 2.7749 (GenMD4) v.s. 2.7847 (MD4). On ImageNet 64x64, we get 3.4233 (GenMD4) v.s. 3.4273 (MD4). We see consistent improvements in test BPD over MD4 throughout training although the improvement is smaller than that of text experiments (1.34 v.s. 1.37 BPC on text8 & 21.80 v.s. 22.13 PPL on OpenWebText validation set).

Simplified GenMD4

We followed the reviewer’s suggestion to perform an additional experiment on text8 that only learns a scalar w for GenMD4 that is shared across all schedules. The result is 1.37 BPD which is almost the same as the result of MD4 and worse than GenMD4 with a vector w (1.34 BPD). We believe this is because the simplified GenMD4 reduces to MD4 with a learnable schedule (we can show this by choosing $\\alpha_t$ in (19) as a scalar schedule times an all-one vector) and MD4 is invariant to masking schedules.