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

TL;DR

This paper introduces LDLM, a joint training approach for latent diffusion language models that improves text generation quality and speed by learning a suitable latent space with a novel training recipe.

Contribution

The paper presents a method for jointly training the latent space, encoder, and diffusion model, significantly enhancing performance and efficiency over existing models.

Findings

LDLM outperforms existing diffusion models on OpenWebText and LM1B datasets.

LDLM is 2 to 13 times faster than comparable models.

A combination of training techniques is crucial for high-quality text generation.

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.