Scaling Diffusion Language Models via Adaptation from Autoregressive Models

Thank you for your time and valuable suggestions. We list our point-by-point responses and hope to address your concerns.

Weak1: qualitative examples and human evaluation

We list the unconditional generation results in Appendix C, which shows better quality with larger decoding T. For human evaluation, thank you for this suggestion. Most of the tasks can be evaluated using automatic metrics, and human evaluation is expensive. We will consider conducting human evaluation for some generation tasks if resource permits.

Weak2: significance testing

All experimental results involving random seed are averaged from the results of 3 separate trials, with a confidence interval of p < 0.01 based on t-test. Experimental results also reveal significant disparities in accuracy among different models/tasks.

Weak3: different T

We discussed the unconditional generation quality change with different T in Fig3, where less T resulted in higher PPL. For downstream tasks like multiple-choices in Table1, smaller T might result in a slight accuracy drop, we list some examples the following.

We thank reviewer rZH1 for valuable suggestions. We are happy to know that you think our experiment is comprehensive and solid. We’d like to discuss the unresolved scientific questions with you.

Weakness 1 What will happen given more computation? Will it converge to local optimum than training from scratch?

Thank you for raising this interesting question. The experiments in transferring GPT2 give us a positive signal that transferring from GPT2 using DD loss will bring us quite good DLMs (compared with SEDD and GPT2, noting that DiffuGPT trains less tokens than SEDD). However, this has not been proven in LLaMA due to the resource limitation. We believe it is an important direction that needs to be examined in the future of the DLM community but this is beyond our current paper. In our paper, we aim to show some positive signals of 7B DLMs with the limited resources. Based on this, we believe the community and us have confidence in the next steps of exploring larger DLMs.

Weakness 2 Loss in Table 3

We agree that ELBO can serve as the reference. We list the loss curve of different base models in Fig 6. Initailized with GPT2-M, we evaluate three loss functions: DD, DD-no shift, and DD-no anneal for fine-tuning GSM data. The corresponding loss and accuracy are shown in the table below. Additionally, results for DiffuGPT and DiffuLLaMa are presented for reference. All results highlight the negative correlation between the loss and accuracy.

Weakness 3 FLOPs of training

KV-cache is usually used during inference time for LLMs. During training, under the same setting, the theoretical computing FLOPs of LLM and DLM should be very close. We use utils in deepspeed to get the FLOPs of AR LLMs and DLMs: For DiffuLLaMA 256tflops, for LLaMA 255tflops per step per gpu. We train 6950 steps for DiffuLLaMA and the total FLOPs are 1.7*10^18, which refers to 15.7K GPU hours on H200.

L373-375: "In tasks that require more extensive global reasoning, such as complex mathematics and coding, DLMs consistently exhibit better performance compared to AR models that rely solely on left-to-right modeling capabilities." This is not compared with continuing pre-training AR.

Thank you for highlighting the controlled comparison between AR and DD. We have added supporting experiments at the end.

To clarify, lines L373-375 explain Table 1, while Table 3 presents a different experimental setup leading to distinct conclusions. In Table 3, we fine-tuned the GPT-2 model directly on GSM data using either AR or DD loss, without further training on the pre-training corpus. The superior performance of DD loss suggests that DD loss better leverages GSM data during fine-tuning. In Table 1, we obtain the adapted DLMs through continual pre-training on unsupervised pre-training corpus. We aim to evaluate their capacities in GSM fine-tuning and other zero-shot tasks like code infilling. The goal here is not to compare how each loss utilizes the pre-training corpus, but rather to benchmark our models with other existing models (AR LMs vs. previous DLMs vs. our DLMs). According to Table1, our DLMs are better than AR LMs in GSM and code. We update our manuscript to avoid misunderstandings.

We understand your concern regarding controlled comparisons with AR models continually pre-trained on the same corpus. To address this, we conduct a continual pre-training of GPT2 on the same corpus under the same settings as DiffuGPT. However, the loss curve flattens after 1k steps, and the zero-shot performance, shown below, indicates no improvement. As analyzed in paper [2], this may be due to the stability gap introduced by continual pre-training, leading to performance degradation. Additionally, since our used corpus is similar to the one used for GPT2’s initial pre-training, continual pre-training may offer limited new knowledge.

Unclear title in Table 2.

Thank you for pointing this out. Actually, the first column is either different models or different prompting methods, because there are different models (DD/AR), different prompting methods (ZS/FS) and different evaluation methods (hit@k). So we will change “Models” to “Settings” here. Stylizing the respective row with hit@k into gray color does not indicate it is the best method, but just for the reference to show the model’s potential upper bound (mentioned in line 434), just like many papers will list the oracle results for reference in the table. We have added clarification to our updated manuscript.

Q1 The order of Section 3.1

Thank you for pointing this out. Continuous time DD is recently discussed in other papers (see citations in the paper). It is close to our model so we list the details in section 3.1 as the preliminary of the whole section 3. We didn’t claim this as our contribution in the introduction. We will consider to move it to section 2.

Q2: Decoding length

We need to pre-define the generation length ahead. We set a max number of new tokens by using a sequence of [MASK] tokens and split at the first [EOS] token to produce the final response. We do this post processing because it is a common practice in DLMs [5][6].

Reference

[1] Scaling Laws for Neural Language Models. OpenAI 2020

[2] Efficient Continual Pre-training by Mitigating the Stability Gap. arxiv 2024

[3] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution. ICML 2024 Best paper

[4] TinyLlama: An Open-Source Small Language Model. arxiv 2024

[5] DiffuSeq: Sequence to Sequence Text Generation With Diffusion Models. ICLR 2023

[6] A Reparameterized Discrete Diffusion Model for Text Generation. COLM 2024

We thank reviewer ZRNS for the detailed review. We’re happy to know you find our method effective. We noticed your concerns are focused on certain statements in the experiment section. Below, we provide a point-by-point rebuttal to clarify any possible misunderstandings and revise our draft accordingly.

L364-368: "by observing different scales of our adapted diffusion models, we can draw the conclusion that scaling diffusion language models results in improved performance" This is not verified by training scratch.

We believe there may have been a misunderstanding regarding our statement, as it was not intended to suggest "training from scratch". Please note that our work does not discuss or focus on training diffusion models from scratch.

To clarify, our paper does not aim to investigate strict scaling laws, such as those observed in AR LLMs (e.g., openai scaling law [1]). Instead, we focus on adaptation and evaluating the adapted models. When we refer to “scaling diffusion models”, we are specifically discussing the performance of adapted models rather than models trained from scratch. It is possible that adaptation retains some of the knowledge and skills from AR LLMs, and the fact is we obtain a larger DLM (7B) which yields superior performance than smaller models. So the conclusion we draw is that larger adapted models perform better than smaller adapted models. In response to your feedback, we have revised lines L364-L368 to ensure our intention is clearer and to avoid any misinterpretation. We hope this revision addresses your concern.

L367-369: "This discrepancy is presumably attributed to the extensive amount of training tokens processed by the DiffuGPT series models, whereas the 7B model has not undergone equivalent training." This is not verified.

Considering the tendentious resources required for training 7B models, we stopped at 65B tokens, so we just assume the possible reason of DiffuLLaMA’s performance in this statement is the insufficient training.

There might be misunderstandings because you think DiffuLLaMA has been pre-trained on more than twice as many tokens (65B vs 30B) than DiffuGPT but actually it is not correct. 65B vs 30B is the amount of data used, but not the amount of training tokens. We should emphasize: (i) DiffuGPT training iterates more than one epoch on prepared 30B tokens (see the footnote in page 6) and in the end it consumes more than 100B tokens. DiffuLLaMA only iterates once using 65B tokens. (ii) 7B scale is much larger than GPT2-small and medium size, and according to scaling law [1] in AR LLMs, 65B tokens are definitely insufficient for 7B models.

We rewrite L367-369 to avoid misunderstanding, please let us know if you have any additional questions.

Regarding the training of DiffuGPT and DiffuLLaMA on different datasets, we mentioned this in Appendix (Line 1043) and provide further clarification here. Our dataset selection aims to align with their respective pre-training objectives. Since we train on a relatively smaller number of tokens and do not intend to introduce new capabilities, we prioritize maintaining continuity with the models' pre-training distributions to minimize distributional shift [2].

For GPT2-based DiffuGPT, we use the FineWeb dataset, which closely resembles OpenWebText (the dataset used in GPT2's pre-training). Considering that SEDD [3] iterates OpenWebText for more than 1 epoch and the total training amount is around 200B tokens, we also iteratively train DiffuGPT on 30B FineWeb data to more than 100B training tokens. In contrast, LLaMA is pre-trained over 1T tokens within one epoch on web and code data. To align with this, we follow TinyLLaMA [4] and use a mixture of SlimPajama and Starcoder data, designed to reflect LLaMA 2's pre-training data. While we acknowledge that these dataset choices may result in performance differences, it is important to note the reasonable overlap between SlimPajama and FineWeb due to their similar data sources. We update this to the dataset section of the appendix and the dataset mixtures for diffusion modeling could be explored in future work.

Weak3: True advantages of DLMs

The advantage of DLMs has been discussed and fairly compared in the previous peer-reviewed work, like in math [4] and RL policy [5]. Additionally, there are concurrent studies under submission to ICLR, including [6] and [7]. These works highlight the unique strengths of diffusion models in search and planning tasks, largely due to their bi-directional modeling capabilities. In our work, we primarily focus on leveraging the pre-trained weights of large language models (LLMs) to efficiently train DLMs. To show its strength, we validate that DiffuLLaMA/DiffuGPT outperforms LLaMA/GPT2 on the GSM8K fine-tuning and infilling setting. To further examine their advantages mentioned in [4] (self-correction) and [6] (global planning), we follow [4] to conduct qualitative analysis and observe similar self-correction ability in DLMs. We show the last few steps of sampling trajectories and can observe that DLMs will refine the intermediate numbers without the restriction of left-to-right:

<<3*15=45>> <<4*45=180>> <<180+300=00>> ####  00</s>
<<3*15=45>> <<4*45=180>> <<180+400=000 #### #### 000</s>
<<3*15=45>> <<4*45=180>> <<180+400=400>> #### 480</s>
<<3*15=45>> <<4*45=180>> <<180+300=500>> #### 580</s>
<<3*15=45>> <<4*45=180>> <<180+300=580>> #### 580</s>
<<3*15=45>> <<4*45=180>> <<180+300=480>> #### 580</s>
<<3*15=45>> <<4*45=180>> <<180+300=480>> #### 580</s>
<<3*15=45>> <<4*45=180>> <<180+300=480>> #### 480</s>
<<3*15=45>> <<4*45=180>> <<180+300=480>> #### 480</s>

For global planning, [8] only shows the potential global planning capacity but not verify it and our model settings are different. Thus we follow [6] to finetune DLMs on counting down (CD) datasets. CD is a mathematical reasoning challenge and is a generalized version of the game of 24 which even advanced models such as GPT-4 struggle with [12]. We compare DiffuGPT with other AR baselines with different model size, showing the global planning advantage of DLMs as demonstrated in [6].

While we acknowledge that the generalization capabilities of DLMs remain an open question, further exploration of their performance across a broader range of downstream tasks—particularly using 7B-scale DLMs—could be a promising direction for future work.

For infilling tasks, we also try to query LLaMA model by prompting “given the <prefix> and <suffix>, please answer the <middle> part”, which includes both prefix and suffix information. However, this approach is no better than performing text completion with solely prefixes, likely because the LLaMA model needs tuning for filling in the middle (like codellama and [9]). Additionally, [9] notes that using autoregressive (AR) models for infilling presents challenges, such as prompting difficulties and repetition. In contrast, DLMs are naturally suited for this task, as they are trained to handle masked inputs [10], which is a key advantage.

Additionally, we conduct a controlled experiment by training both AR and DLMs on 100M tokens from the Starcoder dataset, using CodeLLaMA as the base model and evaluating performance on HumanEval infilling. We finetune CodeLLaMA autoregressively with FIM [9] in both suffix-prefix-middle (SPM) and prefix-suffix-middle (PSM) formats. Our results in the following table show that DiffuCodeLLaMA outperforms PSM, suggesting that prompt format affects AR models but not DLMs. Furthermore, we believe that training on more than 100M tokens, which is relatively small, could enhance performance.

Reference

[1] DiffusionBERT: Improving generative masked language models with diffusion models. ACL 2023

[2] Diffusion language models can perform many tasks with scaling and instruction-finetuning. Arxiv 2023

[3] Fine-Tuning by Curriculum Learning for Non-Autoregressive Neural Machine Translation. AAAI 2020

[4] Diffusion of Thoughts: Chain-of-Thought Reasoning in Diffusion Language Models. NeurIPS 2024

[5] Diffusion Policy Visuomotor Policy Learning via Action Diffusion. RSS 2023

[6] Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning. https://openreview.net/forum?id=NRYgUzSPZz

[7] Implicit Search via Discrete Diffusion: A Study on Chess. https://openreview.net/forum?id=A9y3LFX4ds

[8] PLANNER: Generating diversified paragraph via latent language diffusion model. NeurIPS 2023

[9] Efficient Training of Language Models to Fill in the Middle. OpenAI

[10] Diffusion Models for Video Prediction and Infilling. TMLR 2022

[11] Diffusion Language Models https://sander.ai/2023/01/09/diffusion-language.html

[12] Stream of search (sos): Learning to search in language. COLM 2024

We sincerely thank Reviewer XhLb for acknowledging our contributions to the DLM research. We wish to address your concerns by providing detailed responses to each of your comments.

Weak1 & 2: The difference with previous work

For the two works DiffusionBert [1] and scaling MLM [2], we discussed them in the related work (line 511). The substantial differences between [1][2] and our work include:

1. Both [1] and [2] primarily focus on mask language models (such as BERT and XLM), while our work is unique in its successful adaptation of casual language models (AR). Compared with MLMs, resources allocated to AR large language models (LLMs) are significantly more today, and the scalability of LLM is more promising.
2. Despite of the benefits of adapting AR LLMs, it is non-trivial to achieve this. As discussed in [2], attempts to develop DLMs from LLaMA were unsuccessful, indicating that adapting AR LLMs to DLMs is more complex than adapting MLMs to DLMs. Theoretically, the gap between causal LLM and DLM is more pronounced and to address this, we bridge this gap in sec 3.2 and propose a series of techniques in sec 3.3. We are the first to successfully adapt AR LLM to a DLM.
3. The evaluations conducted in [1] and [2] do not compare to state-of-the art AR LLMs (they mostly focus on test set PPL and specific tasks like translation). Our work develops a comprehensive evaluation across diverse tasks including QA/multiple-choices/code/math, facilitating a consistent assessment alongside popular LLMs. We also evaluate in zero-shot, few-shot and fine-tuning settings.

For paper [3], thanks for pointing this related work out. After reading this paper, we share similar ideas about leveraging existing AR models but the technical details differ: they mix the input sequence and the attention mask of AT/NAT, while we directly shift the input sequence and use the intermediate state between the two attention masks. Besides, diffusion models and NAT are easier to train on labeled data (like seq2seq tasks- machine translation). However, training DLMs on unlabeled data is more challenging due to the sample efficiency for discrete diffusion training [11].

Overall, thank you for pointing out the need for thorough discussion and we have added more details in the introduction and related work. Please check our updated manuscript.

Thank you for suggestions and sorry for the late reply! MDLM is a great work! This work is already added in our camera ready version!