PaperlayerView exact revision on arXiv

Source-linked AI summary

How to Train Your Latent Diffusion Language Model Jointly With the Latent Space

Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov, Nikita Gushchin, Ilya Koziev, Alexander Korotin, Dmitry Vetrov

arXiv:2605.07933v1cs.CL

TL;DR

Text diffusion can generate sequences through parallel iterative refinement, but constructing a suitable latent space remains a central challenge. LDLM jointly trains its latent encoder, diffusion model, and decoder, and reports a better quality-diversity trade-off with faster sampling than competing diffusion baselines on LM1B and OpenWebText. The authors identify limitations in the extreme few-step regime and in the extent of recipe and architecture tuning.

  • Problem

    Constructing a suitable latent space is a major challenge for continuous text diffusion, and prior transformer-based latent diffusion methods keep pretrained encoders frozen during diffusion training.

    Source: §1 · p. 2

  • Method

    LDLM jointly trains a latent encoder, diffusion model, and decoder so the diffusion objective can shape the latent space.

    Source: §1 · p. 2

  • Results

    On LM1B and OpenWebText, LDLM achieves a better quality-diversity trade-off while sampling faster than competing diffusion baselines.

    Source: §9 · p. 9

  • Takeaways & Limitations

    Jointly learned contextual latents outperform alternative latent encoding approaches for text diffusion.

    Source: §1 · p. 2

    The extreme few-step regime remains challenging, and the training recipe and encoder/decoder architectures were not exhaustively tuned.

    Source: G Limitations · p. 24

Abstract

from arXiv · show

Latent diffusion models offer an attractive alternative to discrete diffusion for non-autoregressive text generation by operating on continuous text representations and denoising entire sequences in parallel. The major challenge in latent diffusion modeling is constructing a suitable latent space. In this work, we present the Latent Diffusion Language Model (LDLM), in which the latent encoder, diffusion model, and decoder are trained jointly. LDLM builds its latent space by reshaping the representations of a pre-trained language model with a trainable encoder, yielding latents that are easy to both denoise and decode into tokens. We show that naive joint training produces a low-quality diffusion model, and propose a simple training recipe consisting of an MSE decoder loss, diffusion-to-encoder warmup, adaptive timestep sampling, and decoder-input noise. Ablations show that each component substantially impacts generation performance. On OpenWebText and LM1B, LDLM achieves better generation performance than existing discrete and continuous diffusion language models while being $2{\text -}13\times$ faster, indicating that jointly learning the latent space is a key step toward making latent diffusion competitive for text generation.

1 Introduction

Diffusion offers parallel iterative refinement as an alternative to autoregressive generation, while continuous latent methods depend strongly on the structure of their latent space. This work addresses the untrained-to-diffusion encoder gap with LDLM, joint training, and a recipe for effective optimization, reporting stronger generation performance and faster sampling than competing diffusion models.

1 Introduction

  • Diffusion generates text through iterative refinement of all positions in parallel, unlike autoregressive generation’s left-to-right, one-token-at-a-time pattern.Autoregressive models cannot correct previous mistakes under that generation pattern.Source: §1 · p. 2
  • Discrete diffusion predicts tokens independently at each denoising step, limiting simultaneous multi-token generation and making high-quality few-step generation difficult.This follows from factorizing the joint token distribution.Source: §1 · p. 2
  • Continuous diffusion decodes latents to tokens only after generation, allowing positions to be refined gradually; transformer-based contextual encoders have improved latent spaces and generation quality.The paper identifies latent-space structure as a key influence on text diffusion quality.Source: §1 · p. 2
  • Joint training of transformer-based encoders with diffusion had not been implemented; existing latent diffusion models keep pretrained encoders frozen during diffusion training.Shallow embeddings can remain suboptimal, while pretrained representations are not optimized for diffusion because their encoders are trained on surrogate tasks.Source: §1 · p. 2
  • LDLM jointly trains the latent encoder, diffusion model, and decoder to shape a latent space directly for diffusion.The approach is motivated by the importance of latent-space structure to generation quality.Source: §1 · p. 2
  • Naive joint training leads to low-diversity generations, so the paper proposes a recipe to make end-to-end latent learning effective.The recipe is described in Sections 5 and 6.Source: §1 · p. 2
  • On LM1B and OpenWebText, LDLM achieves higher performance than recent discrete and continuous diffusion language models, while staying 2-13× faster.Source: §1 · p. 2

2 Related work

Related work spans discrete diffusion, continuous diffusion on token-level representations, and latent diffusion built from contextual encoder outputs. Prior latent methods keep the encoder independent of diffusion training, whereas this work lets the diffusion objective shape its latent space.

2 Related work

  • Discrete diffusion commonly uses uniform [41] or masking/absorbing [40] [46] terminal noise, but factorized reverse steps make joint multi-token updates difficult [21].Masking-based methods also cannot naturally revise unmasked tokens and discard predictions for masked positions that are not selected for unmasking.Source: §2 · pp. 2–3
  • Token-level continuous diffusion uses fixed encodings [22] [36] [26] [39] or learned embeddings [27] [13] [16], but these representations lack contextual semantic structure.Joint embedding-diffusion optimization can also be unstable and may require regularization to avoid collapse or norm explosion.Source: §2 · p. 3
  • Prior latent diffusion methods use contextual outputs from a pretrained text encoder that remains frozen during diffusion training, while this work jointly trains the encoder with diffusion.PLANNER [57] and LD4LG [28] use a lower-dimensional bottleneck, and COSMOS improves smoothness and robustness through masking and noise injection during autoencoder training.Source: §2 · p. 3

3 Preliminaries

Continuous diffusion gradually corrupts clean continuous samples with Gaussian noise and trains a network to recover the clean sample from a noisy observation. For text, a sequence is first encoded into a continuous latent representation, which is denoised and then decoded back into tokens.

3 Preliminaries

  • Gaussian diffusion corrupts a clean sample z0 with noise to form zt, and trains a neural network to recover z0 from the noisy observation.The noise schedule decreases from ᾱ0 = 1 to ᾱ1 = 0.Source: §3 · p. 3
  • For continuous text diffusion, a token sequence is mapped to a continuous latent representation, denoised in that space, and decoded back into tokens.The encoder maps a fixed-length sequence into R^n×d, and the decoder maps the latent back into the vocabulary sequence.Source: §3 · p. 3

4 Method

LDLM jointly learns its latent space with a diffusion model, using a frozen pre-trained token encoder and trainable latent encoder and decoder. It trains diffusion over the learned latents and decodes generated latents back into text.

4.1 Latent autoencoder

  • A frozen pre-trained token encoder supplies contextual states that a trainable latent encoder reshapes into representations suited to decoding and diffusion.The latent decoder maps perturbed latents back to hidden states, which the token decoder converts to tokens; the hidden-state alignment loss is the only reconstruction signal affecting the latent encoder.Source: §4.1 · pp. 3–4

4.2 Diffusion model

  • The diffusion model learns to recover clean latents from noisy versions, and its objective directly shapes the trainable latent space.The denoiser may use self-conditioning on a previous estimate of the clean latent.Source: §4.2 · p. 4

4.3 Overall objective and sampling

  • At inference, Gaussian noise is reverse-diffused with an Euler-Maruyama solver, and the resulting clean latent is decoded without decoder-input corruption.Decoder-input noise is used during training only.Source: §4.3 · p. 4

5 Joint training recipe

LDLM's joint-training recipe combines hidden-state reconstruction with a gradual diffusion signal to the encoder, adaptive timestep sampling, and noise at the decoder input.

5.1 Pre-trained encoder and hidden state reconstruction

  • MSE hidden-state reconstruction avoids strongly separating nearby latents while preventing latent collapse, a distinction the authors find crucial for robust diffusion latents.The encoder operates on contextual hidden states from a frozen pre-trained token encoder, which provide a favorable starting point for latent learning.Source: §5.1 · p. 5

5.2 Diffusion-to-encoder warmup

  • Jointly applying reconstruction and diffusion losses from the start makes meaningful encoder representations difficult to learn, so diffusion gradients to the encoder are gradually introduced after reconstruction-only warmup.The diffusion model itself trains from the beginning; the encoder first has time to build a decodable latent space.Source: §5.2 · p. 5

5.3 Adaptive timestep sampling

  • Because the latent space evolves during training, LDLM adapts timestep sampling, seeking a loss curve linear in the sampled timestep by sampling proportionally to dL/du.Since dL/du is unavailable in closed form, the method approximates it using 100 intervals and exponential moving averages of losses at their edges; probabilities are updated every 5,000 iterations.Source: §5.3 · pp. 5–6

5.4 Decoder noise σdec

  • Gaussian noise at the latent decoder input encourages efficient latent representations, improves decoder robustness to diffusion errors, and helps stabilize input variance across timesteps.The authors treat σdec as a hyperparameter and study its effect in ablations.Source: §5.4 · p. 6

6 Ablation study

The ablations show that reconstruction loss, encoder choice, diffusion-to-encoder warmup, timestep sampling, and decoder-input noise each shape generation quality. In particular, the recipe must balance a decodable, smooth latent space against diffusion’s denoising demands.

Reconstruction loss

  • Hidden-state MSE yields better generation quality than token-level cross-entropy, while adding cross-entropy to MSE degrades performance.MSE also produces smoother interpolations, supporting its use to shape the latent space.Source: §6 · pp. 6–7

Token encoder

  • Contextual GPT-2 hidden states substantially outperform shallow or scratch-trained token embeddings under the shared cross-entropy comparison.The authors attribute the benefit to contextual features that capture sequence-level structure.Source: §6 · p. 7

Diffusion-to-encoder warmup

  • Without warmup, diffusion rapidly shrinks latent magnitude, impeding meaningful encoder features and degrading generation quality.At least 50k warmup steps improve decoder quality and overall generation; longer warmup can delay diffusion-driven shaping and leave fewer iterations for convergence.Source: §6 · p. 7

Adaptive timestep sampling

  • Adaptive timestep sampling substantially improves generative perplexity and Mauve while leaving entropy essentially unchanged.It makes diffusion-loss growth linear with timestep, unlike the notably worse denoising-difficulty distribution under fixed noise scheduling.Source: §6 · p. 8

Decoder input noise

  • Increasing decoder-input noise initially improves generation by smoothing the latent space and making the decoder more robust to diffusion errors.The effect is non-monotonic: at σdec = 4.0, low generative perplexity accompanies sharply lower Mauve and entropy, indicating collapsed generations; the authors use σdec = 3 thereafter.Source: §6 · p. 8, Figure 4 · p. 8

7 Latent space comparison

The latent-space comparison evaluates progressively richer representations under shared diffusion and decoder architectures. LDLM’s jointly tuned latents perform best among the compared approaches.

7 Latent space comparison

  • The comparison holds the GPT-2 backbone, tokenizer, diffusion model, and decoder architecture fixed, varying the latent-space construction.The alternatives include shallow embeddings, frozen embeddings, pretrained-encoder hidden states, and autoencoder latents with smoothing and robustness objectives.Source: §7 · p. 8
  • LDLM latents achieve the best results among the compared latent spaces on OWT-128.Shallow embeddings perform poorly, contextual GPT-2 embeddings improve generation, and smoothing those embeddings improves it further.Source: Table 2 · p. 8, §7 · p. 9
  • The authors argue that contextual information in latents simplifies denoising, and that joint training directly shapes how information is organized for diffusion.Their comparison concludes that the most diffusable latent space requires tuning it jointly with the diffusion model.Source: §7 · p. 9

8 Main results

On OWT and LM1B, LDLM achieves the strongest quality-diversity trade-off among the compared diffusion methods and samples faster. Its sampling speed advantage is measured on OWT-1024 under the stated evaluation setup.

8 Main results

  • On both OWT and LM1B, LDLM has the strongest Gen. PPL–entropy Pareto frontier and comes close to real-text quality.At comparable entropy, it achieves noticeably lower generative perplexity than the baselines.Source: §8 · p. 9
  • LDLM sampling runs 2–13× faster than all competing methods on OWT-1024.Measurements use batch size 16, one NVIDIA A100 GPU, and 1024 NFEs; LDLM training speed is limited by invoking the autoencoder at every step.Source: Table 3 · p. 9, §8 · p. 9
  • The efficiency comparison measures training and sampling time on OWT-1024, with sampling using 1024 NFEs.LDLM discards the token decoder during training and tunes it separately afterward to avoid unnecessary overhead.Source: Table 3 · p. 9, §8 · p. 9

9 Conclusion

LDLM jointly learns its latent space with the diffusion model, and its training recipe makes this approach effective. On LM1B and OpenWebText, it achieves a better quality-diversity trade-off while sampling faster than competing diffusion baselines.

9 Conclusion

  • On LM1B and OpenWebText, LDLM achieves a better quality-diversity trade-off while sampling faster than competing diffusion baselines.The conclusion suggests that shaping the latent space directly with the denoising objective is essential for latent text diffusion models.Source: §9 · pp. 9–10

A.1 ELBO interpretation of the joint objective

  • The ELBO derivation uses a fixed-noise relaxation because a deterministic posterior lacks finite KL divergence to a continuous prior, then simplifies the objective for implementation.The implemented objective drops constants, takes the deterministic-encoder limit, and replaces exact ELBO weighting with the main-text x0-prediction loss.Source: §A.1 · pp. 16–17

A.2 Diffusion-to-encoder warmup

  • During warmup, the diffusion target remains unchanged while a schedule gradually increases the diffusion gradient passed to the latent encoder.The experiments use k = 10 and c = 0.8; the coefficient rises toward 1 over Swu steps to limit influence from the initially poorly fitted prior.Source: §A.2 · p. 18, Figure 6 · p. 18

B.1 Effect of GPT-2 layer index

  • LDLM outperforms both baselines at every GPT-2 layer, with the final layer yielding its lowest generative perplexity and Mauve near its best configuration.The COSMOS-style latent space is more layer-sensitive, performing best with the third-to-last layer and substantially worse in Mauve with the final layer; TEncDM changes only moderately across layers.Source: §B.1 · p. 18

B.2 Post-training a smaller token decoder

  • After fixing the latent encoder and diffusion model, a post-trained token decoder with two or more layers reaches generation quality comparable to the reference decoder.The smaller decoder is trained on noisy latents using token-level cross-entropy, without hidden-state reconstruction.Source: LDLM (Ours) · p. 19

LDLM (Ours)

  • A large decoder is mainly useful during joint training for hidden-state reconstruction; afterward, direct token-prediction training can substantially reduce its size without large quality loss.A one-layer post-trained decoder is clearly weaker, while adding layers beyond two brings only moderate generative-perplexity improvement.Source: §B.2 · p. 18, LDLM (Ours) · p. 19

B.3 Pretrained encoder generalization

  • BERT-base produces strong generation quality under the GPT-2 training recipe without BERT-specific hyperparameter tuning, suggesting the approach is not tied to GPT-2 representations.The best hyperparameters may depend on the pretrained encoder.Source: §B.3 · pp. 19–20

B.4 Latent encoder architecture

The latent encoder architecture affects generation quality, but the proposed training recipe also works with bidirectional self-attention. The Perceiver Resampler performs better in generative perplexity and diversity.

B.4 Latent encoder architecture

  • The bidirectional self-attention encoder gives reasonable generation quality, but performs worse than the Perceiver Resampler in generative perplexity and diversity.The authors suggest that aggregating information into trainable latent vectors through cross-attention better fits this setup than repeatedly updating pretrained hidden states with self-attention.Source: §B.4 · p. 20

C Implementation details

The implementation uses a frozen GPT-2 encoder, Perceiver Resamplers, and a Diffusion Transformer, with training and sampling settings varying by experiment. Detailed benchmark tables cover LM1B and OpenWebText across sampling budgets.

C Implementation details

  • The main LM1B and OpenWebText runs use 1,000,000 iterations, while ablations use 500,000; the dataset-specific encoder warmup is 25k steps for LM1B and 50k for OpenWebText.All models use AdamW; the training recipe also includes hidden-state MSE reconstruction, decoder-input noise, and adaptive timestep sampling.Source: C Implementation details · pp. 20–21
  • The default architecture uses GPT-2 small hidden states from the third-to-last layer, 6-layer Perceiver Resamplers, and a 12-layer Diffusion Transformer.Sampling uses an Euler-Maruyama solver with a tangent noise schedule; the main OpenWebText run uses sequence length 1024.Source: C Implementation details · p. 21
  • At inference, iterative sampling evaluates only the denoiser; the latent and token decoders, including the vocabulary projection, run once after denoising.The reported timing comparison measures sampling on OpenWebText at sequence length 1024, batch size 16, and 1024 denoising steps using the same hardware setup for all methods.Source: C Implementation details · p. 21

D Detailed unconditional generation results

  • The reported LM1B and OpenWebText benchmarks compare six baselines with LDLM across twelve sampling budgets from 2 to 4096 denoising steps.The tables report generative perplexity, token-level entropy, n-gram diversity, and Mauve; their sequence lengths are 128 for LM1B and 1024 for OpenWebText.Source: D Detailed unconditional generation… · p. 21

E About metrics

The paper distinguishes within-sequence entropy from corpus-level n-gram diversity and provides unconditional generation examples. It also notes limits in few-step sampling and tuning, and discusses asset use and societal risks.

E About metrics

  • Entropy measures within-sequence token variability, whereas n-gram diversity measures corpus-level variation across generated samples.The metrics can diverge, so reporting both distinguishes repetitive sequences from repeated patterns across samples.Source: E About metrics · pp. 21–22

F Existing assets and licenses

G Limitations

  • Extreme few-step generation remains challenging, and the sampler uses neither distillation nor specialized solver optimization.The authors also have not exhaustively tuned schedules, noise levels, or encoder and decoder architectures.Source: G Limitations · p. 24

H Societal impact

  • The authors note that improved quality and lower sampling cost could make misleading, low-quality, or spam-like text easier to produce at scale.They say responsible release should consider misuse evaluation and safeguards appropriate to the released artifacts.Source: H Societal impact · p. 24

I Generation samples

  • The paper presents unconditional-generation examples for LM1B and OpenWebText, with samples annotated by generative perplexity and token entropy.The appendix includes LM1B samples at varying training steps and OpenWebText samples at NFEs of 128, 256, 512, and 1024.Source: I Generation samples · p. 24, Figure 7 · p. 25, Figure 8 · p. 26
Loading 2605.07933v1…