Simple Guidance Mechanisms for Discrete Diffusion Models

We address the reviewer’s concern about performance on large vocabularies with classifier-free experiments on NLP datasets which use the bert-base-uncased tokenizer (vocab. size = 30,522).

Classifier-free Guidance: Part of speech generation

We find that D-CFG with Diffusion models (MDLM and UDLM) is able to outperform AR models with large vocabularies with comparable sample quality (generative perplexity) and much higher control scores (F1).

Additionally, we see in the Table below, that only UDLM seems to respond to higher $\gamma$ tuning compared to MDLM and AR:

The normalization in our discrete guidance methods is not a computational bottleneck. Rather it amounts to simply applying the .softmax(dim=-1) to the un-normalized log probabilities (along the vocabulary dimension) and is therefore not prohibitive, even in large vocabulary regimes.

Details about this experiment: We train models on the E2E dataset, which consists of 50k restaurant reviews (Novikova et al. 2017). The goal is to generate sentences that follow a part-of-speech template, which acts as the guidance label. Quality is measured by generative perplexity and controllability is measured as the F1 score between the POS tagging of the generated sequences and that provided as the control target.

We thank the reviewer for their detailed comments. Below we respond to the specific concerns raised in detail, the high level responses are as follows:

• Comparison to Flow Matching works: UDLM and Gat et al. use different objectives (theirs is not an ELBO) and samplers; Stark et al. is significantly different from our work: it models continuous-states not discrete-states
• Comparison to Nisonoff et al.: Our guidance formula is substantially different from Nisonoff et al. (see below); both use the Taylor trick from Vignac et al. 2022
• We report experiments on large vocabularies; we also report new experiments where novelty is a concern (e.g., vs. Gat et al.)

In the table below, we summarize the key differences between UDLM and other works that examine discrete corruption with uniform noise.

Comparison to Gat et. al

The main difference is the training objective: we use a principled ELBO, while Gat et al. simply train a BERT-style denoiser. Our variational bound is novel, tight, and performant (see Table 2 and 3). Being probabilistic enables applying UDLM to language modeling, which involves a principled comparison across models (e.g., AR transformers) using perplexity. There is no comparison in Gat et al. to baselines using true perplexity (or its lower bound); they only report generative perplexity.

To underscore the difference, we re-train our model with uniform noise using the simple cross-entropy loss used in Gat et al. and find that training with the UDLM objective leads to better language modeling performance (experiment below reflects the QM9 and Species10 datasets):

Our method also uses a different sampler, although in theory one could extract a rate matrix from our model and use it with Gat’s sampler (and vice versa).

Comparison to Stark et. al

We emphasize that Stark et al. 2024 defines a flow over continuous representations of the data. This is distinct from our approach, which directly defines the corruption and denoising processes on the discrete data representations. Another difference is model parameterization: velocity in Stark et al. vs. mean parameterization in our work.

Stark et al. 2024 can be seen as an analogous work to those that perform continuous diffusion for discrete data, e.g. Han et al. 2022, but for flow models. Of note, Stark et al. 2024 also explore classifier-free and classifier-based guidance for flows but rely on the fact that the velocity fields they learn reflect the scores obtained in continuous diffusion modeling. Thus the guidance methods of Stark et al. 2024 are more similar to the previous definitions of classifier-free and classifier-based guidance from the continuous diffusion literature.

“Uniform prior is widely known and exists in the literature”

We clarify that our claim is not to have developed discrete diffusion with uniform noise, but rather the novelty comes from: (1) the optimized continuous time ELBO with which we train UDLM (Section 3.2); (2) the application of uniform noise diffusion to controllable generation.

We stress that the novelty of Equations 10 and 11 is significant in that it tightens the ELBO by taking $T \rightarrow \infty$ and allows us to avoid computing values that can be analytically set to zero (in terms of $\mathcal{L}\_{prior}$ and $\mathcal{L}\_{recons}$). The importance of this improved objective is reflected in our language modeling results, which show that UDLM outperforms previous uniform-noise discrete data methods (as shown in the bottom 3 rows of Tables 2 and 3, which we reproduce below).

Guidance methodology vs. Nisonoff et al. 2024

First, Nisonoff et al. 2024 is an unpublished manuscript that is concurrent with ours.

Second, the method of Nisonoff et al. 2024 is different from ours. The guidance formula is different, see table below. In particular, the method of Nisonoff et al. 2024 requires sophisticated CTMC theory and several pages of mathematical derivation. In contrast, our guidance mechanisms involve only elementary probability and can be defined in under a page; however, they outperform all published baselines (Diffusion+NOS, AR+FUDGE etc.).

Finally, please note that the Taylor approximation is not core of either our method nor that of Nisonoff et al. In fact, both methods cite earlier usage in Vignac et al. 2022. In both methods, it circumvents performing $L\cdot N$ passes through the classifier model, where $L$ is the sequence length and $N$ is the vocabulary size.

Dear authors,

I appreciate your response to my rebuttal response. I wish to continue to respond in good faith, and I hope the authors will understand my issue with the suggestion that Gat et al. 2024 is a concurrent work.

Very respectfully, this claim is not warranted. MDLM, MD4, and DFM are largely the same paper disguised under slightly different notation and under the guise of Flow Matching vs. Discrete Diffusion. In the discrete setting (discrete state space) there is no difference between MDLM and DFM. The end objective is a cross-entropy based denoiser---they lead to the same ELBO optimization for the masked diffusion case. The main technical difference in DFM is the path construction, which allows uniform paths. Given that you compare with MDLM which will also be published at NeurIPS 2024 this year, I maintain insufficient novelty between this work and DFM.

If you truly wish to convince me that uniform paths have substantial value then I would expect a language modeling benchmark of size LM1B or larger where UDLM clearly outperforms MDLM and DFM. Alternatively, scale to a larger image benchmark e.g. imagenet 128x128, and show UDLM beats MDLM. A skim through the rebuttal experiment findings does not suggest a better Gen PPL or a better PPL upper bound. Maybe I missed your result, if so please point me to it.

If either of these results are obtained, I shall happily increase my score to 6. I sincerely hope you do succeed, because I would love to update my, perhaps insufficient, understanding of UDLMs.

For the language models we trained on LM1B, below we present Gen PPL, which demonstrate that UDLM outperforms MDLM especially for smaller $T$ where the value of UDLM is most apparent (since it can recover from mistakes).

These results are significant because one of the key potential advantages of diffusion vs. AR is parallel sampling and the reduction of forward passes through the model by setting $T$ < $L$ (the sequence length). In this setting we see that UDLM significantly outperforms MDLM.

(Note these outputs come from the models we re-trained where the LM1B sequences were wrapped as opposed to pad filled)

In terms of comparing our work to MDLM/MD4/DFM more broadly, it is important to note none of these works discuss adapting guidance to discrete diffusion. Our work is novel in this respect. Moreover, these works focus only on absorbing state diffusion. Our work investigates uniform noise and for the first time finds a training objective that renders this model competitive with the absorbing state counterparts, and we also find settings where UDLM outperforms these previous works, such as the LM1B results just shared, the CIFAR10 results previously discussed, and the extensive guidance experiments we present in our work.

As promised above, please see below for our updated numbers regarding LM1B PPL when sequences are right-padded to fill the context size of 128 tokens (which is more in-line with the data processing of previous works)

Notice that UDLM:

• even outperforms SEDD Absorbing state,
• is the best uniform noise discrete diffusion model,
• and narrows the gap between comparable models that use absorbing state vs. uniform relative to previous work.

Dear Reviewer,

At your suggestion, we have scaled our image generation experiment to ImageNet (128x128), keeping the setup analogous to our CIFAR-10 experiment. Similar to our previous findings on CIFAR-10, we find that UDLM outperforms MDLM in the context of guidance or fast sampling, and otherwise attains similar performance.

Improved Conditional Generation:

UDLM responds better to increasing the guidance strength parameter $\gamma$, which controls the strength of conditioning on a class label. This improves over the FID of MDLM.

The above settings include a standard conditional diffusion model ($\gamma=1$), and one with an increased level of guidance ($\gamma=2$) conditioned on a class label. Guidance that is too strong ($\gamma=3$) eventually reduces quality, as in prior work on diffusion (Ho & Salimans 2022). We find that being able to revise generated tokens in UDLM improves controllability (and hence quality) over MDLM.

Improved Faster Sampling

UDLM maintains sample quality when we reduce the number of sampling steps $T$ by 2x and 4x, while the quality of MDLM degrades.

We use $\gamma= 1$ in the above table (i.e., we sample from a conditional diffusion model without additional guidance) to study the benefits of UDLM for fast sampling independently of guidance. We again observe that being able to revise generated tokens over the course of sampling allows UDLM to produce high quality images using a fraction of the steps of MDLM.

Details about this experiments: For both MDLM and UDLM we train a 117M parameter DiT model for 400k iterations using the exact same hyperparameters. We trained on a subset of ImageNet that consisted of 100 randomly sampled classes (to limit dataset size to 128k images and finish training by the rebuttal deadline). Note that this strategy has been used by previous works, e.g., SupContrast (Khosla et al., 2020). We patched the image into 16x16 blocks (i.e., each image consists of 16x16 tokens) that were quantized to a vocabulary of size 16k (as in Sun et al 2024). We evaluated FID using 10k samples (due to computational and time considerations). Samples were generated using our discrete classifier guidance algorithm by first sampling a random class id, and then sampling an image conditioned on that id.

We thank the reviewer again for their continued engagement with our work and hope that these latest results address the previously raised concerns.

References

Ho, Jonathan, and Tim Salimans. "Classifier-free diffusion guidance." arXiv preprint arXiv:2207.12598 (2022).

Khosla, Prannay, et al. "Supervised contrastive learning." Advances in neural information processing systems 33 (2020): 18661-18673.

Sun, Peize, et al. "Autoregressive Model Beats Diffusion: Llama for Scalable Image Generation." arXiv preprint arXiv:2406.06525 (2024).

We thank the reviewer for their detailed feedback. Please see below for the answers to the specific questions raised.

Concern 1 Which aspects of the work are novel?

We clarify that our claim is not to have developed discrete diffusion with uniform noise, but rather the novelty comes from the following:

1. We present novel guidance mechanisms for discrete diffusion
   • Our derivation for discrete classifier-based guidance is novel and finds a simple yet principled way to get around the issue that the classifier model need not factorize the same way as the generative decoding model
2. We introduced UDLM (analogy of MDLM (Sahoo et al., 2024) for uniform noise), which is more guideable. Specific novel components include a refined ELBO with terms analytically set to zero and a continuous time limit.

Concern 2 Missing citation to Zheng et al., 2023 and Zhao. et al., 2024

We thank the reviewer for pointing out Zheng et al., 2023 and Zhao. et al., 2024. These are indeed highly relevant in terms of the formulation for the interpolating discrete diffusion, and we cite them in our revised manuscript along with Sahoo et al. 2024 when introducing this class of models.

Comparisons to other uniform discrete diffusion models were done using the metric of perplexity on standard language benchmarks (e.g., text8 and LM1b). Neither Zheng et al. nor Zhao et al. report this metric nor do they test on these standard benchmarks. Furthermore, Zheng et al. does not admit a true variational bound and therefore cannot be compared.

Over the coming week, we will attempt to reproduce Zhao et al. in our experiments and will update here once we have the results.

Below we list differences between our work and these references.

Comparison to Zheng et al. 2023

Most importantly, we train with a principled lower bound. In contrast, the use of “teacher forcing” in Zheng et al. (see Section 4.2 and Appendix D in their work) allows them to derive a loss term that amounts to a simple weighted cross entropy term, but this loss is not to a true variational bound. Additionally, our loss is improved by analyzing the continuous time case. Zheng et al. only examine discrete time.

Comparison to Zhao et al. 2024

The main difference between Zhao et al. and our work is the parameterization (see below for more details). Additionally, Zhao et al. train models with an “augmented” loss where they add a cross-entropy term to the ELBO. Our models directly optimize the ELBO that we derive.

The difference in parameterization is the following: We use mean parameterization, meaning that our network outputs a distribution over the “clean” data given a noisy input. We then plug this estimate, which we denote as $\mathbf{x}\_\theta$ into the pre-defined noise process $q_t(\mathbf{z} \mid \mathbf{x} = \mathbf{x}_\theta)$. We can use this parameterization to derive an equivalence between our ELBO and the one from SEDD (Lou et al. 2023) – please see the new Appendix B in the updated pdf we posted to OpenReview. To minimize the SEDD ELBO, one could use the following parameterization for the score: $s\_\theta(\mathbf{z})\_{\mathbf{z}’} := \frac{q\_t(\mathbf{z}’ \mid \mathbf{x} = \mathbf{x}\_\theta)}{q\_t(\mathbf{z} \mid \mathbf{x} = \mathbf{x}\_\theta)}$

In contrast Zhao et al., which also use a model to output a distribution over the clean data, parameterize the score as follows: $s\_\theta(\mathbf{z})\_{\mathbf{z}’} := \sum_{x}\frac{q_t(\mathbf{z}’ \mid \mathbf{x})}{q_t(\mathbf{z} \mid \mathbf{x})} \mathbf{x}^\top\mathbf{x}_\theta$

We thank the reviewer for recognizing the contributions of our work and address specific questions and comments below

Concern 1: Clarifying why UDLM improves guidance

Absorbing state discrete diffusion models suffer from a drawback that is shared with AR models: once a token is predicted, it cannot be unchanged (see for example Equations 6 and 7 in Sahoo et al. 2024).

In contrast, when using uniform noise as the limiting distribution, even if the model samples a particular token at some time $t$, that token can change as the sequence is denoised at times $s < t$. We therefore hypothesize that discrete diffusion models trained with uniform noise will be more “guidable” than absorbing state counterparts, since mistakes made during the decoding process can be corrected. This hypothesis is made more clear in the experiment below.

We examine what happens when we decode with smaller time steps $T$. MDLM needs to decode multiple tokens per denoising step and these tokens cannot be changed at later denoising iterations. Thus the model can potentially be “locked-in” to mistakes. In contrast, UDLM can update any given token on every iteration. In this experiment for CIFAR10 generation, we find that MDLM’s performance deteriorates with smaller $T$, whereas UDLM is robust to this setting:

Details about this experiment: We generate 50k images sampled from a conditional model (which is equivalent to using $\gamma=1$ and D-CFG). F1 is the accuracy of a separate classifier in identifying the class label used for the conditional generation.

Concern 2: Which equation is the ELBO in Section 3.2?

Equation 10 represents the continuous time (i.e., $T \rightarrow \infty$) ELBO for uniform noise. We derive this by showing how $\mathcal{L}\_{recon}$ and $\mathcal{L}\_{prior}$ can be analytically set to zero and analyzing $\mathcal{L}\_{diffusion}$ in the infinite time limit. Equation 11 adapts this ELBO for sequences of multiple tokens and is what we use to train UDLM.

We highlight the difference between an the unsimplified diffusion loss term from the standard ELBO and the one in our objective below:

Unsimplified

$T\cdot\frac{N\alpha\_t \mathbf{x}\_i + \alpha\_{t|s} - \alpha\_t + (\alpha\_s - \alpha\_t)\mathbf{x}\_i + \frac{(\alpha\_s - \alpha\_t)(1- \alpha\_s)}{\alpha\_s N}}{N \alpha\_t \mathbf{x}\_i + 1 - \alpha\_t} \cdot \log\Bigg[\frac{\frac{N \alpha\_t \mathbf{x}\_i + \alpha\_{t|s} - \alpha\_t + (\alpha\_s - \alpha\_t)\mathbf{x}\_i + \frac{(\alpha\_s - \alpha\_t)(1 - \alpha\_s)}{\alpha\_s N}}{N \alpha\_t \mathbf{x}\_i + 1 - \alpha\_t}}{\frac{N \alpha\_t(\mathbf{x}\_\theta)_i + \alpha\_{t|s} - \alpha\_t + (\alpha\_s - \alpha\_t)(\mathbf{x}\_\theta)_i + \frac{(\alpha\_s - \alpha\_t)(1 - \alpha\_s)}{\alpha\_s N}}{N \alpha\_t (\mathbf{x}\_\theta)_i + 1 - \alpha\_t}}\Bigg] + T\cdot \sum\_{j \in [N]: (\mathbf{z}\_t)\_j = 0} \frac{(\alpha\_s - \alpha\_t)\mathbf{x}\_j + \frac{(\alpha\_s - \alpha\_t)(1-\alpha\_s)}{N \alpha\_s}}{N \alpha\_t\mathbf{x}\_i + 1 - \alpha\_t}\log\Bigg[\frac{\frac{(\alpha\_s - \alpha\_t)\mathbf{x}\_j + \frac{(\alpha\_s -\alpha\_t)(1-\alpha\_s)}{N \alpha\_s}}{N \alpha_t\mathbf{x}\_i + 1 -\alpha\_t}}{\frac{(\alpha\_s-\alpha\_t)(\mathbf{x}\_\theta)\_j +\frac{(\alpha\_s - \alpha\_t)(1-\alpha\_t)}{N \alpha\_s}}{N \alpha\_t(\mathbf{x}\_\theta)\_i + 1 - \alpha\_t}}\Bigg]$

Ours

$\frac{\alpha\_t'}{N\alpha\_t}\Bigg[\frac{N}{N\alpha\_t\mathbf{x}\_i + 1 - \alpha\_t} - \frac{N}{N\alpha\_t(\mathbf{x}\_\theta)\_i + 1 - \alpha\_t}\Bigg] +\frac{-\alpha_t'}{N\alpha\_t}\sum\_{j \in [N] : (\mathbf{z}_t)_j = 0} \Bigg(\frac{N\alpha_t\mathbf{x}\_j + 1 -\alpha\_t}{N\alpha\_t\mathbf{x}\_i + 1 -\alpha\_t}\Bigg)\log\Bigg[\Bigg(\frac{N\alpha\_t(\mathbf{x}\_\theta)\_i + 1 - \alpha\_t}{N\alpha\_t(\mathbf{x}\_\theta)\_j + 1 - \alpha\_t}\Bigg)\Bigg(\frac{N\alpha_t\\mathbf{x}\_j + 1 - \alpha\_t}{N\alpha\_t\mathbf{x}\_i + 1 - \alpha\_t}\Bigg)\Bigg]$

Concern 3: D-CBG is mentioned in the intro before being defined

Thank you for bringing this to our attention. We have corrected this in our revised manuscript.

References

Ou, Jingyang, et al. "Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data." arXiv preprint arXiv:2406.03736 (2024).

Sahoo, Subham Sekhar, et al. "Simple and Effective Masked Diffusion Language Models." arXiv preprint arXiv:2406.07524 (2024).

Shi, Jiaxin, et al. "Simplified and Generalized Masked Diffusion for Discrete Data." arXiv preprint arXiv:2406.04329 (2024).

Thank the author for the response. My concerns have been addressed so I would maintain my score to support this paper.

We thank the reviewer for their feedback and for recognizing the contributions of our work. In the comments below we respond to the specific concerns and questions raised by the reviewer.

Concern: Does UDLM provide benefits over MDLM (especially with larger vocabularies)?

Our new results demonstrate that UDLM outperforms MDLM in the context of guidance or fast sampling, and otherwise attains similar performance in many settings.

Benefits of UDLM for Guidance

ImageNet

We scaled our image generation experiment to ImageNet (128x128), which was quantized to a large vocabulary of 16k tokens.

UDLM responds better to increasing the guidance strength parameter $\gamma$, which controls the strength of conditioning on a class label. This improves over the FID of MDLM.

The above settings include a standard conditional diffusion model ($\gamma=1$), and one with an increased level of guidance ($\gamma=2$) conditioned on a class label. Guidance that is too strong ($\gamma=3$) eventually reduces quality, as in prior work on diffusion (Ho & Salimans 2022). We find that being able to revise generated tokens in UDLM improves controllability (and hence quality) over MDLM.

NLP

We added an NLP experiment with a large vocabulary (~30k) that demonstrates the benefits of UDLM and our guidance methods vs AR. The goal is to generate sentences that follow a part-of-speech template, which acts as the guidance label. We show that diffusion models with our guidance outperform AR. Moreover, UDLM is competitive with MDLM on generation quality and responds better to higher guidance weights.

Benefits of UDLM in Faster Sampling Settings

ImageNet

UDLM maintains sample quality when we reduce the number of sampling steps $T$ by 2x and 4x, while the quality of MDLM degrades.

We use $\gamma= 1$ in the above table (i.e., we sample from a conditional diffusion model without additional guidance) to study the benefits of UDLM for fast sampling independently of guidance. We again observe that being able to revise generated tokens over the course of sampling allows UDLM to produce high quality images using a fraction of the steps of MDLM.

CIFAR10

For CIFAR10, as with ImageNet, with smaller $T$ where the diffusion models predict multiple pixels in parallel we find a setting where MDLM’s performance deteriorates, whereas UDLM is more robust:

Concern: Is UDLM novel relative to DFM/MDLM and other work?

Contributions of our work

1. We present novel guidance mechanisms for discrete diffusion. This is a core focus of our work and is novel relative to both DFM (Gat et al. 2024) no MDLM (Sahoo et al. 2024)
2. We introduced UDLM (analogy of MDLM for uniform noise), which is more guideable. Specific novel components include a refined ELBO with terms analytically set to zero and a continuous time limit.

Training with UDLM ELBO outperforms DFM

We first make clear that DFM is concurrent work with ours. Moreover, we find that using the training objective proposed in DFM leads to worse language modeling performance (as measured by PPL, the standard quality metric):