PaperlayerView exact revision on arXiv

Source-linked AI summary

Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution

Aaron Lou, Chenlin Meng, Stefano Ermon

arXiv:2310.16834v3stat.MLcs.CLcs.LG

TL;DR

Diffusion models have lagged behind autoregressive models on discrete data such as language, despite score matching's success in continuous diffusion. The paper introduces score entropy for learning discrete data-distribution ratios and builds SEDD. SEDD outperforms prior language diffusion models, competes with autoregressive models, and supports unannealed generation, compute-quality trade-offs, and controllable infilling.

  • Problem

    Language diffusion models have not matched autoregressive models on likelihood, sampling speed, or sample quality without heavy annealing and empirical alterations.

  • Method

    SEDD parameterizes reverse discrete diffusion with data-distribution ratios learned using score entropy, a loss analogous to score matching.

  • Results

    SEDD outperforms existing language diffusion models and is competitive with same-size autoregressive models, including beating GPT-2 on zero-shot perplexity tasks.

  • Takeaways & Limitations

    SEDD provides a discrete diffusion framework supporting faithful unannealed text generation, compute-quality trade-offs, and controllable arbitrary-position infilling.

  • Takeaways & Limitations

    Mean-prediction alternatives learn reverse densities indirectly, while implicit score entropy is impractical for large-scale tasks because uniform sampling introduces additional variance.

Abstract

from arXiv · show

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\times$ better generative perplexity than un-annealed GPT-2), can trade compute and quality (similar quality with $32\times$ fewer network evaluations), and enables controllable infilling (matching nucleus sampling quality while enabling other strategies besides left to right prompting).

1. Introduction

SEDD applies score entropy to discrete diffusion, addressing limitations that have kept language diffusion models behind autoregressive modeling. It achieves strong language-modeling performance, faithful unannealed generation, compute-quality trade-offs, and controllable infilling.

  • Autoregressive models remain the only competitive discrete-data modeling approach, but sequential sampling is slow, difficult to control, and often degrades without annealing.
  • Existing language diffusion models do not rival autoregressive models on likelihood, sampling speed, or sample quality without heavy annealing and empirical alterations.
  • SEDD parameterizes reverse discrete diffusion with data-distribution ratios learned through score entropy, a loss analogous to continuous score matching.
  • SEDD outperforms existing language diffusion models by large margins and is competitive with same-size autoregressive models, including beating GPT-2 on zero-shot perplexity tasks.
  • SEDD supports high-quality unconditional samples, compute-quality trade-offs, arbitrary-position prompting, and infilling comparable to autoregressive nucleus sampling by MAUVE.

2. Preliminaries

Discrete diffusion evolves probability mass over a finite support through a continuous-time Markov process and reverses it using concrete-score ratios. Existing training strategies have important drawbacks, motivating score-entropy-based objectives.

  • Discrete diffusion models represent distributions as positive probability-mass vectors over a finite support and evolve them with a continuous-time Markov process.
  • Diffusion matrices preserve total probability mass through nonnegative off-diagonal entries and zero-sum columns, while simple schedules can approach a limiting base distribution.
  • The process can be simulated with small Euler steps by sampling transitions from the diffusion matrix's transition densities.
  • The reverse process uses ratios of perturbed probabilities, called the concrete score, as the discrete analogue of the continuous score function.
  • Discrete diffusion training methods have mixed empirical results, motivating renewed examination of how to learn the reverse-process ratios.
  • Mean prediction learns reverse densities indirectly, making the problem harder and requiring a continuous-time approximation that largely underperforms empirically.
  • Ratio matching departs from standard score matching and requires specialized, expensive architectures, while concrete score matching uses an ℓ2 loss with problematic behavior for nonpositive ratios.

3. Score Entropy Discrete Diffusion Models

Score entropy learns positive data-distribution ratios for discrete diffusion, yielding a consistent, tractable objective that supports likelihood evaluation and scalable sequence modeling.

  • Score Entropy: Score entropy learns the collected concrete score, extending score-matching ideas to positive ratios that evolve under discrete diffusion.The loss is based on a Bregman divergence and generalizes cross entropy to positive values.
  • Score Entropy Properties: Under infinite samples and model capacity, the optimal network recovers p(y)/p(x) for every pair, with zero score entropy.This consistency result assumes full support and positive weights.
  • Score Entropy Properties: Gradient rescaling acts as a natural log-barrier that keeps predicted probability ratios nonnegative and addresses problematic gradients in concrete score matching.For unit weights, gradients are scaled by the inverse predicted ratio relative to concrete score matching.
  • Scalable Training: A denoising score entropy variant makes training computationally tractable and scalable by using one network evaluation and tractable diffusion transition densities.The implicit formulation becomes impractical in high dimensions because uniform sampling introduces additional variance.
  • Likelihood Training: Score entropy defines an ELBO that enables likelihood-based training and evaluation using likelihood values and perplexity for language modeling.The bound is expressed through diffusion-weighted denoising score entropy.
  • Practical Implementation: Structured token transition matrices avoid storing dense GPT-2-scale operators, making training roughly as fast and memory-efficient as standard autoregressive training.The sequence state factorizes and tokens are perturbed independently, so only Hamming-distance-one ratios need modeling.

4. Simulating Reverse Diffusion with Concrete Scores

The paper uses concrete score ratios to improve reverse-diffusion simulation, replacing inefficient single-position updates with simultaneous token updates and enabling arbitrary-position conditioning.

  • Reverse Diffusion: Reverse diffusion samples follow transition densities derived from the columns of the time-dependent transition matrix.The process can be simulated with small Euler steps and random transition sampling.
  • τ-Leaping: τ-leaping updates all token positions simultaneously, avoiding the inefficiency of Euler simulation when only one position can change per step.This strategy applies independent token transitions in parallel.
  • Tweedie Denoising: With perfect score learning, the Tweedie τ-leaping update minimizes KL divergence to the true reverse transition among independent simultaneous token-update strategies.The result extends discrete Tweedie denoising while enforcing τ-leaping independence.
  • Controllable Sampling: Because the concrete score represents probability information, Bayes’ rule provides conditional scores from an unconditional score for infilling at arbitrary positions.Conditional sampling changes only the values at the conditioned positions during τ-leaping updates.

5. Experiments

Experiments evaluate SEDD across likelihood-based language modeling, unconditional generation, and conditional generation. SEDD substantially improves over prior diffusion models, matches or exceeds autoregressive baselines in several settings, and trades sampling compute for quality.

  • Core language modeling: SEDD substantially improves upon D3PM on text8 while achieving the best result among non-autoregressive models.It is beaten overall only by an autoregressive transformer and a discrete flow using an autoregressive backbone.
  • GPT-2 zero-shot tasks: SEDD Absorb beats GPT-2 on a majority of zero-shot perplexity tasks and outperforms prior continuous and discrete diffusion baselines.The evaluation covers LAMBADA, WikiText2, PTB, WikiText103, and One Billion Words.
  • Unconditional generation: 6-8× better generative perplexity than GPT-2 at 2048 sampling steps, while matching GPT-2 quality with 32× fewer network evaluations.SEDD's sampling steps form a predictable log-log linear Pareto frontier between compute and generative perplexity.
  • Conditional generation: SEDD is highly competitive for conditional generation, beats GPT-2 and SSD-LM with standard prompting, and supports infilling without distribution annealing.It does not explicitly encode left-to-right prompting as an architectural inductive bias.

6. Related Work

Related work develops continuous and discrete diffusion approaches for language, but SEDD distinguishes itself through a principled score-entropy objective and stronger empirical performance. It also exposes trade-offs involving sampling speed, memory, controllability, and remaining gaps with modern language models.

  • Continuous Diffusion Models for Text Data: Continuous language diffusion embeds tokens in a latent space and dequantizes by selecting the nearest neighbor.Prior work improved these models through several empirical components and alternative loss functions.
  • Discrete Diffusion Models: Discrete diffusion has been applied especially outside language, although some language applications demonstrate strong sequence-to-sequence performance and more efficient generation.Across prior work, discrete diffusion has tended to reduce network evaluations relative to continuous diffusion.
  • SEDD vs Prior Work: SEDD uses denoising score entropy as a principled, scalable objective, avoiding the intractability and high variance of implicit score entropy in prior methods.The paper attributes large language-task improvements to these objective-level differences.
  • SEDD vs Prior Work: SEDD achieves better perplexity and generation results than continuous diffusion models without empirically driven heuristics.The authors suggest adapting successful continuous-diffusion designs as future work.
  • SEDD vs Autoregressive Sampling Iterations: SEDD challenges autoregressive models with competitive perplexity and generation quality, including beating GPT-2 and nucleus sampling, while remaining behind modern large language models.Its sampling comparison involves distinct procedures and implementation considerations.
  • SEDD vs Autoregressive Sampling Iterations: SEDD matches autoregressive inference time at around 100 sampling steps and can increase batch size roughly 4–6 times by removing KV-cache memory.These observations come from an unoptimized codebase and the existing Hugging Face transformers library.

7. Conclusion

The paper introduces efficiently trained SEDD models that use score entropy to parameterize discrete diffusion. SEDD beats previous language diffusion models and rivals autoregressive models in perplexity and quality.

  • SEDD parameterizes a discrete diffusion model with the concrete score and trains it efficiently using a novel score entropy loss.
  • SEDD beats previous language diffusion models and rivals autoregressive models for both perplexity and quality.
  • The framework is proposed as a basis for developing alternatives to the modern autoregressive language-modeling paradigm.

Impact Statement

The impact statement characterizes the work as advancing natural-language generation while identifying no specific danger beyond existing ethical concerns.

  • The paper advances natural-language generation and is described as largely theoretical and not conducted at a scale posing a specific danger.The statement separately acknowledges existing concerns including bias, toxicity, and fake content.

A. Proof of Main Results

The appendix proves properties of the proposed score-entropy objectives and their diffusion-process bounds. The arguments use categorical identities, convexity, Bayes’ rule, cross-entropy, path-measure KL divergence, and Dynkin’s formula.

  • With infinite samples, the score-entropy loss is analyzed through an equivalent optimization objective.
  • The loss is minimized when the learned score ratio satisfies sθ(x)y = p(y) / p(x), with zero loss at the optimum.
  • A categorical analogue of the divergence theorem converts weighted score terms into an equivalent objective after removing constants independent of θ.
  • The denoising score-entropy result follows by marginalizing the explicitly identified score term while treating the remaining factors as unchanged or constant.
  • The main bound is formulated using path-measure KL divergence between the true reverse noising process and the learned reverse process.
  • Dynkin’s formula provides an explicit expression for the conditional path-measure KL term, analogous to how Girsanov’s theorem computes measure changes for standard SDEs.
  • The reverse process analysis simplifies after parameterizing reverse rate matrices with sθ, while Bayes’ rule establishes the stated forward and conditional transition identities.
  • Under the factorization assumption, the objective reduces to cross-entropy terms for each coordinate, and the τ-leaping condition assumes no changes in other dimensions.

B. Algorithms for Training and Inference

SEDD trains a reverse discrete diffusion process by learning data-distribution ratios with score entropy, then samples through transition densities from noisy states. Its algorithms support both unconditional generation and conditional infilling from arbitrary prompt positions.

  • Training: Score Entropy Training samples clean data and noise levels, constructs noisy states, and optimizes a score-entropy objective for the network.The training loop uses a data distribution, noise schedule, token transition matrix, and time interval.
  • Inference: Unconditional sampling starts from the stationary distribution of the token transition matrix and iteratively constructs transition densities while decreasing the noise time.The sampler requires a network, noise schedule, transition matrix, time step, and step size.
  • Model parameterization: The implementation normalizes transition probabilities and uses positivity-preserving network postprocessing when forming the learned score representation.The model exponentiates outputs and normalizes conditional probabilities when necessary.
  • Inference: Conditional sampling fixes prompted positions to their corresponding tokens, then applies the same reverse process to generate the remaining positions.This procedure accepts arbitrary prompt spaces and token assignments rather than requiring only left-to-right conditioning.
  • Noise and priors: Noise schedules and priors are chosen so terminal-distribution mismatch and approximation error are negligible for the modeled data distribution.Uniform transitions are scaled by 1/N with a uniform base distribution, while absorbing transitions use a MASK-based prior with slight leakage.
  • Architecture: SEDD models match corresponding GPT-2 hidden dimensions, block counts, and attention heads while using separate input and output embedding matrices.The reported models contain approximately 90M and 320M non-embedding parameters for small and medium sizes.

C.3. Training Details

The experiments use standardized optimization and hardware settings, compare against established autoregressive and diffusion baselines, and evaluate likelihood, sampling quality, and conditional generation with specified protocols.

  • Optimization: All models use batch size 512, learning rate 3 × 10^-4, gradient clipping at 1, 2000-iteration warmup, and 0.9999 exponential moving averages.Training uses either 8 A100 80GB GPUs or 16 A100 40GB GPUs, with gradient accumulation when needed.
  • Optimization: The study does not perform hyperparameter or architecture search, instead selecting settings for convenience or lifting them from prior training recipes.The architecture comes from DDiT with rotary embeddings, while the learning rate and EMA follow common recipes.
  • Baselines: Baselines include autoregressive flows and transformers alongside multinomial, Bayesian, and discrete diffusion models from prior work.Additional comparisons use Diffusion-LM, BERT-mouth, DiffusionBert, PLAID, and retrained D3PM models matched for size and training specifications.
  • Evaluation: Likelihoods are estimated with 1000-timestep Monte Carlo sampling, using invertible tokenizers and test-set evaluation except for WikiText02.WikiText02 results use the training set because its test set is shared with WikiText103.
  • Evaluation: Generative samples use the Tweedie denoiser, which performs 1–4 perplexity points better than Euler sampling, with 1000 samples generated per model.Conditional MAUVE evaluation generates five samples per ground-truth sample and conditions on 50-token contexts.
  • Evaluation: Figure 2 reports generative perplexity for SEDD Uniform.The figure is part of the sampling-quality evaluation.

D.1. Ablation of Concrete Score Matching

Replacing score entropy with a concrete score-matching objective substantially worsens training and likelihood performance at GPT-2 scale.

  • Ablation: The concrete score-matching ablation produced 3–4× higher likelihood loss than score entropy.The ablation replaced the score-entropy term with an ℓ2-based loss while retaining scaling by Q_t(x, y).
  • Ablation: The corresponding perplexity was 10,000× higher, and the alternative objective generally did not train well.This result supports the reported empirical advantage of score entropy in the tested GPT-2-scale experiments.

D.2. Further Evaluation of Generative Perplexity

Further evaluation finds that SEDD Uniform has a weaker sampling tradeoff than SEDD Absorb, while analytic sampling generally outperforms Euler and trained baselines fall between the two SEDD variants.

  • Generative perplexity: SEDD Uniform does not exhibit the same linear generative-perplexity tradeoff curve as SEDD Absorb, likely because generation quality forms a bottleneck.The passage attributes this difference as a likely explanation rather than an established causal mechanism.
  • Sampling schemes: Analytic Tweedie sampling generally outperforms Euler sampling and is a major factor for the Uniform model’s performance.The evaluation compares analytic sampling with Euler sampling based on reverse diffusion.
  • Baselines: Trained baselines perform substantially worse than SEDD Absorb but slightly better than SEDD Uniform.The comparison includes baselines from Austin et al. and Gulrajani and Hashimoto.

D.3. Additional Samples

This section presents additional unconditional and conditional sample figures for GPT-2 and SEDD variants, alongside assorted generated text excerpts.

  • Additional figures show unconditional GPT-2 Small analytic samples and SEDD-Uniform Small samples.
  • The section includes additional generated-text excerpts, including passages attributed to news, commentary, and other source-like material.
  • Other figures show unconditional SEDD-Absorbing Small, GPT-2 Medium analytic, and SEDD-Absorbing Medium samples.
Loading 2310.16834v3…