Source-linked AI summary
Cosmos: Compressed and Smooth Latent Space for Text Diffusion Modeling
Viacheslav Meshchaninov, Egor Chimbulatov, Alexander Shabalin, Aleksandr Abramov, Dmitry Vetrov
TL;DR
Autoregressive generation is sequential and can struggle with global coherence, while diffusion for text must contend with high-dimensional representations. Cosmos trains an autoencoder to create a compact, smooth latent space for diffusion; across diverse tasks, it reports comparable or better generation quality with faster sampling than token-level diffusion.
Problem
Autoregressive generation is sequential and can struggle to maintain long-range coherence, while compressing text representations for diffusion remains an open question.
Source: §1 · pp. 1–2, §2 · p. 2Method
Results
Across diverse tasks, Cosmos matches or outperforms token-level diffusion and autoregressive baselines, with up to 2× faster sampling than conventional token-level diffusion.
Source: §1 · p. 2, §9 · p. 10Takeaways & Limitations
The results position latent-space diffusion as a promising alternative to token-level models for text generation.
Source: §1 · p. 2Experiments use relatively small backbones of roughly 130M parameters, and latent dimensionality is fixed across experiments.
Source: E Limitations · p. 18
Abstract
from arXiv · showhide
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}
1 Introduction
Autoregressive generation is sequential and can struggle with long-range coherence, while text diffusion faces the challenge of high-dimensional representations. Cosmos learns a compact, smooth latent space for diffusion and reports generation quality that rivals or sometimes surpasses token-level baselines.
1 Introduction
- Autoregressive decoding is sequential, and token-by-token decisions can hinder long-range logical or narrative coherence.The introduction also links local log-likelihood optimization to hallucination.Source: §1 · pp. 1–2
- Cosmos asks how far textual representations can be compressed while retaining or exceeding the generative fidelity of token-level representations.The paper frames this question against latent diffusion’s use of compact representations in vision.Source: §1 · p. 2
- Cosmos combines a compressed text autoencoder with diffusion in latent space, and reports that its training objectives produce a manifold conducive to high-quality synthesis.Naive token-reconstruction loss alone yields brittle latent geometry that hampers diffusion.Source: §1 · p. 2
- On standard text-generation benchmarks, Cosmos offers up to 2× faster sampling than token-level diffusion while matching or slightly exceeding its quality and diversity.Modestly increasing the number of latent vectors also lets it surpass token-level diffusion and autoregressive models on unconditional generation.Source: §1 · p. 2
2 Related work
Prior latent-diffusion approaches either rely on autoregressive decoders or retain full-length text representations. Cosmos targets the latter gap by compressing BERT-level representations while preserving their suitability for diffusion.
2 Related work
- Earlier systems LD4LG [19] and PLANNER [47] use autoregressive decoders, while TEncDM [36] applies diffusion directly to full-length BERT representations.The paper notes that reliance on powerful autoregressive decoders makes it difficult to isolate latent diffusion’s intrinsic generative capacity.Source: §2 · p. 2
- LD4LG [19] and PLANNER [47] use diffusion to refine compressed codes that are then decoded by autoregressive models.Both systems are described as improving controllability.Source: §2 · p. 2
- Cosmos shrinks BERT-level representations by up to 8× while preserving, and sometimes enhancing, their suitability for latent-space diffusion.This addresses the compression question left open by full-length TEncDM [36].Source: §2 · p. 2
3 Preliminary
Diffusion learns to reverse a progressive noising process, while latent diffusion first encodes data into a compact representation and performs diffusion there. The paper applies this framework to contextualized text representations.
3 Preliminary
- A forward process progressively corrupts a clean sample with Gaussian noise, and a learned denoiser is trained to invert that corruption.The noise schedule decreases from α_0 = 1 to approximately α_1 = 0; inference iteratively transforms pure noise into a sample.Source: §3 · p. 3
3.1 Latent diffusion
- Latent diffusion encodes data into a low-dimensional representation, models diffusion in that space, and decodes the resulting latent.The latent-space network can be shallower and faster than one operating in the original space while achieving comparable perceptual quality.Source: §3.1 · p. 3
- For text, the framework compresses contextualized representations and trains diffusion directly within the resulting latent space.The paper adopts this framework for textual data.Source: §3.1 · p. 3
4 Methodology
Cosmos encodes contextualized text into fixed-length latents, trains diffusion in that space, and decodes latents back into text. Its autoencoder combines token reconstruction with semantic alignment and perturbation-based training to produce compact representations suited to diffusion.
4.1 Overview
- A frozen pretrained text encoder supplies contextual token states to a fixed-length compressor, and diffusion operates on normalized latents before decompression and token prediction.The Perceiver Resampler compressor uses cross-attention to gather information across the input sequence.Source: §4.1 · pp. 3–4
4.2.1 How compact can text latents be?
- On WIKIPEDIA, 16 latents reconstruct 512-token sequences with 100% token-level accuracy, a 32× compression relative to the initial BERT hidden states.The compressor uses 12 transformer blocks; Figure 2 plots reconstruction accuracy against latent sequence length N.Source: §4.2.1 · p. 4, Figure 2 · p. 4
4.2 Learning a compact text latent space
- High token-reconstruction accuracy alone does not establish that the latent manifold is suitable for generative modeling.The paper motivates additional robustness and smoothness training because an unsuitable manifold can prevent Gaussian diffusion from sampling high-quality text.Source: §4.2.1 · p. 4, §4.2.2 · p. 4
4.2.2 Learning robust compact text representations
- Token-level cross-entropy trains reconstruction, while an MSE penalty between frozen encoder activations and their reconstructions preserves contextual semantics.The paper also uses perturb-and-recover training: augmented activations are reconstructed toward the original representations, with random masking and Gaussian noise; latent feature masking encourages robustness to feature removal.Source: §4.2.2 · p. 5
5 Latent-space properties that facilitate diffusion training
The section probes whether smooth, perturbation-resilient latent representations support diffusion beyond training examples. Interpolation and noise tests show smoother latents and a smaller train–inference gap when both decoder and diffusion robustness are addressed.
5.1 Smoothness of the latent manifold
- The proposed autoencoder's latent manifold is smoother than the cross-entropy baseline's: perplexity rises more slowly between training examples, despite similar endpoint reconstruction.The interpolation test averages GPT-2 perplexity over 1,000 random endpoint pairs.Source: §5.1 · pp. 5–6
5.2 Reducing the train–inference mismatch
- MSE loss on decompressor output and latent masking contribute most to decoder stability, enabling almost perfect reconstruction under noise with σ = 1.The tests add Gaussian noise to encoded real-text latents and score decoded outputs against the originals with BLEU.Source: §5.2 · p. 6, Figure 4 · p. 6
- The latent-space diffusion model is more stable when Gaussian noise is injected during a late reverse-process step.Stability is assessed by the MSE between the final estimates from the original and perturbed trajectories.Source: §5.2 · p. 6, Figure 5 · p. 6
- Each proposed refinement reduces the measured train–inference mismatch, indicating that the training procedure strengthens robustness to diffusion perturbations.The mismatch is measured as MSE between sampled latents and the encoder representation of their decoded text.Source: §5.2 · p. 7, Table 1 · p. 7
6 Empirical study of latent representations for diffusion modeling
The experiments assess how autoencoder training objectives and latent sequence length affect diffusion generation. Robustness-oriented training supports quality under compression, while the scaling results show a trade-off between compression and generation quality.
6.1 Empirical analysis of autoencoder training regimes
- Adding BERT-activation alignment raises MAUVE to 0.767, comparable to TEncDM's 0.762 on full-length BERT representations.The passage reports that the MSE penalty on predicted versus reference BERT activations alone boosts all quality metrics by about 50%.Source: §6.1 · p. 7, Table 2 · p. 7
- The reported result shows MAUVE beyond 0.83 and perplexity at 30.2, while diversity remains essentially unchanged for the described text-generation configuration.Source: §6.1 · p. 8
6.2 Impact of scaling the number of latent vectors
- The robustness-oriented autoencoder retains high generation quality at 8× shorter latent sequence length, whereas reducing N rapidly degrades the cross-entropy compressor.The comparison is shown in Figure 6.Source: §6.2 · p. 8, Figure 6 · p. 8
- At N = 16 versus N = 128, MAUVE falls from 0.940 to 0.836 and perplexity rises from 26.3 to 30.2; N = 32 gives only minor quality degradation and is recommended at about 4× compression.The section holds the embedding dimension fixed to isolate latent diffusability from changes in model capacity or structure.Source: §6.2 · p. 8, Table 3 · p. 8
7 Comparison across generative paradigms
COSMOS is compared with diffusion and autoregressive baselines on story generation, summarization, detoxification, and question generation. It matches or slightly exceeds baselines on longer-context tasks while reducing inference time, and its larger latent representation performs strongly on ROCStories.
7 Comparison across generative paradigms
- On ROCStories, COSMOSN = 128 reaches MAUVE 0.940, near the human reference's 0.953 and above GPT-2's 0.789.The passage notes that COSMOSN = 128 lags GPT-2 in perplexity.Source: §7 · p. 9
- On longer-context tasks, COSMOSN = 128 matches or slightly exceeds diffusion and autoregressive baselines while substantially reducing inference time.The evaluation covers summarization and question generation, using 512-token input contexts for these tasks.Source: §7 · pp. 8–9, Table 4 · p. 9
8 Experiments on OpenWebText
OpenWebText experiments test diffusion-model scaling and compare COSMOS configurations on 512-token generation.
8 Experiments on OpenWebText
- The OpenWebText experiments evaluate diffusion models trained for large-scale text generation.The section introduces the evaluation on OpenWebText.Source: §8 · p. 10
8.1 Impact of Diffusion Model Depth
- Increasing diffusion-model depth improves generation quality and diversity while the autoencoder and latent-space configuration remain fixed.The experiments compare 12-, 24-, and 48-layer models, using an autoencoder with 128 latent vectors for 128-token sequences.Source: §8.1 · p. 10, Table 5 · p. 10
8.2 Comparison on Long-Sequence Generation
- MAUVE 0.792 vs. 0.286 and perplexity 55.0 vs. 228.3: COSMOSN=512 outperforms GIDD on both reported metrics.COSMOSN=256 degrades relative to the uncompressed model but still outperforms full-length TEncDM.Source: §8.2 · p. 10, Table 6 · p. 10
9 Conclusion
COSMOS performs text generation in a compact, smooth latent space tailored to diffusion, with generation quality comparable to or better than token-level diffusion and autoregressive baselines across diverse tasks.
9 Conclusion
- COSMOS shifts text generation from token-level representations to a compact latent space tailored for diffusion.Its compressed, smooth representations are learned with autoencoder objectives.Source: §9 · p. 10
- Across diverse tasks, COSMOS matches or outperforms traditional token-level diffusion and autoregressive baselines.Source: §9 · p. 10
- The reported results position latent diffusion as a promising direction for fast, high-quality language models.Source: §9 · p. 10
A Dataset descriptions
The appendix describes datasets, related diffusion approaches, COSMOS’s latent diffusion pipeline, and evaluation configurations.
A Dataset descriptions
- ROCStories provides five-sentence unconditional stories, while Wikipedia supplies 8.6 million sequences for large-scale experiments.ROCStories has 88,161 training instances; the Wikipedia subset holds out 50,000 sequences for validation.Source: A Dataset descriptions · p. 15
- XSum supports abstractive summarization, SQuAD2.0 is reversed into question generation, and ParaDetox pairs toxic comments with neutral rewrites.Inputs are truncated to 512 tokens for XSum and SQuAD2.0, and to 40 tokens for ParaDetox.Source: A Dataset descriptions · p. 15
B.1 Gaussian Diffusion on Embeddings
- Gaussian diffusion on embeddings includes DiffuSeq [10], SeqDiffuSeq [45], and AR-Diffusion [43], with varied noising, architectures, and sequential conditioning.Source: §B.1 · p. 15
B.2 Simplex-Based Diffusion
- TESS [20] applies non-autoregressive diffusion to the logit simplex and uses self-conditioning.The passage reports strong performance across natural-language generation tasks while requiring fewer diffusion steps.Source: §B.2 · p. 15
B.3 Masked Diffusion
B.4 Latent Diffusion
- Latent-diffusion baselines include LD4LG [19] and PLANNER [47], while TEncDM [36] models contextual language-model encodings.Source: §B.4 · p. 16
C.1 Latent diffusion pipeline
- COSMOS denoises noisy latents toward clean representations, uses self-conditioning, and conditions sequence-to-sequence generation on encoded source latents.Conditional source features enter through cross-attention; the denoiser is initialized from unconditional-task weights.Source: §C.1 · p. 16
C.2 Training and inference configuration
- Training uses eight NVIDIA A100 GPUs; inference timing uses batches of 512 on one A100 with bfloat16 and averages five runs.SEDD uses 32 diffusion steps for the stated comparison setting.Source: §C.2 · p. 16, §C.3 · p. 16
D.1 Marginal utility of variational prior in text autoencoder
The appendix finds little benefit from KL regularization, reports decoder robustness to Gaussian noise, and summarizes extended task metrics.
D.1 Marginal utility of variational prior in text autoencoder
- 33.6 → 32.6: β = 0.01 yields only a marginal perplexity gain over β = 0, while β = 1 causes collapse.At β = 1, perplexity explodes and MAUVE falls.Source: §D.1 · p. 17, Table 8 · p. 17
- The decoder maintains reasonable performance when sampling from the prior with σ = 1, supporting the authors’ account of limited KL gains.Source: §D.1 · p. 17
- The final model omits the KL term and relies on latent masking, which the authors report gives larger and more robust improvements.Source: §D.1 · p. 17
D.2 Additional details of comparison across generative paradigms
- Extended conditional-generation comparisons report ROUGE for XSum and SQuAD2.0 and J-score for ParaDetox.The passage defines J-score as the product of style accuracy, fluency, and content preservation; Table 9 compares autoregressive and diffusion baselines.Source: §D.2 · p. 18, Table 9 · p. 18
E Limitations
The section outlines areas left for future work, including joint autoencoder–diffusion training, latent-dimension selection, architecture scaling, and reporting error bars, and discusses societal risks. It also states that LLMs were not important, original, or non-standard components of the core methods.
E Limitations
- The experiments use roughly 130M-parameter backbones, hold latent dimensionality fixed, and leave architecture scaling and a systematic dimensionality sweep to future work.Changing latent dimensionality would require redesigning the diffusion backbone; scaling is described as an orthogonal engineering effort expected to reinforce the observed trends.Source: E Limitations · p. 18
- Jointly training the autoencoder and diffusion model remains open; the authors suggest it may significantly improve training efficiency.Source: E Limitations · p. 18
F Societal Impact
- The authors state that latent continuous diffusion does not bring significant new societal risks beyond those associated with current language models.Source: F Societal Impact · p. 18
2. Limitations
- The paper does not report error bars for comparison results across all tasks, following prior work.Source: §2 · p. 25
- The authors state that LLMs were not used in the described research scenarios or as important, original, or non-standard components of the core method.They mark the LLM-use question as not applicable.Source: §2 · p. 29