PaperlayerView exact revision on arXiv

Source-linked AI summary

Likelihood-Based Diffusion Language Models

Ishaan Gulrajani, Tatsunori B. Hashimoto

arXiv:2305.18619v1cs.CLcs.LG

TL;DR

Diffusion language models lack demonstrated nontrivial likelihoods on standard benchmarks, limiting direct comparison with autoregressive models. This paper develops Plaid through algorithmic improvements and scaling-law analysis, and releases Plaid 1B, which outperforms GPT-2 124M in zero-shot likelihood across six benchmarks while generating fluent, controllable text.

  • Problem

    Diffusion language models had not demonstrated nontrivial likelihoods on standard benchmarks, hindering precise comparison with autoregressive models.

  • Method

    The paper develops Plaid through algorithmic improvements for likelihood-based training and scaling-law analysis to derive compute-optimal training regimes.

  • Results

    Plaid 1B outperforms GPT-2 124M in zero-shot likelihood across six standard benchmarks and generates fluent, controllable text.

  • Takeaways & Limitations

    The work establishes a competitive likelihood-based diffusion language model while narrowing the compute-efficiency gap with autoregressive models to 64×.

  • Takeaways & Limitations

    The method’s theoretical equivalence relies on limiting conditions in which σ2(0) approaches 0, σ2(1) approaches infinity, and T approaches infinity.

Abstract

from arXiv · show

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.

1 Introduction

This work explores likelihood-based diffusion language models as an alternative to autoregressive modeling, targeting likelihoods exceeding GPT-2 124M. It introduces Plaid, studies its scaling laws for compute-optimal training, and releases Plaid 1B with fluent, controllable generation.

  • Autoregressive language models maximize likelihood through left-to-right factorization, enabling strong downstream performance and complex behaviors such as in-context learning.
  • Autoregressive generation produces tokens one at a time, creating challenges for long-range planning and controllable generation.
  • The paper explores continuous diffusion models because prior language diffusion work used non-likelihood-based objectives, limiting evaluation on standard likelihood benchmarks.
  • The authors propose Plaid, an algorithmic framework for likelihood-based diffusion language models, and validate its design choices through compute-matched ablations.
  • Plaid training improves log-likelihood predictably with more compute, while its compute-optimal training recipe differs substantially from the usual autoregressive rule.
  • Plaid 1B, pretrained on OpenWebText2, outperforms GPT-2 124M in zero-shot likelihood across six standard benchmarks and produces fluent, controllable text.

2 Variational Diffusion Models for language

The paper formulates likelihood-based diffusion language models using the Variational Diffusion Models framework, embedding discrete text into continuous vectors and modeling progressive Gaussian noising with a learned reverse process. Likelihood is optimized and evaluated through a variational lower bound, while the noise schedule is learned to improve estimation and training.

  • Forward process: The model maps discrete token sequences into continuous embeddings and applies a diffusion process that progressively adds Gaussian noise.An invertible token-wise embedding function transforms each token, while the forward process is a Markov chain over latent variables.
  • Reverse process: A learned Markov reverse process denoises latent states backward in time using a neural network that estimates the clean embedded sequence.The reverse transitions are defined from Gaussian posteriors evaluated at the denoiser’s predicted clean embedding.
  • Likelihood objective: Likelihood training and evaluation use a variational lower bound consisting of prior, reconstruction, and diffusion terms, estimated with Monte Carlo in continuous time.The bound upper-bounds −log pθ(x), and the resulting continuous-time likelihood bound is used for both training and evaluation.
  • Noise schedule: The noise schedule σ2(t) is learned by gradient descent because the VLB is differentiable through reparameterization and invariant to schedule values away from the endpoints in continuous time.The endpoints maximize the VLB, while intermediate schedule values minimize Monte Carlo estimate variance, a proxy for gradient covariance trace.

3 The Plaid framework

Plaid is a diffusion-language-model framework built around maximum-likelihood algorithmic improvements. Its design makes embeddings jointly learnable, avoids memorizing embedding vectors, and supports generalization across sequence lengths.

  • Plaid (Perplexity-based LAnguage Inverse Diffusion) is a framework for diffusion language models developed through algorithmic improvements to the basic setup.
  • Embedding optimization: The Plaid loss bounds discrete-data log-likelihood, enabling joint optimization of model weights and embeddings without additional constraints.Prior reconstruction losses can admit the degenerate solution WEmbed = 0 and x̂θ(zt) = 0.
  • Categorical reparameterization: Plaid reparameterizes denoiser outputs as softmax-weighted averages of embeddings, replacing the need for the model to memorize embedding vectors at high precision.The reparameterization interprets the neural network as learning a posterior over each discrete token conditioned on zt.
  • Categorical reparameterization: A closed-form Gaussian prior term constrains token predictions to be plausible given zt, while the neural network estimates relationships between different tokens.The prior is linearly annealed over the first 5000 training steps.
  • Stochastic sequence length: Truncating 3% of training examples to random lengths enables generalization across shorter sequence lengths without impacting full-length performance.

4 Ablation experiments

Compute-matched ablations validate Plaid’s algorithmic components and show that likelihood training produces competitive sample quality. Plaid also surpasses a reimplemented CDCD baseline in likelihood, including at half the compute budget.

  • Likelihood weight schedule: Crowdworkers preferred likelihood-trained samples over both heuristic alternatives, with Weighting A’s win rate 0.449 and Weighting B’s win rate 0.457.Weighting A: p = 0.001, 95% CI [0.417, 0.482]; Weighting B: p = 0.005, 95% CI [0.425, 0.490].
  • Algorithmic components: Each algorithmic component described in Section 3 leads to improved likelihoods in a compute-matched ablation study.The study trains Plaid models on OpenWebText2 and compares held-out log-likelihood bounds.
  • Compute-matched setup: The reference full-method model uses a 16 × 384 Transformer with 28M non-embedding parameters and 1.12 × 10^18 non-embedding FLOPs.It is trained for 92K steps with batch size 256 and sequence length 256; ablations preserve the non-embedding FLOP budget as closely as possible.
  • Comparison to prior work: Even Plaid’s half-compute-budget version surpasses the reimplemented CDCD model in likelihood.CDCD follows the same configuration, with tuned noise schedule endpoints, embedding dimension, and embedding weight initialization.

5 Scaling laws for Plaid

Plaid’s compute-optimal likelihood follows a power law with nearly the same slope as autoregressive models, validating its scalability while revealing substantial efficiency and model-size differences. Plaid requires about 64× more compute and should use models about 4× smaller, trained 4× longer, for compute-optimal scaling.

  • Scaling-law methodology: IsoFLOP analyses fit quadratic loss curves over log model size to estimate compute-optimal loss and model size across FLOP budgets.The analyses compare Plaid and autoregressive models over budgets from 10^16 to 10^19 FLOPs.
  • Likelihood scaling: Across many orders of magnitude, Plaid’s compute-optimal log-likelihood follows a power law whose slope is almost exactly the same as autoregressive models.The result validates Plaid’s scalability and suggests larger-scale training can yield strong improvements.
  • Likelihood scaling: 64× more compute is required for Plaid models to match autoregressive equivalents, regardless of scale.The paper identifies this as a constant-factor efficiency gap that future diffusion-model engineering may reduce.
  • Parameter scaling: 4× smaller models are compute-optimal for Plaid, which should therefore be trained for 4× longer than autoregressive models.The differing compute-optimal settings indicate that autoregressive scaling laws can incur substantial loss when used to select Plaid model sizes.

6 Plaid 1B

Plaid 1B is a 1.3B-parameter diffusion language model trained at unprecedented scale that consistently outperforms 124M-parameter GPT-2 in likelihood. Despite unconditional training, it generates fluent text and follows diverse zero-shot conditioning structures.

  • Model and training: Plaid 1B has 1.3B parameters, was trained for 314B tokens on OpenWebText2, and used 2.5 × 1021 FLOPs.The model is Transformer-based and was released as a large Plaid model.
  • Likelihood evaluation: Plaid 1B was evaluated zero-shot on six benchmark datasets used by Radford et al.The suite includes Penn Treebank, enwik8, text8, WikiText2, WikiText103, and the One Billion Word corpus.
  • Likelihood results: Plaid 1B consistently outperforms the 124M parameter GPT-2 model in likelihood, reaching perplexities on par with a small modern autoregressive model.The comparison uses Plaid 1B and 124M parameter GPT-2 likelihoods.
  • Unconditional generation: With sequence length 1024, Plaid 1B generates fluent, on-topic unconditional text over several hundred words.The reported samples are selected examples, with random un-picked samples provided separately.
  • Zero-shot control: Token guidance enables Plaid 1B to follow diverse zero-shot conditioning structures despite purely unconditional training.The demonstrated structures include span conditioning, lexical constraints, token-level weights, and negation.

7 Related work

This work differs from prior diffusion language models by targeting strong likelihood performance and benchmarking against autoregressive models. It extends Variational Diffusion Models from image generation to language modeling while focusing on full-sequence training rather than constrained or hybrid settings.

  • 7 Related work: The central departure from prior diffusion language models is targeting strong likelihood performance and benchmarking against well-known autoregressive models.Prior work had not attempted this except for an appendix result from Li et al.
  • 7 Related work: CDCD is the most comparable general-purpose diffusion language model, but standard likelihood benchmarks are unavailable for precisely comparing it with autoregressive models.Its reported results either underperform an autoregressive baseline or use saturated metrics lacking statistical power to distinguish models.
  • 7 Related work: Other diffusion language models study constrained generation or sequence-to-sequence tasks, while hybrid approaches combine pretrained autoregressive models with diffusion.Concurrent work finetunes OPT 13B into a model that is autoregressive over 25-token blocks and uses diffusion within blocks; this work instead focuses on training a full model.
  • 7 Related work: The method builds on advances in diffusion models for images by adopting and extending the Variational Diffusion Models framework to language modeling.This provides the work’s connection to image diffusion research while adapting the framework to a language-modeling setting.

8 Conclusion · A VLB and heuristic weight schedules

The paper presents Plaid 1B as a competitive likelihood-based diffusion language model, combining algorithmic improvements with scaling-law analysis and achieving GPT-2 124M-level likelihood. It also reports fluent unconditional and zero-shot conditional samples, while examining VLB and heuristic weight schedules in ablations.

  • 8 Conclusion: Plaid 1B matches GPT-2 124M in likelihood.This result combines several algorithmic improvements with scaling-law analysis.
  • 8 Conclusion: The work takes first steps toward a competitive likelihood-based diffusion language model.
  • 8 Conclusion: Several algorithmic improvements contribute to Plaid 1B’s likelihood performance.
  • 8 Conclusion: Scaling-law analysis is part of the approach used to build Plaid 1B.
  • 8 Conclusion: Maximizing likelihood does not substantially harm sample quality in the reported ablations.
  • 8 Conclusion: Plaid 1B produces fluent samples in unconditional settings.
  • 8 Conclusion: Plaid 1B also produces fluent samples in zero-shot conditional settings.
  • A VLB and heuristic weight schedules: Ablation experiments use both VLB weight schedules and heuristic weight schedules.

B Experiment details · B.1 Dataset

The experiments use a filtered OpenWebText2 subset and a tokenizer trained specifically on its training split. Documents labeled non-English are removed, and the data is tokenized with a 32K-token BPE tokenizer.

  • B Experiment details: All models are trained on a subset of OpenWebText2 unless otherwise noted.
  • B.1 Dataset: The OpenWebText2 subset is filtered to remove documents labeled as non-English.
  • B.1 Dataset: The dataset is tokenized using a BPE tokenizer.
  • B.1 Dataset: The tokenizer has a 32K-token vocabulary.
  • B.1 Dataset: The BPE tokenizer is trained on the OpenWebText2 training split.
  • B Experiment details: The stated dataset and tokenization procedure apply unless otherwise noted.

B.2 Architecture · B.3 Optimization

The models use pre-activation Transformers with RMSNorm and GeLU, alongside efficiency-oriented kernels and a 16-layer default configuration. Optimization uses AdamW with µTransfer-derived learning rates, scheduled regularization, and batch sizes tailored to experiment type.

  • B.2 Architecture: Plaid models use standard pre-activation Transformers with RMSNorm normalization layers and GeLU nonlinearities.
  • B.2 Architecture: 16 Transformer layers are used for Plaid models unless otherwise noted, based on approximate optimality at their scale.
  • B.2 Architecture: FlashAttention and other fused kernels are used wherever applicable to improve implementation efficiency.
  • B.2 Architecture: Training uses sequence length 256 for all experiments except Plaid 1B.
  • B.3 Optimization: All models are optimized with AdamW and parameter-specific learning rates derived by µTransfer from a learning rate of 1.4 × 10−3 at width 256.
  • B.3 Optimization: Each parameter’s weight decay is set to 4×10−5 η, where η denotes that parameter’s learning rate.
  • B.3 Optimization: Learning rate and weight decay receive linear warmup over the first 2500 steps, followed by linear decay to zero over training.
  • B.3 Optimization: Batch size is 256 for algorithm ablations and 128 for scaling law experiments.

B.4 Plaid 1B Training · C IsoFLOP profiles

Plaid 1B uses a 1.3B-parameter Transformer denoiser trained for 1.2M steps on 314B tokens. The accompanying IsoFLOP profiles compare autoregressive and diffusion models.

  • B.4 Plaid 1B Training: Plaid 1B uses a 1.3B-parameter Transformer denoiser with 24 layers and width 2048.Its vocabulary contains 32K tokens.
  • B.4 Plaid 1B Training: 2 × 10−3 is the base µTransfer learning rate used at width 256.
  • B.4 Plaid 1B Training: 1.2M steps are used for training with batch size 256 and sequence length 1024.
  • B.4 Plaid 1B Training: 314B tokens are processed during Plaid 1B training.
  • B.4 Plaid 1B Training: The denoiser operates over a vocabulary of 32K tokens.
  • C IsoFLOP profiles: Figure 5 presents IsoFLOP profiles for autoregressive models and diffusion models.Autoregressive profiles appear on the left, while diffusion profiles appear on the right.

D Plaid 1B Random Samples … D.7 Negation: “Donald” anywhere and “Trump” nowhere

Plaid 1B produces varied text across unconditional, prefix-completion, infilling, weighted-word, and lexical-constraint settings. The supplied samples include coherent passages as well as malformed or unrelated continuations.

  • D.1 Unconditional: The unconditional samples span news, commentary, reviews, scientific writing, and political reporting, including passages that mention Trump.Examples include economic commentary, product reviews, climate and finance text, and impeachment-related reporting.
  • D Plaid 1B Random Samples: Across the supplied Plaid 1B random samples, outputs cover multiple controlled-generation tasks and show both topical adaptation and substantial variation in continuation quality.The evidence includes coherent domain-specific passages alongside garbled or weakly related text.
  • D.2 Prefix completion: “Generative models of text are very versatile: they can be used”: Prefix completion preserves the opening phrase but produces continuations ranging from relevant descriptions to repetitive, malformed, or semantically unrelated text.The outputs include discussions of classification, images, data processing, and models, with variable coherence.
  • D.3 Infilling: “A year ago in Paris, [...] Wow, what a great day!”: Infilling completes the Paris prompt with diverse contexts, including gaming, tourism, speeches, politics, parliamentary proceedings, sports, and personal anecdotes.Several continuations retain the requested ending, “Wow, what a great day!”
  • D.4 Word-level weights: “Let’s talk about law and medicine.” with 5× weight on “law”: With 5× weight on “law,” completions remain centered on law and medicine but vary across arbitration, firearms regulation, emergency powers, ethics, cybersecurity, and criminal justice.The outputs repeatedly preserve the prompt while extending it into legal and institutional topics.
  • D.5 Word-level weights: “Let’s talk about law and medicine.” with 5× weight on “medicine”: With 5× weight on “medicine,” completions shift toward healthcare, medical research, cancer prevention, emergency-room ethics, genetics, and brain health.The samples include passages about health data, colorectal cancer, congestive heart failure, and appetite-related brain function.
  • D.6 Lexical constraints: “Donald” anywhere: The “Donald” lexical constraint yields passages containing Donald in names, captions, entertainment reporting, biographies, and political reporting.Examples include Donald Taylor, Donald Pleasence, Donald Stoddard, and Donald Trump.
Loading 2305.18619v1…