💬 NLP

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

This paper introduces the Latent Diffusion Language Model (LDLM), which jointly trains a latent encoder, diffusion model, and decoder to create a high-quality continuous latent space, achieving superior text generation performance and significantly faster inference speeds compared to existing discrete and continuous diffusion models.

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

Published 2026-05-11

📖 4 min read☕ Coffee break read

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to write a story. The old way (called "autoregressive") is like a student writing a sentence one word at a time, left to right. If they make a mistake early on, they can't go back and fix it without rewriting the whole thing. They also can only write one word at a time, which is slow.

The new method in this paper, called LDLM, is like a sculptor working on a block of marble. Instead of carving one word at a time, the sculptor starts with a rough, noisy block of stone (a "noisy" version of the story) and gradually chips away the noise to reveal the final shape. They can look at the whole block at once and refine it in parallel, fixing mistakes anywhere in the story instantly.

However, there's a catch: To sculpt effectively, you need the right kind of marble. If the stone is too hard or too soft, the sculptor can't do their job. In AI terms, this "marble" is called the latent space—a hidden, mathematical representation of the text that the model works with before turning it back into words.

The Problem: The "Frozen" Blueprint

Previous attempts at this "sculpting" method used a pre-made blueprint for the marble. They took a smart, pre-trained language model (like a dictionary that already knows how words fit together) and froze it in place. They tried to sculpt the text using this fixed blueprint.

The authors of this paper realized this was like trying to sculpt a statue using a blueprint designed for a different type of stone. The pre-trained model wasn't optimized for the "sculpting" process (diffusion), so the results were clunky, slow, or low-quality.

The Solution: Joint Training (The "Team-Up")

The authors built a new system called LDLM (Latent Diffusion Language Model). Instead of using a frozen blueprint, they created a three-person team that learns together from scratch:

The Encoder: A translator that turns words into the "marble" (latent space).
The Sculptor (Diffusion Model): The part that removes the noise to refine the story.
The Decoder: A translator that turns the refined marble back into readable words.

The magic is that all three learn together. As the Sculptor gets better at removing noise, it tells the Encoder, "Hey, I need the marble to be shaped this way to work better." The Encoder then adjusts its shape to help the Sculptor. It's a feedback loop where the whole team evolves to fit each other perfectly.

The "Recipe" for Success

The authors found that just putting these three together didn't work immediately; the team would argue and fail to learn. They had to follow a specific "recipe" to get them to cooperate:

The Warm-Up: At the very beginning, they let the Encoder and Decoder practice translating words without the Sculptor interfering. This gives the Encoder a chance to build a stable foundation before the Sculptor starts trying to shape it. It's like letting a musician tune their instrument before the band starts playing.
The "Noise" Trick: During training, they intentionally add a little bit of static (noise) to the Decoder's input. This forces the Encoder to be robust and store information efficiently, rather than relying on perfect, fragile signals. It's like training a runner with a heavy backpack so that when they run without it, they feel light and fast.
Smart Timing: They don't train the model at a constant speed. They adjust when they practice based on how hard the task is at that moment, ensuring the model learns the hardest parts just as well as the easy parts.

The Results: Faster and Smarter

When they tested this new team on large datasets of text (like news articles and books):

Quality: The stories generated were more coherent and diverse than previous methods.
Speed: Because the model works on the "marble" (the hidden math) rather than trying to pick individual words from a giant dictionary at every single step, it is 2 to 13 times faster than the competition.
Efficiency: It achieves a better balance between writing high-quality text and keeping the text varied (not repetitive).

Summary

In short, this paper introduces a new way to teach AI to write by having the "translator," the "sculptor," and the "refiner" learn together as a single unit, rather than using a pre-made, frozen tool. By following a specific training recipe, they created a system that writes better text, much faster than before.

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

Problem Statement

Autoregressive models currently dominate text generation but are constrained by their left-to-right sequential nature, which prevents parallel generation and correction of previous mistakes. Diffusion language models offer an alternative by generating text through iterative refinement in parallel. However, existing approaches face significant hurdles:

Discrete Diffusion: Operates directly on token space but suffers from factorized joint distributions, predicting tokens independently at each step. This limits the ability to generate multiple tokens simultaneously, making high-quality few-step generation difficult.
Continuous Diffusion: Operates in a latent space, allowing gradual refinement of all positions. While superior to discrete methods, current continuous approaches rely on frozen pre-trained encoders (e.g., Transformer-based encoders) to construct the latent space. These encoders are trained on surrogate tasks (like next-token prediction) and are not optimized for the specific needs of the diffusion denoising process. Conversely, shallow token embeddings (one-hot or fixed vectors) yield suboptimal latent spaces even when tuned for diffusion.

The core challenge is constructing a latent space that is simultaneously decodable (can be mapped back to tokens) and diffusion-friendly (easy to denoise), without the mismatch caused by separating the encoder training from the diffusion training.

Methodology: LDLM

The authors propose the Latent Diffusion Language Model (LDLM), a framework where the latent encoder, diffusion model, and decoder are trained jointly. This allows the diffusion objective to directly shape the latent space.

Architecture

Two-Stage Encoding:
- Frozen Token Encoder: Input tokens $w$ are first processed by a frozen pre-trained token encoder (GPT-2) to extract contextual hidden states $h$ .
- Trainable Latent Encoder: A trainable encoder $E_\theta^z$ maps these hidden states $h$ into the diffusion latent space $z_0$ . This reshapes the representation to be suitable for diffusion while maintaining decodability.
Two-Stage Decoding:
- Latent Decoder: A trainable decoder $D_\theta^h$ maps the latent $z_0$ (perturbed by noise during training) back to hidden states $\hat{h}$ .
- Token Decoder: A token decoder $D_\theta^w$ converts $\hat{h}$ into tokens $w$ .
Training Objectives:
The model minimizes a combined loss: $L(\theta) = L_{diff} + L_h + L_w$ $L (θ) = L_{d i f f} + L_{h} + L_{w}$ .
- $L_{diff}$ : Standard continuous diffusion loss (MSE) to recover clean latents from noisy observations.
- $L_h$ (Hidden-State MSE): Reconstructs the frozen encoder's hidden states $h$ from the latent space. Crucially, this uses Mean Squared Error (MSE) rather than Cross-Entropy (CE).
- $L_w$ (Token CE): Reconstructs tokens from the hidden states. Gradients are stopped at the reconstructed hidden states to prevent affecting the latent encoder.

Key Training Recipe (The "Four-Part" Solution)

The authors identify that naive joint training fails (producing low-diversity generations) and propose a specific recipe to make it work:

MSE Decoder Loss: Unlike prior work using Cross-Entropy, LDLM uses MSE to reconstruct hidden states. CE forces latents to be well-separated to distinguish tokens, which conflicts with the smooth geometry required for diffusion. MSE is less strict, allowing latents to be close if their underlying hidden states are similar, preventing latent collapse while maintaining a smooth manifold.
Diffusion-to-Encoder Warmup: At the start of training, the reconstruction loss (encouraging large latent magnitudes) and diffusion loss (encouraging small magnitudes to trivialize denoising) pull the encoder in opposite directions. The authors implement a warmup where the diffusion gradient is scaled by a coefficient $\gamma(s)$ that starts near 0 and increases to 1. This allows the encoder to first learn a decodable latent space before the diffusion objective shapes it.
Adaptive Timestep Sampling: Since the latent space evolves during training, a fixed noise schedule is suboptimal. The authors dynamically adjust the timestep sampling distribution to ensure the denoising loss grows linearly with the timestep, making all timesteps contribute equally to uncertainty reduction.
Decoder-Input Noise: Gaussian noise ( $\sigma_{dec}\epsilon$ $σ_{d ec} ϵ$ ) is injected into the latent decoder input during training. This serves three purposes:
- Encourages the encoder to store information efficiently (preventing unused dimensions from corrupting the signal).
- Acts as data augmentation, making the decoder robust to diffusion errors.
- Stabilizes the variance of latent coordinates, improving diffusion parameterization.

Results

The authors evaluate LDLM on OpenWebText (OWT) and LM1B, comparing it against state-of-the-art discrete (MDLM, Duo) and continuous (CANDI, FLM) diffusion baselines.

Quality-Diversity Trade-off: LDLM achieves the best Pareto frontier between Generative Perplexity (Gen. PPL) and Entropy. It significantly outperforms baselines in Gen. PPL at comparable entropy levels, approaching the quality of real text.
Latent Space Comparison: Ablation studies show that LDLM's jointly learned latents outperform:
- Shallow embeddings (trained from scratch or frozen).
- Contextual embeddings from frozen encoders (TEncDM style).
- Latents from autoencoders trained with explicit smoothing objectives (COSMOS style).
  This confirms that the latent space must be tuned specifically for the diffusion objective.
Computational Efficiency: LDLM is 2–13× faster than competing methods during sampling. Unlike token-level diffusion models that project to the vocabulary at every step, LDLM operates in a lower-dimensional latent space throughout the generation process, only projecting to tokens once at the end.

Significance and Claims

The paper claims that jointly learning the latent space is a critical, previously unaddressed step toward making latent diffusion competitive for text generation.

Overcoming Naive Joint Training: The work demonstrates that simply training an encoder and diffusion model together is insufficient; specific techniques (MSE loss, warmup, adaptive sampling, decoder noise) are required to stabilize the training and achieve high quality.
Superiority of Contextual Latents: The results indicate that high-level sequence information (captured by the frozen GPT-2 encoder) is essential, but it must be reshaped by a trainable encoder specifically for diffusion.
Efficiency: By avoiding per-step vocabulary projections, LDLM offers a substantial speedup, addressing a major bottleneck in diffusion-based text generation.

The authors conclude that shaping the latent space directly with the denoising objective is essential for building effective latent text diffusion models, achieving state-of-the-art generation quality while maintaining superior inference speed.