Simplified and Generalized Masked Diffusion for Discrete Data

Summary

This paper proposes a new framework for masked diffusion models for generative modeling of discrete data. Masked diffusion models offer an alternative to autoregressive models for discrete data but have faced challenges due to complex formulations and unclear relationships between different approaches. This paper presents a simplified and generalized framework to address these issues, enhancing the performance and training of masked diffusion models. The key contributions includes: 1. Simplification of Model Formulation: The paper establishes properties for the forward process and its time reversal using elementary arguments, provides a simple expression for the Evidence Lower Bound (ELBO), demonstrating it as a weighted integral over time of cross-entropy losses, and shows invariance properties similar to continuous space diffusions. 2. Re-derivation of Training Objectives: The paper demonstrates how various previously proposed discrete diffusion training objectives can be derived from the ELBO objective by altering parameterization, relaxing constraints, or modifying loss weighting. 3. Performance Improvements: The paper demonstrates state-of-the-art likelihood and zero-shot transfer results on text and image tasks using the proposed ELBO objective. 4. Generalized Masked Diffusion Model: The paper proposes a generalized masked diffusion model that allows state-dependent masking schedules, further improving predictive performance on test likelihoods.

Strengths

1. The paper makes a notable contribution to the field of generative modeling for discrete data by introducing a simplified and generalized framework for masked diffusion models. 2. The quality of the paper is reflected in the thoroughness of its methodology and the robustness of its experimental validation. 3. The paper is well-written and clearly structured, making it accessible to both experts and those new to the field. 4. The significance of the paper lies in its potential to impact a wide range of applications in generative modeling for discrete data.

Weaknesses

1. While the paper provides a robust theoretical foundation, there could be more emphasis on practical applicability. The paper could benefit from additional practical guidelines for implementing the proposed framework, such as more detailed pseudocode and specific implementation challenges. 2. The experimental results presented are strong, but the range of tasks and datasets could be expanded, such as VQ-Diffusion [1] for token-based text-to-image. 3. I am unfamiliar with diffusion models for text generation. For image generation, the paper has reported likelihood results, missing some other common metrics, such as FID and IS. [1] Vector Quantized Diffusion Model for Text-to-Image Synthesis

Questions

Please see weaknesses.

Rating

6

Confidence

3

Soundness

4

Presentation

4

Contribution

4

Limitations

The authors have adequately described the limitations and potential negative societal impact of their work.

Summary

The paper simplifies the mathematical formula for the absorbing state diffusion process. By doing so, the authors derive a continuous-time ELBO for masked diffusion models. Their method, MD4, achieves better perplexity scores than SEDD on text8 and zero-shot perplexity on numerous datasets.

Strengths

Simplifies the complex mathematical formulations for the absorbing state diffusion for D3PM.

Weaknesses

Weaknesses: 1. Weak empirical results 1. The zero shot numbers for D3PM in Table 1 look fishy. There are only 2 differences between Md4 and Absorbing State D3PM: 1. Mathematical simplification. In the discrete case (Eqn. 6), even though MD4 features a Simplified functional form for the ELBO, it shouldn't give it any performance benefits in terms of perplexity since it is mathematically equivalent to D3PM. 2. The improvement in ELBO could be because of the continuous time formulation. However, VDM [1] has shown that for gaussian diffusion, improvement from discrete (T=1000) to continuous time (T = $\infty$) barely improves the likelihood by less than 1%. For this reason, I request the authors to perform eval on an already trained model and report the perplexity numbers on text8 or OWT using Eqn (6) with T=100, 1000, 10000. If the numbers reported for D3PM in Table 1 are indeed correct, and if the entire improvement is coming from the continuous time formulation, then the discrete time MD4 should get a number that's comparable to D3PM's zero shot ppl numbers. Questions: How did they retrain D3PM? Did they use the same transformer backbone as MD4? Did they use the same model size and data pre-processing scheme? Did they use uniform state or absorbing state diffusion process? The authors need to clarify this. 2. CIFAR10 Experiments. The AR baselines use old transformer models hence the comparison isn't quite fair. Current SOTA diffusion models on Imagenet 32 achieve a NLL of 2.55 [2] which is far better than the absorbing state diffusion models. So, I'm unsure about the takeaway from Table 3. In the conclusion section, the authors claim that "… on text and image data, the resulting masked diffusions outperform existing discrete and continuous diffusion models …" which is factually incorrect given that their method largely underperforms against gaussian diffusion [1, 2]. 2. Limited evaluation of GenMD4. The authors mention that GenMD4 performs poorly on zero-shot tasks. I request the authors to quantify this poor performance by providing 1. Validation ppl numbers on OWT 2. zero-shot ppl numbers. [1] Kingma, D., Salimans, T., Poole, B. and Ho, J., 2021. Variational diffusion models. *Advances in neural information processing systems*, *34*, pp.21696-21707. [2] Sahoo, S., Gokaslan, A., Sa, C., Kuleshov, V., 2024. Diffusion Models With Learned Adaptive Noise. arXiv:2406.07524

Questions

1. Clarification on D3PM experiments in Table 1 as mentioned in the "weaknesses" section in the reviews. 2. Why did the authors decrease the dropout to 0.02 for OWT experiments and not set it to 0? Diffusion models are heavily regularized due to the randomness in the input to the model and oftentimes don't require additional regularization such as dropout. Hence, an intuitive or an empirical explanation would be helpful.

Rating

6

Confidence

5

Soundness

2

Presentation

2

Contribution

3

Limitations

Yes

Summary

The paper proposes a streamlined and generalized framework for masked diffusion models, addressing the complexities and inefficiencies of existing models, including those based on Score Entropy Discrete Diffusion (SEDD). It introduces a continuous-time variational objective for masked diffusion models, simplifying the evidence lower bound (ELBO) to a weighted integral of cross-entropy losses. Additionally, the paper presents state-dependent masking schedules, enhancing the flexibility and performance of these models. The proposed methods demonstrate state-of-the-art results in text and image tasks, significantly improving likelihood and zero-shot transfer performance.

Strengths

- The paper offers a novel theoretical formulation of the continuous-time variational objective for masked diffusion models, simplifying the training process and ensuring consistency between forward and reverse processes. - The introduction of state-dependent masking schedules provides a more adaptable approach, catering to the specific characteristics of the data and improving model performance. - The proposed methods achieve state-of-the-art performance in both text and image generative tasks, significantly enhancing likelihood and zero-shot transfer capabilities. - By reducing the ELBO to a weighted integral of cross-entropy losses, the paper makes the training and understanding of masked diffusion models more accessible and potentially more stable. - The paper includes comprehensive experimental validation on various datasets, demonstrating the robustness and superiority of the proposed methods.

Weaknesses

- Despite the theoretical simplifications, the practical implementation of state-dependent masking schedules can still be complex and computationally demanding. Specifically, obtaining the starting x_T is challenging, and since the sampling process lacks stochasticity, sampling cannot be done from the completely masked state. - The state-dependent models have a tendency to overfit to dataset statistics, which can limit their effectiveness in zero-shot transfer tasks. - While the paper demonstrates superior performance, a more detailed comparative analysis with other state-of-the-art methods, particularly regarding computational efficiency and training times, would provide a clearer picture of the advantages.

Questions

- Could the authors provide more insights into the practical challenges faced during the implementation of the state-dependent masking schedules? - How does the proposed model ensure consistency between the forward and reverse processes, and how does this impact training stability compared to SEDD? - Could the authors provide a detailed and separate description of the training and sampling algorithms, similar to what is provided in the Appendix of the SEDD paper, to better and more easily understand the proposed method? - How sensitive is the proposed method to hyperparameter choices? Do multiple runs with the same hyperparameters yield consistent performance?

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

None

Summary

Summary: This paper introduces a framework for masked diffusions that consolidates previous research on the topic and organizes it into a cohesive structure. The authors also present a generalized model within this framework, which enables the use of state-dependent masking schedules and optimization of scheduler parameters.

Strengths

1. The GenMD4 framework offers a valuable approach to optimize the forward process. In earlier studies, forward processes were typically manually designed and set within the model. However, GenMD4 adjusts the forward distribution to align with the estimated distribution, thereby improving the forward process. This innovation may serve as a source of inspiration for developing more effective forward processes. 2. This paper summarizes previous formulations of masked diffusion models and establishes the connections between them.

Weaknesses

1. In line 90. The handling of $p(x _0|x _{t(1)})$ could be enhanced. Assuming $p(x _0|x _{t(1)}) \propto q(x _{t(1)} | x _0)$ is equivalent to assuming that $q(x _0)$ is uniformly distributed. In reality, it should be treated the same as other $p(x _s|x _t)$. 2. In line 114. When discussing multidimensional data, it is not straightforward to assume that the backward process factorizes across tokens. This is because the distribution $p(x _0)$ does not factorize across tokens. Achieving factorization necessitates a small time step dt, which may not be easily observable. Additionally, in the previous single-token scenario, dt dose not need to be small, indicating that one step is sufficient to model the distribution $p(x _0 | x _1)$. This aspect is crucial for multidimensional data and should be emphasized in a fundamental paper like this. 3. In append F. The presence of a non-zero $\alpha _1$ may result in the "medium brightness problem" [1]. However, there is no singularity when $\alpha _1$ is zero if log-SNR is not introduced, and the time interval can be extended to [0, 1]. 4. In append G2. When applied to masked diffusion, $R_{kj}$ is zero when $ j \ne k$ and $j \ne m$. Given that $R_{kk} + R_{km} = 0$, $\tilde{q}$ can only take on one value (m), resulting in no additional variance. 5. In image experiments, MD4 employs masked noise, while $\tau$LDR uses Gaussian noise. We recommend conducting experiments with the same noise scheduler to demonstrate conclusively that MD4 is superior. If the goal of this paper is solely to establish that masked noise outperforms Gaussian noise, we recommend explicitly stating this claim. Additionally, we advise detailing the sampling method, as variations in methodology can influence the quality of generated samples. [1] Common Diffusion Noise Schedules and Sample Steps are Flawed, Lin et al., 2024

Questions

1. GenMD4 has not been tested on image datasets. Could you please share the results of GenMD4 when applied to image datasets? 2. Since introducing GenMD4 results in additional variance, what if all tokens share the same w (referred to as "simplified-GenMD4")? This would result in less variance. Given that GenMD4's performance is close to MD4, can simplified-GenMD4 achieve the same BPC?

Rating

7

Confidence

5

Soundness

3

Presentation

4

Contribution

3

Limitations

The method is only applied to masked diffusions.

Correction for Typos

Dear Reviewer, We noticed a few typos in the original rebuttal and would like to correct them: * It should be $\\alpha_{t(1)} = 1 - 10^{-4}$ in the response to "The handling of 𝑝(𝑥0|𝑥𝑡(1)) could be enhanced" * The equation in the response to "Appendix G2: Variance of Campbell et al. (2022)" should be $$ \beta(t) \sum_{i: x_{t, i} \neq m} \log (x_{t, i}^\top \mu(x_t(i\to m))_i) $$

Update

>The AR baselines use old transformer models hence the comparison isn't fair. By newer models, I was referring to newer transformer architecture variations like the Llama 2 flavor of transformers or transformers, transformers with RoPE etc. The comparison to SEDD is fair comparison as noted, but it would be good to see more AR. Also non-AR non-sequence-based diffusion baselines should still be listed in Table 3 for the image models as the performance is nowhere near SOTA for non-sequence based non-AR models and it may confuse readers who are not familiar with generative models. Comparison to SOTA Gaussian diffusion models should be drawn, even though they are unfavorable to this class of the methods. I would encourage the author to study the numerical stability of their objective more in the final paper and it could make the submission much stronger. I have raised my score 6 to reflect that some of my concerns have been addressed.