SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We plot both of our scaling laws in Figure 2. Our first finding is that over many orders of magnitude, the compute-optimal log-likelihood of Plaid models closely matches a power law function of the compute. Surprisingly, we find that the slopes of both the autoregressive and diffusion scaling laws are almost exactly the same. These results validate Plaid's scalability and suggest that we can obtain strong improvements by training at larger scale.
Regardless of scale, Plaid models require a constant factor of about  $64 \times$  more compute to match their autoregressive equivalents. While this factor is large, our work represents the very first attempt at efficient diffusion model training and focused engineering effort on constant-factor improvements to diffusion models may enable them to perform similarly to autoregressive models in the future.
Table 2: Plaid 1B outperforms GPT-2 124M in zero-shot likelihood across six benchmark datasets from Radford et al. [29]. Our GPT-2 numbers differ from the originals due to striding and detokenization (see Section 6.1).
PTB (PPL)	enwik8 (BPC)	text8 (BPC)	WikiText2 (PPL)	WikiText103 (PPL)	1BW (PPL)
Plaid 1B (ours)	74.33	1.18	1.12	29.42	28.28	77.64
GPT-2 124M	87.97	1.24	1.22	35.01	35.92	87.85
GPT-2 345M	64.92	1.09	1.11	26.80	26.13	67.34
GPT-2 762M	53.42	1.04	1.06	23.30	22.24	59.48
GPT-2 1.5B	47.59	1.00	1.02	21.33	20.13	54.09
Table 3: Chosen unconditional samples from Plaid 1B demonstrate fluent syntax and long-range coherence. See Appendix E for un-picked random samples.
New research rolled out at an annual scientist meeting finds that the industry will need to recover between 4,000 and 10,000 tons every year of fracked and produced oil and gas fossil reserves in order to do that, according to an analysis done by lead author Dr. Ernesto Monteiro of the University of Alberta in Canada. \n\ nThe eye-watering figure represents the total amount of oil and gas - shale and natural gas produced, extracted and sold - will likely need to be recovered in coming years to meet the carbon mitigation goals. ...[698 words omitted]... the team concluded in a report published in an academic journal prepared for the annual meeting of the National Academies of Sciences.	The Barcelona Golf Course doesn’t look like a golf course, but it is an oasis for gardening in this busy city. \n\ nIt’s a massive course spread over 120 acres with fairways that are split right down the middle, unlike the designs on most golf courses. The course enjoys the stunning view of the skyline above it. \nA giant oak tree almost 40 meters in diameter serves as one main highlight to the golf course’s design. …[414 words omitted]… The new golf course is accessible on the area’s busy streets with shops and restaurants, so the community can enjoy all the leisure activities in the green space. \n\ nThe team uprooted the previously existing Pérez Tree, to make room for the new trees to complement Gillet Park.
5.3 Compute-optimal training recipe
Our next goal is to understand how to optimally use a given compute budget  $C$  to maximize the held-out likelihood of a model. Specifically, we must choose between training a large model for fewer iterations or training a small model for longer. For this, we leverage our parameter scaling law  $N^{*}(C)$  which predicts the optimal model size given a compute budget.
We plot both of our parameter scaling laws in Figure 3 and again find that the trends have nearly the same slope but differ by a constant factor. Specifically, compute-optimal Plaid models should be about  $4 \times$  smaller (and therefore trained for  $4 \times$  longer) than compute-optimal autoregressive models. The large gap in compute-optimal settings suggests that selecting model sizes based on existing scaling laws [15, 13], which were developed for autoregressive models, could incur a substantial loss in the effective compute budget.
6 Plaid 1B
To demonstrate the scalability of Plaid models and achieve our goal of outperforming an autoregressive model in likelihoods, we train, evaluate, and release a large Plaid model called Plaid 1B. Plaid 1B is a Transformer-based Plaid model with 1.3B parameters, trained for 314B tokens on OpenWebText2 [7]. In total, Plaid 1B was trained for  $2.5 \times 10^{21}$  FLOPs, which to our knowledge equals the largest purely diffusion-based language model trained in prior work [6]. We give further training details in Appendix C.
6.1 Likelihood evaluation
We evaluate Plaid 1B's likelihood in a zero-shot setting on a suite of six benchmark datasets originally used in Radford et al. [29]: Penn Treebank [25], enwik8 and text8 [14], WikiText2 and WikiText103 [26], and the One Billion Word corpus [2].
Table 4: Chosen conditional samples from Plaid 1B in different zero-shot control settings. Highlighted spans are prompts. See Appendix E for un-picked random samples.
Prefix completion:
Generative models of text are very versatile: they can be used as a data classification model and also incorporated into multiple data processing engines. In this article, we present two new neural memory models capable of processing terabytes of data and the neural networks and computational techniques that are used in those models.
Infilling:
A year ago in Paris, prior to the tournament, I went to Elijah's to eat and get drunk. Everyone in the venue was seventeen. I was there for a few minutes and then I went back to the event. Wow, what a great day! So relaxed and too happy. I do not think I was always like that.
Token-level weights  $(5\times$  weight on "law"):
Let's talk about law and medicine. n\n\nIn her dissent, Justice Ron Sen, a veteran administrative law judge, points out that the decision "ignores the fact that the original separation agreement was reached by binding arbitration" that responded to "the legitimate ethical concerns of the university administration," which is what lies "at the heart of law and medicine."
Token-level weights  $(5\times$  weight on "medicine"):
Let's talk about law and medicine. In part because of advancements in technology, personal information about medical and drug use is spreading. Healthcare professionals across the nation rely on this personal data to make decisions about drug prescriptions and clinical trials and monitor people at immediate risk of serious or chronic diseases.
Lexical constraints ("Donald" anywhere):
Also facing legal challenges is Donald Trump's executive order banning immigration from seven Muslim-majority countries that is facing a temporary halt, with nothing scheduled to go into effect. Two federal judges have ruled that such an order violates the establishment clause.
Composition and negation ("Donald" anywhere and "Trump" nowhere):
A month later, with little time to spare, the government hired Donald V. Davis, a former senior aide to Senator Tom Mondale of Minnesota and former Chief Security Operations Officer at the White House, to lead tactical centers.
Radford et al. [29] use sliding windows of size 32 in their likelihood computation. As a non-autoregressive model, Plaid doesn't support sliding-window likelihood evaluations, so we use non-overlapping 1024-token sequences when computing likelihoods. Following Radford et al. [29], we use heuristic invertible detokenizers for PTB, 1BW, and WikiText to minimize the effect of tokenization artifacts on the perplexity results. For a fair comparison, we also recompute GPT-2 likelihoods using the same protocol, resulting in different numbers than Radford et al. [29].
In Table 2 we observe that Plaid 1B consistently outperforms the 124M parameter GPT-2 model, demonstrating that diffusion models are capable of scaling to perplexities on par with a small modern autoregressive model.
6.2 Unconditional samples
We generate from Plaid 1B by starting from  $z_{1} \sim \mathcal{N}(0, \sigma^{2}(1)I)$ , performing ancestral sampling of  $p(z_{t - (1 / T)}|z_t)$  for  $T = 4096$  steps, and finally arg max  $p_{\theta}(x|z_0)$ . Following Dieleman et al. [6], we sample using a score temperature of  $\tau = 0.9$ , which in our formulation corresponds to adding  $\frac{1 - \tau}{\tau} (\hat{x}_{\theta}(z_t) - z_t)$  to  $\hat{x}_{\theta}(z_t)$  at each step.
We generate unconditional samples with sequence length 1024 and present chosen samples in Table 3. We observe that the model is capable of generating fluent text and remaining on-topic over several hundred words. We provide random un-picked samples in Appendix E.
6.3 Zero-shot control
Although Plaid models are trained in a purely unconditional fashion, we present a zero-shot control technique called token guidance which allows us to implement a number of conditioning structures at generation time. We begin with classifier guidance, a technique which allows diffusion models to generate samples conditioned on an arbitrary attribute  $y$ . Classifier guidance first trains a probabilistic classifier  $p(y|z_t)$  of  $y$  given noisy latents  $z_t$ , and then biases the diffusion model's sampling steps by
2We choose the splitting boundaries using the Plaid tokenizer, yielding sequences slightly shorter than 1024 tokens under the GPT-2 tokenizer.
a guidance term derived from the gradient of the classifier probability  $\nabla_{z_t}\log p(y|z_t)$ . Now, recall from Section 3.2 that our denoiser  $\hat{x}_{\theta}(z_t)$  is parameterized in terms of a model  $f_{\theta}(z_t)$  which learns the distribution over the token  $x^{(i)}$  at each position  $i$  given  $z_{t}$ . We can therefore implement many different conditioning structures via classifier guidance on probabilities derived from  $f_{\theta}$  itself. We give a few examples:
Conditioning on a span: We perform guidance on the joint probability of the span under the factorial model  $p(x^{(a:b)}|z_t) \propto \prod_{i=a}^b p(x^{(i)}|z_t)$ , where  $f_\theta$  estimates each factor in the product. This lets us implement prefix completion and infilling as special cases. Lexical constraints: In order to condition on the presence of a token without specifying its location, we perform guidance on the token's probability under the unigram distribution  $p(x^{(\text{any})}|z_t) \propto \sum_i p(x^{(i)}|z_t)$ , where  $f_\theta$  estimates each term in the sum. Token-level weights: We can emphasize a specific conditioning token by multiplying the corresponding guidance term by a scalar weight. Negation: We condition on the negation of an attribute  $y$  by performing guidance on the complement probability  $1 - p(y|z_t)$ .
Using Plaid 1B and token guidance, we generate samples under various zero-shot control settings. We present chosen samples in Table 4 and random samples in Appendix E. Despite being trained unconditionally, Plaid 1B is able to follow diverse conditioning structures.
We contribute to a growing body of work on diffusion-based language models [21, 3, 9, 33, 10, 6, 8, 24, 37, 23, 36, 11]. Our biggest departure from those works is that we aim for strong likelihood performance, which to our knowledge has not been attempted in any prior work except for an appendix result from Li et al. [21]. We therefore benchmark against well-known autoregressive models instead of prior diffusion language models.
The work most comparable to ours is CDCD [6], which is also a strong general-purpose diffusion language model. However, without the ability to use standard likelihood-based benchmarks [25, 14, 26], it is difficult to say precisely where CDCD stands in comparison to autoregressive models: in every result, either CDCD underperforms the autoregressive baseline, or the evaluation metric saturates and lacks the statistical power to distinguish the models. Many of the other works above share similar difficulties. In contrast, our likelihood-based approach enables unambiguous comparisons to widely-known models.
Other diffusion language model works consider more constrained settings like controllable generation [21] or sequence-to-sequence tasks [3, 9, 8, 37, 23, 36], or propose hybrid approaches involving pretrained autoregressive models [24, 11]. Particularly, in concurrent work, Han et al. [11] finetune an OPT 13B [38] model into a hybrid model which is autoregressive over 25-token blocks and uses diffusion within blocks. Compared to their work, we focus on the more general setting of training a fully diffusion-based language model from scratch.
Finally, our work builds on recent advances in diffusion models for images [32, 12, 30, 18, 5]. Most notably, we adopt the framework of Variational Diffusion Models [18] and extend it to language modeling.
In this work, we have taken the first steps toward a competitive likelihood-based diffusion language model. We built Plaid 1B, which matches GPT-2 124M in likelihood by combining several algorithmic improvements and a scaling law analysis. Our ablations show that maximizing likelihood does not substantially harm sample quality, and we show samples are fluent in both unconditional and zero-shot conditional settings. Despite this progress, substantial work remains: Plaid narrows the compute-efficiency gap between diffusion and autoregressive language models to  $64 \times$ , and we view this gap as a tractable and exciting open problem that may be addressed with further research.
We thank Sander Dieleman, Lisa Li, and John Thickstun for valuable conversations. This work was supported by an Open Philanthropy grant.
[1] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[2] C. Chelba, T. Mikolov, M. Schuster, Q. Ge, T. Brants, P. Koehn, and T. Robinson. One billion word benchmark for measuring progress in statistical language modeling. arXiv preprint arXiv:1312.3005, 2013.
[3] T. Chen, R. Zhang, and G. E. Hinton. Analog bits: Generating discrete data using diffusion models with self-conditioning. arXiv preprint arXiv:2208.04202, 2022.
[4] Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 35:16344-16359, 2022.
[5] P. Dhariwal and A. Q. Nichol. Diffusion models beat GANs on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
[6] S. Dieleman, L. Sartran, A. Roshannai, N. Savinov, Y. Ganin, P. H. Richemond, A. Doucet, R. Strudel, C. Dyer, C. Durkan, C. Hawthorne, R. Leblond, W. Grathwohl, and J. Adler. Continuous diffusion for categorical data. arXiv preprint arXiv:2211.15089, 2022.
[7] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy. The Pile: An 800gb dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027, 2020.
[8] Zhujin Gao, Junliang Guo, Xu Tan, Yongxin Zhu, Fang Zhang, Jiang Bian, and Linli Xu. Diffformer: Empowering diffusion model on embedding space for text generation. arXiv preprint arXiv:2212.09412, 2022.
[9] Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and LingPeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. arXiv preprint arXiv:2210.08933, 2022.
[10] Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. Ssd-lm: Semi-autoregressive simplex-based diffusion language model for text generation and modular control. arXiv preprint arXiv:2210.17432, 2022.
[11] Xiaochuang Han, Sachin Kumar, Yulia Tsvetkov, and Marjan Ghazvininejad. Ssd-2: Scaling and inference-time fusion of diffusion language models, 2023.
[12] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[13] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, O. Vinyals, J. W. Rae, and L. Sifre. An empirical analysis of compute-optimal large language model training. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[14] Marcus Hutter. The human knowledge compression contest. http://prize.hutter1.net, 2012.
[15] J. Kaplan, S. McCandlish, T. Henighan, T. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361, 2020.
[16] T. Karras, M. Aittala, T. Aila, and S. Laine. Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[17] N. S. Keskar, B. McCann, L. R. Varshney, C. Xiong, and R. Socher. CTRL: A Conditional Transformer Language Model for Controllable Generation. arXiv preprint arXiv:1909.05858, 2019.
[18] D. Kingma, T. Salimans, B. Poole, and J. Ho. Variational diffusion models. Advances in Neural Information Processing Systems (NeurIPS), 2021.
[19] B. Krause, A. D. Gotmare, B. McCann, N. S. Keskar, S. Joty, R. Socher, and N. F. Rajani. GeDi: Generative Discriminator Guided Sequence Generation. arXiv preprint arXiv:2009.06367, 2020.
[20] Yoav Levine, Noam Wies, Or Sharir, Hofit Bata, and Amnon Shashua. The depth-to-width interplay in self-attention. arXiv preprint arXiv:2006.12467, 2020.
[21] X. L. Li, J. Thickstun, I. Gulrajani, P. Liang, and T. Hashimoto. Diffusion-LM improves controllable text generation. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
[22] C.-C. Lin, A. Jaech, X. Li, M. Gormley, and J. Eisner. Limitations of Autoregressive Models and their alternatives. In North American Association for Computational Linguistics (NAACL), 2020.
[23] Zhenghao Lin, Yeyun Gong, Yelong Shen, Tong Wu, Zhihao Fan, Chen Lin, Weizhu Chen, and Nan Duan. Genie: Large scale pre-training for text generation with diffusion model. arXiv preprint arXiv:2212.11685, 2022.
[24] Justin Lovelace, Varsha Kishore, Chao Wan, Eliot Shekhtman, and Kilian Weinberger. Latent diffusion for language generation. arXiv preprint arXiv:2212.09462, 2022.
[25] M. Marcus, B. Santorini, M. A. Marcinkiewicz, and A. Taylor. Treebank-3, 1999.
[26] S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models. arXiv preprint arXiv:1609.07843, 2016.
[27] N. Mostafazadeh, N. Chambers, X. He, D. Parikh, D. Batra, L. Vanderwende, P. Kohli, and J. Allen. A corpus and cloze evaluation for deeper understanding of commonsense stories. In North American Association for Computational Linguistics (NAACL), 2016.
[28] OpenAI. Gpt-4 technical report, 2023.
[29] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. OpenAI blog, 1(8):9, 2019.
[30] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
[31] C. Sahara, W. Chan, S. Saxena, L. Li, J. Whang, E. Denton, S. K. S. Ghasemipour, B. K. Ayan, S. S. Mahdavi, R. G. Lopes, T. Salimans, J. Ho, D. J. Fleet, and M. Norouzi. Photorealistic text-to-image diffusion models with deep language understanding. arXiv preprint arXiv:2205.11487, 2022.
[32] J. Sohl-Dickstein, E. Weiss, N. Maheswaranathan, and S. Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning (ICML), 2015.
[33] Robin Strudel, Corentin Tallec, Florent Alché, Yilun Du, Yaroslav Ganin, Arthur Mensch, Will Grathwohl, Nikolay Savinov, Sander Dieleman, Laurent Sifre, et al. Self-conditioned embedding diffusion for text generation. arXiv preprint arXiv:2211.04236, 2022.
[34] B. Wang and A. Komatsuzaki. GPT-J-6B: A 6 billion parameter autoregressive language model. https://github.com/kingoflolz/mesh-transformer-jax, 2021.
[35] Greg Yang, Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs v: Tuning large neural networks via zero-shot hyperparameter transfer. arXiv preprint arXiv:2203.03466, 2022.
[36] Jiasheng Ye, Zaixiang Zheng, Yu Bao, Lihua Qian, and Mingxuan Wang. Dinoiser: Diffused conditional sequence learning by manipulating noises. arXiv preprint arXiv:2302.10025, 2023.
[37] Hongyi Yuan, Zheng Yuan, Chuanqi Tan, Fei Huang, and Songfang Huang. Seqdiffuseq: Text diffusion with encoder-decoder transformers. arXiv preprint arXiv:2212.10325, 2022.
[38] Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, Todor Mihaylov, Myle Ott, Sam Shleifer, Kurt Shuster, Daniel Simig, Punit Singh Koura, Anjali Sridhar, Tianlu Wang, and Luke Zettlemoyer. Opt: Open pre-trained transformer language models, 2022.
Plaid 1B Sample Quality We compare unconditional samples from Plaid 1B and GPT-2 through a human study. We generate samples of length 128 from both GPT-2 and Plaid 1B ( $n = 1023$  samples each), and repeatedly ask Mechanical Turk crowdworkers to choose the most coherent sample from a pair (one GPT-2, one Plaid, blinded and randomly ordered). The crowdworkers found the models comparable, with the  $95\%$  CI for the win rate of Plaid ranging from 0.47 to 0.55, with a mean win rate of 0.51.
Sampling Timesteps When generating samples, we attempt to approximate the infinite-timestepuse a naive sampler (vanilla ancestral sampling) paired with a very large number of timesteps than necessary (4096) to approximate the continuous-time limit. To validate that our choice of 4096 steps is sufficient to approximate the continuous-time limit (with respect to human judgement), we run a human binary choice preference study of unconditional Plaid 1B samples generated with 2048 and 4096 steps ( $n = 1024$  samples each). We find that crowdworkers do not strongly prefer either choice, with the 4096-step samples having a win rate of 0.51 (CI: [0.47, 0.53]).
B VLB and heuristic weight schedules
Figure 4: VLB weight schedule and and heuristic weight schedules used in ablation experiments.
Unless otherwise noted, all models in this work are trained on a subset of OpenText2 [7] which we filter to remove documents labeled as non-English. The data is tokenized using a 32K-token BPE tokenizer which we train on the OpenText2 training split.
C.2 Architecture
We use standard pre-activation Transformers models with RMSNorm normalization layers and GeLU nonlinearities throughout. Unless otherwise noted, all Plaid models have 16 Transformer layers, which we found to be approximately optimal for our scale. We scale autoregressive model depth approximately following [20]. For efficiency, our implementation uses FlashAttention [4] and other fused kernels wherever applicable. We train at sequence length 256 for all experiments except Plaid 1B.
C.3 Optimization
We optimize all models using AdamW with parameter-specific learning rates derived by  $\mu$  Transfer [35] based on a learning rate of  $1.4 \times 10^{-3}$  at width 256. Each parameter's weight decay is set to  $\frac{4 \times 10^{-5}}{\eta}$  where  $\eta$  is that parameter's learning rate. We use a linear warmup on the learning rate and weight decay over the first 2500 steps, followed by a linear decay to zero over training. We train at batch size 256 for algorithm ablations and 128 for scaling law experiments. All of our small runs take less than 24 hours on a single A100. We perform all forward and backward computations in double precision except for the Transformer layers themselves, which happen in bf10at16 mixed precision. This comes at a negligible extra cost since the Transformer layers dominate the overall cost. Our learning rate and precision choices are optional: when implemented carefully, our method trains stably and performs well when single-precision floats and a single learning rate for all parameters is used.
C.4 Plaid 1B training