Cosmos: Compressed and Smooth Latent Space for Text Diffusion Modeling

TL;DR

Cosmos introduces a novel compressed, smooth latent space for text diffusion, enabling faster, high-quality text generation across multiple tasks by combining autoencoding and semantic alignment.

Contribution

It presents a new latent space for diffusion models that reduces dimensionality and improves efficiency, with a training method that aligns autoencoder representations with pretrained language models.

Findings

Compression of text representations by 8x without quality loss

Cosmos surpasses baselines with longer latent sequences

Achieves over 2x faster inference while maintaining quality

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 $8\times$ 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 $2\times$ faster inference. Code is released at \href{https://github.com/MeshchaninovViacheslav/cosmos}{GitHub}