Generalized Interpolating Discrete Diffusion

We thank the reviewer for their feedback and insightful questions. We especially feel that 2) deserves careful discussion since we consider the theoretical contributions a key part of our work. We would like to encourage the reviewer to share some additional detail on their concerns.

1. We agree with the reviewer that a discussion of [4] needs to be added. While [4], and also [2] (App. A.2.6), experiment with a BERT-like combination of masking + uniform noise, these approaches are only discrete-time and do not investigate what improvements/challenges arise from it.
2. We agree that there are some similarities between our theoretical contributions and the suggested references [1, 2] and would like to highlight the important differences. The theory behind GIDD is a strict generalization of masked diffusion (MD) in [1] (see Cor. 3.8). Unlike MD, GIDD allows for a time-variable mixing distribution $\pi_t$, which makes MD a special case where $\pi_t=\mathbf{m}$ is constant w.r.t. time. Another example would be setting $\pi_t=\mathbf{1}/|V|$ where $|V|$ is the vocabulary size, which simplifies to uniform diffusion from [2]. As for [2], this paper covers the most general case of discrete diffusion and forms the foundation of much subsequent work, including [1] and ours. In this framework (and in discrete diffusion in general), complete knowledge of the transition matrix is paramount and necessary both for computing the ELBO as well as for sampling. In this regard, our theoretical contribution is to solve the inverse problem of finding the (closed-form) Markovian transitions given only the marginal forward transitions for a special family of linearly interpolating diffusion processes (see Prop. 3.3). In addition, while [2] only covers discrete-time, the GIDD ELBO and Markov transitions are stated in continuous-time by applying the tools from [3]. We hope that this sufficiently highlights our theoretical contributions and that we were able to address the concerns of the reviewer.
3. We agree with the reviewer that the confidence-weighting in the self-correction step is similar to the MaskGIT sampling algorithm. However, ours is a fixed-point iteration that resamples one token at a time and is only applied after denoising, whereas MaskGIT sampling is an adaptive masked diffusion sampler prioritizing high-confidence tokens. Indeed, applying MaskGIT sampling to GIDD ($p_u=0.0$; base) finds that this is prone to collapse, with generated samples being low-diversity and consisting mostly of repeated tokens. This manifests in a low generative entropy (via Gemma-2-9B) of 1.31 compared to the baseline of 3.13. MaskGIT sampling is unfortunately not applicable to $p_u>0$ models since those require token replacements in addition to unmasking. We thank the reviewer for pointing out this similarity and will discuss it in an updated version.
4. Preliminary experiments showed that setting the ELBO weights to 1 works for $p_u=0.0$ but is challenging for $p_u>0$, and in light of the clipping strategy working much better, we did not investigate this any further. We hypothesize that the weight ratio between uniform and noise-free tokens is important for learning the correct posterior: What is the distribution of the correct token, given that it is incorrect with only a small probability?
5. Regarding overfitting, we would like to point out that for $p_u=0.0$, MDM and GIDD are equivalent, so it is likely that if GIDD is overfitting, then the baselines are also overfitting. One potential cause could be the way long sequences are handled in our data loader (random cropping instead of splitting), which leads to more frequent repetition of short sequences, potentially leading to overfitting. The 1.1B MDM baseline is included purely for context as training 1.1B GIDD models was unfortunately not feasible given the available resources.
6. Discouraged by the higher loss and PPL of $p_u>0$ models, we did not experiment much with noise levels above $p_u=0.2$. Generally speaking, higher values are more challenging as the SNR drop gets more concentrated on early steps.
7. Unfortunately, we are not entirely sure which entropy the reviewer is referring to. Since PPL is just the exponent of cross-entropy, the comparison would not qualitatively change by reporting cross-entropy.
8. Since self-correction is run for at most 128 iterations (subject to early stopping) and the samples are generated in 128 denoising steps, the worst-case computational overhead is double. However, even if self-correction incurs some overhead, it is generally possible to improve sample quality while keeping the inference compute budget constant by reducing the number of denoising steps in favor of additional post-hoc correction steps.

• [1] Shi et al., 2024. https://arxiv.org/abs/2406.04329
• [2] Austin et al., 2021. https://arxiv.org/abs/2107.03006
• [3] Campbell et al., 2022. https://arxiv.org/abs/2205.14987
• [4] Gu et al., 2021. https://arxiv.org/abs/2111.14822

We thank the reviewer for their insightful review and constructive feedback. In the following, we would like to respond to the reviewer’s comments and suggestions and provide additional experimental results to further bolster the claim of self-correction in the presented models.

1. We mostly agree with the reviewer on the comment regarding overfitting. The difference is indeed small, and performance is almost random on some datasets. Given that the difference is consistent, albeit small, perhaps the appropriate characterization would be to say that it is likely run-to-run variance with some likelihood of overfitting. As mentioned in our response to Reviewer LeK6, we see a potential cause of overfitting in the way long sequences are handled in our data loader, which leads to more frequent repetition of short sequences, potentially causing overfitting.
2. We agree with the reviewer that this would be a good addition and would help quantify the diversity loss during self-correction. Given the mode-seeking nature of the self-correction step, some decrease in diversity is to be expected, but this should be restricted to a reasonable amount. Indeed, we find a decrease in entropy for $p_u>0$ models, with the decrease correlating linearly with the number of changed tokens (see Table 1 below). However, the decrease is moderate and does not indicate any collapse. This is also supported by the LLM-based evaluation (see 3.). For further context, we also provide qualitative self-correction examples in Appendix K, which give some intuition on the nature and extent of the reduced diversity.
3. Unfortunately, conducting a user study comes with its own host of challenges and is beyond the scope of this project. However, we have conducted an LLM-based evaluation of the quality of generated samples via GPT-4o using single-answer grading [1] (prompt omitted due to char. limit). The samples are graded on a scale from 1 to 10 in terms of clarity, grammaticality, factuality, writing style, and creativity. We would like to emphasize that these absolute numbers are highly dependent on the judge model’s calibration and should therefore be taken with a grain of salt. Nevertheless, we find consistent improvements at high significance levels, with $p_u=0.2$ exhibiting both the largest improvement and the highest scores overall (see Table 2 below). We report the self-correction setting with the largest effect for each noise level.
4. Given that our models are comparatively small and not trained on instruction following or conditional generation, evaluation on GSM8k is unfortunately not possible. Even if we try to do pseudo-conditional generation by forcing the logits of prompt tokens, the model does not generate meaningful continuations.

Regarding the reviewer’s questions:

1. The “number of tokens changed” refers to the number of tokens that differs from the initial sequence after convergence/termination of the self-correction step. While this roughly correlates with the number of inference steps, it is not a one-to-one correspondence since self-correction may, at times, oscillate between two or more similar states (e.g. multiple options that are equally correct/incorrect). We therefore deem the number of changed tokens after convergence to be a more meaningful metric compared to simply the number of self-correction steps.
2. Indeed, due to the continuous-time nature of the GIDD diffusion process, the number of denoising steps can be set arbitrarily high. Due to our chosen parameterization, where the model predicts the fully noise-free data which is then re-noised up to the appropriate level, uniform noise is actually continually injected at low levels, so tokens continue to change throughout. Future work can explore different approaches that avoid this, which may lead to improvements.

[1] Zheng et al., 2023. https://arxiv.org/pdf/2306.05685

---

Table 1

Table 2

Significance levels: * $>2\sigma$ difference, ** $>5\sigma$ difference.

We thank the reviewer for their insightful and detailed review and for acknowledging the theoretical contributions of GIDD. In the following, we would like to shed some additional light on the empirical evaluation of the proposed model and, hopefully, address the reviewer’s concerns.

Broadly speaking, our evaluation (and evaluation of language models in general) consists of two parts. The first part is likelihood-based evaluation, which tests the model’s ability to recognize well-formed, high-likelihood sentences, and includes PPL on the test set as well as downstream multiple-choice benchmark performance. The second part is to test the model’s ability to generate well-formed, high-likelihood sentences and for unconditional generation, this is most commonly measured with generative-PPL. This ability is equally, if not more, important because it is how language models are usually used in practice. And while the two abilities may correlate in general, sampling from the model can sometimes introduce challenges not present when simply evaluating likelihood. For example, teacher-forcing is used both when training autoregressive models and when evaluating their likelihood, but not when generating samples, which makes them prone to self-induced mistakes [1]. The same holds true for mask-only diffusion models. The point of our experiments is to show that while, indeed, hybrid diffusion models find it more challenging to evaluate the likelihood of given samples, they have self-correction capabilities and are more robust/less prone to self-induced errors at sampling time, thus generating higher-quality samples overall and especially so for low inference-compute budgets (see Sec. 5.4, L435 ff.; Fig 6, App. D). To further bolster this point, we provide an additional LLM-based evaluation in our response to Reviewer wojY, where we quantify the improvements in sample quality during the self-correction step in terms of clarity, grammaticality, factuality, writing style, and creativity. Perhaps the distinction between recognition and generation should be highlighted and discussed more prominently in the paper, which we will be happy to do in an updated version.

Regarding the second weakness: The reason we cannot naively apply self-correction to the multiple-choice benchmarks is because these benchmarks rely on likelihood-based answer selection. Specifically, for a given question (or prompt) with a set of possible answers (or continuations), the one with the highest likelihood under the model is selected. Since this process does not involve sampling from the model, it is a priori not possible to utilize the self-correction capabilities. As a result, benchmark scores of $p_u > 0$ models are hampered in correlation with their worse likelihood (see Tab. 5, App. B). Future work may aim to close the gap by using different criteria for selecting the correct answers (e.g. self-accuracy instead of likelihood), but unfortunately, this was beyond the scope of this project. Alternatively, future work may extend the proposed models to conditional generation, which will allow evaluation on generative benchmarks like GSM8k or HumanEval, where the generative strength and self-correction capabilities of hybrid models may bring substantial improvements over mask-only diffusion models.

[1] Bachmann & Nagarajan, 2024. https://arxiv.org/abs/2403.06963