Continuous text diffusion

Overview

Text is discrete, so continuous diffusion needs three extra pieces: a learned embedding that maps each word to a vector, a noising process applied to those vectors, and a rounding step that maps the denoised vectors back to words 2205.14217 §1 2205.14217 §3.3 2604.11748 §3.1. S-FLM normalises the embeddings to unit length, so they sit on a sphere, and it interpolates along the sphere for its noising step 2605.11125 §3 2605.11125 eq. 10. Generation starts from noise, one vector per position, and denoises all positions together. The noise is Gaussian in the first two papers and a uniform point on the sphere in S-FLM 2205.14217 §4.1 2604.11748 alg. 2 2605.11125 alg. 2.

The invariant all variants share is the last step: the final words are the nearest or most probable vocabulary entries for the denoiser's output 2205.14217 §4.1 2604.11748 alg. 2 2605.11125 eq. 13. The denoiser itself differs. Diffusion-LM predicts the clean embedding x_0 directly 2205.14217 §4.2. LangFlow predicts a probability for each vocabulary item, and its embedding estimate is the probability-weighted average of the embedding table 2604.11748 §3.1. S-FLM also predicts a posterior over the vocabulary, and uses it both for the velocity of each step and for the final argmax 2605.11125 §3.1 2605.11125 eq. 15.

The wiki has three compiled cards on this topic: Diffusion-LM, LangFlow and S-FLM.

Variants

Non-autoregressive generation with a clamped x_0 predictor

The base model embeds the words, learns the embedding jointly with the denoiser, and trains with a mean-squared loss on x_0 plus a rounding loss 2205.14217 §4.1 2205.14217 app. E. Predicting x_0 directly at every step stayed stable as the embedding dimension grew, where predicting the noise collapsed 2205.14217 §7.4. At sampling time each predicted x_0 is snapped to its nearest word embedding before the next step is drawn 2205.14217 §4.2. The paper uses a square-root noise schedule, which it argues is more robust for text than linear or cosine schedules 2205.14217 app. A.

Reported decoding cost: the paper uses 200 steps for end-to-end generation and says this is about 7x slower than autoregressive decoding; on ROCStories it keeps 2000 steps because 200 hurt quality 2205.14217 §6.1. Card: Diffusion-LM Improves Controllable Text Generation.

Classifier-guided control on the latents

The latents are steered during sampling. At each step the method takes a few gradient updates on the latent, using the control classifier's loss plus a term that keeps the latent near the sampler's own step 2205.14217 §5.1. The paper takes three updates per diffusion step and reduces the schedule from 2000 to 200 steps to pay for them 2205.14217 §5.1. Five control tasks use this path: semantic content, parts of speech, syntax tree, syntax spans and length 2205.14217 §6.2. Two of them are composed in the paper's compositional test 2205.14217 §7.2.

Author-reported success on the classifier-guided tasks, with fluency as lm-score (lower is better), against the baselines the paper compares to 2205.14217 tab. 2:

The FT-search row is a fine-tuned GPT-2 oracle with search, which the paper includes as a reference point 2205.14217 tab. 2 2205.14217 §6.3. The table also lists a fine-tuned sampling baseline (FT-sample), which is omitted here for space.

Sentence infilling with minimum Bayes risk selection

For infilling, the method draws several samples and keeps the one with the lowest average loss against the others, using negative BLEU as the loss 2205.14217 §5.2. On aNLG the paper reports BLEU-4 7.1 and CIDEr 30.7, ahead of the COLD and DELOREAN rows. Those two rows are copied from their papers, not re-run, so the setups may differ 2205.14217 tab. 5. Human scores there are given for Diffusion-LM and an autoregressive model trained for infilling only 2205.14217 tab. 5. The number of samples is not stated 2205.14217 §5.2.

Cross-entropy training with a learned noise schedule (LangFlow)

LangFlow keeps the continuous latent but trains the denoiser with token cross-entropy instead of a squared error on the embedding 2604.11748 §3.1 2604.11748 alg. 1. Noise is added on a variance-preserving path, indexed by the log noise-to-signal ratio gamma 2604.11748 §3.1. The noise level is drawn from a learned Gumbel-shaped proposal. Its parameters are fitted to the observed loss with a stop-gradient, so the schedule does not change the main network 2604.11748 §4.1 2604.11748 alg. 1. The paper's diagnosis is that a uniform schedule over the usual time variable spends more than half its steps where the loss is almost zero 2604.11748 §4.1.

Self-conditioning is part of the method. In training, with probability 0.25, a first pass runs without the self-conditioning input, and its embedding prediction, with gradients stopped, is fed back as a second input 2604.11748 §5 2604.11748 alg. 1. At sampling time it is always on 2604.11748 §4.2 2604.11748 alg. 2. On LM1B the paper reports that self-conditioning moves LangFlow's PPL from 49.0 to 30.0, while it moves MDLM's PPL the other way, from 31.0 to 32.7 2604.11748 tab. 1.

Sampling integrates an ODE with an Euler method from high to low noise level, on a grid set by the quantiles of the learned proposal. The final tokens are the argmax of the network output at the lowest noise level 2604.11748 alg. 2 2604.11748 app. "Numerical Solver". The paper's generative-perplexity runs use 128 steps on LM1B and 1024 on OWT 2604.11748 §5. Quality falls with fewer steps: on LM1B the appendix sweep gives Gen. PPL 92.2 at 128 steps and 179.6 at 16 2604.11748 app. "NFE comparison on LM1B". Card: LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling.

Author-reported language-modelling results, from the same paper's main table. Gen. PPL is measured with GPT-2 Large over 1024 samples; PPL for the diffusion models is an upper bound 2604.11748 §5 2604.11748 tab. 2:

On zero-shot perplexity, LangFlow is below the autoregressive Transformer on four of seven benchmarks (PTB, Lambada, PubMed and Arxiv), and above it on the other three 2604.11748 tab. 3.

Hyperspherical flow with cross-entropy and a truncated schedule (S-FLM)

S-FLM keeps each token's embedding as a unit vector, and trains with the same unweighted token cross-entropy as LangFlow 2605.11125 §3 2605.11125 eq. 14. The noisy latent is a spherical interpolation between a uniform point on the sphere and the clean embedding, at a signal level set by the schedule. The gradient flows through this interpolation into the embedding table, so the table is learned jointly 2605.11125 §3 2605.11125 alg. 1.

Sampling starts from a uniform point at each position and moves each point along the sphere with a geodesic Euler step. The velocity is the posterior-weighted sum of log maps toward every normalised embedding (exact velocity). With top-k, only the k most likely entries are kept and renormalised. Top-1 keeps a single log map toward the most likely token, which the paper describes as a greedy mode-following approximation 2605.11125 §3.1 2605.11125 alg. 2. The final tokens are the argmax of the posterior at the last time 2605.11125 eq. 13.

Training uses only a truncated part of the schedule range. Its upper bound comes from a formula in the vocabulary size and the embedding dimension, which gives 0.121 for about 50k tokens and d = 768 2605.11125 §3.2 2605.11125 eq. 17 2605.11125 tab. 4. The adaptive part refits the schedule every 50 steps, from recent pairs of time and loss 2605.11125 §3.2 2605.11125 app. B.2. The backbone, called S-arch, is a transformer whose activations stay on the sphere, after nGPT 2605.11125 §3.3.

Ablation on GSM8K, accuracy (%) after 250k training steps. Each row adds one component to the row above it 2605.11125 tab. 2:

The same table lists AR with greedy decoding at 63.3, MDLM (1024 steps, no temperature scaling) at 18.0, Duo at 17.2, and the prior continuous flow models CANDI at 0.2 and FLM at 0.3 2605.11125 tab. 2. On Sudoku, plain S-FLM is below MDLM at all three difficulty levels. With the truncated and adaptive schedule it is above MDLM and FLM at all three, and Duo is higher on easy and hard puzzles 2605.11125 tab. 1. On OpenWebText, S-FLM is roughly on par with prior flow models, and Duo is best overall 2605.11125 fig. 4 2605.11125 §4.3. Card: Language Modeling with Hyperspherical Flows.

When to use which

Budget note: the Diffusion-LM paper trains on one GPU for about 5 hours per 200K iterations on one A100 2205.14217 app. B. The control path needs a separate classifier per task 2205.14217 app. D. LangFlow reports about 292 GPU hours per GPU, on 4 RTX 6000 for LM1B and 32 A100 for OWT 2604.11748 app. "Training Details". S-FLM with the adaptive schedule and the S-arch backbone trains at 2.6 steps per second on TinyGSM, against 4.2 for the same schedule with the DiT backbone, and takes 213.7 GPU hours against 132.3 2605.11125 tab. 3.

Code: the Diffusion-LM release includes its training loss and control scripts. The LangFlow release has sampling and evaluation scripts only; its training code is listed as not yet released, so a LangFlow training loop has to be written from its Algorithm 1 2604.11748 §5 README.md. The S-FLM release has training and sampling scripts, but no LICENSE file, so the card cites paths only. Its GSM8K sampling script defaults to 32 steps, while the table uses 1024 sfm_sphere_arch_truncated_adaptive_no_renorm.sh 2605.11125 tab. 2.

Implementation sketch

Training: embed the words, add a small fixed noise, noise to a random step, predict x_0, and add a squared error, a prior term and a rounding loss 2205.14217 §4.1 2205.14217 app. A 2205.14217 app. E.

# One training step for the x_0-parametrised objective  [§4.1] [App E]
# Input:  token_ids (B, n); one random step t per sequence
# Output: scalar loss

def training_step(token_ids):
    # Embed the words and add the fixed starting noise  [App A]
    clean = embed(token_ids)                      # (B, n, d)
    x_zero = clean + 0.1 * normal_like(clean)
    # Diffuse x_0 to x_t at the sampled step  [§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; 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 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: start from noise, predict x_0, snap it to the nearest word embedding, draw the previous step from the posterior, and round at the end 2205.14217 §4.2 2205.14217 app. E.

# Sampling with the clamping trick  [§4.2] [App E]
# Input:  shape (B, n, d); number of diffusion steps
# 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  [§4.2]
        x_zero_hat = snap_to_nearest_embedding(model(x_t, t))
        # Draw x_{t-1} from the posterior  [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))

For control, add the guided update inside the loop: a few gradient steps on x_{t-1} against the classifier loss, plus a fluency penalty 2205.14217 §5.1 2205.14217 app. B. The paper's fluency term is lambda log p(x_{t-1} | x_t); the released code uses a squared penalty around the sampler's mean instead (see Contradictions).

The LangFlow variant replaces the squared error with token cross-entropy and the x_0 predictor with a probability head. Its sampler is an Euler ODE step on the gamma grid 2604.11748 alg. 1 2604.11748 alg. 2:

# One Euler step of the LangFlow sampler  [alg. 2]
# Input:  z at level gamma_k; embedding table E (V, D); self-conditioning z_hat
# Output: z at level gamma_{k+1}, and the new self-conditioning estimate

def euler_step(z, gamma_k, gamma_next, z_hat, E):
    # Token probabilities, given the previous embedding estimate  [§4.2]
    probs = model(z, gamma_k, z_hat)              # (B, L, V)
    # Embedding estimate: average of E under the predicted probabilities  [§3.1]
    z_hat = probability_weighted_mean(probs, E)   # (B, L, D)
    # Euler update in z / sigma  [alg. 2]
    z = euler_update(z, z_hat, gamma_k, gamma_next)
    return z, z_hat

The S-FLM variant works on the sphere instead. Each embedding is normalised to unit length, the noisy latent is a spherical interpolation, and each step moves the latent along a geodesic by the velocity of the posterior-weighted log maps 2605.11125 §3 2605.11125 alg. 2. The velocity needs a sum over the vocabulary at every position and step. Top-k cuts that sum to k terms 2605.11125 §3.1.

Contradictions

Between papers:

• Squared error on the embedding. Diffusion-LM trains with a squared error on x_0 plus a rounding loss, and reports that predicting x_0 stayed stable as the embedding dimension grew 2205.14217 §4.1 2205.14217 §7.4. LangFlow reports that regressing the clean embedding with mean squared error collapses the embedding table, with nearest-neighbour distances near zero, and uses cross-entropy for this reason [arxiv:2604.11748 app. "Embedding Collapse under Mean Square Error"] 2604.11748 §3.1. That collapse result is for Plaid's objective; the card does not say whether it applies to Diffusion-LM. S-FLM also trains with cross-entropy 2605.11125 §3 2605.11125 eq. 14.
• Noise schedule. Diffusion-LM argues for a square-root schedule over linear or cosine ones 2205.14217 app. A. LangFlow argues that a uniform schedule over the usual time variable wastes steps, and learns its schedule 2604.11748 §4.1. S-FLM truncates its range and refits an adaptive schedule from the loss 2605.11125 §3.2. Its comparisons do not include LangFlow 2605.11125 tab. 1 2605.11125 tab. 2. The papers compare schedules on different time variables and do not test each other's choice.
• Do continuous models match discrete ones? LangFlow reports LM1B PPL 30.0 against 31.0 for MDLM, and OWT PPL 24.6 against 23.2 for MDLM 2604.11748 tab. 2. On GSM8K at T = 0.1, S-FLM with top-1 velocity reaches about 18%, while MDLM is about 33% and Duo about 36% 2605.11125 §4.2 2605.11125 fig. 3. The two papers measure different things, perplexity against exact-answer accuracy at a low temperature, so the cards do not reconcile them.
• Likelihood numbers. Diffusion-LM reports a variational bound on E2E and ROCStories, where its bound is 2.28 on E2E against 1.77 for the autoregressive Transformer 2205.14217 §7. LangFlow reports an ODE upper bound on LM1B and OWT 2604.11748 §5 2604.11748 thm. 3.1. The datasets and bounds differ, so the numbers are not comparable across the two papers.

Within a paper:

• LangFlow LM1B Gen. PPL. The paper gives 81.5 in its self-conditioning ablation 2604.11748 tab. 1 and 92.2 in its main table 2604.11748 tab. 2. The card follows the appendix sweep, which supports 92.2 2604.11748 app. "NFE comparison on LM1B".

Within S-FLM's own sources:

• Truncation side. The text says training uses only the higher noise levels 2605.11125 §3.2. The upper bound of the released adaptive scripts is the 0.121 value from eq. 17 sfm_sphere_arch_truncated_adaptive_no_renorm.sh. The fixed DiT script instead sets a lower bound of 0.869 and leaves the upper bound at 1 sfm_sphere_dit_truncated_fixed_no_renorm.sh. The ablation table labels 0.879 as alpha*(0.1) 2605.11125 tab. 7 2605.11125 app. C.6. The table does not state its convention, so the card does not say which side produced the GSM8K results.
• Adaptive density. The text says the density is proportional to the absolute slope of the loss 2605.11125 §3.2. Algorithm 4 and the code clamp the slope at zero, so negative slopes are dropped rather than flipped 2605.11125 alg. 4 noise_schedules.py.

Within the Diffusion-LM paper's own sources, the released code departs from the text in several places:

• Control learning rates. Table 6 gives syntax-tree learning rate 0.2, syntax-span 0.15 and parts-of-speech 0.05 2205.14217 tab. 6. The released scripts pass different values for these tasks infill.py infill.py. The paper does not say which values produced Table 2.
• Fluency term. The text's lambda log p(x_{t-1} | x_t) is replaced in code by a quadratic penalty around the DDIM step mean infill_util.py.
• Last steps. The code skips classifier updates when t < 10; the paper does not describe this cut-off infill_util.py.
• Prior term. The paper says the prior term is needed only when the last step is not pure noise 2205.14217 app. E. The code always includes it, and the paper does not say whether this matters for its schedule gaussian_diffusion.py.

Pitfalls

• Classifiers on latents. Section 5.1 says the classifier is trained on the diffusion latents 2205.14217 §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 2205.14217 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 2205.14217 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} 2205.14217 app. B.
• Speed numbers. The paper says the controller is 1.5x slower than FUDGE 2205.14217 §6.1. The appendix times give 80 s against 50 s, which is 1.6x 2205.14217 app. C. Sampling needs O(2000) model calls, and control needs 3 classifier calls per step, so 600 model calls in total 2205.14217 app. C.
• Decoding is the slow part. Autoregressive decoding costs O(n) model calls. Diffusion-LM costs one call per diffusion step 2205.14217 app. C. Decoding on ROCStories at 200 steps hurt quality, so it uses 2000 2205.14217 §6.1 2205.14217 app. C.
• Clamping is needed. Without it the model does not commit to a word at each step 2205.14217 §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 2205.14217 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 2205.14217 §4.2 2205.14217 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 2205.14217 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 2205.14217 §6.2. 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 2205.14217 §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 2205.14217 fig. 1 2205.14217 fig. 2 2205.14217 fig. 3 2205.14217 fig. 4.
• Exact velocity cost. The exact velocity sums over the whole vocabulary at each position and step, so the per-step cost grows with the vocabulary size. Top-k reduces the sum to k terms 2605.11125 §3.1.
• Truncation bound is a heuristic. Eq. 17 comes from a model with uniform random embeddings and an approximation of the angle. The authors say it should be treated as a heuristic, not as a guaranteed bound 2605.11125 app. C.2 2605.11125 §7.
• Single run per row. The gaps between rows of the S-FLM GSM8K table come from one training run each. Error bars cover the test set only 2605.11125 app. C.5.
• Sampling steps set the cost. GSM8K results use 1024 steps for diffusion models 2605.11125 tab. 2. Results at 32 steps will not match the table.

Papers

• Diffusion-LM Improves Controllable Text Generation: introduces continuous diffusion over word embeddings for text, with a clamped x_0 predictor and classifier-guided control on the latents.
• LangFlow: Continuous Diffusion Rivals Discrete in Language Modeling: a continuous diffusion language model over token embeddings, trained with token cross-entropy, a learned Gumbel noise schedule and self-conditioning, sampled with an Euler ODE solver.
• Language Modeling with Hyperspherical Flows: a flow over unit-norm token embeddings on the sphere, trained with cross-entropy on a truncated, adaptive schedule, with an S-arch backbone whose activations stay on the sphere.