[Paper Review] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution

Aaron Lou, Chenlin Meng+1 more|arXiv (Cornell University)|Oct 25, 2023

Topic Modeling4 citations

TL;DR

This paper introduces Score Entropy Discrete Diffusion (SEDD), a novel discrete score matching loss that enables stable, scalable training of diffusion models on natural language data. By estimating ratios of the data distribution via a principled score entropy objective, SEDD achieves GPT-2-level perplexity while enabling compute-quality tradeoffs, faithful distribution learning (4× better than GPT-2), and arbitrary infilling beyond autoregressive generation.

ABSTRACT

Despite their groundbreaking performance for many generative modeling tasks, diffusion models have fallen short on discrete data domains such as natural language. Crucially, standard diffusion models rely on the well-established theory of score matching, but efforts to generalize this to discrete structures have not yielded the same empirical gains. In this work, we bridge this gap by proposing score entropy, a novel loss that naturally extends score matching to discrete spaces, integrates seamlessly to build discrete diffusion models, and significantly boosts performance. Experimentally, we test our Score Entropy Discrete Diffusion models (SEDD) on standard language modeling tasks. For comparable model sizes, SEDD beats existing language diffusion paradigms (reducing perplexity by $25$-$75$\%) and is competitive with autoregressive models, in particular outperforming GPT-2. Furthermore, compared to autoregressive mdoels, SEDD generates faithful text without requiring distribution annealing techniques like temperature scaling (around $6$-$8 imes$ better generative perplexity than un-annealed GPT-2), can trade compute and quality (similar quality with $32 imes$ fewer network evaluations), and enables controllable infilling (matching nucleus sampling quality while enabling other strategies besides left to right prompting).

Motivation & Objective

• To address the lack of effective score matching methods for discrete diffusion models in natural language generation.
• To develop a stable, scalable training objective that enables high-quality discrete diffusion modeling comparable to autoregressive models like GPT-2.
• To enable algorithmic advantages such as compute-quality tradeoffs and arbitrary infilling beyond standard left-to-right autoregressive generation.
• To demonstrate that non-autoregressive diffusion models can match autoregressive performance at scale on language modeling tasks.

Proposed method

• Proposes score entropy, a novel discrete score matching loss that estimates the ratio of perturbed data distributions, forming a variational lower bound (ELBO) for maximum likelihood training.
• Derives a denoising variant of the score entropy loss to enable efficient optimization without requiring explicit density estimation.
• Introduces a score-based ancestral sampling method and a general infilling procedure that allows conditioning on arbitrary masked positions in the sequence.
• Scales the SEDD model to GPT-2 model sizes using architectural improvements and efficient training techniques.
• Uses a continuous-time SDE framework adapted to discrete spaces, modeling the score function as a neural network that predicts gradients of the log-perturbed density.
• Employs a denoising objective to approximate the score function without direct access to the true data density, improving training stability.

Figure 1: The graphs of $\mathcal{L}_{\rm CSM}$ versus $\mathcal{L}_{\rm SE}$ for a ground truth score of $0.2$ . The score entropy loss respects nonnegativity.

Experimental results

Research questions

• RQ1 Can a discrete score matching loss be designed that is both stable and scalable for language modeling, enabling diffusion models to match autoregressive baselines?
• RQ2 Does the proposed score entropy loss lead to better distributional fidelity compared to standard autoregressive models like GPT-2?
• RQ3 Can SEDD models achieve competitive perplexity while enabling new capabilities such as arbitrary infilling and compute-quality tradeoffs?
• RQ4 How does the performance of SEDD compare to GPT-2 in terms of likelihood, sample quality, and inference efficiency?

Key findings

• SEDD achieves perplexity scores within +10% of GPT-2 on similarly sized models, with some configurations even outperforming the autoregressive baseline.
• SEDD models learn a more faithful sequence distribution, achieving approximately 4× better distributional fidelity than GPT-2 when evaluated with large language models using ancestral sampling.
• SEDD enables a compute-quality tradeoff, requiring only 16× fewer network evaluations than GPT-2 to match its generation quality.
• SEDD supports arbitrary infilling, allowing conditioning on any masked positions in the sequence, unlike standard autoregressive models restricted to left-to-right generation.
• The score entropy loss is stable and forms a valid ELBO, enabling principled maximum likelihood training for discrete diffusion models.
• The denoising variant of the score entropy loss enables efficient optimization without requiring explicit density estimation or complex sampling schemes.

(a) Generative Perplexity vs. Sampling Iterations.