Diffusion-LM Improves Controllable Text Generation

In plain words

Most language models write one word after another, left to right. That makes it hard to ask for a sentence with a particular shape, such as a given grammar tree or a fixed number of words, because the model cannot look ahead to the end while it writes the start.

Diffusion-LM takes a different route. It begins with a sequence of random noise, one noisy vector per word, and removes the noise a little at a time until each vector settles on a word. Every step sees the whole sentence at once, so the sentence is shaped as a whole rather than word by word. The intermediate, still-noisy versions are continuous, which means you can nudge them. A small classifier that checks for the wanted property can push those nudges, and the sentence bends toward the property while staying fluent.

The cost is speed and polish. Generating takes many denoising steps, and the text it produces is less fluent than that of a standard left-to-right model. The paper argues that the control gained is worth it for constraints that are hard to express in a left-to-right model.

TL;DR

• Diffusion-LM is a non-autoregressive language model. It denoises a sequence of Gaussian vectors into learned word embeddings, with the embedding table trained jointly with the denoiser §4.1.
• Gradient updates on the continuous latents, guided by a classifier, control the output. Six control tasks are reported; four use a classifier, and the length and infilling tasks are classifier-free §6.2.
• Author-reported: on classifier-guided control it beats PPLM and FUDGE on all five tasks and matches or beats the fine-tuning oracle on syntax and spans §7.1 tab. 2.
• Use it when you need global or fine-grained control and can accept a slow sampler and higher perplexity. Its likelihood is worse than that of an equivalent autoregressive Transformer §7 §8.

Problem

Left-to-right autoregressive models cannot easily condition on the right context. Controls that involve the whole sentence, such as a parse tree or a span constraint, need lookahead or specialised decoding §1 §2. Plug-and-play methods steer a frozen autoregressive model with a classifier, which works for coarse attributes such as sentiment and topic but, the paper argues, fails on fine-grained structure §1 §2. The paper asks whether a non-autoregressive model built on continuous diffusion can support plug-and-play control on structured targets §1.

Continuous diffusion works well for images and audio, but text is discrete, so it needs an embedding step, a rounding step back to words, and a way to learn the embeddings §1 §3.3.

Method

Overview. Each word w_i is mapped by a learned embedding Emb(w_i) to a vector of dimension d. A sequence of length n becomes a vector of size n d §4.1. A forward process adds Gaussian noise to these vectors until the last step is close to pure noise §3.3. A Transformer denoiser, shared across all steps, learns to reverse the process §4.1 §6.1. Generation starts from noise and denoises step by step. A final rounding step picks the most probable word for each position §4.2.

The paper has no numbered algorithm; the steps below are in prose and in the pseudocode blocks that follow §4 §5.

Forward process and embedding (§4.1). The first noising step is written as q_phi(x_0 | w) = N(Emb(w), sigma_0 I), with sigma_0 the initial noise level §4.1. The paper's appendix sets the initial standard deviation to 0.1 [App A]. The embedding is learned end to end, so q_phi has trainable parameters and the reparametrisation trick is used §4.1.

Noise schedule (App A). The paper introduces a square-root schedule, alphabar_t = 1 - sqrt(t/T + s), with s = 1e-4 and T = 2000 [App A]. It argues that standard linear and cosine schedules are not robust for text, since small noise near t = 0 does not change a word's nearest neighbour [App A].

Training objective (§4.1, App E). The starting point is the usual diffusion variational bound, equation (1) §3.3. The paper replaces it with a mean-squared objective that predicts x_0 directly at every step (the x_0-parametrisation). The final end-to-end form is given unnumbered in Appendix E. It has three parts: a squared-error term at each step, a prior term on the last step, and a rounding term, the negative log-likelihood of the words given x_0 [App E]. Equation (2) is the end-to-end objective as first stated §4.1. Equation (3) is the per-step simplification of the KL terms, which the appendix uses to derive the end-to-end form [App E]. The rounding network is a softmax over the vocabulary for each position §4.1.

Clamping trick (§4.2). The paper found that the denoiser does not commit to a single word embedding at each step. Its fix is to parametrise the denoiser by x_0, and at sampling time to snap each predicted x_0 to the nearest word embedding before computing the step §4.2. The paper says this is applied at all steps, and that applying it everywhere does not hurt quality much §4.2 footnote. Sampling then draws x_{t-1} from the posterior given the predicted x_0 and x_t [App E].

Controllable generation (§5.1). Control works on the latents x_{0:T}. At each step the method runs gradient updates on x_{t-1}, with the objective lambda log p(x_{t-1} | x_t) + log p(c | x_{t-1}). The first term keeps the text fluent, and the second is the control classifier §5.1. The paper takes three Adagrad steps per diffusion step, and reduces the number of diffusion steps from 2000 to 200 to pay for them §5.1. The classifier is applied to the latent variables §5.1.

Classifiers (App D). Semantic content: a GPT-2-small language model predicts the field-value pair, and the log-probability of the value is used [App D]. Parts-of-speech: a BERT-base tagger returns per-word tag probabilities, and the log-probabilities are summed [App D]. Syntax tree and spans: a constituency parser, with the log-probability of each labelled or unlabelled span summed or taken for the target span [App D]. Length is classifier-free for Diffusion-LM §6.2 footnote.

Minimum Bayes risk decoding (§5.2). For infilling, which needs a single output, the method draws a set of samples and keeps the one with the lowest average loss against the others, using negative BLEU as the loss §5.2.

Infilling (§6.2, §7.3). Given a left and a right sentence, generate the middle sentence. The paper reports this as classifier-free with MBR decoding §5.2 §7.3.

# Training step for the end-to-end objective, x_0-parametrised  [§4.1] [App E]
# Input:  token_ids (B, n) with n = 64; one random step t per sequence
# Output: scalar loss

def training_step(token_ids):
    # Embed the words, then add the fixed starting noise  [§4.1] [App A]
    clean = embed(token_ids)                       # (B, n, d)
    x_zero = clean + 0.1 * normal_like(clean)
    # Sample a step and diffuse x_0 to x_t  [§3.3]
    t = sample_step()
    x_t = diffuse(x_zero, t)
    # Predict x_0 directly at every step  [§4.2]
    x_zero_hat = model(x_t, t)
    # Squared error to x_0; to the clean embedding at the first step  [App E]
    if t == 1:
        mse = squared_error(clean, x_zero_hat)
    else:
        mse = squared_error(x_zero, x_zero_hat)
    # Prior term on the last step, and rounding loss on the words  [App E]
    prior = squared_error(last_step_mean(x_zero), 0.0)
    rounding = cross_entropy(round_logits(x_zero), token_ids)
    return mse + prior + rounding

# Sampling with the clamping trick  [§4.2] [App E]
# Input:  shape (B, n, d); the number of diffusion steps T
# Output: word ids (B, n)

def sample(shape, num_steps):
    # Start from pure Gaussian noise  [§3.3]
    x_t = normal(shape)                            # (B, n, d)
    for t in steps_from(num_steps, down_to=1):
        # Predict x_0 and snap each vector to its nearest word embedding  [§4.2]
        x_zero_hat = snap_to_nearest_embedding(model(x_t, t))
        # Sample x_{t-1} from q(x_{t-1} | predicted x_0, x_t)  [App E]
        x_t = posterior_sample(x_zero_hat, x_t, t)
    # Round each final vector to its most probable word  [§4.1]
    return argmax(round_logits(x_t))

# One guided denoising step for a classifier control  [§5.1]
# Input:  x_t (B, n, d); step t; classifier; fluency weight lam; Adagrad lr
# Output: x_{t-1} (B, n, d)

def guided_step(x_t, t, classifier, lam, lr):
    # Usual reverse step: mean and std of x_{t-1} given x_t  [§5.1] [App E]
    mean, std = reverse_step(x_t, t)
    x = mean + std * normal_like(mean)
    # Three Adagrad steps per diffusion step  [§5.1]
    for _ in range(3):
        # Control loss from the classifier  [§5.1] [App D]
        control = classifier_loss(x)
        # Fluency: the paper writes lam * log p(x_{t-1} | x_t); the code
        # uses a squared pull toward the mean scaled by std  [code]
        fluency = lam * squared_error(x, mean) / std
        x = adagrad_update(x, control + fluency, lr)
    return x

# Minimum Bayes risk selection among samples  [§5.2]
# Input:  samples (S sequences); loss, e.g. negative BLEU
# Output: one sequence

def mbr_select(samples, loss):
    # Average loss of each sample against all samples  [§5.2]
    risks = [mean_loss(a, samples, loss) for a in samples]
    # Keep the sample with the lowest expected risk  [§5.2]
    return samples[argmin(risks)]

Hyperparameters

Values are as stated in the paper. "Not stated" means the paper and the appendix do not give the value.

Per-task control settings, from the paper's table of selected values tab. 6:

Baselines [App B]: PPLM uses learning rate 0.04, KL scale 0.01, and tunes the number of gradient updates per token over {10, 30}. FUDGE tunes its trade-off over {16, 8, 4, 2}, with pre-condition top-K 200 and post top-K 10.

Results

All numbers are author-reported. Setup: five classifier-guided control tasks, each with 200 control targets from the validation split and 50 samples per target §6.2. Success is exact-match or parse-based, depending on the task §6.2. Fluency is the perplexity under a fine-tuned GPT-2 (lm-score, lower is better) §6.2 footnote. Baselines are PPLM, FUDGE and fine-tuned GPT-2 (FT), with GPT-2 small as the autoregressive base §6.3.

Classifier-guided control, success (ctrl, higher is better) and lm-score (lower is better), as reported in the paper tab. 2:

Composing two controls, semantic content with syntax or with parts of speech, as reported in the paper tab. 4 §7.2:

Sentence infilling on aNLG, with human evaluation from the GENIE leaderboard tab. 5 §6.2 §7.3:

The COLD and DELOREAN rows are copied from their papers, not re-run tab. 5 caption.

Language-model likelihood, reported as the variational bound (exact likelihood is intractable) §7 §4.1: Diffusion-LM scores 2.28 on E2E against 1.77 for the autoregressive Transformer; on ROCStories it scores 3.88 against 3.05. Scaling the model and data brings ROCStories to 3.10 §7. The paper treats the likelihood gap as a cost of the approach §8.

Ablations §7.4 fig. 4: learned embeddings beat random embeddings, most clearly on ROCStories; and predicting x_0 directly stays stable as the embedding dimension grows, while predicting the noise collapses at larger dimensions. Figure 4 is the only evidence and gives no table of numbers in the text.

Paper vs code

Official code: https://github.com/XiangLi1999/Diffusion-LM, linked from a footnote on the paper's first page, at commit 759889d58ef38e2eed41a8c34db8032e072826f4. The owner is an author. Code is Apache-2.0. Code citations are relative to the repository root.

• Training length on ROCStories. Paper: 800K iterations [App B]. The released command uses --lr_anneal_steps 400000 README.md. The E2E command uses 200000, which matches the paper README.md.
• Training objective. The paper's final objective is the x_0-simple form with a squared error at each step, a prior term and a rounding term [App E]. The released e2e mode runs training_losses_e2e. It computes an unweighted squared error to x_0, with the clean embedding at the first step; it adds the rounding cross-entropy and the prior term gaussian_diffusion.py script_util.py. The paper says the prior term is only needed when the schedule does not reach pure noise [App E]. The code always includes it, and the paper does not say whether that matters for its schedule.
• e2e-simple is a different path. The code's e2e-simple branch computes only the cross-entropy term; its squared error is commented out gaussian_diffusion.py gaussian_diffusion.py. The paper does not describe this branch. The released scripts use e2e README.md.
• Noise schedule. The paper writes alphabar_t = 1 - sqrt(t/T + s), s = 1e-4 [App A]. The code uses 1 - sqrt(t + 0.0001) with t = i/T gaussian_diffusion.py gaussian_diffusion.py. These agree. The 200-step control schedule is built from the same formula, which is the same as the paper's subsampling alphabar_{10t} §5.1 [App C]. This follows from the formula depending only on t/T. The initial noise std of 0.1 matches sqrt(1 - alphabar_0) gaussian_diffusion.py.
• Control fluency term. The paper's term is lambda log p(x_{t-1} | x_t) §5.1. The code adds coef * (mean - x)^2 / sigma to the classifier loss, where mean and sigma come from the DDIM step infill_util.py infill_util.py gaussian_diffusion.py. This is a Gaussian-style penalty around the sampler's mean, not the diffusion model's likelihood.
• Control step count. The paper takes three Adagrad steps per diffusion step §5.1. The code takes three steps only while t >= 10, so the last ten steps get no control infill_util.py infill_util.py.
• Sampler for control. The paper describes sampling from the posterior at each step [App E]. The released control script uses the DDIM sampler with eta = 1.0 infill.py infill.py README.md.
• Syntax-tree learning rate. Paper: 0.2 tab. 6. Code passes 0.1 with trade-off 0.0005 infill.py.
• Syntax-spans learning rate. Paper: 0.15 tab. 6. Code passes 0.1 with trade-off 0.1 infill.py.
• Parts-of-speech settings. Paper: trade-off 0.0005, learning rate 0.05 tab. 6. Code passes learning rate 0.1 infill.py and hard-codes the trade-off to 0.0001 with the comment "prev default" infill_util.py.
• Semantic and length settings. These match the paper: trade-off 0.01 with learning rate 0.1 for semantic content infill.py infill_util.py, and trade-off 0.01 with learning rate 0.1 for length infill.py.
• Optimiser schedule and EMA. The paper's learning-rate schedule is linear decay from 1e-4 [App B]. The code decays linearly to zero over lr_anneal_steps train_util.py. The defaults are learning rate 1e-4, weight decay 0, and an EMA rate of 0.9999 train.py. The released decoding command uses the EMA weights README.md. The paper does not mention EMA.
• Batch size and dropout. The defaults are batch size 64 and dropout 0.1 run_train.py run_train.py, as in the paper [App B].
• Clamping. The code snaps the predicted x_0 to its nearest word embedding using a nearest-neighbour lookup, and this is used in the control sampler test_util.py infill.py. This matches the paper's clamping trick §4.2.

Pitfalls

• Classifiers on latents. Section 5.1 says the classifier is trained on the diffusion latents §5.1. Appendix D says the semantic classifier is a language model trained on text, with log p(c | x_t) computed from per-token log-probabilities of the value [App D]. The paper does not say how the classifiers are trained on latents. The released tree-control path needs a separately trained classifier, and its directory is hard-coded infill.py infill.py README.md.
• Fluency is not the paper's term. An implementation that follows the paper's lambda log p(x_{t-1} | x_t) will not reproduce the code. Use the quadratic penalty if you want the released behaviour infill_util.py.
• The last ten steps have no control. The code skips classifier updates when t < 10. The paper does not describe this cut-off, so a re-implementation that follows the text will differ at the end of sampling infill_util.py.
• Fine-grained hyperparameters are not reproducible from the paper alone. Table 6 gives values that differ from the released scripts for syntax tree, spans and parts of speech. The paper does not say which values produced Table 2 tab. 6 infill.py.
• FUDGE's grid. Table 6 lists FUDGE trade-off 20 for three tasks, but the text's tuning grid is {16, 8, 4, 2} [App B].
• Speed numbers. The paper says the controller is 1.5x slower than FUDGE §6.1. The appendix times give 80 s against 50 s, which is 1.6x [App C]. Sampling needs O(2000) model calls, and control needs 3 classifier calls per step, so 600 model calls in total [App C].
• Decoding is the slow part. Autoregressive decoding costs O(n) model calls. Diffusion-LM costs one call per diffusion step [App C]. Decoding on ROCStories at 200 steps hurt quality, so it uses 2000 §6.1 [App C].
• Clamping is needed. Without it the model does not commit to a word at each step §4.2. Implementations must include the snapping step at sampling time, not only in training.
• Prior term for other schedules. The prior term is needed only when the last step is not pure noise. Otherwise the embedding can grow large norms and make the task degenerate [App E]. The paper's choice of schedule avoids this; a different schedule needs the term.
• Missing ablation appendix. Section 4.2 and the Hyperparameters appendix point to an ablation appendix that is not in the source §4.2 [App B]. Ablation numbers are not available beyond Figure 4.
• Timestep sampling and MBR size are not stated. The paper does not say how training timesteps are drawn. The code uses a uniform sampler train_util.py train.py. The number of MBR samples is not stated; the infilling script defaults to 50 infill.py.
• Baselines are not re-run. The infilling baselines COLD and DELOREAN are numbers from their papers tab. 5. Compare them with care, since the evaluation setups may differ.
• lm-scores are not comparable with other papers. The paper uses a fine-tuned GPT-2 as the teacher because its models output UNK tokens that GPT lacks §6.2 footnote. lm-scores from papers that use GPT are not directly comparable.
• Likelihood is worse. Diffusion-LM has higher perplexity than the autoregressive baseline and converges more slowly §8. The paper argues that controllability is the main benefit.
• Figures. Described here, not reproduced. Figure 1 shows the denoising process, Figure 2 the graphical model, Figure 3 a t-SNE plot of learned embeddings grouped by part of speech, and Figure 4 the ablation fig. 1 fig. 2 fig. 3 fig. 4.

Lineage

• Builds on: diffusion models for continuous domains, and the simplified mean-squared objective used for images §3.3. Also the plug-and-play control idea of steering a frozen language model with a classifier §2.
• Compared against: PPLM and FUDGE for plug-and-play control; a fine-tuned GPT-2 oracle (FT); COLD and DELOREAN for infilling; and an autoregressive model trained for infilling §6.3 §6.4.
• Superseded by: not stated in this wiki.