💬 NLP

TextLDM: Language Modeling with Continuous Latent Diffusion

The paper introduces TextLDM, a language modeling framework that adapts the visual Diffusion Transformer (DiT) architecture to text generation by mapping discrete tokens to continuous latents via a VAE enhanced with Representation Alignment (REPA), achieving performance comparable to GPT-2 and advancing the goal of unified multimodal diffusion models.

Original authors: Jiaxiu Jiang, Jingjing Ren, Wenbo Li, Bo Wang, Haoze Sun, Yijun Yang, Jianhui Liu, Yanbing Zhang, Shenghe Zheng, Yuan Zhang, Haoyang Huang, Nan Duan, Wangmeng Zuo

Published 2026-05-11

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jiaxiu Jiang, Jingjing Ren, Wenbo Li, Bo Wang, Haoze Sun, Yijun Yang, Jianhui Liu, Yanbing Zhang, Shenghe Zheng, Yuan Zhang, Haoyang Huang, Nan Duan, Wangmeng Zuo

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 how to write a story. For a long time, the standard way to do this has been like a typewriter: the robot types one letter, then the next, then the next, in a strict line. This is called "Autoregressive" (AR) modeling. It works well, but it's slow because it can't write the whole sentence at once; it has to build it brick by brick.

On the other hand, when teaching robots to draw pictures, scientists recently discovered a better way. Instead of drawing line by line, they start with a canvas full of static noise (like TV snow) and slowly "denoise" it until a clear image appears. This is called Diffusion.

This paper, TextLDM, asks a simple question: Can we teach the robot to write stories using the same "denoising" method we use for drawing pictures?

Here is how they did it, explained through analogies:

1. The Problem: Words are Discrete, Noise is Continuous

The main hurdle is that words are like Lego bricks. You can't have half a brick; it's either there or it isn't. But the "denoising" method for pictures works with smooth, continuous colors (like mixing paint).

If you try to turn Lego bricks directly into smooth paint, the result is usually a blurry mess. Previous attempts to do this for text failed because the "smooth versions" of words they created were too messy for the robot to understand how to turn them back into good sentences.

2. The Solution: The "Translator" (TextVAE)

The authors built a special Translator (called a TextVAE).

The Job: It takes a sentence made of Lego bricks (discrete words) and translates it into a smooth, continuous liquid (latent vectors).
The Trick: They found that just translating the words wasn't enough. The liquid needed to be "smart." So, they added a Mentor (a frozen, pre-trained language model called Qwen3).
The Alignment (REPA): Imagine the Translator is a student and the Mentor is a master teacher. The student tries to translate the words, but the teacher constantly checks: "Is this smooth liquid representation capturing the true meaning of the word, just like I would?" This process, called Representation Alignment (REPA), forces the Translator to create a "liquid language" that is perfectly structured for the denoising process.

3. The Generator: The "Denoising Sculptor" (TextDiT)

Once the Translator is ready, the actual writing happens.

Instead of typing word-by-word, the robot starts with a bucket of random noise (like static on a TV).
It uses a Sculptor (a Diffusion Transformer, or DiT) to slowly chip away the noise, guided by the "liquid language" the Translator made.
The Magic: Because it's working in this smooth, continuous space, the robot can generate the entire story at once, rather than one word at a time.

4. Why This Matters (The Results)

The paper tested this new method on four different writing challenges (from simple children's stories to complex encyclopedia entries).

Speed: Because it generates the whole story in parallel (like painting a whole canvas at once) rather than sequentially (like typing), it can be much more efficient for long texts.
Quality: The new method (TextLDM) beat all previous "diffusion" text models. In fact, it performed just as well as the best "typewriter" models (GPT-2) of similar size.
The Key Insight: The paper proves that the exact "recipe" used to make amazing AI images (Smooth Translation + Smart Mentor + Denoising Sculptor) works perfectly for text too, provided you fix the "Translator" with the right alignment.

Summary

Think of it this way: Before, writing with AI was like molding clay one tiny piece at a time. This paper shows that if you first turn the clay into a perfectly smooth, intelligent liquid, you can pour the whole shape at once and get a result that is just as good, but potentially faster and more flexible. They didn't invent a new clay; they just found a better way to smooth it out before shaping it.

Technical Summary: TextLDM

Problem Statement

The field of multimodal modeling seeks a unified framework capable of seamlessly generating both visual and textual content. While Diffusion Transformers (DiT) trained with flow matching in a VAE latent space have successfully unified image and video generation, language modeling has remained dominated by the autoregressive (AR) paradigm. Existing attempts to unify understanding and generation often rely on discrete token-based paradigms or discrete diffusion, leaving a methodological gap. The central challenge in applying continuous latent diffusion to text is not the diffusion backbone itself, but the quality of the continuous text representations. A VAE trained solely for token reconstruction can achieve high fidelity yet produce latent spaces poorly suited for conditional denoising, leading to suboptimal generation quality.

Methodology: TextLDM

The authors propose TextLDM, a framework that transfers the visual latent diffusion recipe to language modeling with minimal architectural modification. The system operates in a two-stage process:

1. TextVAE: Continuous Latent Representation

A Transformer-based Variational Autoencoder (TextVAE) maps discrete tokens to continuous latent vectors.

Architecture: Unlike prior approaches that compress token sequences into shorter latent sequences, TextVAE maintains a one-to-one mapping where each token corresponds to a continuous latent vector. The encoder ( $E_\phi$ ) and decoder ( $D_\psi$ ) are both Transformers.
Representation Alignment (REPA): To address the bottleneck of latent effectiveness, the authors introduce REPA. The VAE encoder is aligned with a frozen pretrained language model (Qwen3-1.7B) via a cosine similarity loss on intermediate representations. This shapes the latent geometry to be more amenable to diffusion-based generation without compromising reconstruction fidelity.
Training Objective: The VAE is trained with a composite loss comprising cross-entropy reconstruction loss, KL divergence regularization, and the REPA alignment loss.

2. TextDiT: Flow Matching in Latent Space

A standard Diffusion Transformer (TextDiT), architecturally identical to its visual counterpart, performs flow matching within the learned latent space.

Conditional Formulation: The model treats language generation as a conditional process, taking clean context latents ( $z_c$ ) and noisy target latents ( $z_t$ ) as input to predict a velocity field.
Flow Matching: The model is optimized using Conditional Flow Matching (CFM) loss, interpolating between Gaussian noise and the target latent.
Scheduling and Guidance: The system adopts the logit-normal timestep schedule (proven effective in Stable Diffusion 3) and utilizes Classifier-Free Guidance (CFG) to improve generation quality.
Inference: Generation is performed in parallel via an Euler ODE solver, avoiding token-level autoregressive decoding. This allows for length-invariant inference efficiency (constant Number of Function Evaluations regardless of sequence length).

Key Contributions

Unified Architecture: TextLDM successfully transfers the visual diffusion recipe (VAE + DiT + flow matching + CFG) to language modeling, taking a step toward unified multimodal architectures. The entire system is trained from scratch without relying on pretrained encoders or decoders for the diffusion process.
Identification of Representation Bottleneck: The authors identify that reconstruction fidelity alone is insufficient for high-quality text diffusion. They demonstrate that Representation Alignment (REPA) is critical for shaping the latent space to support downstream generation quality.
State-of-the-Art Performance: TextLDM achieves state-of-the-art results among diffusion language models and matches autoregressive baselines (GPT-2) under identical settings.
Efficiency: The framework offers a more constant-time generation profile compared to the linear growth of function evaluations in autoregressive models.

Experimental Results

The models were trained from scratch on OpenWebText2 and evaluated on text continuation across four benchmarks: One Billion Words, TinyStories, Wikipedia, and WikiSource.

Performance vs. Diffusion Baselines: TextLDM substantially outperforms prior continuous (e.g., SSD-LM) and discrete (e.g., Block Diffusion) diffusion language models across all metrics (ROUGE, BERTScore, MAUVE).
Performance vs. Autoregressive Baselines:
- On TinyStories and One Billion Words, TextLDM matches or exceeds GPT-2 models of similar or larger size on ROUGE metrics.
- On out-of-distribution benchmarks (Wikipedia, WikiSource), the 768M TextLDM variant outperforms all GPT-2 baselines.
- A minor gap remains on BERTScore for some benchmarks, where AR models retain an advantage.
Scaling Behavior: TextLDM demonstrates consistent scaling behavior; performance improves significantly as the DiT size increases from 114M to 768M parameters.
Training Efficiency: Training dynamics show that TextLDM matches or exceeds GPT-2-medium performance within a comparable number of training steps on most benchmarks, a significant improvement over prior continuous diffusion models that required substantially more compute.
Ablation Insights:
- REPA: Removing REPA causes a substantial drop in generation quality, confirming its necessity for latent space structure.
- Reconstruction vs. Generation: All VAE configurations achieved near-perfect token reconstruction (>97.5%), yet generation quality varied significantly based on REPA and latent dimension, proving that latent geometry, not reconstruction, drives performance.
- Hyperparameters: The logit-normal timestep schedule (std=1.5) and a CFG scale of 7 yielded optimal results.

Significance and Claims

The paper claims that TextLDM establishes that the "visual DiT recipe" transfers effectively to language. By proving that a unified architecture can handle both visual synthesis and text generation, the work provides a concrete foundation for future unified diffusion architectures for multimodal generation and understanding. The authors emphasize that the key to success was not just the diffusion backbone, but the specific adaptation of the latent space via Representation Alignment, which allows discrete text to be effectively modeled in a continuous diffusion framework.

The authors remain modest regarding limitations, noting that the two-stage training pipeline adds complexity and that reconstruction accuracy on out-of-domain data (e.g., Wikipedia) is slightly lower than in-domain data, which may propagate errors. They suggest that expanding the training corpus with more diverse data could mitigate these issues.