Likelihood-Based Diffusion Language Models

TL;DR

This work pioneers likelihood-based diffusion language modeling by introducing Plaid, a framework with learned embeddings, a categorical reparameterization, an output prior, and self-conditioning. It provides a comprehensive scaling-law analysis (IsoFLOP) and demonstrates Plaid 1B achieving zero-shot likelihood competitive with GPT-2 124M across six benchmarks, along with fluent unconditional and zero-shot conditioned generation. The results illuminate how diffusion models scale with compute in language tasks, quantify the compute-efficiency gap (about 64x relative to autoregressive models), and propose compute-optimal training guidelines to bridge that gap. Overall, Plaid shows diffusion models can attain nontrivial likelihoods on standard benchmarks and offers a concrete release (Plaid 1B) that outperforms a small autoregressive baseline on likelihood while enabling controllable generation.

Abstract

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.

Figures (5)

• 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 \ref{['tab:likelihood']}).
• 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.
• Figure 4: VLB weight schedule and and heuristic weight schedules used in ablation experiments.
• Figure 5: IsoFLOP profiles for autoregressive models (left) and diffusion models (right).