TextLDM: Language Modeling with Continuous Latent Diffusion

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

Source-linked AI summary

TextLDM: Language Modeling with Continuous Latent Diffusion

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

arXiv:2605.07748v1May 8, 2026cs.CL

TL;DR

Text generation lacks a demonstrated transfer of the visual latent-diffusion recipe, particularly because reconstruction quality does not ensure useful continuous text representations. TextLDM combines a Transformer VAE, REPA, and a standard DiT with flow matching, outperforming prior diffusion language models and matching GPT-2 under the same settings.

Problem
Continuous text representations remain a bottleneck because near-perfect reconstruction does not ensure representations suitable for conditional diffusion denoising.
Method
TextLDM maps tokens to continuous latents with a Transformer VAE, aligns them with a frozen language model using REPA, and applies a standard DiT with flow matching.
Results
TextLDM substantially outperforms prior diffusion language models and matches GPT-2 and autoregressive baselines under identical settings.
Takeaways & Limitations
The visual diffusion recipe transfers effectively to language modeling with minimal architectural modification, supporting a path toward unified multimodal diffusion architectures.
Takeaways & Limitations
TextLDM requires a two-stage VAE-then-DiT training pipeline, and lower out-of-domain TextVAE reconstruction accuracy may limit generation quality.

Abstract

from arXiv · show

Diffusion Transformers (DiT) trained with flow matching in a VAE latent space have unified visual generation across images and videos. A natural next step toward a single architecture for both generation (visual synthesis) and understanding (text generation) is to apply this framework to language modeling. We propose TextLDM, which transfers the visual latent diffusion recipe to text generation with minimal architectural modification. A Transformer-based VAE maps discrete tokens to continuous latents, enhanced by Representation Alignment (REPA) with a frozen pretrained language model to produce representations effective for conditional denoising. A standard DiT then performs flow matching in this latent space, identical in architecture to its visual counterpart. The central challenge we address is obtaining high-quality continuous text representations: we find that reconstruction fidelity alone is insufficient, and that aligning latent features with a pretrained language model via REPA is critical for downstream generation quality. Trained from scratch on OpenWebText2, TextLDM substantially outperforms prior diffusion language models and matches GPT-2 under the same settings. Our results establish that the visual DiT recipe transfers effectively to language, taking a concrete step toward unified diffusion architectures for multimodal generation and understanding.

1 Introduction

TextLDM transfers the visual latent-diffusion recipe to language modeling through a Transformer VAE and DiT with minimal architectural modification. The paper identifies latent representation effectiveness as the key bottleneck, addresses it with REPA, and reports strong text-continuation performance from scratch.

Method: TextLDM maps discrete tokens into continuous latents with a Transformer-based VAE and applies flow matching using a standard DiT.The TextDiT is architecturally identical to its visual counterpart.
Challenge: Representation effectiveness, rather than reconstruction fidelity alone, is the central bottleneck for conditional denoising in latent text diffusion.VAEs with virtually identical reconstruction accuracy can produce substantially different generation quality.
Representation Alignment: REPA aligns TextVAE encoder features with a frozen Qwen3-1.7B language model, improving downstream generation quality without affecting reconstruction.The alignment shapes latent geometry to better support diffusion-based generation.
Results: TextLDM is trained from scratch on OpenWebText2 and substantially outperforms prior continuous and discrete diffusion language models across four text-continuation benchmarks.It matches GPT-2 baselines under the same settings.
Ablations: Ablations show that visual diffusion components, including logit-normal scheduling and classifier-free guidance, transfer effectively to language modeling.The transferred recipe includes VAE, DiT, flow matching, and CFG, without pretrained encoders or decoders.

2 Related Work

Prior work established flow matching, VAE, and DiT as a standard visual-generation recipe, while diffusion language models explored continuous and discrete formulations. TextLDM builds on these directions with a from-scratch, non-autoregressive text VAE enhanced by language-model representation alignment.

Diffusion Models for Visual Generation: Flow matching, VAE, and DiT have become a standard recipe for scalable visual generation in latent spaces.This line of work unified diffusion models with flow matching and extended them to latent representations.
Diffusion Language Models: Diffusion language models comprise continuous approaches operating in embedding or simplex spaces and discrete approaches.Latent-diffusion variants such as LD4LG and COSMOS use pretrained encoders or compressed latent spaces.
Variational Autoencoders for Text: TextVAE differs from prior text VAEs by training from scratch with a non-autoregressive decoder and adapting REPA to align its encoder with a frozen language model.Prior text VAEs typically rely on pretrained components or autoregressive decoders.

3 Method

TextLDM uses a two-stage architecture that compresses each text token into a continuous latent and models generation in that space with a Flow Matching Diffusion Transformer. Its TextVAE combines reconstruction, KL regularization, and representation alignment with a frozen Qwen3-1.7B, while conditional diffusion generates target segments in parallel.

Overall framework: TextLDM consists of a Transformer TextVAE and a Flow Matching Diffusion Transformer operating in the learned latent space.The TextVAE encodes tokens into continuous representations, after which the DiT models generative dynamics.
TextVAE: Each token maps one-to-one to a latent vector, preserving the original sequence length rather than compressing tokens into a shorter latent sequence.The encoder produces one latent vector z_i ∈ R^d for every token x_i.
TextVAE: The Transformer decoder reconstructs all tokens in parallel from the latent sequence, while random truncation exposes it to varying portions and lengths.The decoder predicts a vocabulary distribution for each position non-autoregressively.
Representation Alignment (REPA): REPA aligns encoder intermediate representations with frozen Qwen3-1.7B hidden states using cosine similarity, with the 3rd-to-last layer outperforming the last layer.A linear projection matches dimensions when needed, and stop-gradient is applied to the language-model representation.
Training Objective: The TextVAE objective combines cross-entropy reconstruction, KL regularization, and REPA alignment, using β = 10^-3 and λ = 1.The KL term regularizes the posterior toward a standard Gaussian prior.
TextDiT: The DiT concatenates clean context latents with noisy target latents and predicts a velocity field under Conditional Flow Matching.Noisy targets are formed by linear interpolation between Gaussian noise and target latents, with timesteps sampled from a logit-normal distribution.
Inference: The model uses classifier-free guidance with p_uncond = 0.1, and generates the entire target segment in parallel instead of decoding token by token.Unconditional training replaces context latents with zero vectors; inference iteratively updates noisy latents before decoding them.

4 Experiments

TextLDM substantially outperforms prior diffusion language models and is competitive with or superior to GPT-2 across four text-continuation benchmarks. Performance improves consistently with DiT scaling, while REPA and selected latent-space configurations are more important than simply increasing VAE capacity.

Main results: TextLDM substantially outperforms SSD-LM and Block Diffusion across all reported metrics and benchmarks, with the 768M model achieving the best results on most metrics.The gains cover ROUGE, BERTScore, and MAUVE, including comparisons at comparable model sizes.
Main results: TextLDM matches or exceeds similarly sized GPT-2 models on TinyStories and One Billion Words, while its 768M variant outperforms all GPT-2 models on Wikipedia and WikiSource.Autoregressive models retain an advantage on some BERTScore results.
Scaling: 32.7 versus 21.6: 768M TextLDM improves MAUVE on WikiSource over 114M, while ROUGE-1 rises from 33.0 to 37.5 there and from 10.3 to 21.4 on One Billion Words.On Wikipedia, ROUGE-1 increases from 27.5 to 38.9, demonstrating consistent scaling across metrics and datasets.
Training efficiency: TextLDM matches or exceeds GPT-2-medium on ROUGE and MAUVE for WikiSource, Wikipedia, and TinyStories within a comparable number of training steps.It lags slightly on One Billion Words, where short-sample frequency gives autoregressive training an advantage.
Ablations: REPA substantially improves every metric and dataset, especially Wikipedia and WikiSource, whereas increasing VAE capacity beyond 350M yields no consistent gains.The 350M VAE achieves the best ROUGE scores on most datasets, suggesting REPA matters more than raw VAE capacity.
Ablations: The strongest ablation configuration uses 64 latent channels, third-to-last-layer REPA, a logit-normal schedule with std=1.5, and CFG=7.These choices outperform alternatives on the reported metrics; reconstruction remains ≥99.6% on TinyStories and One Billion Words and ≥97.5% on Wikipedia and WikiSource.

5 Limitation

TextLDM’s two-stage VAE-then-DiT training adds complexity, while lower TextVAE reconstruction accuracy on out-of-domain data may propagate errors and limit generation quality.

Training complexity: The VAE-then-DiT training pipeline is more complex than end-to-end autoregressive training.The two stages are trained separately.
Out-of-domain reconstruction: ∼97.5% reconstruction accuracy on Wikipedia and WikiSource may propagate errors and limit DiT generation quality on those domains.The TextVAE performs worse on out-of-domain samples; more diverse training data is expected to mitigate this drop.

6 Conclusion

TextLDM is a continuous-latent diffusion framework for language modeling that combines a REPA-trained TextVAE with a standard Flow Matching Diffusion Transformer. It achieves state-of-the-art results among diffusion language models while matching autoregressive baselines, supporting the transfer of the visual-generation recipe to language.

Framework: TextLDM operates entirely in a continuous latent space for language modeling.The framework is presented as a latent diffusion approach rather than direct discrete-token diffusion.
Architecture: TextLDM trains a TextVAE with Representation Alignment (REPA) and a standard Diffusion Transformer with Flow Matching.These components constitute the proposed latent diffusion framework.
Results: TextLDM achieves state-of-the-art results among diffusion language models while matching autoregressive baselines.The conclusion identifies performance against both diffusion-language-model peers and autoregressive baselines.
Key finding: The paper identifies the exact recipe proven in visual generation as a key finding for language modeling.The passage explicitly frames this visual-generation recipe as central to TextLDM’s conclusion.

A Implementation Details

TextLDM uses Transformer-based TextVAE and a standard DiT with RoPE positional encoding. Training uses AdamW-based optimization, specified step counts, KL warmup, classifier-free guidance dropout, and multi-GPU execution.

A Implementation Details: The TextVAE encoder and decoder are standard pre-norm Transformers with LayerNorm and RoPE positional encoding.The encoder outputs per-position mean and log-variance vectors, while the decoder maps sampled latents to vocabulary logits.
A Implementation Details: The DiT follows the standard DiT architecture, concatenates clean-context and noisy-target latents, and uses RoPE without timestep embeddings.The absence of timestep embeddings is consistent with LLaDA and RADD.
A Implementation Details: TextVAE training runs for 200K steps with AdamW at learning rate 1e-4 and weight decay 0.01, while DiT training runs for 1M ablation steps or 2M main-result steps.The KL weight β follows a warmup schedule.
A Implementation Details: Classifier-free guidance training uses an unconditional dropout rate of p_uncond = 0.1 on 8× NVIDIA H200 GPUs.Experiments use approximately 100K tokens per GPU per mini-batch.

B Broader Impacts

TextLDM advances diffusion-based language modeling toward a shared framework for language and vision generation, while retaining the potential risks associated with text-generating systems.

Potential benefits: Unifying language and vision generation under a shared diffusion framework could simplify multimodal model development.The passage also suggests this could lower the barrier to building controllable generation systems.
Potential risks: TextLDM could potentially be used to generate misleading or harmful text, like all language models.The passage frames this as a societal implication common to text generation systems.

C Future Work

Future work focuses on testing TextLDM’s scaling behavior at larger sizes and corpora, while developing fairer evaluation methods for language understanding.

Scaling laws: Future work should test whether TextLDM’s favorable scaling trends continue with substantially larger models and training corpora.This is identified as an important step toward practical diffusion language models.
Fair evaluation on language understanding benchmarks: Future work should establish fair evaluation on language-understanding benchmarks whose likelihood-based scoring may disadvantage diffusion language models.MMLU compares per-token log-probabilities of candidate answers, creating an evaluation challenge for current methods.

D Qualitative Examples: Step-by-Step Denoising

TextLDM’s qualitative Wikipedia continuations become progressively more coherent and factually accurate as diffusion denoising proceeds from steps 10 through 50.

Progressive Denoising: At diffusion steps 10, 20, 30, 40, and 50, TextLDM generates progressively refined continuations from a conditioning prefix.The examples visualize the continuation at each denoising stage.
Progressive Denoising: The qualitative examples show coherence emerging progressively as the number of denoising steps increases.The denoising process is illustrated on Wikipedia text continuations.
Progressive Denoising: Factual accuracy also improves progressively across the displayed denoising steps.Each example is conditioned on a prefix labeled “Cond.”