Diffusion on Language Model Encodings for Protein Sequence Generation

Thank you for your review and positive assessment of our work. We appreciate your recognition of our systematic ablation studies and approach to developing protein-specific diffusion parameterizations rather than simply adopting techniques from the image domain. Below, we address each point raised.

C1, C2. Formatting and typos.

Thank you for pointing out the formatting issues and typos. We will fix the formatting inconsistencies, including proper capitalization of titles and section headings.

Q1. Line 98: I'm wondering whether the normalize is applied over the length dimension $s$ or the hidden dimension $d$. From the last sentence, it seems that the normalization is over the length dimension, which seems weird for generation.

We apologize for the unclear description. The normalization is applied over the hidden dimension $d$ rather than the sequence length dimension $s$. Specifically, for each component $d_i$, we precompute the mean and variance over the training data and then apply normalization using these statistics to achieve zero mean and unit variance. We will clarify this in the revised manuscript.

Q2. Do the authors also use ESM-2 decoder for diffusion on CHEAP and SaProt representations?

We use the same approach across all encoder architectures. For CHEAP and SaProt representations, we fine-tune their corresponding pretrained decoders alongside the diffusion model, similar to our approach with ESM-2. We found that low-dimensional embeddings (e.g., ESM-2 8M, $d$=320) are less robust to small perturbations than higher-dimensional ones (ESM2/SaProt 650M, $d$=1280, CHEAP, $d$=1024). Fine-tuning the decoders helps minimize these effects during diffusion generation.

Q3. What's the rationale behind the heurisitc approach that the reconstruction loss should exhibit an approximately linear increase over diffusion time?

The key issue with standard noise schedules (linear, cosine) is that they corrupt data very gradually at small timesteps. This means the model spends significant training resources on nearly trivial denoising tasks where input and target are almost identical. Our approach ensures the difficulty of the denoising task increases steadily with each timestep. By designing a noise schedule where reconstruction loss grows linearly with diffusion time, we create learning problems of incrementally increasing difficulty, allowing the model to make consistent progress throughout training rather than facing an abrupt jump in task complexity.

Q4. Typically, the quality of generated samples from diffusion models can be controlled by adjusting the temperature. I'm wondering how the authors set the temperature values in different experiments and for different baselines.

For baseline comparisons, we use the sampling parameters recommended by the original authors of each method. For autoregressive models like ProGen2 and ProLLAMA, which show suboptimal quality and collapse to highly repetitive sequences on default settings, we performed grid searches to identify optimal temperature and top-p values. DiMA has two analogous parameters for navigating the quality-diversity trade-off: the number of generation steps and self-conditioning rate. The dependencies of quality, diversity and novelty on these parameters are shown in (figures 2 and 3 on https://tinyurl.com/icml25re).

---

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 thorough review and constructive comments. We appreciate your recognition of our work's comprehensive experimental design and clear methodology.

W1. Limited theory behind the proposed method, making is a straightforward application practice of well-established method

While latent diffusion models have been established in other domains, our work demonstrates that naive application of continuous diffusion to protein sequences yields poor results. Our ablation studies in Table 1 show that standard diffusion implementations struggle with protein generation, achieving only ~60 pLDDT compared to DiMA's 83.3 pLDDT. We demonstrate that domain-specific engineering is essential for effective latent diffusion models.

The challenge of applying diffusion to proteins is further evidenced by existing approaches requiring custom modifications. For example, discrete model DPLM [1] requires specialized initialization and sampling strategies [1, Appendix D.1], MultiFlow requires distillation steps, and so on.

Our work contributes by establishing a principled framework for latent diffusion that generalizes across different representation spaces (from 8M to 3B parameter sequence-based and multimodal encoders) and tasks with a single architecture. As shown in our fold-conditioning experiments, this enables DiMA to achieve stronger performance than specialized structure-generation models like RFDiffusion (TM-score 0.93 vs 0.48).

W2. Not comprehensive discussion of potential limitations or failure cases

Thank you for this suggestion. In our revision, we will add a dedicated discussion section addressing current challenges and future opportunities. Structure generation in our approach relies on heavyweight folding models like ESMFold. We see significant potential for enhancing performance through joint training of the diffusion model with lightweight structure prediction components. This approach could both improve computational efficiency and structural awareness of the model.

W3. Not strong performance compared to existing methods which may weaken the contribution of this work

As shown in Figure 2, DiMA, MultiFlow, RFDiffusion, and DPLM each demonstrate Pareto-optimal performance on the quality-diversity tradeoff, with each dominating in different aspects. Considering model sizes, DiMA shows remarkable performance, using two orders of magnitude fewer parameters than e.g. DPLM. Table 8 demonstrates that only DiMA and MultiFlow achieve balanced metrics on this tradeoff, with the notable advantage that DiMA trains exclusively on sequences while MultiFlow requires both sequence and structure data, making DiMA more scalable due to its independence from limited 3D structural data.

Q1. How does DiMA compare in speed (training/inference) to discrete diffusion models such as ESM3, DPLM?

We analyzed inference speed by generating 100 sequences of length 250 with 5 repetitions (for details see Figure 4 at https://tinyurl.com/icml25re). DiMA-35M is ~10x faster than DPLM-150M, which is slowed by using a larger model at each generation step. Structure models are notoriously computationally demanding. We have not measured training speed, but we expect DiMA to be faster as it requires only forward encoder passes, whereas discrete models need backward encoder passes as well.

Q2. Can you train or fine-tune DIMA with structure-based tasks and data?

Yes, DiMA incorporates structural information in two ways:

1. We use local structural tokens for each amino acid passed with the sequence into the SaProt encoder. This allows explicit use of 3D structure in the scaffolding task, achieving strong results (details in section 3.6.1, appendix E.8, and Figure 3).

2. We employ DiMA with the CHEAP encoder based on ESMFold, enabling dual decoding (sequence and structure) from sequence alone. This approach outperforms RFDiffusion in fold-specific generation using explicit structure guidance (details in sections 3.6.3 and appendix E.9, and Figure 4).

C1. Line 187-202. the long list of bullet points is not reader-friendly. Please consider re-organize it.

Thank you for this feedback. We will reorganize this section in the revision to improve readability.

C2. A discussion on how to scale up DiMA (maybe by combining it with structural models like AlphaFold) would be valuable.

Thank you for this suggestion. We are currently exploring the integration of structural generation models with our latent diffusion approach for co-generation. In the revision, we will add a discussion section highlighting future directions and the potential benefits of combining DiMA with structural models.

---

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.

[1] https://arxiv.org/abs/2402.18567

Thank you for your review and the valuable suggestions. We appreciate your recognition of DiMA's thorough evaluation and benchmarking approach. Below, we address each point raised.

Concurrent work that the authors will be interested in.

Regarding the work by Lu et al. (2024), we became aware of PLAID and evaluated it after our submission, just before receiving the reviews. We evaluated PLAID-100M (https://github.com/amyxlu/plaid) using the same protocol, we apply to all the models in our work. Our analysis shows that DiMA substantially outperforms PLAID, especially in terms of protein quality:

We will include these results in the final manuscript, specifically we will update Section 3.5 ('Comparison with Large Pretrained Models') and Figure 2.

Can the authors comment on the 4JHW and 5YUI PDB IDs as failure modes? What causes DiMA and other methods to struggle with these particular cases?

These cases highlight limitations in the current benchmark rather than specific model shortcomings. Our analysis reveals that even when using reference sequences, the resulting structures after ESMFold prediction fail to meet the benchmark success criteria. For 4JHW, the reference sequence yields motif RMSD exceeding 6.0Å with pLDDT around 30, far below the thresholds for success (RMSD ≤ 1.0Å, pLDDT ≥ 70). Similarly, 5YUI produces motif RMSD above 3.0Å. The detailed results of our analysis on these challenging cases are presented in Table 1 at [https://tinyurl.com/icml25re].

Recent work [1] has recognized these benchmark limitations and is developing improved evaluation protocols, specifically addressing them in Table 4 of their preprint.

It would be more informative to include a graphic overview of the model and approach as Fig 1, rather than the noise schedule.

We appreciate your suggestion about including a graphic overview of our model approach. In the final version, we are commited to add a figure to illustrate DiMA's architecture and workflow.

Distance metrics like FD-seq and other metrics should be introduced and defined in the main text.

We will also ensure that distance metrics like FD-seq and other evaluation metrics are properly introduced and defined in the main text to improve clarity.

---

Thank you again for your constructive feedback, which will help us improve the final manuscript.

[1] https://arxiv.org/pdf/2502.12479

We sincerely thank you for your thoughtful and detailed feedback. We appreciate the time you have taken to review our work and your positive comments about our paper. We aim to address your concerns and questions below.

W1. On technical novelty.

While diffusion models are not new, our work shows that naive application of continuous diffusion to proteins yields poor results (~60 vs 83.3 pLDDT, Table 1). Our systematic ablations quantify the impact of each component. As reviewer Y6i7 notes, we contribute by 'exploring training and inference strategies instead of directly copying parameterization from image domain'. This enables DiMA to achieve strong performance across diverse representation spaces (from 8M to 3B parameter sequence-based and multimodal encoders) and multiple tasks with a single architecture - outperforming specialized models like RFDiffusion in fold conditioning (TM-score 0.93 vs 0.48).

C1. On the meaning of "sd-10".

Thank you for noting the unclear meaning of "sd-10" in Figure 1. This refers to our noise schedule based on the formula $\alpha_t = \frac{1}{1 + d^2tan^2(\frac{\pi t}{2})}$, where d=10 determines the rate of information decay. We will rename this to "tan-10" in the revised manuscript for better clarity.

C2. On padding omitting.

The "w/o length sampling" entry in Table 1 corresponds to the model trained with padding and without length sampling. We will clarify this connection in the revision.

C3. On sequence length determination.

For fair comparison across all methods, we apply the same length sampling strategy during inference: sequence lengths are sampled from the empirical distribution of the training set.

The difference noted on line 236 refers to protein domains (distinct structural/functional units within proteins), not overall sequence length. While we control overall sequence length distribution across methods, domain length distributions emerge naturally from each model's generation process. Our analysis shows DiMA generates domain length distributions closely matching natural proteins, whereas DPLM skews toward longer domains (Figure 18C).

Q1. On the necessity of continuous vs. discrete diffusion.

Thank you for this important question. While discrete diffusion may seem more intuitive for sequences, continuous representations have proven highly effective in protein domains including representation learning (ESM, ProtT5), structure prediction (AlphaFold, ESMFold), and backbone generation. Recent work (CHEAP [1]) confirms continuous representations capture richer protein features.

Continuous diffusion offers several advantages:

1. Direct application of established score-based techniques like classifier and classifier-free guidance without requiring discrete approximations
2. Seamless integration with multimodal representations (CHEAP, SaProt) that jointly capture sequence and structure
3. More stable and efficient training compared to discrete spaces
4. Fine-grained optimization of diffusion parameters, as shown in our ablation studies

Our experiments demonstrate that this approach not only achieves strong performance but also enables structure-aware generation and fold-conditioning that are challenging in purely discrete frameworks.

We believe both approaches have merits, and our work contributes to understanding how continuous diffusion can be effectively applied to protein design. We will expand this discussion in our final manuscript.

Q2. On sequence length preference.

Thank you for this question. We conducted additional experiments and our results show that DiMA maintains consistent performance across all sequences length ranges, with quality metrics closely tracking the natural protein distribution (figure 1 at https://tinyurl.com/icml25re). DiMA achieves more stable diversity than DPLM-3B, which shows a drop in diversity for longer sequences.

Q3. On control over diversity and novelty.

Thank you for raising this point on balancing quality, diversity, and novelty. Our framework offers two knobs to control this trade-off without changing the encoder architecture (figures 2, 3 at https://tinyurl.com/icml25re):

1. The number of sampling steps. Reducing the number of steps increases diversity and novelty and maintains reasonable quality.
2. The self-conditioning rate parameter (w), which controls how much the model relies on its previous predictions during sampling. Lower values of w (0.6-0.8) yield higher diversity and novelty with a modest quality trade-off.

Q4. On the code release.

We are currently preparing our codebase for public release and plan to make it available upon publication.

---

We are grateful for your detailed feedback and will incorporate your suggestions in the final version of our manuscript. Please let us know if any questions still need clarification.

[1] https://www.biorxiv.org/content/10.1101/2024.08.06.606920v2