SlowGuess/NeurIPS_2023_OCR huggingface dataset:
igul222@gmail.com
thashim@stanford.edu
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.
Large language models lie at the center of recent advances in artificial intelligence. Shared across nearly all such language models is a common recipe: learn a model that maximizes data likelihoods using an autoregressive, left-to-right factorization. Maximum-likelihood pretraining has been a remarkably successful paradigm, leading to models that perform well on a range of downstream tasks and display complex behaviors like in-context learning [1, 28].
Thus far, autoregressive modeling has been a core part of this process due to its computational efficiency and empirical performance. However, this choice carries drawbacks. Autoregressive models generate tokens one at a time, making it difficult to perform long-range planning or controllable generation [17, 19, 21]. In addition, certain sequence distributions may be fundamentally more difficult to model autoregressively [22].
Given the importance of language modeling, these potential drawbacks motivate us to explore alternatives to the autoregressive approach. As a promising candidate, we turn to continuous diffusion models [32, 12], which have achieved state-of-the-art results in image modeling [5, 30, 31]. In language, prior works on diffusion models exist [e.g. 21, 10, 6], but these optimize non-likelihood-based objectives. Without the ability to use standard likelihood-based benchmarks [25, 14, 26], it is difficult to say precisely how these models compare to autoregressive models (see Section 7 for a discussion). Somewhat concerningly, there is no work showing that it is possible for diffusion language models to achieve any nontrivial likelihoods on standard benchmarks.
In this work, we explore the limits of likelihood-based diffusion language models. Our goal is to train and release a diffusion model which achieves better likelihoods than GPT-2 124M [29], which we consider the smallest widely-adopted autoregressive model today. To achieve this goal, we first
<sup>1</sup>We release our code and pretrained models at https://github.com/igul222/plaid.
Figure 1: Plaid models scale predictably across five orders of magnitude. Our largest model, Plaid 1B, outperforms GPT-2 124M in zero-shot likelihood (see Table 2).
develop an algorithmic framework, then study its scaling laws to enable compute-optimal training, and finally train a large model called Plaid 1B.
1. We explore the design space of likelihood-based diffusion language models and propose an algorithmic framework called Plaid. We validate the design choices of Plaid through compute-matched ablations.
2. We study the scaling laws of Plaid training. Our analysis shows that the log-likelihood of Plaid models improves predictably with more compute. We derive a recipe for compute-optimal training which differs substantially from the usual autoregressive rule.
3. We train and release Plaid 1B, a large diffusion language model pretrained on OpenWebText2 [7]. Plaid 1B outperforms GPT-2 124M in zero-shot likelihood across six standard benchmarks. We demonstrate Plaid 1B's ability to perform fluent and controllable text generation.
2 Variational Diffusion Models for language
In this background section, we formally define continuous diffusion models over text sequences, adopting the Variational Diffusion Models (VDM) framework [18] which is a natural fit for likelihood-based training (see Karras et al. [16] for a survey on other formalisms). For brevity, we simplify some details in our exposition and refer the reader to Kingma et al. [18] for details.
Consistent with prior work (e.g. Li et al. [21]), our basic approach will be to map discrete text sequences into a continuous space with a token-wise embedding function and then construct a diffusion model on the embedded data.
2.1 Forward diffusion process
Consider a sequence of tokens  $x = [x^{(1)}, \ldots, x^{(L)}]$  drawn from the data distribution  $q(x)$ . We transform  $x$  into a sequence  $\tilde{x}$  of embedding vectors using an invertible token-wise embedding function  $\operatorname{Embed}(\cdot)$ , such that  $\tilde{x}^{(i)} := \operatorname{Embed}(x^{(i)})$ .
The forward process is a Markov chain over latent variables  $z_{t}$  from  $t = 0$  to  $t = 1$  which progressively adds Gaussian noise to  $\tilde{x}$ . Let  $\sigma^2 (t)$  be some monotonic function that specifies the total noise
added by time  $t$ . We then define the forward process distribution  $q$  with  $T$  discrete timesteps as
$$
q (x, z) := q (x) q \left(z _ {0} \mid x\right) \prod_ {i = 1} ^ {T} q \left(z _ {i / T} \mid z _ {(i - 1) / T}\right) \tag {1}
$$
where  $q(z_0|x) \coloneqq \mathcal{N}(\tilde{x},\sigma^2 (0))$  and  $q(z_{t}|z_{s}) \coloneqq \mathcal{N}(z_{s},\sigma^{2}(t) - \sigma^{2}(s))$ . It follows from this that  $q(z_{s}|z_{t},\tilde{x})$  is also Gaussian, which will be useful later.
2.2 Reverse generative process
We can approximate the forward process distribution  $q$  by a learned Markov reverse process where time runs backward from  $t = 1$  to  $t = 0$ . The reverse process with  $T$  timesteps is defined via the decomposition
$$
p _ {\theta} (x, z) := p \left(z _ {1}\right) \left(\prod_ {i = 1} ^ {T} p _ {\theta} \left(z _ {(i - 1) / T} \mid z _ {i / T}\right)\right) p (x \mid z _ {0}). \tag {2}
$$
Let  $z_{t}^{(i)}$  denote the portion of  $z_{t}$  at sequence position  $i$ . Then we choose  $p(z_{1}) := \mathcal{N}(0, \sigma^{2}(1)I)$  and  $p(x|z_{0}) := \prod_{i} p(x^{(i)}|z_{0}^{(i)})$  where  $p(x^{(i)}|z_{0}^{(i)}) \propto q(z_{0}^{(i)}|x_{i})$ . Recalling that  $q(z_{(i-1)/T}|z_{i/T}, \tilde{x})$  is Gaussian, for the remaining factors we choose  $p_{\theta}(z_{(i-1)/T}|z_{i/T}) := q(z_{(i-1)/T}|z_{i/T}, \tilde{x} = \hat{x}_{\theta}(z_{i/T}))$  where  $\hat{x}_{\theta}(z_{t})$  is a denoiser neural network that approximates  $\mathbb{E}_q[\tilde{x}|z_t]$ . Finally, our generative model is given by the marginal distribution  $p_{\theta}(x) = \int_z p_{\theta}(x,z)$ . If  $\hat{x}_{\theta}$  is optimal, then the forward and reverse processes express the same joint distribution as  $\sigma^2(0) \to 0$ ,  $\sigma^2(1) \to \infty$ , and  $T \to \infty$ .
2.3 Likelihood bound
To optimize and evaluate the likelihood, we can write a variational lower bound (VLB) for the log-likelihood as
$$
- \log p _ {\theta} (x) \leq - \operatorname {V L B} (x) := D _ {\mathrm {K L}} \left(q \left(z _ {1} | x\right) \| p \left(z _ {1}\right)\right) + \mathbb {E} _ {q \left(z _ {0} | x\right)} \left[ - \log p \left(x \mid z _ {0}\right) \right] + \mathcal {L} _ {T} \tag {3}
$$
$$
\mathcal {L} _ {T} := \sum_ {i = 1} ^ {T} \mathbb {E} _ {q \left(z _ {i / T} | x\right)} \left[ D _ {\mathrm {K L}} \left(q \left(z _ {(i - 1) / T} \mid z _ {i / T}, x\right) \| p _ {\theta} \left(z _ {(i - 1) / T} \mid z _ {i / T}\right)\right) \right]. \tag {4}
$$
In the  $T\to \infty$  limit,  $\mathcal{L}_T$  simplifies to
$$
\mathcal {L} _ {\infty} = - \frac {1}{2} \mathbb {E} _ {t \sim U [ 0, 1 ], z _ {t} \sim q (z _ {t} | x)} [ \mathrm {S N R} ^ {\prime} (t) \| \tilde {x} - \hat {x} _ {\theta} (z _ {t}) \| _ {2} ^ {2} ] \tag {5}
$$
where  $\mathrm{SNR}'(t) \coloneqq \frac{d}{dt} \frac{1}{\sigma^2(t)}$ . We use Monte-Carlo estimates of the resulting continuous-time likelihood bound to train and evaluate our model.
2.4 Learned noise schedule
A crucial hyperparameter in diffusion models is the noise schedule  $\sigma^2 (t)$ , which specifies how much noise to add at each time in the diffusion process. In our setting, the VLB is differentiable with respect to  $\sigma^2 (t)$  via the reparameterization trick. Moreover, the VLB is invariant to the value of  $\sigma^2 (t)$  except at  $t = 0$  and  $t = 1$  in the continuous-time limit.
We can therefore parameterize  $\sigma^2(t)$  as a scalar-to-scalar neural network and learn it by gradient descent. We train the endpoints  $\sigma^2(0)$  and  $\sigma^2(1)$  to maximize the VLB, and the schedule in between the endpoints to minimize the variance of the Monte-Carlo estimate of the VLB. Minimizing the loss variance is a proxy for minimizing the gradient covariance trace, which generally speeds up learning. See Kingma et al. [18] for further implementation details about this training procedure.
3 The Plaid framework
In this section, we present a series of algorithmic improvements to the basic setup described in Section 2. The result is a framework for diffusion language models which we refer to as Plaid (Perplexity-based LLanguage Inverse Diffusion).
3.1 Learned embeddings
In an autoregressive language model, the embedding operation is simply the first layer of the neural network and thus can be treated as just another part of the network. This is not true of embeddings in diffusion language models, which play a more fundamental role: they determine the order in which different tokens get generated. Tokens whose embeddings are far apart become distinguishable early in the reverse process, whereas nearby embeddings are distinguishable only later, at low noise levels.
Despite the importance of embeddings in diffusion language models, the loss functions used in prior work [21, 6] lead to ill-posed problems when optimized over  $W_{\text{Embed}}$ : for example, if our objective is  $L_2$  reconstruction, then collapsing the embeddings by setting  $W_{\text{Embed}} = 0$  and  $\hat{x}_{\theta}(z_t) = 0$  yields a degenerate solution with zero loss. Prior work addresses this with workarounds like choosing  $W_{\text{Embed}}$  by hand [3, 33] or using heuristic regularizers [21] or constraints [6].
In contrast, the Plaid loss function is a bound on the log-likelihood of the discrete data, which is a meaningful objective over both the model weights and embeddings. We therefore optimize the embedding matrix  $W_{\mathrm{Embed}}$  jointly with the rest of the model without additional constraints.
3.2 Categorical reparameterization
When optimally trained,  $\hat{x}_{\theta}(z_{t})$  learns to approximate a conditional expectation  $\mathbb{E}[\tilde{x} | z_t]$  over sequences of word embeddings  $\tilde{x}$ . At low noise levels, some or all of the embeddings in  $\tilde{x}$  are deterministic given  $z_{t}$ , so an optimal  $\hat{x}_{\theta}(z_{t})$  should output these exactly. However, doing so requires memorizing embedding vectors to high precision somewhere inside the model parameters, which is a poor use of capacity.
Instead of forcing the model to memorize the embedding vectors, we reparameterize  $\hat{x}_{\theta}(z_t)$  as an average of embeddings weighted by a softmax over tokens. More formally, let  $f_{\theta}(z_t)$  be a neural network which outputs logits and define  $\hat{x}$  as an average over embeddings  $\hat{x}_{\theta}^{(i)}(z_t) := W_{\mathrm{EmbedSoftmax}}(f_{\theta}^{(i)}(z_t))$ . We can interpret  $f$  as learning a posterior over each discrete token  $x^{(i)}$  given  $z_t$ . This relates to methods proposed in prior work, but these either require proxy objectives [21, 6] or consider image models [3].
3.3 Output prior
When we interpret  $f_{\theta}$  as a posterior over tokens, the optimal value of  $f_{\theta}(z_t)$  is  $\log q(x^{(i)}|z_t) + Z$ , which decomposes as  $\log q(z_t^{(i)}|x^{(i)}) + \log q(x^{(i)}|z_t^{(\neq i)}) + Z$  where  $z_t^{(\neq i)} := \{z_t^{(j)} : j \neq i\}$ . We view the first term as a prior constraining the model's predictions to those which are plausible given  $z_t$ , while the second term models relationships between different tokens.
To make it easier to model  $f_{\theta}$ , we compute the first term in closed form as the log-density of a Gaussian  $\mathcal{N}(\tilde{x}^{(i)},\sigma^2 (t)I)$  and add it to the output. This leaves the neural network with only the easier task of estimating  $\log p(x^{(i)}|z_t^{(\neq i)})$ . Empirically, we found it helpful to linearly anneal in this prior over the first 5000 steps of training.
3.4 Learned conditional likelihood
Recall that our loss function (3) includes a conditional likelihood term  $\log p(x|z_0)$ . We are free to choose  $p$  however we wish, and in Section 2 we chose a position-wise factorial model  $p(x|z_0) \coloneqq \prod_i p(x^{(i)}|z_0^{(i)})$ , with a simple fixed distribution for each factor. This choice is optimal for sufficiently small  $\sigma^2(0)$ , but using a more powerful model allows  $\sigma^2(0)$  to take a larger value, effectively truncating the reverse process and therefore making it simpler to learn.
Here we leverage the fact that, after applying the categorical reparameterization (Section 3.2), our neural network  $f_{\theta}(z_t)$  can be interpreted as learning the logits for  $q(x^{(i)}|z_t)$  at all positions  $i$ . We therefore choose to keep  $p(x|z_0)$  as a factorial model, but define each factor  $p(x_i|z_0^{(i)})$  using the more powerful learned model  $\mathrm{softmax}(f_{\theta}^{(i)}(z_t))$ .
Implementing this change naively requires two evaluations of  $f_{\theta}$  for each minibatch example during training, corresponding to the two terms of (3)  $\mathcal{L}_{\infty}$  and  $\log p_{\theta}(x|z_0)$ . We instead split each minibatch,
Table 1: Compute-matched ablations of algorithmic components on OpenText2.
NLL bound (val.)
Our full method	3.89
Our full method (0.5× compute)	4.01
No learned noise schedule	4.17
No learned embeddings	4.54
No categorical reparameterization	4.25
No output prior	3.95
No learned conditional likelihood	4.03
No self-conditioning	3.98
CDCD [6] (our reimplementation)	4.23
using some examples to compute  $\mathcal{L}_{\infty}$  and the rest to compute  $\log p_{\theta}(x|z_0)$ . We allocate examples between the two terms according to the ratio  $\sqrt{\frac{\operatorname{Var}(\mathcal{L}_{\infty})}{\operatorname{Var}(\log p_{\theta}(x|z_0))}}$ , where we compute the variances using running estimates of each term's first and second moments. This minimizes the variance of the full loss (3).
3.5 Self-conditioning
Self-conditioning [3] is a technique which improves the performance of diffusion language models. The core idea is to reparameterize the denoiser  $\hat{x}_{\theta}(z_{t})$  as the fixed point  $y_{\infty}$  of a recurrence  $y_0 \coloneqq 0, y_{i+1} \coloneqq \hat{x}_{\theta}'(z_{t}, y_{i})$  where  $\hat{x}_{\theta}'$  is a neural network which now takes two inputs instead of one. During training, we approximate the fixed point  $y_{\infty}$  by randomly unrolling the recurrence to either  $y_1$  (with probability 0.75) or  $y_2$  (otherwise). When we unroll to  $y_2$  during training, we zero the gradients with respect to  $y_1$ , the noise schedule, and the embeddings. During held-out likelihood evaluation, we always unroll to  $y_2$ . During sampling, instead of solving the recurrence from scratch at each step of the diffusion chain, we compute  $\hat{x}_{\theta}'(z_1, 0)$  for the first step and  $\hat{x}_{\theta}'(z_t, \hat{x}_{\theta}'(z_{t + (1 / T)}, \ldots))$  for subsequent steps.
3.6 Other details
We perform all forward and backward computations in double precision except for the Transformer layers themselves, which happen in bffloat16 mixed precision. This comes at a negligible extra cost since the Transformer layers dominate the overall cost.
Architecture choices We condition  $\hat{x}_{\theta}(z_{t})$  on the timestep  $t$  by adding a sinusoidal encoding of  $t$  to the Transformer's residual stream before the first layer. Before feeding  $z_{t}$  into the Transformer, we rescale it by a factor of  $\sqrt{1 + \sigma^2(t)}$  which makes each input dimension approximately unit-variance. Whereas autoregressive Transformers are relatively insensitive to aspect ratio [15], we find that Plaid performance increases significantly with Transformer depth up to about 16 layers. We also find that performance is sensitive to the choice of embedding dimension, with small values performing best. In all experiments, we use embedding dimension 16.
Stochastic sequence length Unlike autoregressive models, diffusion language models can only operate on sequences of exactly the same length as those seen during training. To enable our model to generalize to shorter sequence lengths, we truncate a small random subset of examples seen during training to random lengths. We observe that truncating even  $3\%$  of examples allows the model to generalize well across lengths without impacting full-length performance. Short-sequence performance does not improve substantially as we increase the number of truncated examples.
4 Ablation experiments
In this section, we validate different aspects of the Plaid framework through compute-matched ablation experiments.
4.1 Validating likelihood-based training
We take a likelihood-based approach in this work for multiple reasons: it has a principled interpretation, it simplifies training and evaluation, and it has yielded strong results in autoregressive models. Here, we validate that the log-likelihood objective can attain competitive sample quality through a human evaluation.
In diffusion models, the log-likelihood bound is an expectation over noise levels of a reconstruction loss weighted by a specific function of the noise level. In contrast, most prior work on diffusion models for language [21, 33] as well as images [12, 5] use heuristic weight schedules. Motivated by the intuition that human perception is more sensitive to coarse structure than fine details, these typically assign more weight to higher noise levels than the likelihood weight schedule.
We train three Plaid models: one with the likelihood weight schedule ("VLB") and two with heuristic weight schedules ("Schedule A" and "Schedule B") which we plot in Appendix B. The models are trained on a large dataset of short children's stories which we constructed by finetuning GPT-J [34] on ROCStories [27]. Because learning embeddings is only straightforward when training against the likelihood bound, all models use fixed embeddings obtained from a previously-trained known-good model.
We repeatedly asked crowdworkers to choose from a pair of model samples, where one sample came from the likelihood-trained model and the other came from a heuristically-trained model. On average, crowdworkers preferred the likelihood-trained model over both alternatives: Weighting A's win rate was 0.449 ( $p = 0.001$ , 95% CI [0.417, 0.482]) and Weighting B's win rate was 0.457 ( $p = 0.005$ , 95% CI [0.425, 0.490]). Because we only consider two alternative weight schedules, we do not claim that the likelihood objective yields optimal sample quality, but our results suggest that it performs at least comparably to other choices.
4.2 Validating algorithmic components
Having validated our likelihood-based approach, we show in this section that each of the algorithmic components described in Section 3 lead to improved likelihoods in a compute-matched ablation study.
We train Plaid models on OpenWebText2 [7] and report their log-likelihood bounds on held-out data in Table 1. Our reference model ("full method") is a  $16 \times 384$  Transformer with 28M non-embedding parameters, trained for 92K steps at batch size 256 and sequence length 256, corresponding to  $1.12 \times 10^{18}$  non-embedding FLOPs. For each ablation model, we stay as close to this configuration as possible while preserving the number of non-embedding FLOPs (we exclude FLOPs from the embedding and output projections because these become negligible at large scale). We observe that ablating each of the components described in Section 3 results in a worse log-likelihood. As a comparison point, we also train a model at half the compute budget ( $5.6 \times 10^{17}$  FLOPs) by halving the model size. See Appendix C for more training details.
Finally, as a comparison to prior work, we reimplement CDCD [6], train it following the same configuration, and report its log-likelihood. We follow the authors' description as faithfully as possible except for the noise schedule endpoints, embedding dimension, and embedding weight initialization, which we tune to maximize log-likelihood. We observe in Table 1 that even the half-compute-budget version of Plaid surpasses our CDCD implementation in likelihood. Note that CDCD was not developed as a likelihood-based model, and the lack of a public implementation means that there are most likely differences between our implementation and the original.
5 Scaling laws for Plaid
Having developed an algorithmic framework for diffusion language models, we now study its scaling properties in order to guide large-scale training of Plaid models. In the case of autoregressive models, the work of Kaplan et al. [15] demonstrates that model log-likelihoods follow a log-linear scaling law: across many orders of magnitude, training with more compute predictably improves likelihood. Using these scaling laws, Kaplan et al. [15] and Hoffmann et al. [13] accurately predict the optimal model size as a function of the given compute budget across many orders of magnitude. Both results together enable effective large-scale training. In this section, we experimentally determine them for Plaid models.
Figure 2: Plaid models improve with compute at a similar rate to autoregressive models, but Plaid is less efficient by a constant factor of  $64\times$
Figure 3: Compute-optimal Plaid models should be  $4\times$  smaller (and trained for  $4\times$  longer) than compute-optimal autoregressive models.
Our main experimental method will be an IsoFLOP analysis [13]. We first fix a set of FLOP budgets  $\{C_1,\dots ,C_K\}$ . For each budget  $C$ , we train models with different sizes  $\{N_{C,1},\ldots ,N_{C,M}\}$  and perform a quadratic fit of the loss  $L$  to  $\log N$ . We plot all the data along with quadratic fits in Appendix D and find that the fits approximate the data well.
The minimum of the quadratic gives us the compute-optimal loss  $L_{C}^{*}$  and corresponding model size  $N_{C}^{*}$  for that budget.
Given the compute-optimal loss  $L_{C_i}^*$  for each FLOP budget  $C_i$ , we fit the parameters of a loss scaling law
$$
\min _ {\alpha , \beta} \sum_ {i} (\log (L _ {C _ {i}} ^ {*}) - \beta \log (C _ {i}) - \alpha) ^ {2}
$$
which can then be used to predict the compute-optimal loss as  $L^{*}(C) = \alpha C^{\beta}$ . We also fit a parameter scaling law  $N^{*}(C)$  in the same fashion from the model sizes  $N_{C_i}^*$ .
We perform IsoFLOP analyses for both Plaid and autoregressive models in order to compare the results. We choose compute budgets log-uniformly between  $10^{16}$  and  $10^{19}$  FLOPs and corresponding model sizes heuristically. We choose learning rates using  $\mu$  Transfer [35] and batch sizes, weight decays, and aspect ratios by well-tuned heuristics. When computing FLOPs, we exclude FLOPs from the embedding layers and output projections. This enables us to use much smaller compute budgets than Hoffmann et al. [13], but it causes our autoregressive scaling law to differ slightly from theirs. We consider this acceptable since we are mainly interested in the differences between our autoregressive and Plaid scaling laws.
5.2 Loss improves predictably with compute