Title: Cosmos: Compressed and Smooth Latent Space for Text Diffusion Modeling
Venue: NeurIPS 2025
Accepted: yes
Keywords: latent diffusion

text generation

language models

autoencoders
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 up to $8\times$ while maintaining generation quality comparable to token‑level diffusion models. Furthermore, increasing the latent sequence length allows \textsc{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 https://github.com/MeshchaninovViacheslav/cosmos.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper studies the use of a learned continuous and compressed latent spaces for text generation. A large part of the paper is focused on how to regularize the latent space to improve its “diffusability” and on evaluations to understand how each of these regularizations affect the learned latent space and the subsequent performance by the latent diffusion model. One of the main benefits of the proposed method comes from running a generative model in a lower dimensional space (typically 4x and 8x smaller), which leads to faster sampling.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The analysis and evaluations of the regularization techniques used are quite thorough and indicate that the regularizations used are quite relevant for the task at hand. I consider this to be quite nice, as learning the right latent space for these tasks (i.e. well suited for a diffusion model to operate in the latent space) is critical.\n- The model gets faster inference without sacrificing performance.\n- The fact that the method builds on continuous space diffusion models enables the use of tools developed in that space, such as distillation and other fast sampling techniques, which are not yet as developed for discrete diffusion models.\n\nWeaknesses:\n- Typical AR language models rely on low temperature sampling, which appears to be somewhat important for good performance. Moving to the continuous latent diffusion framework, it is unclear how this could be done. While there are techniques for low temperature sampling for diffusion models in continuous spaces, these are quite heuristic, and the final effect may be quite different to the one obtained when doing low temperature sampling for LLMs (the fact that the model operates on a latent space may exacerbate this).\n- I don’t find the empirical evaluation of the final diffusion models entirely convincing regarding the performance of the method at scale. Models are on the smaller side (100M-200M parameters), trained on small-to-medium size datasets (ROCStories, XSum, etc.), and on fairly small latent spaces (16 x dim). For longer sequences, the 16 should be increased? Or would this approach be used in an autoregressive manner generating blocks?\n- While I don’t consider this necessary, it could be very nice to have. One strength I mentioned above involves the use of distillation techniques that are only applicable in continuous space diffusion models. It would be great to see how these techniques interplay with the proposed method. This would be orthogonal to the paper’s contributions, but is aligned with the whole direction of increased efficiency. Maybe such techniques could give another boost in speed, which would be very nice.\n\nQUESTIONS:\nIn the diffusion model evaluations, what is the decoded sequence length (that is, L)? How exactly can the method be used to generate longer outputs than the ones obtained by decoding the fixed length latent variables to a fixed output size. (Related to the next question.)\n\nI understand in this work the method is always trained on inputs of certain size L=512, and the latent space has L / c, where c is the compression factor >1. You could train this on sequences of varying lengths I suppose? By using not a fixed size latent space, but a latent space of size (L / c) x d, where L changes every batch.\n\nI was not totally clear on how you are introducing context (prompt) in the method, for instance, for short or long context tasks. \n\nWhy are other baselines, for instance MDLM, not considered? Is it due to the scale of the networks used? (This work focuses on architectures in the 100M-200M range).\n\nLIMITATIONS:\nThe paper includes a limitation section. See weaknesses above for some questions regarding other potential limitations of the method.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces COSMOS, an approach for text generation that leverages a compressed and smoothed latent space specifically tailored for diffusion models. Using a combination of a frozen language encoder, a lightweight compressor-decompressor architecture, and robustness-promoting training objectives, the method aims to achieve efficient, high-quality text generation with latent spaces significantly smaller than token-level representations. COSMOS is evaluated on a range of text generation tasks (such as story generation, summarization, question generation, and detoxification) and compared against both autoregressive and diffusion baselines. The experiments demonstrate that COSMOS achieves comparable or superior text quality while providing substantial speed-ups in inference.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**\n1. Applying diffusion models to the field of text generation is an interesting and reasonable approach that can alleviate some of the issues caused by sequential decoding. \n2. The paper presents extensive ablation studies and evaluates the proposed method, COSMOS, on a variety of downstream tasks. The improvements over baselines provide compelling evidence for the method’s effectiveness.\n\n**Weakness**\n1. The paper lacks a clear and detailed comparison with related work, particularly with methods such as LD4LG. A more thorough discussion of conceptual and empirical differences would strengthen the contribution.\n2. All evaluated examples are relatively short in length. To better demonstrate the scalability and robustness of the proposed method, experiments on longer text sequences are recommended.\n3. Although the introduction highlights exposure bias as a major motivation, the experimental validation of this point is indirect. More explicit analysis would be beneficial.\n\nQUESTIONS:\nQ1. What is the fundamental difference between your method and LD4LG?  \nQ2. Public datasets and methodologies must include citations or sources.  \nQ3. In Section 4.2.2, there are two consecutive ending punctuation marks in the last line. Please fix this issue.  \nQ4. In Section 6.2, you state: \"It achieves an 8× reduction in latent sequence length while surpassing the quality of uncompressed representations.\" However, according to Figure 6, this phenomenon only holds for the comparison between N=128 and N=16, not universally. Recommend revising the wording to reflect this nuance.  \nQ5. Reproducible code and implementation details.  \nQ6. The evaluation metrics lack human evaluation results under few-shot settings.  \nQ7. Related work deserves a more detailed discussion.  \nQ8. Long-context experiments.  \nQ9. Direct experimental validation of exposure bias mitigation.\n\nLIMITATIONS:\nRefer to weakness", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work explores latent diffusion for language generation. It focuses primarily on the design of the latent diffusion autoencoder. Two primary approaches exist from past work. The first approach involves compressing activations from a language model and then decoding them autoregressively (LD4LG and PLANNER). The second approach learns a diffusion model directly in the activation space of a pretrained language model and decodes them with a learnable non-autoregressive decoder. This work develops an autoencoder that both compresses the activations of a language model and decodes them non-autoregressively. It closely follows the architecture of the LD4LG autoencoder. They introduce a number of carefully ablated choices that improve the effectiveness of the latent space for diffusion modeling and decoding. They validate their approach across a number of tasks (unconditional generation, summarization, detoxification, and question generation) with a wide range of baselines.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\nLatent diffusion is a relatively underexplored approach for language generation compared to autoregressive LMs and discrete diffusion LMs and work in this space is welcome.\n\nThe proposed design choices (e.g. dropout, noise augmentation, intermediate MSE regularization) in this work are well-motivated and the ablation studies are comprehensive. Prior work in the image domain has demonstrated the importance of the autoencoder design, so exploring it for language is natural. The analysis of the robustness of the latent space is interesting and insightful.\n\nThe baseline selection is comprehensive and covers a range of different approaches (autoregressive LMs, word embedding diffusion, discrete diffusion, latent diffusion). The authors validate their approach on a reasonable selection of tasks. Their approach consistently compares favorably to the set of baselines.\n\n\n**Weaknesses**\n\nThis work develops dataset-specific language diffusion models instead of trying to develop a general pretrained language model trained on web-text as is the current language modeling paradigm. I think that the authors’ setting is more computationally tractable and is reasonable given the early stage of the paradigm, but it is a limitation. It is an open question whether such approaches can scale to that setting.\n\nFor the unconditional language generation evaluation, the work does not quantify memorization of the training set. All of the metrics can simply be maximized by reproducing samples from the training set, so the memorization can be quantified in some way. Even just using a simple n-gram overlap metric such as was used by prior work (LD4LG). \n\nOne weakness of non-autoregressive approaches is that variable-length generation is not as natural as with autoregressive models, although I do not see discussion of this in the paper. My understanding is that the number of latents in the decoder must be specified. How is this done during inference?\n\nThe authors mention that they only focus on compressing the sequence length because reducing the dimensionality of the latents requires modifying the diffusion architecture. However, you can generally account for differences in input dimensions by simply introducing a linear projection at the input/output layer that represents a negligible change in capacity. This restriction limits the scope of their study.\n\nAs written, the “Latent-space augmentation” sounds like standard dropout. If it is the exact same operation, then just introducing it as dropout would be clearer.\n\nI would expect things such as the choice of text encoder to have a significant impact on downstream generation quality, but this choice is not explored in this work.\n\nQUESTIONS:\nHow is variable-length generation handled? This the main confusion about the paper that I would like to see addresses.\n\nCan you expand on the challenge of compressing the dimensionality? Adding a linear projection at the input/output layer should enable changing the latent dimensionality without significantly altering the diffusion model capacity. I think it's fine to limit the scope of the study, but do not think this is a legitimate challenge.\n\nIs “Latent-space augmentation” equivalent to dropout or is there some implementation difference?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes COSMOS (COmpressed and SMOth latent Space), a latent diffusion model for text generation. COSMOS learns a compressed and smooth latent representation through carefully designed autoencoder objectives, enabling high-quality generation while reducing inference time.\n\nThe authors carefully demonstrate how their approach achieves desired latent-space properties for diffusion learning (Section 5). They also ablate and show how CE loss and MSE loss were important to learning latent space representations for COSMOS.\n\nSTRENGTHS AND WEAKNESSES:\nThe paper adopts latent diffusion framework for textual data by learning contextualised text representations into the latent space, and trains a diffusion model that operates within this space. The proposed approach is very interesting, the authors have done a great job (specially in Section 5 and 6) to demonstrate how to effectively learn such models for text. However, I am afraid that the paper in its current form and with a limited set of experiments will not be very impactful. \n\nFor example, the paper will benefit from a more thorough related work and a detailed and honest discussion on its limitations, instead of directly claiming that the model is able to match or outperform autoregressive baselines across diverse tasks. \n\n> Our empirical results show that, COSMOS matches or outperforms traditional token-level diffusion and autoregressive baselines across diverse tasks. Our findings challenge the dominance of token-level models and highlight latent diffusion as a promising direction for building fast, high-quality language models.\n\n\nIn Table 4, it would be very important to report on significant differences among various models. Also it would be great to show how model predictions look both for Table 2 and Table 4. \n\nIt would be important to explain why only GPT2 and GPT Neo were chosen for autoregressive models. What are the SOTA performances on these tasks? \n\nI think the experiments were done with 100-200M model parameters. What are the limitations of scaling COSMOS to much larger models?\n\nQUESTIONS:\nPlease see my comments in the Strength and Weakness section. \n\nWhat is the complexity of the compressor and the decompressor? Are they quadratic in the sequence length? \n\nMinor: Appendix were referred, but I could not find relevant appendix in the paper.\n\nLIMITATIONS:\nPlease see my comments in the Strength and Weakness sections. I believe the paper will benefit from a honest discussion on its adaptability to real worlds tasks and the scaling limitations.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-22"}