Cosmos: Compressed and Smooth Latent Space for Text Diffusion Modeling
Viacheslav Meshchaninov, Egor Chimbulatov, Alexander Shabalin, Aleksandr Abramov, Dmitry P. Vetrov
Abstract
Autoregressive language models dominate modern text generation, yet their sequential nature introduces fundamental limitations: decoding is slow, and maintaining global coherence remains challenging. Diffusion models offer a promising alternative by enabling parallel generation and flexible control; however, their application to text generation is hindered by the high dimensionality of token-level representations. We introduce Cosmos, a novel approach to text generation that operates entirely in a compressed, smooth latent space tailored specifically for diffusion. This space is learned using an autoencoder trained simultaneously for token-level reconstruction and alignment with frozen activations from a pretrained language encoder, providing robust semantic grounding and enabling effective perturbation-based augmentations. Empirically, we demonstrate that text representations can be compressed by while maintaining generation quality comparable to token-level diffusion models. Furthermore, increasing the latent sequence length allows Cosmos to surpass both diffusion-based and autoregressive baselines. We evaluate Cosmos on four diverse generative tasks including story generation, question generation, summarization, and detoxification and compare it with various generative paradigms. Cosmos achieves comparable or superior generation quality while offering more than faster inference. Code is released at GitHub
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext c052ddff-9813-4719-a77f-15446cd688beCited by top-tier papers3
- Guided Star-Shaped Masked DiffusionViacheslav Meshchaninov, Egor Shibaev, Artem Makoian, Ivan Klimov et al.ICML 2026
- One-step Optimal Transport via Regularized Distribution Matching DistillationDenis Rakitin, Ivan Shchekotov, Viacheslav Meshchaninov, Dmitry VetrovICML 2026
- Esoteric Language Models: A Family of Any-Order Diffusion LLMsSubham Sekhar Sahoo, Zhihan Yang, Yash Akhauri, Johnna Liu et al.ICML 2026
Builds on24
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- Denoising Diffusion Probabilistic ModelsJonathan Ho, Ajay Jain, Pieter AbbeelNeurIPS 2020 · 35,902 citations
- High-Resolution Image Synthesis with Latent Diffusion ModelsRobin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser et al.CVPR 2022 · 13,123 citations
- BERTScore: Evaluating Text Generation with BERTTianyi Zhang, Varsha Kishore, Felix Wu, Kilian Q. Weinberger et al.ICLR 2020 · 8,443 citations
- Flamingo: a Visual Language Model for Few-Shot LearningJean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech et al.NeurIPS 2022 · 6,707 citations
Related papers
- Latent Diffusion for Language GenerationJustin Lovelace, Varsha Kishore, Chao Wan, Eliot Shekhtman et al.NeurIPS 2023 · 177 citations
- TEncDM: Understanding the Properties of the Diffusion Model in the Space of Language Model EncodingsAlexander Shabalin, Viacheslav Meshchaninov, Egor Chimbulatov, Vladislav Lapikov et al.AAAI 2025 · 11 citations
- AR-Diffusion: Auto-Regressive Diffusion Model for Text GenerationTong Wu, Zhihao Fan, Xiao Liu, Hai-Tao Zheng et al.NeurIPS 2023 · 170 citations
- Beyond Autoregression: Fast LLMs via Self-Distillation Through TimeJustin Deschenaux, Caglar GulcehreICLR 2025
- PLANNER: Generating Diversified Paragraph via Latent Language Diffusion ModelYizhe Zhang, Jiatao Gu, Zhuofeng Wu, Shuangfei Zhai et al.NeurIPS 2023 · 65 citations
