Research map / Image generation / Autoregressive and latent generation
Compressed and Smooth Latent Space for Text Diffusion Modeling
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Cosmos, a frozen-BERT-based autoencoder that compresses token representations into a compact, smooth latent space for diffusion text generation, reaching up to 8x compression with competitive or better quality and 2x faster inference.
Problem
Token-level representations are high-dimensional, making diffusion-based text generation slow and hard to train; latent diffusion for text needs a latent space that is both compact and diffusable.
Approach
Train a Perceiver-resampler compressor/decompressor over frozen BERT hidden states with CE and MSE-on-encoder-activation losses plus activation-space perturbations (masking, Gaussian noise) and latent feature dropout, so the latent manifold is smooth and robust enough for Gaussian diffusion.
Key results
- On ROCStories, Cosmos N=128 reaches MAUVE 0.940 versus 0.789 for GPT-2 and 0.762 for TEncDM.
- Cosmos N=128 on XSum achieves BERTScore 0.704, versus 0.703 for TEncDM, with inference time 35.1s vs 29.6s at L=128 and 36.6s vs 180s at L=512.
- On OpenWebText 512-token generation, Cosmos N=512 gets MAUVE 0.702 and PPL 55.0 vs GIDD's 0.286 and 228.3.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add MSE loss on encoder activations (ablation) | Cosmos autoencoder · ROCStories | MAUVE: 0.294_{.036} -> 0.464_{.024} | Helped |
| Add random masking of encoder activations at rate 0.3 (ablation) | Cosmos autoencoder with MSE activation loss · ROCStories | MAUVE: 0.464_{.024} -> 0.586_{.015} | Helped |
| Add Gaussian noise augmentation with δ=0.7 (ablation) | Cosmos autoencoder with activation MSE and random masking · ROCStories | MAUVE: 0.586_{.015} -> 0.767_{.011} | Helped |
| Add latent feature dropout at rate 0.4 (ablation) | Cosmos autoencoder with activation MSE and activation augmentation · ROCStories | MAUVE: 0.767_{.011} -> 0.836_{.009} | Helped |
| Increase latent sequence length from 16 to 128 (ablation) | Cosmos · ROCStories | MAUVE: 0.836_{.009} -> 0.940_{.011} | Helped |
| Compress from 512 to 256 latent vectors (ablation) | Cosmos · OpenWebText, 512-token generation | MAUVE: 0.702 -> 0.263 | Hurt |
| Increase diffusion model depth from 12 to 48 layers (ablation) | Cosmos · OpenWebText, 128-token generation | MAUVE: 0.849 -> 0.923 | Helped |
| Add a variational KL prior with β=1 (ablation) | Cosmos autoencoder · ROCStories | MAUVE: 0.767 -> 0.011 | Hurt |
| Add a variational KL prior with β=0.01 (ablation) | Cosmos autoencoder · ROCStories | PPL: 33.6 -> 32.6 | Helped |
| Use Cosmos N=128 instead of TEncDM | ROCStories · ROCStories | MAUVE: 0.762 -> 0.940 | Helped |
| Use Cosmos N=128 instead of GPT-2 | ROCStories · ROCStories | MAUVE: 0.789 -> 0.940 | Helped |
| Use Cosmos N=128 instead of GPT-2 | Cosmos · XSum | BERTScore: 0.690 -> 0.704 | Helped |
Surprising findings
- Compressing BERT representations to 16 latents with the proposed autoencoder matched or exceeded uncompressed BERT-representation diffusion quality, and N=128 exceeded the uncompressed source-level MAUVE on ROCStories (0.940 vs 0.953 source is close but the scaling trend is non-monotonic in compress
- A variational KL prior gave no clear benefit and collapsed at beta=1.
Limitations
- Autoencoder and diffusion trained separately; joint training left for future work.
- Latent feature dimension d fixed at 768; not systematically swept.
- Relatively small backbones (~130M parameters).
Experimental setup
Datasets: ROCStories Wikipedia (ROOTS subset) XSum SQuAD2.0 ParaDetox OpenWebText
Benchmarks: ROCStories XSum ParaDetox SQuAD2.0 OpenWebText
Metrics: MAUVE perplexity (GPT-2 Large) diversity memorization BERTScore BLEU-4 inference time
Baselines: GPT-2 GPT-Neo AR-Diffusion DiffuSeq SeqDiffuSeq TESS SEDD LD4LG TEncDM GIDD
Builds on: TEncDM latent diffusion PLANNER Perceiver Resampler
Closest papers by meaning
- Fast and Fluent Diffusion Language Models via Convolutional Decoding and Rejective Fine-tuning (NeurIPS 2025 · similarity 0.88)
- Energy-Based Diffusion Language Models for Text Generation (ICLR 2025 · similarity 0.88)
- Continuous Latent Diffusion Language Model (NeurIPS 2026 · similarity 0.88)
- How to Train Your Latent Diffusion Language Model Jointly With the Latent Space (NeurIPS 2026 · similarity 0.88)
- Scaling Diffusion Language Models via Adaptation from Autoregressive Models (ICLR 2025 · similarity 0.87)
- TEncDM: Understanding the Properties of the Diffusion Model in the Space of Language Model Encodings (AAAI 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.