PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion-LM Improves Controllable Text Generation

Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, Tatsunori B. Hashimoto

arXiv:2205.14217v1cs.CLcs.AIcs.LG

TL;DR

Controllable generation needs methods that steer language models toward complex fine-grained requirements without retraining, beyond simple attributes such as sentiment. Diffusion-LM uses continuous diffusion to denoise Gaussian vectors into word vectors and applies gradient-based control through hierarchical latent representations, achieving strong results across six tasks while retaining important efficiency trade-offs.

  • Problem

    Existing controllable-generation methods have made limited progress on complex fine-grained controls without retraining, such as syntactic structure.

  • Method

    Diffusion-LM denoises Gaussian vectors into word vectors, then applies gradient updates to continuous latent variables while balancing fluency and control satisfaction.

  • Results

    Across six fine-grained control tasks, Diffusion-LM almost doubles prior methods’ control success rate and matches or outperforms fine-tuning baselines.

  • Takeaways & Limitations

    Continuous hierarchical latents enable complex controls, including composed semantic and syntactic constraints and span-anchored tasks such as infilling.

  • Takeaways & Limitations

    The constructed Diffusion-LMs have higher perplexity, substantially slower decoding, and slower training convergence.

Abstract

from arXiv · show

Controlling the behavior of language models (LMs) without re-training is a major open problem in natural language generation. While recent works have demonstrated successes on controlling simple sentence attributes (e.g., sentiment), there has been little progress on complex, fine-grained controls (e.g., syntactic structure). To address this challenge, we develop a new non-autoregressive language model based on continuous diffusions that we call Diffusion-LM. Building upon the recent successes of diffusion models in continuous domains, Diffusion-LM iteratively denoises a sequence of Gaussian vectors into word vectors, yielding a sequence of intermediate latent variables. The continuous, hierarchical nature of these intermediate variables enables a simple gradient-based algorithm to perform complex, controllable generation tasks. We demonstrate successful control of Diffusion-LM for six challenging fine-grained control tasks, significantly outperforming prior work.

1 Introduction

Diffusion-LM addresses complex controllable generation with a continuous diffusion language model whose hierarchical latents support gradient-based steering. Across six control targets, it substantially improves over prior plug-and-play methods and remains competitive with fine-tuning.

  • Contribution: Diffusion-LM denoises Gaussian noise vectors into word vectors, creating continuous latent representations for controllable generation.The model incrementally denoises a sequence and uses the resulting hierarchy of latents for control.
  • Motivation: Continuous diffusion models had succeeded in vision and audio but had not been applied to text because text is inherently discrete.
  • Method: Gradient updates on Diffusion-LM’s continuous latents balance fluency against target structural and semantic controls.Fluency is parametrized by Diffusion-LM, while control requirements are parametrized by a classifier.
  • Evaluation: The approach evaluates six control targets spanning fine-grained attributes such as semantic content and complex structures such as parse trees.
  • Results: Diffusion-LM almost doubles prior plug-and-play success rates and matches or outperforms fine-tuning on the evaluated classifier-guided tasks.It also composes semantic and syntactic controls and supports span-anchored tasks such as length control and infilling.

2 Related Work

Prior controllable-generation work mainly steers frozen autoregressive language models or uses discrete diffusion processes. Their generation-order constraints limit control over global properties, motivating continuous latent representations for flexible steering.

  • Diffusion Models for Text: Continuous diffusion models have produced strong results for images and audio, while prior text diffusion work used discrete state spaces.Diffusion-LM focuses on continuous diffusion for text and introduces continuous latent representations.
  • Autoregressive and Non-autoregressive LMs: Autoregressive language models generate left to right, limiting flexibility for controls involving both left and right contexts.Infilling and syntactic structure control are examples of settings with global or right-context constraints.
  • Plug-and-Play Controllable Generation: Plug-and-play methods keep the language model frozen and steer outputs with potential functions that jointly target control satisfaction and fluency.
  • Plug-and-Play Controllable Generation: PPLM performs gradient ascent on autoregressive hidden activations but remains left-to-right and fails on more complex controls described by the paper.

3 Problem Statement and Background

The paper formulates controllable generation as sampling text that satisfies a control while remaining fluent, and reviews diffusion models as latent-variable Markov chains. Diffusion training adds noise in a fixed forward process and learns a reverse denoising process.

  • Problem Statement and Background: Controllable generation samples a word sequence w from p(w | c) so that it satisfies a target control c.Controls may specify a syntax tree or sentiment label.
  • Problem Statement and Background: In plug-and-play generation, p_lm(w) promotes fluency while p(c | w) promotes control satisfaction, yielding p(w | c) ∝ p_lm(w) · p(c | w).
  • Language Modeling Background: Autoregressive language modeling factors the sequence left to right by repeatedly predicting the next token from the preceding partial sequence.
  • Diffusion Models for Continuous Domains: A continuous diffusion model represents data as a Markov chain x_T … x_0 and incrementally denoises Gaussian latent variables toward the target data distribution.Each reverse transition is modeled as p_θ(x_t−1 | x_t) = N(x_t−1; μ_θ(x_t,t), Σ_θ(x_t,t)).
  • Diffusion Models for Continuous Domains: The forward diffusion process adds Gaussian noise through transitions q(x_t | x_t−1), with β_t controlling the noise added at step t.The forward process has no trainable parameters and supports learning a reverse process that reconstructs the data.
  • Diffusion Models for Continuous Domains: Diffusion models are trained with a variational likelihood objective, while a reweighted mean-squared-error surrogate can stabilize training and improve sample quality.The surrogate compares the neural network’s predicted reverse-process mean with the posterior mean.

4 Diffusion-LM: Continuous Diffusion Language Modeling

Diffusion-LM adapts continuous diffusion to discrete text by learning word embeddings, adding rounding, and redesigning training so denoising predictions commit precisely to words. A clamping trick further reduces rounding errors during decoding.

  • Model construction: Diffusion-LM maps discrete words into continuous vectors and adds embedding and rounding mechanisms to standard diffusion.The embedding maps each word to a vector in R^d, while rounding maps predicted vectors back to words.
  • Model construction: End-to-end training jointly learns the diffusion model and word embeddings, which outperform fixed random embeddings on generation and control tasks.Fixed embeddings were also helpful for held-out perplexity, but trainable embeddings performed best for control and generation.
  • Rounding: Rounding selects the most probable word independently at each position, but predicted x0 vectors may fail to commit to a single word.This mismatch motivates additional decoding-time techniques to make vector predictions more precise.
  • Training objective: The reparameterized objective predicts x0 directly at every diffusion step, encouraging predictions to center precisely on word embeddings.This objective was introduced because the original formulation placed insufficient emphasis on modeling x0’s commitment to a single word embedding.
  • Rounding: The clamping trick replaces each predicted x0 vector with its nearest word-embedding sequence before sampling the next latent state.It forces intermediate predictions to commit to words and reduces rounding errors, although applying it near the noisiest steps may be suboptimal.

5 Decoding and Controllable Generation with Diffusion-LM

Diffusion-LM performs controllable generation by steering continuous diffusion latents with gradient updates, rather than controlling discrete text directly. Fluency regularization, multiple updates, and MBR decoding support fluent or high-quality outputs.

  • Controllable generation: Controllable generation operates on Diffusion-LM’s continuous latent sequence x0:T and rounds the resulting latents into text.This provides a plug-and-play control procedure inspired by a Bayesian formulation over conditioned diffusion trajectories.
  • Controllable generation: Each diffusion step combines the model transition score with a classifier score, using a gradient update on xt−1.The transition term is parameterized by Diffusion-LM, while the control term is parameterized by a neural-network classifier.
  • Controllable generation: Fluency regularization balances the diffusion transition term λ log p(xt−1 | xt) against the control term log p(c | xt−1).The authors report that this regularizer was instrumental for generating fluent text.
  • Controllable generation: Multiple gradient steps improve control quality, while downsampling diffusion steps from 2000 to 200 reduces computation without substantially hurting sample quality.The method uses 3 Adagrad updates per diffusion step and decodes over 200 steps after downsampling.
  • Decoding: For tasks requiring a single high-quality sequence, MBR decoding selects the sample with minimum expected loss from a set of Diffusion-LM samples.Low-quality samples can be penalized because they are dissimilar from the remaining samples under the loss function.

6 Experimental Setup

The experiments train Diffusion-LM on E2E and ROCStories, then evaluate six control tasks spanning semantic, syntactic, length, and infilling requirements. Evaluation combines task-specific success metrics with teacher-LM fluency scores and comparisons to plug-and-play and fine-tuning baselines.

  • Datasets: Diffusion-LM is trained on E2E restaurant reviews and ROCStories five-sentence narratives, with ROCStories presenting the larger vocabulary and more diverse semantic content.E2E contains 50K examples labeled by eight fields, while ROCStories contains 98K stories and an 11K-word vocabulary.
  • Model configuration: The model uses an 80M-parameter Transformer with sequence length n = 64 and diffusion training length T = 2000.Embedding dimensions are d = 16 for E2E and d = 128 for ROCStories; decoding uses dataset-specific step counts.
  • Control tasks: The evaluation covers six controls: semantic content, parts of speech, syntax trees, syntax spans, length, and infilling.The first four use classifiers, while length and infilling are classifier-free for Diffusion-LM.
  • Evaluation: Fluency is measured by teacher-LM perplexity, reported as lm-score, where lower values indicate better sample quality.Task success is defined separately using exact match, POS agreement, parse F1, span match, length tolerance, or infilling evaluations.
  • Baselines: The first five control tasks are compared with PPLM, FUDGE, and a fine-tuning oracle, while infilling is compared with three specialized baseline methods.The fine-tuning oracle is not plug-and-play; infilling baselines include DELOREAN, COLD, and an autoregressive infilling model.

7 Main Results

Diffusion-LM achieves strong controllability across classifier-guided tasks, including structured syntax and compositional controls, while also supporting infilling. Its advantages are linked to non-autoregressive generation and continuous coarse-to-fine representations, though language-modeling likelihood remains weaker than an autoregressive Transformer.

  • Language Modeling: Diffusion-LM’s variational NLL bound underperforms an equivalent autoregressive Transformer, although scaling model and dataset size partially narrows the ROCStories gap.Reported E2E values are 2.28 versus 1.77, and ROCStories values are 3.88 versus 3.05; scaling changes ROCStories from 3.88 to 3.10.
  • Classifier-Guided Control: Diffusion-LM achieves high control success and fluency across five classifier-guided tasks, outperforming PPLM and FUDGE and exceeding fine-tuning on parse-tree and span control.The reported metrics are control success (ctrl ↑) and fluency (lm ↓).
  • Mechanisms: The model’s non-autoregressive process supports future planning, while coarse-to-fine representations let classifiers influence whole sequences and individual tokens.These properties are associated with performance on span, length, part-of-speech, and syntax-tree controls.
  • Compositional Control: Diffusion-LM composes semantic and syntactic controls with high success for both components, outperforming FUDGE and fine-tuned product-of-experts models on structured syntax.The composed controls include semantic content with syntax trees or parts of speech, with higher control success at some fluency cost.
  • Infilling: Diffusion-LM significantly outperforms COLD and DELOREAN on sentence infilling and matches a specialized autoregressive model trained from scratch.Human evaluation found no statistically significant improvement for either method despite slightly better automatic scores.

8 Conclusion and Limitations

Diffusion-LM enables fine-grained controllable generation through continuous diffusions, with learned embeddings and reparametrization improving sample quality. Its main drawbacks are higher perplexity, slower decoding, and slower training.

  • Conclusion: Diffusion-LM is a controllable language model based on continuous diffusions that supports six fine-grained control tasks.
  • Conclusion: Learned embeddings and reparametrization substantially improve sample quality under the lm-score evaluation.
  • Limitations: Diffusion-LM has higher perplexity, substantially slower decoding, and slower training convergence than the desired deployment profile.
  • Design choices: The model uses a square-root noise schedule that starts with higher noise, increases it rapidly early, and slows later injection.

C Decoding Speed

Diffusion-LM decoding iterates through many diffusion steps, making it slower than autoregressive decoding and other plug-and-play methods. Skipping steps can preserve simple-task quality but harms harder language-modeling tasks.

  • Decoding speed: Diffusion-LM sampling requires 2000 diffusion steps, producing O(2000) model calls rather than O(n) autoregressive decoding.
  • Decoding speed: Decoding 50 sequences of length 64 takes around 1 minute in the reported medium-length sequence regime.
  • Step skipping: Naively downsampling 2000 diffusion steps to 200 does not hurt E2E sample quality but hurts quality on harder language-modeling tasks.
  • Comparisons: PPLM takes around 80 minutes and FUDGE 50 seconds to generate 50 samples under their reported sampling setups.

E End-to-end Objective Derivations

The end-to-end objective adapts diffusion training to discrete text by jointly learning embeddings and denoising behavior. Its derivation simplifies Gaussian diffusion terms while accounting for terminal-noise and rounding considerations.

  • Objective derivation: The simplified objective matches the model’s predicted mean for x_t−1 to the true Gaussian posterior mean.
  • Objective construction: Diffusion-LM jointly learns the diffusion model and embedding parameters through an end-to-end training objective.
  • Terminal condition: When ᾱ_T is not zero, a terminal regularization term prevents embeddings from adopting excessively large norms.
  • Diffusion model: The model receives (x_t, t) and predicts the distribution of x_t−1 using one Transformer shared across diffusion steps.
  • Sampling: With x_0-parametrization, each step estimates x_0 and samples x_t−1 from the Gaussian posterior conditioned on that estimate and x_t.

F Log-Likelihood Models and Results

The paper separately investigates log-likelihood using modified diffusion training procedures, but those improvements do not translate into better generation quality in its experiments.

  • Training setup: The log-likelihood investigation departs from the original training procedure, including one-hot inputs and continuous-time diffusion objectives.
  • Findings: Log-likelihood improvements did not translate into better generation quality, so the paper focuses on the original method afterward.
  • Evaluation: The modified likelihood models and baseline autoregressive Transformers are evaluated on E2E and ROCStories with small and medium model sizes.

H Additional Ablation Studies

Additional ablations favor learned embeddings, the square-root noise schedule, and Transformer architecture for Diffusion-LM. The supplied qualitative tables also cover syntax, POS, length, and semantic-content controls, while the discussion notes both beneficial and harmful uses of fine-grained controllability.

  • Learned v.s. Random Embeddings: Learned embeddings outperform random embeddings on both ROCStories and E2E.The passage reports improvements of xx percent and xx percent, respectively, without specifying numeric values.
  • Noise Schedule: The square-root schedule delivers consistently good and stable performance across dimension and parametrization choices.It is substantially more robust than cosine and linear schedules under alternative parametrizations such as ϵ.
  • Transformer v.s. U-Net: Transformer architecture outperforms U-Net architecture for language modeling.The comparison adapts U-Net’s 2D convolutions to 1D convolutions suitable for text data.
  • Broader Implications: Fine-grained controllability may support toxicity mitigation and truthfulness, but may also enable targeted disinformation.The discussion also frames watermarking with fluency and distinguishability constraints as a controllable-generation problem.
  • Qualitative Outputs: Qualitative evaluations examine syntax-span, POS-sequence, length, and semantic-content control outputs.The tables mark correct and failed spans or target-length violations to support visual inspection of control behavior.
Loading 2205.14217v1…