Diffusion-LM Improves Controllable Text Generation
Hold the whole sentence as continuous numbers, and every word stays editable to the last step.
A left-to-right model freezes each word the moment it writes it. Diffusion-LM keeps all 64 positions live at every step, so a gradient can enforce a target parse tree or an exact length.
Explaining the paperDiffusion-LM Improves Controllable Text GenerationAsk a language model for a sentence whose syntactic parse is exactly (S (NP * *) (VP * (NP * *))). Before this paper, the strongest plug-and-play method matched the target parse with an F1 of 17.9 out of 100.
Fine-tuning is the obvious way to control generation: collect pairs of (requirement, text), train the model on them, and it learns to obey. That works. It is also expensive, and it leaves you with one model per requirement, so if you ask for text that is both about Japanese food and shaped like a particular sentence, your only option is to multiply the two fine-tuned models together, which the paper shows composes badly.
The alternative keeps the language model frozen and steers it at decoding time with a separate scorer, which is called plug-and-play control: the language model supplies fluency, an external classifier supplies the requirement, and you sample from the product. Nothing is retrained, and two classifiers compose by adding their scores. For coarse attributes like sentiment or topic this works; for anything positional it has not.
Diffusion-LM changes what is being steered. Instead of a sequence of discrete token choices, generation becomes a trajectory of continuous vectors, refined from Gaussian noise down to word embeddings over two thousand steps. Continuous means differentiable, so a classifier can be attached to those vectors and its gradient can move them directly. Getting there needs an embedding step, a rounding step, a different loss parametrization, and a noise schedule designed for text instead of pixels.
Steering a model you cannot retrain
In Bayesian terms: you have a language model trained on unlabeled text, and for each task a classifier trained on a much smaller labeled set. What you want is a sample from the posterior:
The first factor pushes toward fluent English, the second toward text that satisfies the requirement, and the product is text that does both. For syntactic control the classifier is a probabilistic parser; for part-of-speech control it is a tagger; for semantic content it is a small model that reads a sentence and predicts which field-value pair it mentions.
The paper runs six control tasks; only the first is an attribute, and the other five pin a requirement to a position:
- Semantic content. Given
food : Japanese, produce a sentence containing that value. Scored by exact string match. - Parts of speech. Given
DET NUM NUM VERB ADJ NOUN, produce a sentence of the same length whose tags, under an off-the-shelf tagger, match position by position. Diffusion-LM answers “The Twenty Two provides Chinese food”. - Syntax tree. Given a full constituency parse written as nested brackets, produce a sentence that parses to it. Scored by F1 against an off-the-shelf parser, so numbers on this task are not success rates.
- Syntax spans. Given
(10, 12, PP), produce a sentence in which words 10 through 12 form a prepositional phrase and nothing else is constrained. - Length. Given a target between 10 and 40, land within two words of it.
- Infilling. Given a first and a last sentence, write the middle one so the three connect.
The strongest plug-and-play baseline, FUDGE, matches the part-of-speech target 27.0% of the time and scores 17.9 F1 on parse trees, against 69.9% on semantic content, the one task not tied to a position.
Why left to right makes this hard
Almost every large language model factors the sentence left to right, one token conditioned on the ones before it. To steer such a model at decoding time you have to reweight each next-token distribution, and Bayes says what the reweighting has to be:
Read the second expression carefully. To know whether the token you are about to write is compatible with a requirement anchored at position 11, you need the probability that the requirement will hold once the sentence is finished, and that means summing over every possible completion. With a vocabulary of eleven thousand words and fifty positions left, that sum has terms, and nobody computes it.
The two baselines take opposite ways out. FUDGE trains a neural network to approximate directly, a future discriminator: feed it a prefix and it predicts whether the finished sentence will satisfy the requirement. FUDGE keeps equation (2) exact apart from that one learned approximation. If the prefix has already drifted so that no completion satisfies the target, the discriminator scores every candidate roughly equally badly, the reweighting flattens out, and generation continues with no steering. The paper puts it as: one mistake anywhere in the prefix makes the discriminator assign low probabilities to all continuations.
PPLM leaves that quantity out and runs gradient ascent on the transformer's per-layer key/value cache to raise a classifier's score, thirty updates per generated token in the configuration used here. Its classifier is a single layer on top of the hidden states averaged over time, and averaging over time throws position away, so the classifier cannot express “a prepositional phrase at words 10 to 12” even in principle. That is why PPLM appears on only one of the five control tasks in the results table.
Both baselines share one structural fact: a left-to-right decoder freezes a position the moment it writes it. Advance the run below and watch the counters under each tape: how many positions are written, how many a gradient could still move, and how many the scorer can see. Run it to the end and the left-to-right tape finally shows the scorer the whole sentence, one step after the last chance to change any of it.
So the paper wants a generative model of text in which the entire sequence exists, in a differentiable form, at every step of generation. Continuous diffusion is such a model, and before this paper no continuous diffusion model had produced text.
The diffusion chain, in one pass
Everything in this section is inherited from image diffusion, so it goes quickly. A denoising diffusion model defines a fixed forward process that destroys data and learns to run it backwards. Each forward step shrinks the current vector slightly and adds Gaussian noise:
The shrink factor is tuned so that total variance stays near one, which makes the far end of the chain a fixed you can start sampling from instead of a signal buried under ever-growing noise. Writing and , composing Gaussians gives an exact shortcut to any step:
So training never simulates the chain: pick a step, jump straight to it, learn to undo one step. The other closed form makes the learning target concrete. If you are allowed to look at the clean , the reverse step is itself Gaussian, with a mean that mixes the clean vector and the noisy one:
Training matches the network's reverse step to that cheating posterior. The bookkeeping that makes this a likelihood objective is the variational bound: the quantity in the paper's equation (1) upper-bounds the negative log-likelihood and is minimized. Dropping the per-step weights the bound prescribes gives the squared-error loss everyone trains instead, which is no longer a bound on anything, and trains better anyway.
The last piece of vocabulary: the network can be asked to predict any one of three things, and they are the same prediction in different coordinates. Given , knowing the clean vector determines the noise and the reverse mean , and vice versa, by invertible affine maps with no learning in them. Which one you attach the loss to changes what the model learns about text, and the section after next is about that.
Words in, words out
Everything above lives in . Text does not. Diffusion-LM connects the chain to words at both ends: an embedding step going in, a rounding step coming out.
Going in, each word gets a vector, and a sentence of words becomes their concatenation. The forward process is then extended by one extra transition from the discrete sentence to the first continuous variable, which adds a little noise to the embedding:
The second factor gets you back out. A softmax head reads each position of and returns a distribution over the vocabulary; in the released code its weight matrix is tied to the embedding table, so reading a word out and writing one in use the same numbers. That head makes a probability mass function over sentences instead of a density over , which is why the paper can report its likelihood in nats per token next to an autoregressive model's.
Concretely, on the restaurant-description corpus the paper uses: 64 positions, 16 dimensions each, a vocabulary of 821 words after preprocessing. So is a 64 by 16 array, 1024 numbers, and the embedding table is 821 by 16, about thirteen thousand parameters trained jointly with the 80M-parameter denoiser. On the harder story corpus the vocabulary is 11,043 and the dimension rises to 128.
The denoiser between the two ends is a plain bidirectional encoder, BERT-base in shape (12 layers, 768 wide, 12 heads) and trained from scratch, not initialised from BERT's weights. A two-layer projection lifts each of the 64 slots from 16 dimensions to 768; learned position embeddings and a sinusoidal encoding of the step are added to every slot; the encoder runs with no causal mask, so every position attends to every other; and a second two-layer projection brings each slot back down to 16. With no causal mask, nothing makes position 3 come before position 11, which is why the lower tape in Figure 1 can revise any slot.
Why train the embeddings rather than take a pretrained set? The authors tried both, along with random Gaussian vectors, and found fixed embeddings worse for generation. The denoiser and the embedding table solve a joint layout problem. The table decides how far apart words sit, the denoiser has to resolve them at every noise level, and a layout inherited from an objective like GloVe, one of the two pretrained sets the paper tried alongside GPT-2's, has no reason to place words at distances the noise schedule can work with. Training them together lets the geometry adapt. The learned embeddings cluster by part of speech, although no term in the loss rewards it.
Training the table introduces a failure mode the paper anticipates. Nothing in the squared-error loss stops the embeddings from growing without bound, and scaling them up would make every denoising step easier, because the words would be further apart relative to the same noise. That would also make the far end of the chain wrong: generation starts from , and if the noised data at step is a much wider cloud than that, you would be sampling from the wrong place. A term in the bound depends on the norm of , and after the same simplification as everywhere else it becomes a squared norm on the embeddings scaled by the final . Whether embeddings would actually blow up without it is a hypothesis, not something the paper measures; the paper does not report norms with and without the term, and a later paper from two of the same authors identifies the opposite degenerate solution and calls the term a heuristic regularizer.
The noise level in equation (6) is fixed by the schedule. The code sets it to the schedule's own first noise level, , which the paper's formula puts at 0.1 and the released discretization actually runs at 0.121. Either value is small. Because the encoder is a Gaussian whose only learned quantity is its mean, the reparametrization trick from variational autoencoders carries gradient straight back into the embedding rows of the words in the sentence.
Coming out, rounding means picking the most probable word at each position. In an ideal run would land exactly on a word's embedding and the choice would be obvious. Drag the noise slider below and you can see the geometry that makes it obvious or not: a cloud of noised copies of one word, coloured by whether each still rounds back to it.
Two things need separating here, because the paper uses one word for both. At training and likelihood time, rounding means the softmax head of equation (6). Exact nearest-neighbour lookup, which the figure draws, appears only at decoding time, in a trick covered in the next section. They agree when a point sits near a single embedding and can disagree when it does not.
The trajectory is now continuous end to end, and therefore differentiable end to end. The price is fixed before you start: the array is 64 positions wide whatever you generate, and producing one sample means running the network once per diffusion step, not once per token.
Making the model commit to a word
Trained the straightforward way, the model does not produce an that commits to a single word. It lands between embeddings, rounding becomes a guess, and text quality suffers. The paper traces this to which of the three equivalent heads the squared error is attached to.
Start from the exact relation. If the network predicts and you build its reverse mean by substituting that prediction into the posterior of equation (5), then both and contain the identical term, so subtracting them cancels it and leaves a scalar multiple of the prediction error:
The same substitution in the noise coordinate gives a second identity, this one weighted by the ratio of surviving signal to noise:
So all three losses are the same quantity, , with three different per-step weights:
The paper's footnote calls the parametrizations “equivalent up to scaling constants”, and per step that is right: depends on neither the parameters nor the data, so at any fixed the two losses have the same minimizer. Across steps it is not right, because depends on . Summing a squared error with a weight that varies over six orders of magnitude is a different objective from summing it flat, and one shared network has to divide its capacity between two thousand steps according to that weight.
Drag the cutoff below to see how differently the three divide it. The readout is the share of each objective's total weight on the reconstruction error that falls in the first steps.
Set the cutoff to 1 and the number is 76.5%. Under the posterior-mean parametrization, three quarters of the objective's total weight on reconstructing lands on a single one of the two thousand steps, the one where the noise is already almost gone. The paper's observation, in numbers: the constraint that commit to a word appears only in the terms with near zero. Switching the loss to
weights all two thousand steps equally, so every step carries the full reconstruction term and the network learns quickly that should sit precisely on an embedding. Predicting weights the first 509 of the 2000 steps more heavily than the flat loss does, and the late steps far less. The paper's ablation is that noise-prediction works for small embedding dimensions and collapses for larger ones, while -prediction holds up across the range.
Having an head also makes a decoding trick expressible. At every reverse step the model produces its guess of the clean sentence; the clamping trick replaces that guess with its nearest word-embedding sequence before it is used to build the next state. In the released code this is exact L2 nearest-neighbour search over the embedding table, applied to the prediction and not to the drawn sample, which makes it the operation Figure 2 draws. Clamping makes every step condition on an actual sentence instead of on a cloud of points sitting between words: the guess is forced to commit, so a near-miss at one step cannot be handed forward as a slightly-off vector and compound down the chain. The authors expect clamping at very high noise to be premature, since the model has not decided on words yet, and report that doing it at every step does not hurt much in practice.
Section 4.2 writes the reverse step as
which is the forward marginal of equation (4) with the clean vector replaced by the prediction. It throws away and injects the full noise of that step. The sampler in the released code, and the prose in the paper's own appendix, use the posterior of equation (5) instead. At step 1001 on this schedule the printed rule puts a coefficient of 0.544 on the prediction where the real sampler puts 0.000932, and adds noise of standard deviation 0.839 where the real sampler adds 0.035.
One training step in the released code:
# one training step, E2E settings: n=64 positions, d=16, T=2000
w = sample_sentence() # 64 token ids
x0_bar = Emb[w] # 64 x 16, learned table
x0 = x0_bar + sigma0 * randn_like(x0_bar) # sigma0 = sqrt(1 - abar_1)
t = randint(0, T) # ONE step, not all 2000
eps = randn_like(x0)
xt = sqrt(abar[t]) * x0 + sqrt(1 - abar[t]) * eps
x0_hat = f_theta(xt, t) # BERT-base trunk, 64 x 16 out
mse = mean((x0_hat - (x0_bar if t == 0 else x0)) ** 2)
round_ = cross_entropy(x0 @ Emb.T, w) # rounding head, weights tied
anchor = mean((sqrt(abar[T - 1]) * x0) ** 2) # keeps ||Emb|| in check
(mse + round_ + anchor).backward() # grad -> theta AND EmbThe loss samples a single per example rather than summing over all two thousand, which is the usual unbiased estimate of the sum; but the rounding cross-entropy and the embedding-norm anchor are added at full weight on every step, so relative to the written objective they carry roughly times more weight than the per-step squared error does. A commented-out line in the same file shows the alternative, scaled by , that was not used. And at the regression target switches from the sampled to the noise-free embedding, which the bound's reconstruction term simplifies to.
A noise schedule built for text
A diffusion model shares one network across all its steps, so the noise schedule decides how much of that network's capacity each denoising problem gets. Applied to text, the schedules built for images hand a large share of that capacity to problems the geometry in Figure 2 has already solved.
Take a word whose nearest neighbour in the embedding table sits a distance away. Add the noise of step . The noised point is closer to the neighbour than to the original precisely when the noise projected onto the line between them exceeds half the gap, and that projection is a one-dimensional Gaussian however many dimensions the embeddings have. So the chance that rounding changes its answer has a closed form:
In (10), is the standard normal CDF: is the probability that a draw from a standard Gaussian falls below . With a gap of 1, that probability stays below one in a million for the first 57 steps of the linear schedule and the first 119 of the cosine schedule, and below 1% until steps 125 and 256 respectively. At those early steps the noised embedding still points unambiguously at its own word, and a model that has learned nothing can round it correctly.
Diffusion-LM replaces both schedules with one that gets to real noise immediately, with . Pick a noise level below and the readout gives the first step at which each schedule reaches it.
The paper also claims the sqrt schedule then slows down, so as not to spend many steps on the hardest, highest-noise problems. The figure shows that second half directly: the sqrt curve reaches every level up to 0.71 before the other two, and from 0.72 upward the linear curve overtakes it. Measured as steps spent in the high-noise regime the claim is threshold-dependent. Against linear it holds everywhere. Against cosine it holds from noise 0.95 upward, 368 steps versus 408, and reverses if you lower the threshold: counting steps above noise 0.90, the sqrt schedule spends 682 to cosine's 579. At the front end, where the gap runs to a factor of twenty, it holds against both.
The schedule and the parametrization address the same problem, and the paper reports that they overlap: once the loss is written in coordinates, the appendix reports the sqrt schedule's advantage is no longer salient. With -prediction, the sqrt schedule still keeps performance steady across embedding dimensions and parametrizations that otherwise fall apart.
The closed form is also invalid at its own endpoint: at evaluates to a negative number. A cap of 0.999 on any single step's noise, inside the released schedule builder, fixes this; it fires at a single index, the last, and leaves .
Control by gradient
The trajectory is a continuous object, so the posterior from equation (1) can be imposed on it one step at a time. Each step's conditional factorizes the same way:
That drops the from the classifier term, and the reduction is exact rather than approximate, given the standard forward model in which the requirement attaches to the clean sentence and the noising process ignores it. Taking logs and differentiating, the normalizer depends only on and so has zero gradient, leaving a sum of two terms:
Both terms are computable. The first is a Gaussian whose mean the diffusion model just produced; the second is a neural classifier.
Those classifiers are ordinary supervised models, one per task: a small autoregressive model that reads a sentence and names its field-value pair, a BERT-shaped tagger for parts of speech, and a Transformer constituency parser shared by the two syntax tasks. Each reads the latent array instead of a string of tokens, is trained on the model's own noisy latents, and is told which step it is looking at, so a half-denoised sentence still gets a usable score.
Prior work on guiding image diffusion stops here and takes one analytic step: shift the Gaussian's mean by the covariance times the classifier gradient, then draw. That shift is a first-order Taylor expansion, exact only in the limit of vanishing step covariance, and it is the step classifier guidance for images takes. Diffusion-LM does something else. It draws from the unshifted Gaussian and then pushes the drawn point uphill on a weighted objective:
multiplies the fluency term, not the classifier term, which is the reverse of how a guidance strength usually enters. Maximizing is invariant to positive rescaling, so dividing equation (13) through by gives the same solution:
which is the mode of the reverse Gaussian times the classifier raised to the power . Smaller means stronger control. The tuned values run from 0.1 down to 0.0005, so the effective exponent on the classifier runs from 10 to 2000, and only would give the true posterior.
Small is needed because the fluency term is a Gaussian log-density over all 1024 numbers in the array, at a per-coordinate standard deviation near 0.1 for most of the run, so its gradient runs to order , about a hundred per coordinate. A classifier log-probability is a handful of nats spread over that same array, and its gradient is correspondingly small. Left at the fluency term would pin the sample near the mean and almost no control would happen. Earlier work on guiding image diffusion leaves this term out altogether, and the paper reports that putting it back is what made the generated text fluent.
The step below is drawn in two of the 1024 coordinates. Turn on one control, then both, and watch where the climb ends up, then slide to each end of its range.
The climb is three gradient steps per diffusion step, with Adagrad rather than plain gradient descent because the authors found it far less sensitive to tuning. Reading the released code explains why. The optimizer is constructed fresh inside the three-step loop, so its accumulator only ever holds the current gradient, and the update reduces to the learning rate times the sign of the gradient. Every one of the 1024 coordinates moves by exactly the learning rate, between 0.05 and 0.2 depending on the task, regardless of how large its gradient is. A step whose size does not depend on the gradient is what you want when the two terms in equation (13) differ in magnitude by orders of magnitude and the balance between them shifts along the chain.
Three gradient steps, each with its own classifier backward pass, are expensive, so the sampler downsamples the chain from 2000 steps to 200 by taking every tenth noise level, which costs little on the simpler corpus. The control is also switched off for the final ten steps, letting the sample settle back onto real embeddings before it is rounded.
Composition falls out of equation (12) with no new machinery: two requirements mean two classifier terms, and gradients add. Nothing is retrained and nothing is renormalized. Composing semantic content with a target parse tree, Diffusion-LM satisfies both at 69.8% and 74.8 F1, against 61.7% and 15.4 for FUDGE. A product of two fine-tuned models does better than FUDGE on the syntactic half, at 29.2, and ties it at 61.7 on the semantic half, but stays nowhere near Diffusion-LM; on the part-of-speech composition it collapses on both constraints at once, to 29.4% and 10.5%.
Because the drawn sample is then moved deterministically uphill, the step is no longer a draw from the model's reverse kernel, and the run as a whole is not sampling from . The paper says so itself, calling it a stochastic decoding method that balances maximizing and sampling, and compares it to nucleus sampling or decoding at a low temperature.
The controlled reverse step, end to end:
# one controlled step of the 200-step reverse chain
x0_hat = clamp_to_nearest_embedding(f_theta(xt, t))
mu = posterior_mean(x0_hat, xt, t) # NOT sqrt(abar)*x0_hat + noise
x = mu + sigma_t * randn_like(mu) # the ordinary reverse draw
if t >= 10: # control is off for the last 10
for _ in range(3): # 3 Adagrad steps, fresh state
obj = lam * logp_gauss(x, mu, sigma_t) + classifier.logp(c, x)
x = x + lr * sign(grad(obj, x)) # fixed-size step per coordinate
return xTwo controls with no classifier
Two of the six tasks need no trained classifier, because the requirement can be written straight into the array. They use different mechanisms, and only the released code shows the difference.
Infilling takes the first and last sentences of a story and asks for the middle. Their words occupy their own slots in the latent array, and after every denoising step the sampler writes their clean embeddings back over those slots. The free slots between them are the only positions the model ever chooses. No gradient is taken and nothing is trained; the anchor is a straight overwrite. For anyone comparing this to image inpainting: the anchored slots are held at their clean embeddings even at high noise, rather than at embeddings noised to match the current step.
Length works differently. The array is 64 slots wide whatever length you ask for. A START embedding is pinned at slot 0 and an END embedding at slot , the other 62 are free, and the same three gradient steps per diffusion step pull the model's prediction at the two pinned slots toward those two embeddings. So it is gradient-based control with a hand-written squared-error objective in place of a learned classifier. Two pinned slots out of 64 are enough to land within two words of the target 99.9% of the time.
Infilling needs a single answer, so the paper draws a set of samples and picks by minimum Bayes risk decoding, a decision rule from statistical machine translation:
with a loss like negative BLEU. This picks the sample most similar to all the others rather than the most probable one, the opposite of what beam search does. It works because of how the failures are shaped: bad samples are bad in idiosyncratic ways, so a bad sample sits far from every other sample and pays a large expected loss, while the consensus answer sits in the middle of the pack. Uniform weights are correct here because the samples are drawn from the model itself, which makes the plain average a Monte Carlo estimate of the posterior-weighted expectation.
What it wins, what it costs
Averaging the five control columns, which mix exact-match rates with a parser F1, gives 90.2 for Diffusion-LM against 43.2 for FUDGE, so the paper's “almost doubles” is if anything modest. The fine-tuning oracle is a stricter comparison: it sees (requirement, text) pairs and trains a model per task, and a plug-and-play method is not supposed to reach it.
The wins are on the tasks that need planning. On syntax spans Diffusion-LM scores 93.8% against FT-search's 54.4%, a 39-point margin over the oracle, and that happens because placing a constituent at positions 10 to 12 requires knowing what will occupy 1 through 9 while you decide. On full parse trees it scores 86.0 F1 against 76.4. Table 3 shows two examples of the repair Figure 1 sets up. In one, Diffusion-LM overshoots an interior span by a word and absorbs the error in the spans that follow, dropping a conjunction to stay aligned. In the other, the fine-tuned model undershoots a span by a word and then misaligns everything after it.
Where planning is not the bottleneck, the oracle stays ahead. On semantic content Diffusion-LM reaches 81.2% against FT-search's 89.9%, and on parts of speech 90.0% against 93.0%. The paper calls performance on these two tasks and on length “similar”, which stretches a bit for semantic content's 8.7-point gap. Fluency splits the same way, and the perplexity column needs care because the rows mix plug-and-play methods with fine-tuned ones. FT-search decodes by beam search and has the lowest perplexity on all five tasks. Among the three plug-and-play methods Diffusion-LM has the lowest on four of five, losing only on parse trees, 3.71 against FUDGE's 3.39. Its worst showing is parts of speech, at 5.16 against 4.72 for FT-sample and 3.31 for FT-search.
Infilling comes out a tie, and the paper says so. Against its own autoregressive baseline, trained from scratch on reordered stories to fill gaps, Diffusion-LM scores 7.1 BLEU-4 to 6.7 and 30.7 CIDEr to 26.9, with identical BERTScore, and the human evaluation goes marginally the other way, 0.37 to 0.39, with overlapping intervals. That tie establishes something narrower than “diffusion beats autoregressive models at filling gaps”: a single Diffusion-LM handles a task the baseline needed a bespoke training setup for.
On likelihood the model clearly loses. The variational bound on Diffusion-LM's negative log-likelihood is 2.28 nats per token against an equivalent autoregressive transformer's 1.77 on the restaurant corpus, and 3.88 against 3.05 on stories. The paper notes that scaling model and data narrows this, quoting 3.88 down to 3.10, and that comparison crosses two different setups: the 3.10 comes from a larger model pretrained on eight million synthetic stories, and the autoregressive baseline in that same regime improves to 2.41. So the gap narrows from 0.83 nats to 0.69. The 3.10 and the 2.41 come from a different model family described in an appendix, trained on one-hot vectors in continuous time with a learned schedule and no rounding step, whose improvements the authors report did not carry over to generation quality.
Speed costs differently depending on what you are doing with the model. Unconditional decoding needs one network pass per diffusion step against one per token, so even the 200-step configuration is about seven times slower than autoregressive decoding, and the full 2000-step chain takes about a minute for fifty sequences of 64 tokens. Controlled generation narrows the gap and, against PPLM, closes it entirely, because the baselines pay their control cost per token as well. Fifty controlled samples take about 80 seconds for Diffusion-LM against about 50 for FUDGE, still behind but by a factor of 1.6 rather than 7, and about 80 minutes for PPLM, which runs 30 gradient updates for every token it emits. Training is modest by 2022 standards, 200,000 iterations at batch 64 on a single GPU, roughly five hours on an A100 for the restaurant corpus.
Representing a sentence as 1024 continuous numbers costs likelihood, costs decoding speed, and fixes the sequence length at 64 before generation starts. In exchange, an external gradient reaches every position at every step, which turns “put a prepositional phrase at words 10 through 12” from a search problem into an optimization problem. Later work from the same group, listed in the footnotes, has since closed much of the likelihood gap.
Questions you might still have
If lambda multiplies the fluency term, why does a smaller lambda mean stronger control?
Because maximizing is invariant to positive rescaling. Dividing the objective through by lambda leaves the fluency log-probability with a coefficient of 1 and the classifier log-probability with a coefficient of 1/lambda, so the tuned values 0.1 down to 0.0005 correspond to raising the classifier to the power 10 up to 2000. Only lambda = 1 recovers the actual posterior.
Diffusion-LM has worse perplexity than the autoregressive model it beats on control. How can both be true?
They measure different things. Likelihood asks how well the model spreads probability over all of the held-out text, and a variational bound on a 2000-step chain is a loose way to answer that. Control success asks whether a decoding procedure can be pushed to satisfy a requirement. Diffusion-LM is worse at the first by half a nat per token on the restaurant corpus and eight tenths of a nat on stories, and better at the second because its generation trajectory is continuous and every position stays editable.
Does the sequence length really have to be fixed at 64?
Yes, for this model. The latent array has one slot per position and its width is chosen before generation starts, so shorter sentences are padded out and 64 is a ceiling. Length control works by pinning an END embedding inside that fixed array rather than by changing its size, which is why the length task needs no classifier of its own.
Is this the same classifier guidance that image diffusion models use?
The two share the idea of adding a classifier gradient and differ in procedure. Image classifier guidance shifts the Gaussian mean analytically by the covariance times the classifier gradient and then draws, which is a first-order Taylor approximation. Diffusion-LM draws from the unshifted Gaussian first and then runs three gradient steps on the drawn point, adds a fluency term the image methods omit, and turns control off for the last ten steps. The result is a deterministic push applied to a random draw, which the paper describes as balancing maximizing and sampling rather than as posterior sampling.
Where does this sit next to DDPM and classifier-free guidance?
DDPM supplies everything in the chain section: the forward process, the closed-form posterior, and the squared-error loss. Diffusion-LM adds the embedding and rounding steps at the two ends, switches the loss to predict the clean vector, and replaces the noise schedule. Classifier-free guidance goes the other way, folding the conditioning into the denoiser so no separate classifier is needed; that trades away the modularity Diffusion-LM is built around, since a classifier-free model has to be retrained for each new requirement.
Footnotes & further reading
- The paper: Li, Thickstun, Gulrajani, Liang, Hashimoto, Diffusion-LM Improves Controllable Text Generation (Stanford, NeurIPS 2022). Code. Every code claim on this page was checked against that repository: the schedule in
gaussian_diffusion.py, the clamping trick intest_util.py, and the control loop inscripts/infill_util.py. - The diffusion machinery is inherited from Ho, Jain and Abbeel, Denoising Diffusion Probabilistic Models (explained here), by way of Nichol and Dhariwal's Improved DDPM, whose codebase Diffusion-LM forks and whose cosine schedule is the comparison in Figure 4.
- Guiding a diffusion sampler with an external signal starts with Sohl-Dickstein et al., Deep Unsupervised Learning using Nonequilibrium Thermodynamics (explained), and the version with a trained classifier is Dhariwal and Nichol, Diffusion Models Beat GANs on Image Synthesis (explained), whose Appendix H contains the exact proof that the classifier can be conditioned on the current state alone. The continuous-time statement is in Song et al., Score-Based Generative Modeling through Stochastic Differential Equations (explained), which is the reference Diffusion-LM cites at this point.
- The baselines: Dathathri et al., Plug and Play Language Models, and Yang and Klein, FUDGE: Controlled Text Generation With Future Discriminators. The infilling comparisons are DELOREAN (Qin et al., 2020) and COLD (Qin et al., 2022), with numbers taken from those papers rather than re-run.
- On the datasets, the paper's one-line descriptions are looser than their sources. E2E (Novikova et al., 2017) is 50,602 instances but only 5,751 distinct meaning representations, about eight references each, and the texts are crowdsourced restaurant descriptions elicited from pictures rather than customer reviews; each representation carries 3 to 8 of the 8 ontology slots. ROCStories (Mostafazadeh et al., 2016) reports about 50,000 stories; the 98,000 figure is the corpus after the 2017 release. The vocabulary sizes of 821 and 11,043 are Diffusion-LM's own preprocessing, not dataset statistics.
- The related-work claim that non-autoregressive approaches “fail for language modeling” cites Ren et al., 2020, which studies machine translation, speech recognition and speech synthesis and runs no language-modeling experiment, so the extrapolation to language modeling is the authors' own. The classic statement of the difficulty a non-autoregressive model faces is Gu et al., Non-Autoregressive Neural Machine Translation; Diffusion-LM escapes it because its per-step kernel factorizes across positions but the marginal over 2000 correlated latents does not.
- Minimum Bayes risk decoding as used here is Kumar and Byrne, 2004, who in turn credit the decision rule to Bickel and Doksum (1977) and its speech-recognition use to Goel and Byrne (2000).
- Where this went next: Gulrajani and Hashimoto, Likelihood-Based Diffusion Language Models, from two of the same authors, closes much of the likelihood gap and revisits the embedding-norm term; Gong et al., DiffuSeq, carries the same continuous approach to conditional sequence-to-sequence tasks.
How could this explainer be improved? Found an error, or something unclear? I read every message.