Scaling up Masked Diffusion Models on Text

Authors: Shen Nie, Fengqi Zhu, Chao Du, Tianyu Pang, Qian Liu, Guangtao Zeng, Min Lin, Chongxuan Li
Primary area: generative models
Keywords: Masked Diffusion Models, Scaling Laws, Conditional Generation, Language Understanding, Reverse Curse, Temporal Quality Degradation

Abstract:
Masked diffusion models (MDMs) have shown promise in language modeling, yet their scalability and effectiveness in core language tasks, such as text generation and language understanding, remain underexplored. This paper establishes the first scaling law for MDMs, demonstrating a scaling rate comparable to autoregressive models (ARMs) and a relatively small compute gap. Motivated by their scalability, we train a family of MDMs with up to 1.1 billion (B) parameters to systematically evaluate their performance against ARMs of comparable or larger sizes. Fully leveraging the probabilistic formulation of MDMs, we propose a simple yet effective *unsupervised classifier-free guidance* that effectively exploits large-scale unpaired data, boosting performance for conditional inference. In language understanding, the 1.1B MDM outperforms the 1.1B TinyLlama model trained on the same data across four of eight zero-shot benchmarks. Notably, it achieves competitive math reasoning ability with the 7B Llama-2 model on the GSM8K dataset. 
In text generation, MDMs with 16 times more pre-training time offer a flexible trade-off against ARMs with the accelerated sampling technique KV-Cache: MDMs match ARMs in performance while being 1.4 times faster during sampling.
Moreover, MDMs address challenging tasks for ARMs by effectively handling bidirectional reasoning and adapting to temporal shifts in data. Notably, a 1.1B MDM breaks the *reverse curse* encountered by much larger ARMs with significantly more data and computation, such as 13B Llama-2 and 175B GPT-3. Our code is available at https://github.com/ML-GSAI/SMDM.

Decision: accept (poster)
Reviewer scores: 6, 8, 6, 6

Reviews and discussion:
<thread>
creation date:
1729084707990

thread_type:
official_review

<thread_object>
creator: reviewer_6chg

<content>
summary:
This paper explores the potential of Masked Diffusion Models (MDMs) in language tasks, highlighting their scalability and performance compared to autoregressive models (ARMs). The authors train models with up to 1.1 billion parameters and demonstrate a scaling law for MDMs. MDMs show competitive results in language understanding and provide a flexible trade-off in conditional generation compared to ARMs. Additionally, the authors find MDMs are robust to temporal shifts in data and overcome the reverse curse problem faced by ARMs.

strengths:
- Scaling MDM: The main contribution of this paper lies in scaling mask discrete diffusion (MDM) and provide the corresponding scaling law.
- Identifying distinctive features of diffusion: The authors highlight how diffusion can effectively address shortcomings like temporal shifts in data and the reverse curse found in ARM. This paves the way for additional research and application opportunities in the realm of diffusion language models.

weaknesses:
- The paper primarily compares MDM and AR, but lacks a clear comparison with other pretrained diffusion language models such as Plaid (https://arxiv.org/abs/2305.18619) and SEDD (https://arxiv.org/abs/2310.16834) under the same setting. 
- The diffusion loss provides an upper bound and should not be directly compared with the loss in AR models. The comparison in Figure 2(a) appear to lack rigor, a similar problem for the perplexity comparison detailed in Table 7.
- About the setup described in Table 6, it appears that MDM fine-tuned on the dataset, while GPT3 and LLaMA not. To mitigate the potential confounding variable of fine-tuning versus prompting, a baseline model like fine-tuned GPT-2 ought to be included for a more comprehensive analysis.
- The proposed unsupervised CFG is similar to calibration in AR (http://arxiv.org/abs/2102.09690,http://arxiv.org/abs/2104.08315), but the relationship is not properly discussed.

Minor:
- typo: line 756 “Greddy sampling“

questions:
- Why different likelihood evaluation methods result in distinct behavior on different tasks? How can we decide on the fly given a new task?
- What's the relationship between 10^20 Flops and tokens? Do both ARM and MDM consume the same amount of tokens?
- What’s the inference timesteps in Table 5? Line 409 “Built upon the unsupervised CFG”, does it means the inference timesteps are doubled? It's better to draw a curve to show the speed advantage of diffusion, which I guess is in generating longer texts.
- What’s the computation resource and GPU hours comparison for MDM and AR?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730652962517

thread_type:
official_review

<thread_object>
creator: reviewer_3nzv

<content>
summary:
The paper focuses on scaling Masked Diffusion Models (MDMs) as an alternative to Autoregressive Models (ARMs) in language modeling.  The study establishes the first scaling law for MDMs, indicating that they can achieve performance on par with ARMs at a smaller computational cost gap than previously assumed. Moreover, MDMs demonstrate competitive performance in zero-shot language understanding tasks and effectively address ARM challenges like "reversal curse" and temporal quality degradation.

strengths:
-   The introduction of scaling law for MDMs trained from scratch suggests that they show the same scaling potential as ARMs.
-   The presentation is clear and well-organized, with only minor typos.
-   Results on reversal curse tasks of MDMs is strikingly better than ARMs. This is an interesting point.

weaknesses:
\textbf{Unfair Comparison:} Line 368-369: The authors extend MDM pre-training time by a factor of 16 to achieve a “meaningful” comparison with ARMs. However, a more valid comparison would involve MDMs and ARMs trained with equivalent compute budgets or trained to convergence for a given model size. Comparing an MDM to an ARM trained with only 1/16 of its FLOPs does not yield a balanced assessment.

\textbf{Missing Citations:} The paper lacks a related work section, making it difficult for readers unfamiliar with text DMs. Notably, a closely related work [1] is omitted. This prior research examines scaling trends of MDMs fine-tuned from a pretrained MLM and evaluates MDMs on complex tasks, aligning closely with this paper's contributions.

\textbf{Weak Benchmarks:} The benchmarks used in section 5 are mostly language understanding tasks, i.e, can be mostly formulated as a one-token-completion task. These NLU tasks are of less interest now, compared to summary benchmarks like gigaword and reasoning benchmarks like GSM8K, which require more complex abilities.

\textbf{Typos:} 

Line 458 RELEVING should be relieving

Line 756 & 777 greddy should be greddy

 

[1] Diffusion Language Models Can Perform Many Tasks with Scaling and Instruction-Finetuning https://arxiv.org/pdf/2308.12219

questions:
It is acceptable that MDMs trained from scratch do not match the performance of ARMs, as suggested by the scaling law in this paper as well as previous text DM work. I think a better way to present these findings should be reporting both the pros and cons of MDMs, e.g., they are more computation-consuming but can perform some tasks that state-of-the-art ARMs cannot do at all.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730654475292

thread_type:
official_review

<thread_object>
creator: reviewer_cmg3

<content>
summary:
The paper scaled up Mask Diffusion Model based on previous work [1], demonstrate that previous work has strong scalability of Masked Diffusion Model for text. The paper also fit the scaling law of MDM, and found that MDM need 16x more training computation to match the validation loss of AutoRegressive model. 

The paper also employed simple yet effective classifier free guidance (CFG) for MDM that improved the performance of downstream tasks. 

Finally the authors evaluated the model on reverse curse, a challenging problem of autoregressive language model, and showed that MDM alleviate the reverse curse greatly. 



[1] Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data

strengths:
1. The paper scaled up the Masked Diffusion Model(MDM) for text, and experiment are performed in a solid to demonstrate the strength (bidirectional modeling of text), and weakness (has to consume more computation to reach the performance of AR language model).


2. The paper provided the research community with understanding of MDM in term of scaling and future research directions of more computationally efficient MDM, which is a significant contribution. 


3. The paper is written in a clear and easy to follow way. 


4. The supplementary materials are sufficient to review or reproduce this work.

weaknesses:
1. The major theoretical basis are given in the previous work [1], this work only performed experiments to scale up, which makes this paper less impressive and novel. 

2. In the reverse curse experiment (Table 6), it seems that MDM performance still decay in reverse direction in two dataset. Given the model encodes the text in a bidirectional way, how we can explain this decay ?

3. There are minor typo problems.

[1] Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data

questions:
1. See `2` of Weakness 

2. For CFG based generation, it seems that we need double computation to compute $\tilde{p}$. So are the results in Table 1 using the same sampling steps ? If so, is it possible that the extra computation contribution to performance improvement ? It could be better to use a doubled sampling step w/o CFG to compare the performance so that we can rule out the hypnosis

3. In Line 756  and Line 777, `Greddy` or `Greedy` ?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730821657981

thread_type:
official_review

<thread_object>
creator: reviewer_5xsm

<content>
summary:
Note: The authors addressed my concerns about disentangling masking and diffusion in their response with an additional experiment. I thus adjusted my score from 5 to 6.

-------

The paper studies scaling up masked diffusion models (MDMs) for text. The first establish scaling laws for MDMs and compare them to autoregressive models. They also propose an unsupervised CFG (classifier-free guidance) technique (based on Ho & Salimans, 2022, Chang et al. 2023) to boost performance with only unlabeled data. 

The results show that:
-ARMs and MDMs demonstrate competitive scaling rates and similar scaling behavior (the compute cost of achieving a particular validation loss z with an MDM is 16x higher than that of achieving that loss z with ARMs but this factor stays constant)

-Unsupervised CFG gives consistent gains over vanilla MDM across all the 8 QA tasks.

-MDM and ARM, when given same compute achieve similar performance (but MDM achieves strong performance when trained for 16x longer).  1.B MDM performs comparably to a 1.5B GPT-2.

-MDM has success breaking the reverse curse.

strengths:
-Exploring diffusion models for text is an interesting direction. The authors' work is well motivated. 

-The gains from Unsupervised CFG seem consistent across the 8 tasks.

weaknesses:
-I believe comparing to a a text-based LLM that uses bidrectional masking/span corruption is more fair than a left-to-right next token prediction ARM model like GPT2. This would disentangle the value of the masking (which can also be done in language models) from the diffiusion aspect.

Right now I think these two factors are conflated. For example are the reverse curve results in Table 6 because the MDM uses masking or because its a diffusion model?

-The novelty of the work is also modest.

questions:
-Please see the Weaknesses above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732033100763

thread_type:
public_comment

<thread_object>
creator: ~niklas_nolte1

<content>
title:
relevant related work

comment:
Cool paper!
Just fyi, there is related work on the reversal curse and MDM https://arxiv.org/abs/2406.05183
</content>
</thread_object>

</thread>
<thread>
creation date:
1732375303487

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Paper Revision

comment:
We thank all reviewers for their constructive feedback and have responded to each reviewer individually. We have also uploaded a **Paper Revision,** including additional results and illustrations. **All changes in revision are marked in blue.**

**For Reviewer 5XSM:**

- Section 7.1 and Table 6: We add the T5 results for the reverse curse experiment and explain why MDMs break the reverse curse.
- Section 1: We emphasize our contribution in the introduction.

**For Reviewer cmg3:**

- Section 1: We emphasize our contribution in the introduction.
- Section 7.1: We add an explanation for decay performance in the reverse direction.
- Algorithm 1 (in Lines 810): We fix the typo.
- Section 6 and Appendix C.2: We report that unsupervised CFG outperforms sampling without CFG under equal sampling computation.

**For Reviewer 3NzV:**

- Section 1 and Table 5: In the introduction and Table 5, we highlight that MDM utilizes 16 times pre-training time for the conditional generation.
- Section 8: We elaborate on the advantages and limitations of MDMs and identify potential directions for further research.
- Section 5 (in Lines 355-358) and Table 12 (in Lines 994-1001): We add the results of GSM8K.
- Appendix B.3: We add the experimental details for the GSM8K results.
- Section 3: We add a discussion about the related work.
- Section 7.2 and Algorithm 1 (in Lines 810): We fix the typos.

**For Reviewer 6Chg:**

- Section 1: We emphasize that our work focuses on exploring scalability rather than proposing a new model.
- Table 6: We detail that all models are fine-tuned on the same dataset in the reverse curse experiment.
- Section 4: We add a discussion about the related works.
- Algorithm 1 (in Lines 810): We fix the typo.
- Section 5 and Appendix C.1: We present the empirical observations about likelihood evaluation methods.
- Table 2 and Table 3: We present the number of training tokens.
- Section 6 and Appendix C.2: We report that unsupervised CFG outperforms sampling without CFG under equal sampling computation.
- Section 3: We emphasize that the relationship between training Flops and data tokens is the same for ARMs and MDMs.

**For public comment by Niklas Nolte:**

- Section 7.1: We add a discussion about the mentioned paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors scale up masked diffusion models on the text and present classifier-free guidance methods for a conditional generation. The findings are that masked diffusion models, when scaled appropriately, can perform well and compete with AR models (though not in a compute-matched way). 

The paper studies a problem that has been posed several times in the literature (does scaling up diffusion work well?) but in a context that has not been studied (MDMs, and with CFG) and shows strong results that would be of interest to the diffusion LM community. As reviewer 3NzV notes, the writing of the paper does not quite make it clear that experiments are not always compute-matched. I would strongly encourage the authors to make such un-matched comparisons explicitly upfront starting in the abstract.