Despite a growing interest in diffusion-based language models, existing work has not shown that these models can attain nontrivial likelihoods on standard language modeling benchmarks. In this work, we take the first steps towards closing the likelihood gap between autoregressive and diffusion-based language models, with the goal of building and releasing a diffusion model which outperforms a small but widely-known autoregressive model. We pursue this goal through algorithmic improvements, scaling laws, and increased compute. On the algorithmic front, we introduce several methodological improvements for the maximum-likelihood training of diffusion language models. We then study scaling laws for our diffusion models and find compute-optimal training regimes which differ substantially from autoregressive models. Using our methods and scaling analysis, we train and release Plaid 1B, a large diffusion language model which outperforms GPT-2 124M in likelihood on benchmark datasets and generates fluent samples in unconditional and zero-shot control settings.
Paper
Similar papers
Peer review
Summary
This paper is concentrated on training a Diffusion LM. The authors chose Variational Lower Bound loss used in Diffusion LM as a core of the proposed framework named Plaid, as well as some improvements on the training procedure and architecture of a model (usage of categorical reparametrization, use of trained categorical mapping for learning conditional likelihood term, and usage of self-conditioning method). Authors evaluated proposed changes in the ablation study compared to reimplemented CDCD measured by NLL, derived scaling laws for their model. They trained the Plaid 1B model, outperforming GPT-2 124M in NLL. Also, the authors provided samples from the Plaid 1B model and showed the ability of the model to perform conditional and unconditional sampling.
Strengths
- Authors proposed improvements on top of VLB loss, which appear to be useful measured by NLL on test data within the ablation study. - Categorical reparametrization is an exciting way to mix strengths of score interpolation objective from CDCD with VBL loss which do not imply heuristic constraints on training (e.g., normalization of the embedding matrix). - The paper is well organized and mainly clear (see weaknesses section)
Weaknesses
Experiments: - Specific weightings A and B in Section 4.1 are too artificial. It will be nice to have an explanation in the paper as to why such weightings were selected. - Information regarding pre-trained embeddings in Section 4.1 is necessary. Stating "all models use fixed embeddings obtained from a previously-trained known-good model" does not provide any helpful information regarding the details of the experiment. - There needs to be more information regarding human resources used for crowd working in Section 4.1, while the authors claimed that Human Subjects are N/A. - The authors stated that results from Section 4.1 suggest that VLB performs at least comparably to other choices. Though, I need to see how such a conclusion could be made from this experiment. As for me, it states that if one tries to replace theoretically justified weighting with a specific hand-crafted one, the quality of samples will reduce. If it was the aim of this experiment, different other parts of VBL could also be changed to understand how important they are. - The evaluation of other experiments is poor. Plaid is compared to CDCD and GPT-2 only by NLL, which does not show any information on the quality of samples. E.g., removing self-conditioning leads to an increase of NLL on 0.09 points – how bad is it? Is this a minor reduction of sample quality or not? It is a critical point for me since if we want to use Plaid for text generation, it is not important what NLL values on test data it achieves. It is much more important how good samples from the model are. For the autoregressive model, it is convenient to expect that lower NLL leads to better samples, but is there such a convenience for Diffusion LMs (even trained with VLB loss)? I strongly believe that not. E.g., if one dramatically reduces the noise scale, making embeddings easily distinguishable, then it is easy to make NLL equal to 0, while samples will be poor. It is necessary to include numerical evaluation of model samples (e.g., text repetition, evaluating perplexity of samples with a third-party language model, and others) - Baselines used for experiments are also limited. E.g., if Plaid is built on top of VLB from Diffusion LM, it is necessary to include information on its performance in the paper. - The list of insights on the performance of Plaid could be more extensive. E.g., sampling from Plaid could be performed with a different number of steps. How to sample quality differs with the change in this number? How does the number of sampling steps affect inference speed? The only numerical information regarding Plaid that is available is NLL on test data, which is not sufficient. Motivation: - I see the leitmotif of the first three sections that VLB allows us to reduce dependency on heuristic design choices (e.g., authors claim other methods to be workarounds (L109)). Though the design choices of other works are heuristic, the VLB framework seems to be much more complex and harder to implement (e.g., usage of double-precision parameters for everything except transformer layers and usage of parameter-specific learning rates indicate that making Plaid work stable required a lot of efforts). Why should a practitioner use a method that requires parameter-specific LR and spend more time selecting the best-performing LR over a simple yet heuristic method? Doing so could be justified by higher sample quality, but in the current state, the paper lacks details on sample quality. - Considering the usage of double-precision and parameter-specific learning rates, stating that the likelihood-based approach "simplifies training" (L167) does not seem to be true. Reproducibility - The paper needs to include information regarding training infrastructure and the time necessary to reproduce the experiments. Also, the authors claimed "Yes" within the "compute" box for submission while not providing information regarding computing. - Self-conditioning used in the code is not common. I understand self-conditioning is only added for specific token positions with some offset (line 194 train.py). What is the motivation for doing so? Are there any other works that used such a scheme? If yes, I do not see citing. If not, I need sufficient information on such a scheme in the paper. - I do not see the trained model with supplementary materials, though authors claim that the model is publicly available Narrative: - The paper feels like a merge of two short papers: authors had proposed architectural improvements for training a Diffusion LM with VLB objective without precise studying limits of the Plaid framework and studied scaling laws of Plaid. These scaling laws could have been an interesting complement to the paper if there were more experiments on Plaid performance and behaviour. But now, scaling laws experiments look like studying the scaling laws of a model we do not understand. While at this point, I voted for the rejection of the paper, I still see potential in this work since some of the architectural choices seem interesting to me. I would happily increase the score if the drawbacks were fixed. First, possible changes include adding more metrics for Plaid since NLL value does not provide information on Plaid performance. Once these metrics are added, including more experiments to provide more insights into Plaid behavior will be highly beneficial. E.g., authors could study what trained posterior $f_{\theta}(z)$ looks like (its statistics, etc.).
Questions
Please, refer to points from weaknesses.
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
1 poor
Presentation
3 good
Contribution
2 fair
Limitations
L320 explicitly states, "Our ablations show that maximizing likelihood does not substantially harm sample quality, " though it is not valid. Ablation experiments were performed with only NLL on test data, which does not include any sample quality information. The only experiment that evaluated sample quality was a comparison with hand-crafted weighting functions.
Summary
This paper formalises variation deffusion models for text and makes several algorithmic contributions to such models. They show that these models are able to model the likelihood of text well, and do several other evaluations to validate the likelihood approach, quantify the computational requirements, perform ablations, and evaluate a method for doing conditional generation.
Strengths
I found this perspective on diffusion enlightening and was very interested to see how to use it to model text. The contributions are substantial and the evaluations are informative.
Weaknesses
This paper covers a lot of material for its length, which makes it challenging to read. Several references are made to Appendices, but I couldn't find any; these would have been useful. The model specifications in sections 3.3 - 3.5 are very terse; perhaps more citations or a longer explanation in an appendix would be helpful to many readers. Section 3.2 seems related to simplex diffusion, so some additional citations would be appropriate. The validation of the likelihood-based approach (sec 4.1) only addresses whether the model puts high probability on good texts, not whether it covers the full distribution of texts accurately. The latter is addressed by the likelihood evaluations themselves, but this is exactly what the model is trained to do. It would be better to include previous measure for generation quality, such as both forward and reverse cross entropy (https://arxiv.org/abs/1804.07972).
Questions
Where are the Appendices? Would it be possible to run a forward-cross-entropy evaluation on the generated texts, to see if they cover the full distribution? Suggestion: The caption of table 2 does not define what the numbers are.
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
4 excellent
Presentation
3 good
Contribution
4 excellent
Limitations
Only briefly mentions computational issues.
Summary
This paper aims to close the gap between autoregressive and diffusion-based language models on standard perplexity-based language modeling benchmarks. To achieve this goal, the paper proposes several algorithmic improvements for maximum-likelihood training, and studies the scaling laws of the diffusion models to find optimal training regimes. Finally, based on the improvements, the paper shows results with a 1B diffusion-based language model that outperforms GPT2 in perplexity with various analyses.
Strengths
Solid experiments and analysis. The derived recipe for training diffusion-based LM that differs substantially from the usual autoregressive model is particularly meaningful. The explored design space and open-sourced model have a good contribution to the community on further developing of diffusion-based LM as an alternative foundation model research.
Weaknesses
The basic algorithm is mostly following VDM (Kingma et, al.), and most of the proposed improvements are also seen from existing literature. Therefore, the novelty from the algorithmic side is somewhat limited. Part of the description is confusing and unclear how exactly the authors do (see Q1)
Questions
In line 111, the paper mentioned, “Plaid loss function is a bound on discrete data”. Why is this the case? Do you mean the objective in Eq.5 where x is discrete or x is one-hot vectors? In line 63, the paper mentions each x will be transformed into embedding vectors with invertible token-wise embedding function Embed(.). Is this invertible function the learnable embeddings mentioned in line 111? Why is it invertible? Why is there no x~ in the final objective in Eq.5? For categorical reparameterization, how do you keep x~ still the original embedding? Will there be some mismatch? How this relates to avoiding the model to memorizing the embedding vectors? I can get the method but find it confusing how this is motivated. Can you also do the same in the target embedding (e.g., label smoothing?)
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
Confidence
4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.
Soundness
4 excellent
Presentation
3 good
Contribution
4 excellent
Limitations
The paper does not seem to discuss the limitations. Since the findings on diffusion-based LM seems to suggest that training diffusion-based LM is much harder and may require more resources than its counterpart autoregressive models, it may be interesting to discuss the future steps or any potential combinations.
Summary
This paper propose a variational diffusion language model which is based on variational diffusion model. The work is related on some works such as combineing vae with language model.
Strengths
The application of variational diffussion model to language model is make sense at the some times, and the experiment results verify the effectivenss of the proposed model. However, the proposed model is litter novel, which is simlar to the application of vae on language models.
Weaknesses
The proposed model is a autoagressive language model, and will add more parameters compared with exciting languge models. At the same time, the sampling time of each token will become expensive. It's noted that the autoagressive language model generate each token is slow, and there are many works try their best to construct non-agressive languge models based on diffussion model. While the results are not perfect, their work can speed the generative process. Thus, in my opinion, this paper do a simple combination work to improve language model's performance, which is not very worth .
Questions
Please see the weaknesses
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
3 good
Presentation
2 fair
Contribution
2 fair
Limitations
The sampeling time is expensive. The parametes of model is large
Summary
The paper explores the likelihood-based training of diffusion language models as an alternative to autoregressive models like GPT-2. Several algorithmic improvements are proposed for maximum likelihood training of diffusion LMs, including learned noise schedule, learned embeddings, categorical reparameterization, output prior, learned conditional likelihood, and self-conditioning. Scaling laws are analyzed to derive a compute-optimal regime for training diffusion LMs that differs substantially from autoregressive LMs. The methods enable training Plaid 1B, a 1 billion parameter diffusion LM that outperforms GPT-2 124M in likelihood on benchmarks. The qualitative analysis shows that Plaid 1B generates fluent unconditional samples and demonstrates controlled generation abilities.
Strengths
- The paper shows that diffusion language models can achieve non-trivial likelihoods on standard language modeling benchmarks, outperforming a widely used autoregressive model like GPT-2 124M. This helps establish diffusion models as a promising alternative to autoregressive models. - The analysis of scaling laws and derivation of a compute-optimal training regime is insightful. Following this recipe likely enabled training a large model like Plaid 1B. - The trained model, if released, helps move the field forward, demonstrating fluent unconditional and controlled generation results.
Weaknesses
- More analysis could be provided on the sample quality of the Plaid 1B model beyond likelihood benchmarks. Samples from the model should be evaluated quantitatively. The paper focuses on likelihood, but other metrics like human evaluation of samples could better highlight the benefits of diffusion models over autoregressive ones. - The ablations studying design choices are limited to a single model scale (1B). Ablations at multiple scales could better validate conclusions. - Lack of experiment support (or theoretical explanations) of design choices, such as pre-training sequence length, word embedding dimension, the ratio of examples to compute conditional loss, etc. - Only CDCD is compared, additional comparison to other diffusion language models (such as Diffusion-LM[1]) would add context about progress in this area. [1] https://arxiv.org/abs/2205.14217
Questions
- Need more detailed derivations from Eq(4) to Eq(5) when $T\rightarrow \infty$. - Why Section 3.1 is titled learned embedding since prior diffusion models also use the learned embeddings? What the meaning of ablation learned embedding in Table 1? Use fixed pretrained embeddings? - How to choose the examples ratio between the two terms in Section 3.4? Why use $\sqrt{\frac{Var(L_{\infty})}{Var(logp_{\theta})}}$? - Why use embedding dimension 16 when training? Would it be too small to represent the information of tokens, with a vocabulary size of 32K tokens? - Plaid 1B is trained using the sequence length of 1024. Why use 1024 and what would happen if we want to use a shorter or longer sequence length, without considering the memory consumption? Some preliminary experiments show that the longer the sequence is trained, the harder for models to converge. Do you encounter this and how to solve it? - Plaid 1B is trained using the sequence length of 1024, does it mean that you chunk the pretraining text data using 1024 and batch them? If then, how to undermine the [EOS] of the sampling sentence? You truncate a small random subset of examples with shorter sequence lengths, but how to batch these samples? Did you use [PAD] token to pad the sequence?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.
Soundness
3 good
Presentation
3 good
Contribution
3 good
Limitations
Authors have not addressed the limitations in the paper.
Summary
This research studies diffusion-inspired language models with a primary focus on narrowing the perplexity gap between autoregressive and diffusion-based LMs. To achieve this goal, the paper systematically explores various design choices within diffusion-based language models, addressing questions such as the best method for learning token embeddings and the impact of self-conditioning on performance, among others. Through this comprehensive investigation into design choices, the paper successfully scales up the diffusion language model to an impressive 1 billion parameters. Experimental results conclusively demonstrate that the 1B diffusion LM outperforms GPT-2 124M in terms of perplexity across diverse benchmark datasets.
Strengths
[Motivation] The focus of this research on further enhancing diffusion-based LM is particularly intriguing, given that considerable effort is often required to bring about fundamental changes in the trends of autoregressive LM. [Experiments] The successful scaling-up of diffusion-based LM is also a strength of this paper.
Weaknesses
[Limited Novelty] While the Plaid framework claims to contribute by exploring the design space of diffusion-based language models and introducing compute-optimal MLE training of language diffusion models, its technical novelty remains limited. Many of its design choices, including self-conditioning, have been extensively studied in the past, as evident in Table 1 and previous work such as CDCD. [Comparison to CDCD] It is difficult to determine whether the Plaid framework produces better samples than CDCD based solely on Table 1. While the table shows that Plaid outperforms CDCD in terms of perplexity, it lacks a direct comparison of the generated texts. Without such a comparison, it is challenging to ascertain which model generates higher quality samples or exhibits more desirable language generation characteristics. Additional evaluation and comparison of the generated texts would provide a more comprehensive assessment of the two models' performance. [Comparison to GPT-2] The performance of Plaid 1B is comparable to GPT-2 124M. However, it is challenging to assess the significance of this achievement, as Plaid requires a significantly larger number of parameters to match the performance of GPT-2 124M. To better understand the importance of this result in the context of diffusion LM research, the paper should provide further elaboration. It would be helpful to highlight more benefits of diffusion-based LM, such as its performance on in-filling tasks and its ability to capture long-range dependencies more effectively than autoregressive models with a left-to-right order. By emphasizing these advantages, the authors can better demonstrate the value of the Plaid framework beyond simple perplexity comparisons with GPT-2.
Questions
#1. (ll. 102-105). I can’t comprehend the rationale behind why the loss functions used in CDCD result in ill-posed problems when optimized over the embedding. Since CDCD aims to be learned by minimizing the cross-entropy loss between the output of diffusion and the input tokens, it’s not an ill-posed problem. #2. (ll. 62). The embedding function should be invertible, but there is no constraint on the objective function. After training, is the embedding matrix full-rank? #3. (ll. 139). In what situations is it beneficial to allow \sigma^2(0) to take a large value? It appears that when the first forward step can utilize a large noise scale, the overall process's length is reduced, leading to faster sampling speed. Is this the intention of the authors? #4. In section 4.1, the paper lacks a clear description of the rationale behind the choice of heuristics on the weighting for the ablation study. Could you provide a detailed explanation for selecting these particular heuristics?
Rating
5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Soundness
2 fair
Presentation
2 fair
Contribution
2 fair
Limitations
The paper does not explicitly outline the limitations of the proposed framework. However, a significant drawback of the Plaid framework is its requirement for a significantly higher number of parameters compared to AR (Autoregressive) language models. This contrasts with the trend observed in diffusion models for image generation tasks. For example, when comparing DALL-E 1 to DALL-E 2 or Parti to Imagen, AR-based image generation models typically necessitate a much larger number of parameters than diffusion-based generative models. It would be beneficial to include an explanation in the paper regarding why this observed trend in parameter efficiency is not consistent with diffusion language models, or at least, why it is not evident in the Plaid framework.
Thank you for the detailed answer. I want to apologize for the inconvenience regarding the significance of NLL in the setup of your model. It indeed provides useful information on the likelihood of the data under Plaid. While you have provided an upper bound on NLL, it is clear that the tightness of this bound depends on a number of samples of $t$. Am I understanding correctly that during the generation, you still could use a different number of generation steps to follow the path from $z_1$ to $z_0$? If so, it seems like this number of steps will affect the final quality of metrics. How does this quality depend on a number of steps? You have provided a human evaluation of the comparison with GPT-2, but what if we use two times fewer steps than was used in this experiment? What if we use two times more steps? I believe that providing such information on the quality of Plaid samples in dynamic would significantly improve the paper so that I will have no more concerns about your submission. Also, while for Plaid, it is valid to estimate NLL, as you pointed out in your answer, for CDCD, this estimation does not provide useful information regarding the performance of CDCD, though, with Table 1, you still had compared Plaid to CDCD purely based on NLL. How was this evaluation performed? If you had a reproduction of CDCD (I understand that full reproduction is not possible), then you could use it to compare samples from Plaid and CDCD using automatic metrics.
Thank you for your response! Regarding step counts: We train and evaluate NLLs using equation 5 (via equation 3) which is the infinite-timestep limit of the NLL bound, so our NLL evaluation results do not depend on a chosen number of steps. When generating samples, we are similarly interested in the infinite-step limit behavior, so we use a naive sampling algorithm (ancestral sampling) combined with a much larger number of steps than we believe to be necessary (4000). To confirm that 4000 steps approximates the infinite-step limit, we will run a human study comparing samples generated with 2000 and 4000 steps and include it in the camera-ready version. Given that many more efficient sampling algorithms for diffusion models exist (e.g. DDIM, 2nd-order ODE solvers, diffusion distillation), we focus in this paper on building the strongest models we can without inference compute constraints. Investigating these sampling algorithms with Plaid would be exciting future work. Regarding CDCD: Even though we don’t use it for training, we can compute an NLL bound for CDCD using equation (3) and, as with the Plaid models, it is a correct bound on the discrete data log-likelihood. We developed, debugged, and tuned the hyperparameters of our CDCD reimplementation against the NLL bound, so we don’t think it would be fair to CDCD to evaluate that model on sample quality. We will update the draft to make this point more clear.
Sorry for misunderstanding. And my other concern is why not use the BLEU evaluate the model's performance.
Increasing my score from BR to BA.
Thank you for your comprehensive response. Most of my initial questions about the paper have been clarified. As the authors recognize the absence of human evaluation and the inefficiencies in the network parameters, the manuscript could be further enhanced from another round of revisions. Nonetheless, since my primary concerns have been addressed, I've adjusted my score to BA (borderline accept). Thank you!
Decision
Accept (poster)