TEncDM: Understanding the Properties of the Diffusion Model in the Space of Language Model Encodings

1 Introduction

• Autoregressive generation can let an early mistake spoil later text and requires one model evaluation for each new token.The introduction motivates exploring diffusion as an alternative generation approach.Source: §1 · p. 1
• Non-autoregressive diffusion generates tokens simultaneously, can revise any sequence part, and uses denoising iterations rather than sequence length to determine evaluation count.The number of iterations can be reduced through distillation.Source: §1 · p. 1
• The paper examines text decoding, diffusion architecture, noise schedules, and self-conditioning to identify design choices for continuous text diffusion.It combines these findings in the Text Encoding Diffusion Model (TEncDM).Source: §1 · pp. 1–2
• TEncDM outperforms other non-autoregressive diffusion models on paraphrasing, summarization, and text simplification.The comparison covers models trained in embedding or encoding spaces.Source: §1 · p. 2
• TEncDM trains diffusion in a latent space built from pre-trained Transformer encoder outputs.The approach uses contextual information extracted from embeddings rather than raw embeddings as latents.Source: §1 · pp. 1–2

2 Problem Statement and Background

• Unconditional generation samples a variable-length token sequence from p(y), while conditional generation samples from p(y|x) to satisfy a condition x.The sequence is represented as y = [y1, . . . , yn].Source: §2 · p. 2
• Gaussian diffusion trains a denoising network to reconstruct the original latent z0 from a noisy latent zt, with noise controlled by a schedule.The forward process adds Gaussian noise, and the schedule decreases αt as time advances.Source: §2 · p. 2
• Sampling starts from Gaussian noise and iteratively applies the denoising network to generate latents.Text diffusion also requires an encoder to map tokens to continuous latents and a decoder to convert generated latents back into text.Source: §2 · p. 2

3 Related Work

• Embedding-based diffusion converts predicted latents into text by selecting the token whose embedding is nearest to each latent vector.Other approaches use categorical noise, reduced latent representations, or the probability simplex.Source: §3 · p. 2
• Self-conditioning adds the model’s previous data-sample estimate as input, but its effectiveness and influence on generation dynamics had not been analyzed.Without self-conditioning, the model is conditioned on zt and timestep t.Source: §3 · p. 2
• Text-domain noise schedules have been adjusted to add more noise near t = 0 because small noise may not change a latent’s nearest text.Prior work recommends this to make denoising more difficult.Source: §3 · p. 2

4 Understanding Text Diffusion

• Encodings are final-layer outputs of a pre-trained language model that provide contextual information absent from token embeddings.The paper argues this context makes encodings more suitable latents for diffusion.Source: §4 · p. 2
• The decoder must handle inaccurate latent samples, and the authors report that an advanced context-dependent decoder improves generation quality.They leave encoder fine-tuning for future work, although they state it could further improve quality.Source: §4 · pp. 2–3
• Self-conditioning increases prediction confidence at each denoising timestep, reducing required generation steps; sample quality diminishes as step count increases.The authors hypothesize that a mismatch between training and generation latents may explain the quality pattern.Source: §4 · p. 3
• Encoding-based diffusion requires more noise than commonly used cosine and square-root schedules provide, and the authors hypothesize contextual information makes denoising easier.The paper reports that these schedules add little noise over many timesteps, weakening the training signal.Source: §4 · p. 3

5 Methodology

• The framework combines a diffusion encoder, diffusion model, and decoder; conditional generation adds an input encoder whose output reaches the model and decoder through cross-attention.Figure 1 distinguishes training (top) from generation (bottom).Source: §5 · p. 3, Figure 1 · p. 3

5.1 Diffusion encoder, Ediff

• The pre-trained Transformer diffusion encoder maps text to same-length latents, with padding used to align shorter sequences.Special-token encodings are replaced with embeddings and padding encodings with zeros; encoder weights remain fixed during training.Source: §5.1 · p. 3

5.2 Decoder, D

• A BERT-type decoder uses context to rectify potential latent-prediction errors; training on corrupted latents aims to bridge the distribution mismatch with generated latents.Corruption expands the decoder’s training domain and is sampled by applying the diffusion process at t ∈ [0, 0.15].Source: §5.2 · pp. 3–4

5.3 Diffusion model, ˆzθ

• The 12-layer diffusion model predicts the original latent from its noisy version, using coordinate-wise normalized latents and the tan-d noise scheduler.The scheduler’s d parameter controls the noise introduction rate and defaults to 9.Source: §5.3 · p. 4
• Self-conditioning training alternates between predictions without a prior estimate and predictions conditioned on a stop-gradient estimate, preparing the model for inference behavior.The two cases are sampled with probability 0.5 each.Source: §5.3 · p. 4

5.4 Generation process

• Generation starts from Gaussian noise and uses an Euler solver for 50 denoising steps by default; conditional generation adds cross-attention context and classifier-free guidance.The guidance coefficient is 0.5.Source: §5.4 · p. 4

6 Datasets

• ROCStories and Wikipedia are the unconditional-generation datasets; ROCStories contains 98k five-sentence stories, and Wikipedia contains over 2 million cleaned articles.The stories capture causal and temporal relations between everyday events.Source: §6 · p. 4
• QQP supplies 144k paraphrased question pairs for paraphrasing, while XSum provides 204k BBC article-summary pairs for summarization.Source: §6 · p. 4
• Wiki-Auto consists of aligned complex and simplified Wikipedia sentences for text simplification.Source: §6 · p. 4

7 Empirical Analysis

• The component analysis uses unconditional ROCStories and Wikipedia and leaves comparisons with other methods to separate evaluations.The authors state that findings transfer to conditional generation in Section 8.Source: §7 · p. 4

7.1 Evaluation Metrics

• Evaluation uses GPT-2 Large perplexity, a diversity metric, memorization measured by training-set 4-grams, and MAUVE for generated-text quality.MAUVE compares generated and reference text distributions; evaluation generates 1,000 texts and repeats the procedure five times.Source: §7.1 · pp. 4–5

7.2 Model setup

• BERT encodings outperform BERT embeddings, and BERT encodings provide the clearest latent-space advantage among the tested encoders on ROCStories.The comparison includes BERT, RoBERTa, BART, and T5 encoders of approximately 100M parameters each.Source: Table 1 · p. 5, §7.2 · p. 5
• The three-layer Transformer decoder with contextual information gives the best result, and corrupting decoder inputs improves quality for both datasets.The MLP without corruption has the lowest text quality by perplexity; MAUVE comparisons are difficult because of substantial variance.Source: Table 2 · p. 5, §7.2 · p. 5
• With self-conditioning, generation quality peaks at 50 steps by MAUVE and then falls; at its peak, it surpasses the model without self-conditioning on MAUVE and perplexity.Source: §7.2 · p. 5, Figure 2 · p. 6
• The tan-9 scheduler is marginally superior to other schedulers on MAUVE, with slightly better memorization and diversity, but lags behind tan-7 on perplexity.The authors report that tan-d adds noise more consistently across timesteps, making denoising more challenging; excessive noise can hurt performance.Source: Table 3 · p. 6, §7.2 · p. 6

8 Seq2Seq Experiments

• The evaluation covers paraphrasing, summarization, and simplification, using ROUGE-L, BERTScore, and BLEU-4 for paraphrasing and simplification and ROUGE-1/2/L for summarization.Source: §8 · pp. 6–7
• TEncDM variants outperform popular embedding-based diffusion methods; with an optimal encoder, TEncDM is comparable to three autoregressive baselines on QQP and XSum and outperforms them on Wiki-Auto.Source: Table 4 · p. 7, §8 · p. 7
• Comparisons include non-autoregressive diffusion and autoregressive baselines; the table marks results sourced from prior papers with ⋆ and newly trained results with †.GENIE results are sourced from another paper because its original evaluation used ground-truth labels to select outputs.Source: §8 · p. 7, Table 4 · p. 7
• No tested encoder leads on every task: T5-base excels on QQP, while RoBERTa-base performs best on XSum and Wiki-Auto.BERT-base, T5-base, and RoBERTa-base achieve comparable performance overall.Source: §8 · p. 7

• TEncDM outperforms recent diffusion models and some classical autoregressive methods in downstream task experiments.Source: §9 · p. 7
• The model diffuses in language-model encodings, whose contextual information helps recover latents compared with embeddings.Source: §9 · p. 7
• Self-conditioning increases prediction magnitude, allowing the number of generation steps to be reduced.Source: §9 · p. 7
• The proposed decoder and noise scheduler boost generation quality, and ablations on ROCStories and Wikipedia demonstrate the impact of the design choices.Source: §9 · p. 7

10 Limitations

• Latent dimensionality depends on sequence length, so training slows significantly as sequence length increases.Source: §10 · p. 7
• The authors avoid joint training of the diffusion encoder, decoder, and denoising model to keep the approach from becoming overcomplicated.Source: §10 · p. 7
• Different diffusion encoders work better for different tasks, making selection of the best encoder task-dependent.Source: §10 · p. 7

A Decoder for embedding-based model

• The decoder improves MAUVE, memorization, and diversity over closest-embedding rounding on the BERT embedding-based model.The rounding method’s low perplexity reflects low diversity, not necessarily higher sample quality.Source: A Decoder for embedding-based model · p. 10, Table 5 · p. 10

B Corruption for decoder training

• Corrupting clean latents trains the decoder to handle potentially inaccurate latents generated by diffusion.Source: B Corruption for decoder training · p. 10
• The forward-process corruption technique performs better than adding random Gaussian noise to decoder inputs.The comparison applies both decoders to the same diffusion-produced latents; Figures 5 and 6 plot generation quality against maximum timestep or noise strength.Source: B Corruption for decoder training · p. 10, Figure 5 · p. 10, Figure 6 · p. 10
• The decoder’s forward-process corruption samples a timestep from 0 to tmax and uses the corresponding zt, while its alternative adds Gaussian noise with strength σ.Source: B Corruption for decoder training · p. 10

C.1 Self-conditioning increases prediction magnitude

• Repeatedly feeding the model its previous prediction increases prediction magnitude while the noisy latent and timestep stay fixed.The experiment samples zt, predicts repeatedly, and changes only the self-conditioning input; Figure 7 plots prediction-magnitude trajectories across timesteps.Source: §C.1 · pp. 10–11, Figure 7 · p. 11

C.2 Wikipedia self-conditioning analysis

• On Wikipedia, MAUVE with self-conditioning rises and then declines, while without it quality increases monotonically; both settings perform similarly at 50 and 100 steps.The with-self-conditioning pattern is similar to that reported for ROCStories.Source: §C.2 · p. 11, Figure 8 · p. 11
• Wikipedia prediction-magnitude trajectories resemble ROCStories trends but show a larger spread in their final magnitudes.Source: §C.2 · p. 11, Figure 9 · p. 11

G Low MAUVE for RoBERTa encoder

• RoBERTa’s low MAUVE is attributed largely to generated-text length mismatch, not inferior text quality; BERT’s generated lengths more closely match source lengths.The length-distribution comparison is shown for TEncDM with RoBERTa and BERT encoders and for source text.Source: G Low MAUVE for RoBERTa encoder · p. 12, Figure 11 · p. 12
• After truncating generated and source texts to 40 tokens, MAUVE increased significantly and the values became closer, supporting its dependence on text length.Table 7 reports the truncated-text comparison.Source: G Low MAUVE for RoBERTa encoder · p. 12, Table 7 · p. 12
• RoBERTa text length decreases as clusters become more balanced, whereas BERT text length stays relatively constant; the authors link RoBERTa’s longer texts in source-poor clusters to lower MAUVE.Figure 13 orders clusters by entropy, with more imbalanced clusters on the left.Source: G Low MAUVE for RoBERTa encoder · p. 13, Figure 13 · p. 13

H Training time comparison

• Adding the diffusion encoder slows training by 15% and self-conditioning by 20%, while the encoder does not affect generation speed.Measurements used one A100 GPU, batches of 512 texts, and sequence length 128.Source: Table 8 · p. 13, H Training time comparison · p. 13

I Scalability

• Scaling the diffusion model alone to 223M parameters is reported as effective on Wikipedia; jointly scaling the encoder and diffusion model is deferred because it requires substantially more compute and training time.The larger model uses 16 transformer heads and 24 layers, compared with 12 heads and 12 layers in the base model.Source: I Scalability · pp. 13–14, Table 9 · p. 14

J Comparison with GPT-2

• TEncDM surpasses GPT-2-small on all reported metrics except perplexity, though the paper calls that comparison unfair.The authors attribute GPT-2’s lower perplexity to greater training-set memorization and lower diversity.Source: Table 11 · p. 14, J Comparison with GPT-2 · p. 14