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

2026-05-08

AI Summary

This work proposes the Latent Diffusion Language Model (LDLM) to support efficient parallel denoising by constructing a high-quality continuous latent space suitable for non-autoregressive text generation. LDLM is the first framework to enable end-to-end joint training of the latent encoder, diffusion model, and decoder, leveraging representations from a pretrained language model to reshape the latent space. The approach incorporates several key training strategies, including MSE-based decoding loss, diffusion-encoder warmup, adaptive timesteps sampling, and noised decoder inputs. Experimental results demonstrate that LDLM achieves superior text generation quality compared to existing discrete and continuous diffusion language models on OpenWebText and LM1B benchmarks, while offering 2× to 13× faster inference speeds.

Abstract

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.

Problem

• latent diffusion
• text generation
• latent space
• non-autoregressive
• language model

Innovation

• latent diffusion
• joint training
• non-autoregressive generation
• continuous text representation
• adaptive timestep sampling