Compressed and Smooth Latent Space for Text Diffusion Modeling

We sincerely thank the reviewer for the positive and encouraging feedback. We are pleased that you found our analysis of the regularization techniques thorough and relevant, and we appreciate your recognition of the efficiency gains and the advantages of building on continuous‑space diffusion models. Below, we address each point raised.

W1. On the analogue of temperature sampling in the continuous latent diffusion framework.

Temperature sampling in autoregressive models allows for increasing or decreasing the confidence of token sampling, which in turn typically improves quality at the expense of diversity, and vice versa. In continuous Gaussian diffusions, there is indeed no direct concept of temperature sampling. However, we can propose a technique with an analogous effect. The work of TEncDM [1] shows that self-conditioning increases the network's confidence in its predictions. Based on this fact, we can control the prediction confidence using the following procedure to obtain the prediction $x_0$:

$x_0^{t} = x_{\theta}(x_t, t, \emptyset) + w \cdot \bigg(x_{\theta}(x_t, t, x_0^{t+1}) - x_{\theta}(x_t, t, \emptyset)\bigg).$

Here, the final prediction $x_0^t$ at step $t$ is an extrapolation from the unconditional prediction $x_{\theta}(x_t, t, \emptyset)$ in the direction of the self-conditioned prediction $x_{\theta}(x_t, t, x_0^{t+1})$, controlled by a weight $w$.

A larger $w$ increases model confidence, leading to higher-quality generations, while a smaller $w$ promotes diversity. Our key finding here is that this trade-off can be finely controlled; by tuning $w$, it is possible to achieve substantial improvements in quality with only a minor decrease in diversity.

This represents a very useful practical feature of our approach. We appreciate you raising this point and will add these results to the final version of the paper.

W2. On empirical evaluation.

Latent diffusion for text generation is a relatively new and under-explored area, as also noted by Reviewer Kj8c. Prior to developing a large‑scale model, it is crucial to validate the underlying hypotheses for training the latent space and the diffusion model at a smaller scale. This work accomplishes precisely that, laying the groundwork we consider essential for successful large‑scale development.

To provide initial evidence of this scalability, we have taken the first steps on the OpenWebText dataset by investigating if quality improves when increasing the diffusion model's size, keeping other components fixed. By increasing the number of layers, we found that generation quality and diversity both improve significantly:

Additionally, we trained an autoencoder using our strategy without compression and a diffusion model with a sequence length of 512 on OpenWebText, for comparison with the discrete diffusion method GIDD, which was trained under identical conditions.

These results are a strong positive indicator for scaling.

Regarding generation length, we used a fixed number of latents for all texts in this work. Adapting the model for longer texts requires separate research as, you correctly noted, there are several viable strategies to explore

W3. On distillation of continuous diffusion model.

We agree that exploring the distillation of latent text diffusion models is a fascinating direction for future work. However, we are cautious about the direct application of techniques from the image domain, as we are not aware of any work that has successfully distilled a continuous Gaussian diffusion model into a one-step generator for text. As you rightly point out, such an acceleration technique is orthogonal to our core contribution of learning a compressed latent space for diffusion. We believe our model will become a good framework for beginning research in this direction.

Q1. On sequence length.

As we wrote in Section 7 in the paragraph "Autoencoder setup across tasks," the decoder sequence length (L) is 128 or 512, depending on the experiment. To generate sequences of arbitrary length, we follow the same procedure as in other works on text diffusion [1, 2], where the decoder reconstructs a sequence up to a maximum length with padding, and the output is then cut off after the first EOS token. Consequently, the decoder in our framework does not produce sequences that exceed the predefined maximum length.

Q2. On using not a fixed size latent space.

We thank you for this insightful suggestion. To keep the core methodology focused, we employ a fixed latent space size in this work. We agree that using a variable number of latents for texts of varying lengths is a compelling idea for future exploration, though it requires additional research beyond the scope of this paper.

Q3. On introducing context in the method.

To encode the context before feeding it into the diffusion model, we first pass it through our autoencoder, which generates a compressed representation. It is this representation that is subsequently supplied to the diffusion model via cross-attention. Depending on the task and context length, we use a model configured for either 512 or 128 tokens for this process.

Q4. On baselines.

Our comparative analysis was designed to benchmark our method against key paradigms in text diffusion: Gaussian diffusion on text embeddings, simplex-based diffusion, masked diffusion, and latent diffusion. From the many masked diffusion models available, we selected SEDD as a prominent and established baseline. For a more direct and controlled comparison, we also benchmark against GIDD (as detailed in our response to W2), as it was trained under the same conditions and text length as our model. We considered MDLM but opted against it because its training on single-line concatenated texts produces fragmented, non-comparable samples. To avoid re-training a method specifically for our setup, we decided to use GIDD, which was trained in the exact same setup as our model.

$\newline$

References:

[1] Shabalin, Alexander, et al. "Tencdm: Understanding the properties of the diffusion model in the space of language model encodings" (2025).

[2] von R{"u}tte, Dimitri, et al. "Generalized interpolating discrete diffusion" (2025).

We sincerely appreciate your constructive feedback, which will greatly assist us in improving the final manuscript. Please let us know if there are any remaining points that require clarification.

We sincerely appreciate your thoughtful feedback and are encouraged by your interest in diffusion models for language generation. We are especially grateful for your positive recognition of our ablation studies and evaluation design.

W1. Q7. On related works.

Thank you for this valuable suggestion. While our related works section focuses primarily on latent diffusion methods for text—given their direct relevance to our approach—we agree that a more thorough comparison with a broader set of related methods would strengthen the manuscript.

In the revised version, we will expand our discussion to provide a clearer account of both conceptual and empirical differences between our method and prior work across the broader landscape of text diffusion models. This will include a more detailed analysis of methodological design choices, evaluation strategies, and task settings, in order to better contextualize our contributions within the existing literature.

W2. Q8. On longer contexts.

As also noted by Reviewer Kj8c, latent diffusion for text generation is a relatively new and under-explored area.Before developing a large-scale model, it is crucial to validate the underlying hypotheses for training the latent space and the diffusion model at a smaller scale. This work accomplishes precisely that, laying what we consider the essential groundwork for successful large-scale development. Having established this foundation, we agree that evaluating our method on longer text sequences is a critical next step in assessing its scalability.

To this end, we have taken the first steps and made additional experiments on OpenWebText with extended generation lengths of up to 512 tokens with 512 latents using our 0.12B parameter model. We compared its performance against the discrete diffusion baseline GIDD [1], which has a comparable model size and was pretrained on the same dataset and sequence length. The results show that our latent diffusion approach maintains coherent and diverse text generation at this longer length, demonstrating its robustness beyond short examples.

These findings suggest that the method scales well to longer sequences, and we will include these results in the final version of the manuscript to provide a more comprehensive evaluation.

W3. Q9. On explicit exposure bias analysis.

Thank you for your insightful comment regarding exposure bias. Since our work investigates diffusion in a continuous latent space, the proposed framework inherently mitigates exposure bias. As the diffusion model updates the entire latent representation simultaneously, it can flexibly revise tokens that are out of place at any position.

While extensive empirical evidence supports this behavior in the image domain, we acknowledge that corresponding evidence in the language domain is more limited. In response to your suggestion, we now provide several randomly selected predictions of our model on ROCStories at different time steps. These examples demonstrate the model’s capacity to revise previously predicted tokens. We appreciate your feedback and will include additional examples in the final version of the manuscript to further illustrate this behavior.

Q1. On the differences from LD4LG.

We appreciate your reference to the LD4LG paper, which we also found to be insightful and well-executed. While our work draws conceptual inspiration from LD4LG, there is a key methodological distinction. Specifically, LD4LG employs autoregressive language generation conditioned on latent representations obtained via diffusion. As a result, it inherits the limitations associated with autoregressive modeling, such as exposure bias and constrained token-level dependency.

Moreover, LD4LG utilizes a fine-tuned BART-base model as the decoder for these latent representations. Given the strength of BART-base as a standalone model, it becomes difficult to isolate and quantify the specific contribution of the diffusion-generated latents to the overall generation quality. This design choice makes the impact of the latent diffusion component less transparent. In contrast, our primary contribution is a methodology for learning a textual space specifically designed for diffusion modeling, which in turn allows for the efficient generation of high-quality text within the autoencoder's latent space.

Q2. On the abscence of citations.

Thank you for bringing this issue to our attention. We have made every effort to properly cite all public datasets and methodologies used in our work, either in the main text or in the Appendix. However, it is possible that we may have inadvertently omitted a citation or failed to provide sufficient attribution in some instances.

If you could kindly point us to specific lines or sections where a citation appears to be missing, we would be grateful. We will be happy to correct any such oversights in the final version of the manuscript.

Q3. On the typographical error.

Thank you for noting this typographical error. We will correct it in the final version.

Q4. On the overly strong statement.

Thank you for the suggestion. We will revise and soften the statement accordingly in the final version.

Q5. On the supplementary materials.

We have included the working code in the supplementary material. Implementation details are provided in Appendix Section B.

Q6. On the lack of human evaluation.

We adopt an evaluation framework that is standard in the continuous language diffusion literature [2], [3]. The aggregated automatic metrics used for our selected tasks are well-established, and prior work has shown that they correlate reliably with human judgments. Additionally, conducting large-scale human evaluation for models with fewer than 0.3B parameters is resource-intensive and, unfortunately, beyond our current budget constraints.

$\newline$

References:

[1] von R{"u}tte, Dimitri, et al. "Generalized interpolating discrete diffusion" (2025).

[2] Lovelace, Justin et al., "Latent Diffusion for Language Generation" (2023).

[3] Shabalin, Alexander, et al., "TEncDM: Understanding the Properties of the Diffusion Model in the Space of Language Model Encodings" (2025).

Thank you for your valuable feedback and suggestions. If our responses have addressed your concerns, we would kindly appreciate your consideration of revising your score. Please do not hesitate to let us know if any points require further clarification.

Thank you again for your review and feedback on our work. As the discussion period is ending shortly, we are writing to respectfully check if our detailed response and the new experiments have addressed your concerns. If they have, we would kindly ask you to consider these improvements in your final assessment.

Thank you for your review and positive assessment of our work. We greatly appreciate that you recognized the motivation behind our design choices, the value of our ablation studies, and the insights provided by our analysis of latent space robustness.

W1. It is an open question whether such approaches can scale to that setting.

We thank the reviewer for highlighting this important point. As you rightly observe, latent diffusion for text is a relatively new and under-explored area. Before developing a large-scale model, it is essential to validate hypotheses regarding the training of the latent space and the diffusion model at a smaller scale. This is precisely what we accomplished in this work. We believe this foundational research is a prerequisite for successful large-scale development.

To that end, since submitting this work, we have taken the first steps toward a general pretrained language model on the OpenWebText dataset. Specifically, we investigated two questions: first, whether generation quality improves when increasing the diffusion model's size while keeping the autoencoder and latent space fixed, and second, how our 0.12B model performs at a length of 512 tokens compared to the GIDD discrete diffusion model pretrained under the same conditions.

To increase the diffusion model's size, we focused on increasing the number of layers. The results for 128-token sequences show that generation quality improves significantly in both quality and diversity.

Furthermore, the comparison with discrete diffusion at a generation length of 512 shows that the latent diffusion paradigm is highly promising for advancing the text generation field.

These results provide positive evidence that our approach can be successfully scaled to longer sequences and larger datasets. The final manuscript will incorporate these new experimental outcomes.

W2. On memorization of the training set.

We thank you for bringing this important point to our attention. We have calculated this metric for all methods involved in the comparison. The SEDD method shows the lowest value; however, its value is lower than the reference, which suggests that the generated texts may be less coherent. Cosmos achieved the values closest to the reference, which indicates that Gaussian latent diffusion is not prone to copying data from the dataset. Following Cosmos are the other two Gaussian latent diffusion methods, TencDM and LD4LG. We will add this metric to the final version of the paper.

W3. Q1. On variable length generation.

We acknowledge that non-autoregressive methods do not generate texts of arbitrary length as naturally as autoregressive models. The procedure in our work is consistent with other studies on text diffusion [1, 2], where the decoder reconstructs a sequence up to a maximum length with padding, and the output is then cut off after the first EOS token. The number of latents is the same for all texts.

It is also worth noting that this challenge is a direct consequence of generating the text holistically in an iterative manner. This provides the key benefit of a full bidirectional context but also the drawback of less natural variable-length generation. A block-based approach could partially resolve this issue in future work.

W4. Q2. On changing the latent dimension.

We agree that adding input and output projection layers is a possible way to adapt the diffusion model to a changed latent dimensionality. However, this would not achieve our goal of accelerating inference. The computational complexity would remain the same because the internal operations would still process sequences of the original length and original dimension. This is exactly why we investigated sequence length compression instead, as this is a direct way to reduce attention complexity and speed up the generation process without changing the capacity of diffusion model.

W5. Q3. On latent-space augmentation.

Our augmentation method differs from classic dropout [3] in a key respect. Standard dropout involves zeroing out a fraction of neurons and then scaling the remaining ones by 1/(1-p). Our technique omits this final scaling step. To prevent any potential confusion for the reader, we deliberately chose not to refer to it as dropout.

W6. The impact of text encoder on downstream generation quality.

We agree that the autoencoder is an important component. However, its specific architecture may be less critical than one might assume for general-purpose generation. This is substantiated by the work of TEncDM, who found that results across different downstream tasks did not vary significantly with the choice of autoencoder. While a highly-tuned autoencoder might be necessary to "win" a benchmark, a general one provides a robust and comparable baseline.

We believe a more dominant factor for downstream performance is the autoencoder's pre-training data, as shown by T5. Since our goal in this paper was not to achieve state-of-the-art results on downstream tasks, but rather to compare different diffusion paradigms for text generation, we did not select a text encoder for a specific task and instead focused on developing the compressor and decompressor for the diffusion model.

$\newline$

References:

[1] von R{"u}tte, Dimitri, et al. "Generalized interpolating discrete diffusion" (2025).

[2] Shabalin, Alexander, et al. "Tencdm: Understanding the properties of the diffusion model in the space of language model encodings" (2025).

[3] Hinton, Geoffrey, et al. "Improving neural networks by preventing co-adaptation of feature detectors" (2012).

[4] Raffel, Colin, et al. "Exploring the limits of transfer learning with a unified text-to-text transformer" (2020).

We look forward to incorporating your feedback in our final version and thank you again for your time and insightful comments.

Thank you for your review and the valuable suggestions. We value your acknowledgement of our work and your interest in the latent‑diffusion training technique we propose for textual data.

We are sorry to hear you could not find the appendix. It is included with the full paper in the supplementary material (file: full_paper.pdf). We would be grateful if you could take a moment to review it, as we believe it addresses a number of the points you have raised. For your convenience, we have also duplicated answers to those questions immediately below.

W1. On related works.

In the related works section, we initially focused on latent diffusion for text, as these works form the direct foundation of our research. We then provide a broader discussion of text diffusion methods in Section 7, where we categorize them into four main groups: gaussian diffusion on text embeddings, simplex-based diffusion, masked diffusion, and latent diffusion. Following your suggestion, we will update the related works section to incorporate these alternative approaches for a more comprehensive overview.

W2. On limitations.

We discuss the limitations of our approach in Appendix D. We acknowledge that the two-stage training process is a significant limitation. Before the diffusion model can be trained, a text autoencoder must be prepared. This step, while enabling the diffusion to operate in a more meaningful latent space, introduces an additional training overhead not found in autoregressive models and complicates the overall workflow. Concurrently, we are not claiming our model outperforms large-scale autoregressive models. Nevertheless, our results demonstrate that it achieves performance competitive with autoregressive models of a similar size, specifically ~0.1B, on several generation tasks, albeit with the added training overhead.

W3. On significant differences among various models.

Thank you for this suggestion. To better contextualize our work, we will add a dedicated discussion in the appendix that explores the significant differences among various diffusion models for text, complementing the expanded related works section.

W4. On model predictions.

Please see Appendix F for sample outputs on each dataset we evaluated.

W5. On the choice of GPT-2.

The primary objective of Table 4 is to compare Cosmos with other diffusion-based approaches for text generation, such as Gaussian diffusion on text embeddings, simplex-based diffusion, and masked diffusion. All these methods select their own text representation, define a text corruption process, and train a model to refine the text using full bidirectional context. Our results validate our approach, showing that operating in a compressed latent space can substantially speed up sampling while maintaining or even improving quality, given a properly trained autoencoder. The autoregressive model was included as it represents the standard paradigm, making it an important benchmark for any alternative text generation approach. GPT-2 and GPT-Neo (an open-source alternative to GPT-3) were chosen as classic representatives of autoregressive modeling.

Furthermore, the goal of this paper was not to develop a state-of-the-art model for specific generative tasks, as separate works focus on developing large, specialized models for such purposes [4, 5]. Instead, our primary contribution is a method for learning a compressed textual space specifically designed for diffusion modeling, which in turn allows for the efficient generation of high-quality text within the autoencoder's latent space.

W6. On Cosmos scaling.

Latent diffusion for text is a relatively new paradigm, which makes the scaling process non-trivial. It requires additional research to determine how to best allocate resources: how much should one increase the autoencoder, the latent space size, or the diffusion model?

In this paper, we have partially answered some of these questions. For example, in Section 6.2, we showed that it is beneficial to increase the dimensionality of the latent space even without increasing the size of the autoencoder. Specifically, we investigated how the quality of generated text changes as the number of layers in the diffusion model's transformer is increased. According to preliminary results at a text length of 128, both quality and diversity of generated text improves significantly. This inspires confidence that the proposed approach can be scaled.

We have also performed an experiment with an increased generation length of 512 tokens with our 0.12B model and compared it with the discrete diffusion model GIDD [3], which is of a similar size and was pretrained on the same dataset and length. This result shows that the latent diffusion paradigm for text generation is highly promising and poised to impact the field as a whole.

We hope our work provides a foundation for further research into latent language diffusion as a promising paradigm.

The final manuscript will incorporate these new experimental outcomes.

Q1. On the complexity of the compressor and the decompressor.

The compressor block is linear with respect to the sequence length $L$, with an attention complexity of O($NLd$). The decompressor block is quadratic with respect to the sequence length, with its attention complexity being O($L^2 d$). Here, $L$ denotes the sequence length, $N$ the number of latents, and $d$ the latent dimensionality.

Crucially, we only use the computationally expensive decompressor once per inference, to map the final latent back to text. The primary efficiency gain stems from running the entire iterative diffusion process in the compressed latent space. Since attention complexity is quadratic, operating on this much shorter sequence of length N significantly accelerates generation compared to models working in the full sequence space.

$\newline$

References:

[1] Sahoo, Subham, et al. "Simple and effective masked diffusion language models" (2024).

[2] Wang, Guanghan, et al. "Remasking discrete diffusion models with inference-time scaling" (2025).

[3] von R{"u}tte, Dimitri, et al. "Generalized interpolating discrete diffusion" (2025).

[4] Liu, Yixin, et al. "BRIO: Bringing order to abstractive summarization" (2022).

[5] Zhang, Tianyi, et al. "Benchmarking large language models for news summarization" (2025).

We thank you for your thoughtful feedback and suggestions. We would be grateful if you would consider raising your score, in case we have addressed your concerns. Please let us know if any aspects still need clarification.

Thank you again for your thoughtful review. As the discussion period is ending in less than 24 hours, we just wanted to kindly follow up on our rebuttal. We would be very grateful to know if our detailed response and the new experiments addressed your concerns. If they did, we would be deeply grateful if you would consider raising your score. Thank you for your time.