Improving Discrete Diffusion with Schedule-Conditioning

Authors: Alan Nawzad Amin, Nate Gruver, Andrew Gordon Wilson
Primary area: generative models
Keywords: discrete diffusion, image generation, language model

Abstract:
Discrete diffusion models, like continuous diffusion models, generate high-quality sequence data by gradually undoing noise applied to datapoints via a Markov process. Gradual generation in theory comes with many conceptual benefits; for example, inductive biases can be incorporated into the noising Markov process. In practice however, the best performing discrete diffusion model is consistently masking, which does not denoise gradually. Here we explain the performance of masking diffusion by noting that it makes use of a fundamental difference between continuous and discrete Markov processes: discrete Markov processes evolve by discontinuous jumps at a fixed rate and, unlike other discrete diffusion models, masking diffusion builds in the known distribution of jump times and only learns where to jump to. We show that we can similarly bake in the known distribution of jump times into any discrete diffusion model; despite their simplicity, our new models -- schedule-conditioned diffusion (SCUD) -- generalize classical discrete diffusion and masking diffusion. By applying SCUD to models with noising processes that incorporate inductive biases on images, text, and protein data, we build diffusion models that outperform masking.

Decision: reject
Reviewer scores: 8, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1729726547748

thread_type:
official_review

<thread_object>
creator: reviewer_ifsu

<content>
summary:
This paper attempts to understand the empirical finding that discrete diffusion models using a masking corruption process outperform other noising processes such as uniform or Gaussian noising processes in which tokens are "mutated." The authors hypothesize that masking diffusion dominates due to having knowledge of when corruptions occur, which is not the case for the non-masking noise processes. With this hypothesis, the authors propose a modification to standard non-masking discrete diffusion models that allows the model to condition on information about the schedule of the noising process. The authors derive a new version of the ELBO used to train discrete diffusion models into one that explicitly marginalizes over different noise schedules. This allows them to propose a new training objective and parameterization of the denoising process that conditions on noise schedules. The authors evaluate their approach on images, text, and proteins and demonstrate that their approach consistently improves over models using the same noising process, but without schedule conditioning.

strengths:
The paper is tackling an original problem, namely understanding why masking discrete diffusion outperforms other noising processes including those that have domain-specific inductive biases. While I have some concerns with the hypothesis (see section Weaknesses), I agree that conditioning on noise schedules should help mitigate some of the advantages that masking diffusion models have. The paper contains a rigorous derivation that shows how to incorporate the noise schedules into the modeling framework. The empirical results demonstrate that conditioning on the noise schedule tends to improve likelihood-based metrics over models that use the same noise process but do not incorporate the noise schedule.

weaknesses:
I am not convinced by the primary hypothesis outlined in the abstract that says "...we propose that masking diffusion dominates due to knowledge of when corruptions occur." Unless I am misunderstanding something, masking diffusion models have no knowledge of when an event occurs, only that an event occurred. To put a different way, the identity of the state (masked or not) tells us whether it has been noised. However, we have no idea when the noising occurred. This seems like an important point for motivating the use of noise schedules as input.
  
The primary weakness for me was the experimental results. In particular, the paper seems to be lacking several key details about the experiments. I will reiterate this in the "Questions" section, but I did not see any explanation for why the authors chose to model Cifar10 with both 128 and 256 pixel values. Based on the lack of explanation, I am left feeling that the B=128 experiments were done simply because this was a setting where the authors found they can get Gaussian SCUD to improve over Masking. However, to me, this makes the results appear cherry-picked. Furthermore, in Figure 2, it appears Gaussian is already outperforming Masking, without SCUD (although perhaps not with statistical significance), diminishing the novelty of the result. There should be some commentary on this since.

All of the experiments appear to be done with different training data than the baselines. Thus, I am not sure any of the numbers are truly valid in terms of comparing between methods. Most concerningly, I didn't see any explanation of this in the text or appendix. Baselines should be redone using the same training data.

For the baselines, it seems a crucial baseline would be to use all of the same hyperparameters as SCUD, but without conditioning. This would use the same noising process (e.g. Gaussian or Uniform), the same architecture, and the same training data but without using any conditioning information (e.g. just pass in a constant $s$ every time).

The authors do not describe how BPD/Perplexity are computed. Since there appears to be an additional variable to marginalize out (the noise schedule), these details are important and need to be explained and justified.

Practically speaking, the field seems to be moving away from discrete diffusion models and more towards discrete flow matching. The latter has a much simpler objective function and training procedure while improving results. However, I am very sympathetic to the fact that field is moving so quickly and therefore do not penalize the authors for this. However, any more discussion about how this can be extended to flow-matching would be welcome and it seems like a straightforward extension?

The writing and presentation could be improved for more clarity. There were a few instances where the writing was too imprecise for me to understand what was meant. For example, around lines 235-236, the authors write "...define pr(x_t^d) as the last event in dimension d...". It was not clear to me what is meant by and "event" and what "last" is referring to. I think what is meant is pr(x_t^d) is the state of x_t before the last noise event?

Another place where the writing could have been more clear is in the "Efficient Loss" section when the authors say "... and then add a weight representing how likely an event is to occur at the instant t." Here I think it would be helpful to write what the term is as I was left confused about which terms in Equation 5 were the weight. I believe the weight term is both the term involving Betas and the multiplication of s_t. As a reader, I was expecting this to be more clearly laid out for me.

questions:
1) Can the authors clarify my point about masking diffusion models not having knowledge of when corruptions occur?

2) Why were experiments done with B=128 and B=256? Was is due to the fact that only for B=128 did SCUD outperform masking?
     
3) Why were different training data used in the baselines? 

4) How were the metrics computed? Specifically, how did the authors handle marginalizing the noise schedule?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730636789955

thread_type:
official_review

<thread_object>
creator: reviewer_3fup

<content>
summary:
This paper propose that masking diffusion succeeds because it leverages information about when corruption events occurred, enabling it to focus on modeling relationships between tokens rather than inferring corruption events. Building on this insight, they introduce a new method called schedule-conditioned diffusion (SCUD), which incorporates corruption schedule information into discrete diffusion models. Experimental results demonstrate that SCUD generalizes masking diffusion and outperforms classical discrete diffusion methods on various datasets.

strengths:
​1. **Innovative Insight**: This paper provides a novel explanation for the success of masking diffusion, linking it to the modeling of corruption schedules, which is a valuable addition to the understanding of discrete diffusion processes.

​2. **Methodology**: The introduction of SCUD is rigorous, with mathematically supported schedule conditioning, which extends masking diffusion and further enhances its performance.

​3. **Empirical Evidence**: Experiments on image, text, and protein data show that SCUD outperforms standard discrete diffusion models, supporting its claims of enhanced generative capability.

weaknesses:
​1. **Lack of Motivation Explanation**: The authors seem to focus heavily on explaining how SCUD works, with less emphasis on why this approach is chosen and what its ultimate goal is. This may make it difficult for readers to follow the authors’ line of reasoning.

​2. **Over-Reliance on Appendix for Proofs**: Many key proofs and details are placed in the appendix, which disrupts the main text’s coherence and could interfere with readers’ logical understanding.

​3. **Inadequate Training/Sampling Procedure Description**: The description of the training and sampling processes is not sufficiently detailed, making it difficult to understand SCUD’s training and sampling mechanics without referring to supplementary materials or  external resources.

questions:
​1.	Could the authors clarify the role of $\Delta t$ in discrete Markov process? This is not clearly defined in the “Background” section, nor are there references provided for further reading.

​2.	Would it be possible to simplify some of the proofs to make them more accessible in the main text?

​3.	Why didn't provide the complete training and inference process in algorithmic form in the paper?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730702971954

thread_type:
official_review

<thread_object>
creator: reviewer_bgrb

<content>
summary:
The paper presents Schedule Conditioned Diffusion (SCUD), a new method to enhance discrete diffusion models for conditional sequence generation. The authors identify that existing structured and state-dependent approaches are often outperformed by the simpler "masking diffusion," due to its effective alignment of forward and backward transition schedules. They introduce a decomposition of the Evidence Lower Bound (ELBO) that addresses this mismatch and demonstrate efficient training strategies for SCUD. The findings indicate that SCUD achieves similar performance to masking diffusion, with the authors releasing their code for further exploration.

strengths:
- The paper is well written, equations are nicely embedded and overall presentation is good (marking will be raised to good once the page limit is achieved).

- The paper effectively addresses limitations of existing discrete diffusion models, specifically the dominance of masking diffusion over more gradual and structured processes. It provides a solid theoretical foundation to argue for the introduction of SCUD.

- Introducing the SCUD model, which conditions on the transition schedule, is a novel approach. This model adapts diffusion frameworks by incorporating structured forward processes, potentially expanding discrete diffusion's applications across data types.

- The paper includes a rigorous theoretical framework, with proofs and mathematical propositions that justify the SCUD approach.

- The experiments span various data types, including images, text, and protein sequences. Results on CIFAR-10, LM1B, and UniRef50 datasets convincingly show that SCUD improves performance compared to other non-masking processes.

- The paper compares SCUD with state-of-the-art discrete diffusion models, showing how SCUD better captures transition schedules and can leverage structured processes.

weaknesses:
# Presentation (Minor)

I will re-adjust my mark for presentation from poor to good once the following urgent concern has been addressed by the authors:

- The text of the main paper exceeds the 10-page limit. Please move your remarks regarding Reproducibility to the appendix, ICLR will likely be strict about enforcing the 10-page limit, exceeding the limit may lead to your work being rejected down the line.


# Content (Major)

- The paper's complexity could limit practical adoption. SCUD requires intricate parameterizations and careful handling of components like the rate parameter $\gamma$ and the transition matrix $K$ in particular, potentially increasing computational cost.

-  While SCUD reduces training samples, the discussion on costs associated with complex matrix operations, schedule conditioning, and increased dimensionality in high-dimensional data is minimal.

- The qualitative analysis of CIFAR-10 images suggests that SCUD models lack clear object formation. Though the focus is on likelihood scores, this limitation in sample quality could affect its utility in image generation tasks.

questions:
In general, I am willing to raise my score, if my concerns and questions are addressed with compelling evidence. 

Building on the aforementioned weaknesses, I pose the following questions:

1. Can you provide O-Notation w.r.t. the data-dimensionality for the added computational cost arising from $K$? Furthermore mentioning the GPU hours of the different methods in your work could help put the computational cost of different methods into perspective.

2. In addition to 1.: Could you provide a quantitative comparison of the computational costs associated with SCUD versus standard discrete diffusion methods in terms of memory usage and processing time?

3. Given the increased complexity SCUD introduces, what specific strategies could one use to manage computational demands when applying SCUD to larger datasets or higher-dimensional inputs?

4. Could you elaborate on how sensitive SCUD is to the selection of the rate parameter $\gamma $? How does this parameter interact with other hyperparameters in practice on datasets other than CIFAR-10 as shown in figure 2?

5. Are there particular domains where SCUD could have inherent advantages over masking diffusion?

6. Although SCUD improves likelihood scores, how do you plan to address the relatively low quality of visual samples generated in CIFAR-10? Could modifications to SCUD enhance sample fidelity?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730851770207

thread_type:
official_review

<thread_object>
creator: reviewer_qq5r

<content>
summary:
The choice of forward process has great performance implications for (discrete) diffusion models. This paper theoretically studies the improved performance of the masking forward process in discrete diffusion models. To this end, the authors introduce the notion of an event schedule, which describe times along the forward process where transitions take place, and separate the "when" of transition from the "where" of transition. 

The authors further derive the training objectives corresponding to conditioning on these event schedules, and apply their proposed method SCUD, to the tasks of image generation, language generation and protein sequence generation. Across the different tasks, their proposed method is able to outperform equivalent forward processes, but without the conditioning using less training examples.

strengths:
1. The main strength of the paper is the firm theoretical footing to understand different forward processes in discrete diffusion, and how they relate to the "when" and "where" of the corresponding transitions. Through the notion of event / schedule conditioning, the paper attempts to disentangle the influence of "when" and "where", which leads to corresponding modifications to the ELBO objective. The authors also emphasize the connections of their method to previous discrete diffusion methods. 

2. The paper is very well written, and makes a strong effort to coherently explain the different moving parts. 

3. To connect the method to practice, the authors also propose different tricks such as reversing multiple events jointly, and an efficient loss formulation for high-dimensional data. Experiments across image, protein and language domain show favorable improvements for the same forward process, conditioned on event schedule. The experiments on proteins are especially interesting, using 2 orders less of training data.

weaknesses:
There are no glaring weaknesses in the paper. But there is some minor feedback:

1. The formal notion of event schedule is only introduced in Proposition 4.2. This should instead be moved to before Proposition 3.1, so the readers already know what the event schedule captures, and the corresponding equations become easier to follow.

2. There are claims in the paper regarding SCUD outperforming masking, but evidence of this is not visible in the experiments. It is unclear whether the SCUD conditioning with Gaussian / Uniform will outperforming existing discrete diffusion (e.g. D3PM) with masking, given equal compute.

questions:
1. In light of the point 2 in weaknesses, could the authors run the D3PM model with masking diffusion with the same number of training samples as SCUD? -- The reviewer hopes this is feasible within the available compute budget of the authors. An alternative would be to also consider small versions of D3PM and SCUD that can be trained for equivalent number of samples.

2. How is B defined in the language and protein experiments?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732472172332

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global comment

comment:
We thank all reviewers for their thoughtful reviews! There is a flurry of work extending the capacities and design space of discrete diffusion models. However, in many settings, the simplest diffusion model, masking diffusion, dominates. Therefore, designs and capacities deviating from masking diffusion, for example, incorporate inductive biases in the forward process, often harm rather than help inference. SCUD explains the surprising dominance of masking diffusion. We use this insight to incorporate inductive biases into the forward process without sacrificing schedule conditioning, thus building better models. SCUD opens the door to moving beyond masking diffusion and leveraging the full design space of discrete diffusion.

Below we discuss how we have incorporated each of your suggestions into the new draft of the paper, adding a discussion on the computational complexity of SCUD, and extending SCUD to flow matching. In addition to these changes we have run our image (B=256) and protein experiments for longer along with our own careful implementations of masking diffusion and classical diffusion baselines. These reimplementations use the same training and architecture as SCUD and have strong performance on their own. Nevertheless, the results (Tab 1, 2) show that SCUD outperforms these baselines and the longer training improved the performance of SCUD. [FigS1](https://anonymous.4open.science/api/repo/host_image-6044/file/comparison.pdf?v=ebfdc962) also shows that the sample quality of SCUD greatly improves with longer training. These models will continue to run until the end of the discussion section, when we will update these numbers.

**Table S1: Discrete diffusion models on image data**
|Method|Forward process|Training samples|BPD|
|-|-|-|-|
|D3PM|Uniform|1.9×10⁸|5.08|
|D3PM|Gaussian|1.9×10⁸|3.44|
|τLDR|Gaussian|2.6×10⁸|3.59|
|MD4|Masking|2.6×10⁸|**2.78**|
|Classical|Gaussian|6.4×10⁷|2.94|
|Classical|Masking|6.4×10⁷|2.90|
|SCUD|Gaussian|6.4×10⁷|*2.86*|


**Table S2: Discrete diffusion models on protein data**
|Method|Forward process|Training tokens|Perplexity|
|-|-|-|-|
|D3PM|Uniform|>3×10¹¹|18.82|
|D3PM|BLOSUM|>3×10¹¹|17.16|
|D3PM|Masking|>3×10¹¹|**14.61**|
|Classical|BLOSUM|8×10⁹|15.39|
|Classical|Masking|8×10⁹|15.56|
|SCUD|BLOSUM|8×10⁹|*15.29*|
</content>
</thread_object>

</thread>


Metareview:
metareview:
The current paper proposes to condition on the transition schedule in the forward-pass of discrete diffusion models. The reviewers acknowledged that the authors' work gives novel theoretical insights into the field and that the overall presentation is coherent. The reviewers' questions regarding the computational cost have been addressed during the rebuttal. In the final discussion, there are still major concerns on the experimental evidence including visual qualities of the generated images and a thorough comparison with masking diffusion under a fair evaluation protocol, and thus the current empirical evaluation is not conclusive. Moreover, Reviewer bGrb suggested to study how fast SCUD converges relative to other methods based on checkpoints of the models over time, which is not included in the current paper. To incorporate these additional results would need more efforts and further validation, and we cannot be certain whether this can be accomplished within the tight timeline of ICLR'25 but hope will be achieved in the next iteration of this work. The rejection is based on careful evaluation and thorough discussion with regret. The authors are encouraged to resubmit the paper after the suggested revisions with improved and more stable empirical results.