DPLM-2: A Multimodal Diffusion Protein Language Model

Thank you so much for your support of our work and for so many insightful and thoughtful comments! We would like to address your concern and questions point-by-point below. Your insightful comments and suggestions have greatly improved our manuscript. Please have a check and we are happy to address any further feedback!

I have carefully read all the responses, and I can say that these responses have addressed all my concerns. This is an excellent paper, and I believe it should be accepted by ICLR. Therefore, I have raised my score to 8.

Q9: Regarding the structural information: Similar to ESM-3, you also use discrete structure tokens to represent structural information. This approach offers scalability benefits for the model, yet it also limits the model’s ability to capture finer structural details accurately. This trade-off represents an important challenge in the current field of multimodal protein language models. I am interested to know whether you would consider any methods to mitigate this issue when applying DPLM-2 to more structure-related tasks, as you described in Section 5.

A9: This is a great question, and we fully agree that it represents a crucial challenge in token-based multimodal protein language models. albeit being the key enabler to multimodal PLMs, Structure tokenization is essentially clustering similar local structured environments, which results in lossy compression and the absence of fine-grained structural variation. We are very aware of this issue. The primary principle of the solution is that we need to "recover" and preserve the high-frequency variation that gets lost during quantization. We propose some potential directions for mitigation:

• Separate structure encodings for DPLM-2. We can introduce different structure encoders for encoding and generation purposes, respectively. For parts of a protein where atomic coordinates are already provided, lossy tokenization may not be necessary. Instead, we can use robust features from powerful structure encoders like GVP-Transformer while continuing to use structure tokens for generating the remaining parts. To achieve this, the model can be trained to alternate between these two types of encodings. A similar approach has been applied successfully in recent vision-language MLLMs [1], as the vision-language community has also recognized that understanding and generation often require different types of representations.
• Modeling continuous structure features with hybrid tokenization. In the structure tokenizer, the vector quantizer module converts encoder features into discrete structure token features, but the residuals—differences between the original and quantized features—are lost, removing fine-grained structural details. To address this, we can using continuous generative modeling, such as diffusion/flow-based models, to learn to recover these residuals. This would work by conditioning on the structure tokens and possibly the final hidden states of DPLM-2. The protein structure generation process would involve first generating discrete structure tokens that capture the overall topology, then using those tokens to generate the missing residuals. These residuals would be added up to the structure token embeddings to recover a more complete and accurate structure representation, closer to the features produced by the structure encoder. This approach could significantly improve structure generation. By combining this idea with hybrid structure encodings, DPLM-2 could not only interpret given structures at atomic accuracy but also generate structures that include the missing fine-grained variations. Similar strategies have shown significant success in visual autoregressive generation with visual tokenizers [2].

[1] Janus: Decoupling visual encoding for unified multimodal understanding and generation. Arxiv 2024.

[2] HART: Efficient Visual Generation with Hybrid Autoregressive Transformer

Q8: Regarding catastrophic forgetting: I would like to confirm whether DPLM-2 in Table 7 still utilizes structural pre-training, given that neither of these models used large-scale sequence pre-training. If structural pre-training was indeed applied to DPLM-2, then there are two differences between DPLM and DPLM-2: one is the difference in model architecture, and the other is that DPLM-2 was pre-trained on structural data. With these two variables present, how can we determine that the performance decline in DPLM-2 is due to catastrophic forgetting?

A8: The DPLM-2 in Table 7 does not utilize structural pretraining. We conduct this experiment because we find DPLM-2 demonstrates a significantly larger performance gap compared with the DPLM, which is used for parameter initialization for DPLM-2. As such, we think this task can be a good testbed for us to figure out the factors causing this degradation. We hypothesize two potential causes for the observed degradation:

1. DPLM-2 needs to accommodate additional structural representations given the same model capacity (parameters), which could negatively impact the representation learning performance.
2. As continuous training on smaller magnitude of structure data, DPLM-2 may experience catastrophic forgetting of the representation power gained during DPLM's large-scale sequence pretraining.

To explore (1), we eliminated pretraining factors by retraining both DPLM and DPLM-2 with random initialization on the SwissProt and PDB datasets for 100K training steps. Additionally, we evaluated performance across all three tasks (HumanPPI, MetalIonBinding & DeepLoc) where DPLM-2 underperformed compared to DPLM. As shown in the table below, when large-scale sequence pretraining is removed, DPLM-2 significantly outperforms DPLM (exp 2 vs exp 1). This indicates that incorporating structural information enhances performance rather than harming it, which rejects the hypothesis (1).

However, when DPLM undergoes large-scale pretraining and DPLM-2 is subsequently trained from the pretrained DPLM, the performance of DPLM-2 on certain tasks diminishes (exp 4 vs exp 3). Given the relatively smaller structure data for DPLM-2 training, this suggests that catastrophic forgetting occurs during DPLM-2's multimodal training, reducing the advantages of large-scale pretraining. To verify and mitigate this, during the course of rebuttal of last week, we have curated additional 1.3M predicted structures from AFDB_rep [1], and trained DPLM-2 on this larger data. The experimental results show that the amount of structure data is indeed a key factor for better multimodal protein representations, leading to significantly improved performance over the original data (exp 5 vs exp 4). In particular, on HumanPPI, enlarging data from 200K to 1.5M helps DPLM-2 attain 2.3% improvement, and also outperforms SaProt, a strong multimodal PLM trained with 40M foldseek tokenized AFDB data.

[1] Clustering predicted structures at the scale of the known protein universe. Nature 2023

Q7: Regarding VQ-VAE: I am curious about the vocabulary usage rate of VQ-VAE when the vocabulary size is 1024. I noticed a significant performance difference between VQ-VAE and LFQ in this case, so I wonder if the vocabulary usage rate might be a contributing factor to this issue.

A7: Thank you for your suggestion. Following your advice, we calculated the codebook utilization, as shown in the table below. We observed that LFQ-based tokenizers consistently achieve nearly 100% codebook utilization, with more evenly distributed code usage, whereas the vanilla VQ-VAE suffers from codebook collapse (63.5%). This suggests that severe codebook collapse limits the vanilla VQ-VAE’s ability to learn meaningful vocabulary, at least in our implementation. We are also aware of the successful use of vanilla VQ-VAE in ESM-3, so there might be significant optimization efforts that matter. Nevertheless, the combination of pre-trained structure encoder (GVP-Transformer) + LFQ + AF2-style structure decoder has been found to be an effective and efficient approach, requiring minimal twists for robust development.

Q6: Regarding repetition: In practice, we observed that diffusion protein language models tend to generate a large number of repetitive amino acids. I noticed that DPLM employed a resampling scheme to address this issue. Therefore, regarding the DPLM-2 model, I would like to know whether it also encounters similar repetition issues. If so, what approach did you adopt to resolve it? If not, I would like to know what prevented this issue in DPLM-2.

A6: This is a great question.

About repetition issues in DPLM. Generative models, especially language models, are good at learning what you provide with them. DPLM trained with UniRef50 contains a great number of sequences with high proportion repetition patterns. As a consequence, maximum-a-posterior (MAP) based sampling methods, such as default mask-predict sampling method for masked discrete diffusion, are likely to firstly recover these high likelihood patterns, and then get stuck in this local optimal and turn out to be low quality generation. This phenomenon, i.e., high likelihood != high generation quality, has been widely noticed in the studies of neural text/sequence generation [1,2]. In the original DPLM, authors address this issue with resampling for repetition patterns given a threshold.

About repetition issues in DPLM-2. In DPLM-2, we train the model on pairs of structure tokens and amino acid tokens. As structure vocabulary is much larger than amino acid and structure data is much more complicated, there are basically no such repetitive structure tokens in the resulting structure tokens training data. Meanwhile, PDB and swissprot data are curated much more carefully than UniRef and AFDB in general (predicted structure from unannotated sequences, incl. UniRef), also resulting in less amino acid repetition in our training data. As a result, DPLM-2 just samples "good-looking" proteins (while amino acid repetitions often lead to long helical or disordered loopy proteins) using straightforward sampling strategy with no need for resampling. We can also further enhance sampling diversity in general using temperature annealing sampling, starting with high temperature (e.g., 2.2) for high stochacity and gradually decreasing towards low temperature (e.g., 0.1) for high likelihood and fidelity, similar to langevin dynamics or simulated annealing.

[1] Is MAP Decoding All You Need? The Inadequacy of the Mode in Neural Machine Translation. Coling 2020

[2[ The Curious Case of Neural Text Degeneration. ICLR 2020.

Q5: Regarding the inverse folding task: Comparing the performance of ESM3 and DPLM-2 in Table 5 reveals a pattern where DPLM-2 often exhibits a higher AAR while frequently showing a lower scTM. Could this possibly be due to biases in the protein structure prediction model (ESMFold) or because the sequences generated by DPLM-2, despite having higher ARR, may deviate further from the target protein’s data distribution?

A5: Thanks for your insightful question!

About scTM: Although the actual gap between DPLM-2 (650M) and ESM-3 is not that large (0.88 vs 0.90 on cameo 2022), we suggest that this difference can be attributed to the lossy structure encoding caused by discrete tokenization. ESM3 introduces an important geometric attention module to encode atomic coordinate of input structure when available, while DPLM-2 currently relies on pure structure tokens. This also leads to the discussion of trade-offs and further directions of structure modeling in multimodal PLM, which exactly corresponds to your last question. Please take a look at our elaborate discussion on this in Q9.

About high AAR: This mainly arises from the sampling strategy. Here we ablated the sampling strategy to study its impact. By default, we followed DPLM's approach, using argmax decoding, which selects the token with the highest probability at each timestep. This method generates sequences with high probabilities, resulting in strong amino acid recovery (AAR). In contrast, we introduced a sampling strategy with annealing temperature ranging from 2.2 to 0.1 to enhance sequence diversity. While this approach lowers AAR, it maintains the same scTM score as argmax decoding. This demonstrates that the temperature annealing strategy generates more diverse sequences. Although these sequences are less similar to the ground truth, they still satisfy the structural requirements, highlighting the trade-off between diversity and similarity to the target sequence.

The pLDDT metric reflects the structural plausibility of sequence. In addition to the pLDDT metric, we also include sequence-related metric that we calculate sequence diversity using mmseq2 clustering at sequence identy = 0.5 for good quality samples with pLDDT > 70. This quality threshold for diversity is inspired by Multiflow, which is more informative by avoiding diverse but messy sequences. We follow the experimental setting in DiMA [1], generating 2048 sequences with length sampled from the length distribution of PDB + SwissProt. The results highlight that DPLM-2 is able to generate structurally plausible and diverse sequences for protein generation. We also find that training data distillation greatly helps Multiflow's sequence quality in terms of pLDDT and diversity.

<End of Q4 >

Q4: Regarding the pLDDT metric: Table 2 shows that MultiFlow (retrained on our training data) and DPLM-2 (co-generation) perform similarly on scTM and scRMSD metrics but diverge significantly on the pLDDT metric. A similar trend is observed when comparing Figures 3A and 3C. In Figure 3A, scTM and scRMSD metrics remain relatively stable as sequence length increases, with scRMSD even increasing for sequences of length 500. However, Figure 3C shows a notable upward trend in pLDDT as sequence length increases (from 100 to 300), without a significant drop in pLDDT even for sequences of length 500. These phenomena indicate that the scTM and scRMSD metrics convey somewhat contradictory information to the pLDDT metric. Generally, a significant drop in pLDDT should correspond with a decrease in scTM and an increase in scRMSD. My question here is whether this inconsistency might indicate that the pLDDT metric is unreliable. I raise this because, in previous experiments, we observed that pLDDT, being predicted by a neural network (e.g., ESM-2) trained on natural protein datasets, sometimes fails to handle model-generated proteins that deviate significantly or exhibit severe irregularities, leading to inflated pLDDT values. This phenomenon is quite common in unconditional generation. Therefore, I believe using a more diverse set of metrics is essential. However, you did not provide evaluation results for MultiFlow (retrained on our training data) on the Novelty and Diversity dimensions, which makes it harder to assess the model’s performance on this part of the task.

A4: Thanks for pointing this out! After checking our code, we find that the results of MultiFlow (retrained on our training data) were obtained incorrectly. This occurred because we leveraged default MultiFlow sampling configuration that sampling proteins in lengths of [70, 100, 200, 300], inconsistent with the experimental settings in our paper that sample in lengths of [100, 200, 300, 400, 500]. We apologize for this mistake and we provide a updated results in the Table below with corrected scTM and scRMSD, where MultiFLow (retrained on our training data) achieves significantly lower scTM and scRMSD results compared with DPLM-2, aligning with the performance gap seen in pLDDT. We believe this is because MultiFlow's weaker ability to generate long proteins, resulting in a significant decline in designability when generating proteins of lengths 400 and 500.

Q3: Regarding structure tokens: Figure 2B shows that most structure tokens are concentrated in alpha helix and beta sheet regions, with lower density in loop regions. However, in line 409 of the paper, you state that "loops are highly flexible." Could this pose a potential issue? More variable regions typically contain richer and more diverse information, which may require more structure tokens to model this diversity effectively. Yet, your experimental results do not align with this. Moreover, your results in Figure 4B also indicate that the model’s designability is notably worse in regions with a higher proportion of loops. Could one possible reason for this issue be that the structure tokens in loop regions are not well-learned, resulting in weaker modeling capability?

A3: Thanks for this great question! This leads to the discussion on how to explain different in silico designablity metrics (scTM & scRMSD) and what information the structure tokens learn. For the former, we have provided our thoughts on this in the previous question and we empirically notice that scTM for the samples indeed drops when proportion of loops become higher but still keep greater than 0.8 overall.

About what information the structure tokens learn. For interpratability, we update a more informative simplex plot of struct token vs second struct in Fig. 2B. As you are pointing out, we can observe a strong correlation between a vast majority of the structure tokens and structured local environments, where a lot of structured tokens concentrate on the alpha helix and beta sheet vertices, while some tokens lie between regions or the loop vertice. There are also a subset of struct tokens having less clear clues to specific secondary structures. This suggests that structure tokens mostly capture clear secondary elements, some may correspond to structured local environments (in bewteen helic and sheet), while others could be high-level abstract entities or just not well-learned entries. On one hand, we agree that "more variable regions contain richer and diverse information" as these regions are of high entropy. On the other hand, the nature of lossy and clustering-like vector-quantization methods is highly likely to eliminate such high-frequency high-entropy structural variations, and only keep low-frequency, low-entropy content. We suggest that this could be the major reason for the not-well learned flexible regions. And this also exactly corresponds to your last question about the trade-offs and further directions of structure modeling in multimodal PLM. Please take a look at our elaborate discussion on this in Q9.

Q2: Regarding scTM and scRMSD metrics: In Table 2, under the unconditional backbone generation task, DPLM-2 achieves the highest scTM metric while also showing a significantly worse scRMSD metric, both in terms of mean and variance, which are at high levels. Is this phenomenon somewhat unusual? Given that these two metrics should have a certain level of correlation.

A2: Thank you for pointing this out. We would like to address your concern as follows: The pLDDT metric reflects the structural plausibility of sequence. In addition to the pLDDT metric, we also include sequence-related metric that we calculate sequence diversity using mmseq2 clustering at sequence identy = 0.5 for good quality samples with pLDDT > 70. This quality threshold for diversity is inspired by Multiflow, which is more informative by avoiding diverse but messy sequences. We follow the experimental setting in DiMA [1], generating 2048 sequences with length sampled from the length distribution of PDB + SwissProt. The results highlight that DPLM-2 is able to generate structurally plausible and diverse sequences for protein generation. We also find that training data distillation greatly helps Multiflow's sequence quality in terms of pLDDT and diversity.

We first need to highlight that the generated samples from DPLM-2 share similar scTM (0.925) and scRMSD (3.9) as native PDB samples, which also exhibit good scTM (0.904) with a little bit higher scRMSD (4.623). Moreover, Additionally, DPLM-2 maintains a balanced structural composition (helix: 0.4, strand: 0.2, coil: 0.45), closely resembling natural distributions. In contrast, for MultiFlow, the officially released model with distillation attains much lower scRMSD (3.2), while the performance of our retrained version (on the same DPLM-2 training set) degrades in both scTM (0.871) and scRMSD (6.58). Lower scRMSD in MultiFlow with distillation, appears to be driven by overrepresentation of structured elements (Figure 4A), i.e., significantly biasing towards proteins with more helices, with less strands and loops (also see Figure 4C). This overrepresentation drives the observed scRMSD improvement but deviates from natural protein diversity.

TM-score emphasizes global topology, while RMSD is sensitive to local structural errors. As such, although scTM and scRMSD are generally correlated, discrepancies can arise. The purpose of TM-score is to solve this sensitivity of RMSD, because RMSD is an average distance of all residue pairs in two structures, a local error (e.g. a misorientation of the tail) will raise a big RMSD value although the global topology is correct. In TM-score, however, the small distance is weighted stronger than the big distance, which makes the score insensitive to the local modeling error.

As shown in Fig 4B, some samples from DPLM-2 with higher loop proportion are more conformationally flexible, hence may show high scTM (>0.9) but worse scRMSD (>2.0), similar to natural protein. However, this does not necessarily indicate a limitation in generation quality but reflects differences in metric sensitivity.

As a result, the in-silico designability of protein generation should be evaluated comprehensively using both scTM and scRMSD, as each metric offers distinct insights and serves different purposes. For users aiming to generate samples with accurate global topology, scTM serves as a reliable indicator, whereas scRMSD may occasionally exclude reasonable structures. Conversely, for applications requiring structurally rigid and stable proteins, such as functional designs (e.g., binder design), scRMSD has been shown to correlate more strongly with in vitro success rates, as suggested by RFDiffusion.

Q1: The model in this paper is trained on a relatively small protein structure dataset, lacking exploration on larger-scale structural data. I believe that extending multimodal protein language models to larger protein sequence and structure datasets could further unlock the potential of this approach.

A1: Thanks for your suggestions! As we stated in the discussion section, we fully agree that the limited numbers of structure data may hinder DPLM-2 from unlocking its full potential. As you suggested, during the last week, we curated a 1.3 million predicted structure dataset from AlphaFoldDB, where predicted structures were representative entries obtained by MMSeq2 sequence clustering as well as Foldseek structure clustering. This data size is an order of magnitude larger than the current structure data we used, and not too much larger to help us perform an efficient proof-of-concept of whether enlarging data size can boost our model in a course of a week.

The following table summarizes our primary explorations on how the size of the data, data mixing strategy as well as training strategies affect the model performance (DPLM-2 650M).

Meanwhile, we find that training on the larger-scale structural data can further enhance the representation learning, please refer to Q8 for more details.

Q1: Missing Details: Key details are missing in certain sections. For instance:

Q1.1: A core contribution is the DPLM-2 model itself, but essential details, such as how sequence and structure tokens are combined and how the structure tokenizer is trained, are not included in the main text. Moving these details from the appendix to the primary text, with clear explanations including the frozen pretrained structure encoder and the sequence prediction head on top of the structure decoder, would significantly improve clarity.

A1.1: Thanks for your valuable suggestions on clarity issues! We have updated our paper by moving the essential details into the main text, and alternating ablation study (section 4.1.3) to the appendix.

Q1.2: The distinct noise level for each modality could be better explained, as this aspect is currently underdeveloped.

A1.2: Thank you for pointing this out. We introduce a distinct scheduler to control the noise level of structure and sequence flexibly during training (section 3.1 in our paper). Different combinations of structures and sequence schedulers (denoted as $t_z$ and $t_s$, respectively) imply training for different applications. Specifically, we mainly focus on: (1) sequence-conditioned structure generation (e.g., folding), (2) structure-conditioned sequence generation (e.g., inverse-folding), (3) sequence generation, (4) structure generation, (5) structure-sequence co-generation.

For conditional generation tasks (e.g., folding and inverse-folding), we set the noise scheduler of the conditioned modality to 0, e.g., no noise in the conditioned modality. Specifically, in the folding task, the $t_s$ is always set to $0$, while in the inverse-folding tasks the $t_z$ is always set to $0$.

In the structure-sequence co-generation task, we keep the $t_z$ and $t_s$ for the same, enhancing the structure-sequence consistency in co-generation. The structure or sequence generation tasks do not depend on another modality, so we set the noise scheduler of another modality to $T$, which is the maximum timestep and means 100% noise in another modality. For example, in structure generation task, the $t_s$ is always set to $T$.

During training, we jointly train the above 5 tasks simultaneously. We divide the training data in a batch into 5 parts according to a preset proportion, and each part is used for a specific task training. In our experiment, the proportion for each task is the same, which is 20%.

Q1.3: In Section 4.2, the authors mention that performance can improve with supervised fine-tuning using a folding objective; however, the paper lacks details on this fine-tuning process.

A1.3: According to the above question, we can further enhance a specific generation task by supervised finetuning (SFT). This involves continuing training for the specific task with a proportion of 100%, while the proportion for other tasks is set to 0%. For example, in Tab. 4, the folding supervised finetuning is performed by continuous training from a pre-trained DPLM-2 with 100% proportion of folding objective using the same training data for additional 50K steps with a constant and smaller learning rate (5e-5 vs 1e-5 for joint pre-training).

Q8: In Section 3.3, you state, “Recent efforts have applied this approach to protein structure coordinates (Van Kempen et al., 2024; Liu et al., 2023; Gao et al., 2024; Lu et al., 2024). This allows language models to better learn the composition of local structural elements. However, how to learn an effective structure tokenizer remains an active research question.” Given the similarity of your method to prior approaches, particularly with the use of LFQ as in Gao et al., 2024, could you elaborate on how your method contributes to addressing this active research question?

A8: Thank you for your question. Our approach addresses the active research question of structure tokenization through a simpler and more practical design compared to Gao et al., 2024:

Simplicity of design and training: We use a strong pretrained structure encoder (from ESM-IF) with an AF2-style decoder (triangular modules + IPA) trained using FAPE loss. There are also no unnecessary modifications to LFQ. In contrast, Gao et al introduced very complicated modifications to LFQ (SoftLFQ). This makes our method fairly easy and straightforward to implement, as well as effective and efficient for training.

Performance: Despite its simplicity, our tokenizer performs competitively in structure reconstruction (recRMSD), demonstrating its ability to effectively capture structural features. Higher reconstruction does not necessarily lead to better generation quality, which is widely observed in VQ-VAE literature while the LFQ paper [2] also clearly elaborated on this. We in this paper have shown that our tokenizer can actually support various generation tasks with decent generation performance, while Gao et al., (2024) mainly focused on evaluation of reconstruction. For generations, they only assessed antibody CDR infilling, which is fairly easy, and general protein generation ability of their tokenizer remians unclear.

[1] Gao et al.: FoldToken: Learning Protein Language via Vector Quantization and Beyond. Arxiv 2024

[2] Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation. ICLR 2024

Q5: While DPLM-2 enables flexible generation, what are the trade-offs in structural invariance when using structure tokens instead of 3D coordinates?

A5: Thank you for your insightful question.

Our approach can be seen as decoupling the learning of the structurally invariant topology to the tokenizer and the language model, and the geometric reasoning on 3D coordinates to the structure decoder parameterized by triangular modules and IPAs from AlphaFold2 with FAPE loss. This shares the similarity to AF2 and ESMfold, where the sequence encoding modules like Evoformer (co-evolution encoding) and ESM (amino-acid encoding) provide invariant features (in the form of single and pair embeddings) to the structure decoder that learns to convert invariant features into 3D coordinates. The AF2-style structure decoder does not enforce strict equivariance to rigid transformations. Instead, it relies on the FAPE loss to ensure structural consistency, which minimizes coordinate errors in a manner that is invariant to global rotations and translations.

As such, we suggest that the primary trade-off when using invariant structure tokens instead of 3D coordinates mainly lies in the potential loss of fine-grained structural details. Structure tokens cluster similar local structures into discrete representations, which inherently introduce quantization errors. This trade-off, on one hand, enables efficient multimodal learning and generation by simplifying the representation space, on the other hand, represents an important challenge in the current field of multimodal protein language models as suggested by the Reviewer FqCq (Q9) and in our discussion section. Future efforts should be made towards better structure modeling to mitigate this trade-off. Some potential directions include introducing separate structure encoders for structure encoding and generation purposes [1], or hybrid tokenization for recovering the lost fine-grain structural variations [2].

[1] Janus: Decoupling visual encoding for unified multimodal understanding and generation. Arxiv 2024.

[2] HART: Efficient Visual Generation with Hybrid Autoregressive Transformer.

Q4: Sequence tokens use a smaller vocabulary than structure tokens. Are the corresponding structural embeddings in DPLM-2 trained from scratch during fine-tuning?

A4: That's correct. The structural token embeddings in DPLM-2 are trained from scratch. To efficiently utilize the evolutionary information from pre-trained sequence-based DPLM, DPLM-2 uses a warm-up strategy (outlined in Section 3.2 of our paper). This approach initializes DPLM-2 with the weights of the sequence-trained DPLM.

More specifically, since DPLM's vocabulary consists only of amino acids, DPLM-2 expands this with discrete structure tokens. The embeddings for these new tokens are initialized using the mean and standard deviation of the learned amino acid embeddings. This embedding initialization keeps the distributional statistics of the embedding space consistent with the pre-trained DPLM, ensuring stable early-stage training (for learning structure-sequence alignment) and reducing the risk of extreme gradients that could cause training instability.

Q3: Mixed Results: The results on certain tasks are not fully convincing. In particular, while DPLM-2 exhibits high scTM in unconditional protein generation, it does not achieve lower scRMSD compared to other methods, indicating potential limitations in generation quality.

A3: Thank you for pointing this out. This is a great question about the evaluation metrics and their interpretation for the in silico designability of protein generation or design. We would like to address your question as follows:

We first need to highlight that the generated samples from DPLM-2 share similar scTM (0.925) and scRMSD (3.9) as native PDB samples, which also exhibit good scTM (0.904) with a little bit higher scRMSD (4.623). Moreover, Additionally, DPLM-2 maintains a balanced structural composition (helix: 0.4, strand: 0.2, coil: 0.45), closely resembling natural distributions. In contrast, for MultiFlow, the officially released model with distillation attains much lower scRMSD (3.2), while the performance of our retrained version (on the same DPLM-2 training set) degrades in both scTM (0.871) and scRMSD (6.58). Lower scRMSD in MultiFlow with distillation, appears to be driven by overrepresentation of structured elements (Figure 4A), i.e., significantly biasing towards proteins with more helices, with less strands and loops (also see Figure 4C). This overrepresentation drives the observed scRMSD improvement but deviates from natural protein diversity.

TM-score emphasizes global topology, while RMSD is sensitive to local structural errors. As such, although scTM and scRMSD are generally correlated, discrepancies can arise. The purpose of TM-score is to solve this sensitivity of RMSD, because RMSD is an average distance of all residue pairs in two structures, a local error (e.g. a misorientation of the tail) will raise a big RMSD value although the global topology is correct. In TM-score, however, the small distance is weighted stronger than the big distance, which makes the score insensitive to the local modeling error.

As shown in Fig 4B, some samples from DPLM-2 with higher loop proportion are more conformationally flexible, hence may show high scTM (>0.9) but worse scRMSD (>2.0), similar to natural protein. However, this does not necessarily indicate a limitation in generation quality but reflects differences in metric sensitivity.

As a result, the in-silico designability of protein generation should be evaluated comprehensively using both scTM and scRMSD, as each metric offers distinct insights and serves different purposes. For users aiming to generate samples with accurate global topology, scTM serves as a reliable indicator, whereas scRMSD may occasionally exclude reasonable structures. Conversely, for applications requiring structurally rigid and stable proteins, such as functional designs (e.g., binder design), scRMSD has been shown to correlate more strongly with in vitro success rates, as suggested by RFDiffusion.

Q2: Clarity and Consistency Issues: Minor inconsistencies reduce clarity, such as inconsistent bolding of best results in Tables 2 and 6. In Table 4, DPLM-2 is presented as performing well in zero-shot folding, yet its RMSD is high compared with ESM3, PVQD, and ESMFold, achieving competitive performance only after fine-tuning (SFT). Additionally, clarifying the presentation of mean and median values could help with data interpretation.

A2: Thank you for pointing out the misleading and we have changed the way of highlighting in our paper: the highest number is bold, and the second highest number is underlined. The current results of conditional generation (Table 4, Table 5) are presented as mean/median values, and we have updated our paper for clearer clarification.

About SFT: The results of PVQD and ESMFold are actually obtained by supervised finetuning on folding-only objective, so it is meaningful and fair to compare to DPLM-2 with folding SFT too.

Thank you so much for your constructive suggestions. We understand your concerns about the clarity issues and missing technical details, potential inconsisency in some result intepretation, and requiring analysis and discussion on structure tokenizer. To address these issues, we have made our major efforts in (1) elaborating on technical details and accordingly updating the manuscript for improved clarity; (2) providing discussion on the implications of in silico designablity metrics; and (3) providing more analysis and discussion on structure tokenization. Your insightful comments and suggestions have greatly improved our manuscript. We sincerely thank you once again and welcome any further feedback!

Q6: In Table 2, what is meant by DPLM-2 (seq → structure) or (structure → seq)? If these indicate modality-by-modality generation, could you clarify how this is implemented?

A6: Thanks for pointing this out, we have updated our paper to improve clarity. The co-generation can be performed in simultaneous generation (co-generation) and cascaded workflow: first generating the structure then the sequence conditioned on generated structure (struct → seq), and the reverse way (seq → struct), without the need of other folding or inverse folding models.

Q7: For the protein representation learning evaluation, it might be useful to include a broader range of baselines, such as GNN-based models, for a more comprehensive comparison.

A7: Thanks for your valuable suggestion! We have accordingly added the results of GearNet [1], as the GNN-based baseline, as you suggested. We have updated the results in Table 6 of our paper.

[1] Protein representation learning by geometric structure pretraining. ICLR 2023

Q9: The experiment on DeepLoc is not described in sufficient detail. Why you chose to use only one tiny dataset? Could you provide experiments that show the described catastrophic forgetting issue, which is of high importance?

A9: Thank you for your question. In our study, we selected the DeepLoc subcellular dataset because it highlights a pronounced performance gap between DPLM-2 and DPLM, which serves as the initialization model for DPLM-2. This dataset provides an ideal testbed for investigating the factors contributing to this performance drop.

We hypothesize two potential causes for the observed degradation:

1. DPLM-2 needs to accommodate additional structural representations given the same model capacity (parameters), which could negatively impact the representation learning performance.
2. As continuous training on smaller magnitude of structure data, DPLM-2 may experience catastrophic forgetting of the representation power gained during DPLM's large-scale sequence pretraining.

To explore (1), we eliminated pretraining factors by retraining both DPLM and DPLM-2 with random initialization on the SwissProt and PDB datasets for 100K training steps. Additionally, we evaluated performance across all three tasks (HumanPPI, MetalIonBinding & DeepLoc) where DPLM-2 underperformed compared to DPLM. As shown in the table below, when large-scale sequence pretraining is removed, DPLM-2 significantly outperforms DPLM (exp 2 vs exp 1). This indicates that incorporating structural information enhances performance rather than harming it, which rejects the hypothesis (1).

However, when DPLM undergoes large-scale pretraining and DPLM-2 is subsequently trained from the pretrained DPLM, the performance of DPLM-2 on certain tasks diminishes (exp 4 vs exp 3). Given the relatively smaller structure data for DPLM-2 training, this suggests that catastrophic forgetting occurs during DPLM-2's multimodal training, reducing the advantages of large-scale pretraining. To verify and mitigate this, during the course of rebuttal of last week, we have curated additional 1.3M predicted structures from AFDB_rep [1], and trained DPLM-2 on this larger data. The experimental results show that the amount of structure data is indeed a key factor for better multimodal protein representations, leading to significantly improved performance over the original data (exp 5 vs exp 4). In particular, on HumanPPI, enlarging data from 200K to 1.5M helps DPLM-2 attain 2.3% improvement, and also outperforms SaProt, a strong multimodal PLM trained with 40M foldseek tokenized AFDB data.

[1] Clustering predicted structures at the scale of the known protein universe. Nature 2023

Q6: The paper lacks analysis on the trained structural tokens. Additional exploration of the interpratability of the tokens themselves and untilization of the codebooks would greatly benefit the paper. Do the tokens correspond to some local environments in structure, or are they just abstract entities?

A6: Thanks for your suggestion. In the Figure 2 of the original submission, we have shown the reconstruction accuracy and an interpretation analysis on the correspondence of structure tokens and local structural elements in terms of secondary structures. As you suggested, we calculate the codebook utilization in the following table. We find that LFQ-based tokenizers always achieve nearly 100% codebook utilization with more evenly distributed code usage, while vanilla VQ-VAE struggles with codebook collapse.

For interpretability, we also update a more informative simplex plot of struct token vs second struct in Fig. 2B. We can observe a strong correlation between a vast majority of the structure tokens and structured local environments, where a lot of structured tokens concentrate on the alpha helix and beta sheet vertices, while some tokens lie between regions or the loop vertice. There are also a subset of struct tokens having less clear clues to specific secondary structures. This suggests that structure tokens mostly capture clear secondary elements, some may correspond to structured local environments (in bewteen helic and sheet), while others could be high-level abstract entities.

Q7: The model does not learn the distribution of protein lengths. Have you tried to overcome this limitation?

A7: Thanks for your question. In our paper, our primary purpose is to conduct fair comparisons with previous models under the similar settings to better assess the strengths and limitations of our models, hence follow MultiFlow in sampling within length intervals. Meanwhile, in many protein design applications users have prior knowledge of the target lengths or the ranges of lengths, or indeed require explicit length control and manipulation. As such, we may not need to directly learn the length distribution.

DPLM-2 is capable of generating proteins from the empirical length distribution. Specifically, we sample 2048 sequences with length sampled from the length distribution of PDB + SwissProt. The table below demonstrates that DPLM-2 can generate highly plausible proteins, which is consistent with sampling with length intervals.

Q8: The ablation results presented in 4.1.3 and table 3 are controversial. Could you please clarify the procedure, how many samples was used for evaluation and so on?

A8: Thanks for pointing this out. In ablation study, we evaluate the effects of sequence pre-training and data augmentation on unconditional protein generation. Specifically, We investigate the effect of sequence pre-training by randomly initializing DPLM-2 instead of using DPLM parameters, while for effect of predicted structures we leverage only PDB structures for training. We conduct experiments on a 150M parameter DPLM-2. For each DPLM-2 variant, we sample 100 examples for each length in 100, 200, 300, 400 and 500. We evaluate designability by scTM and diversity by the number of difference clusters for each length. The results of Table 3 in our paper demonstrate the effectiveness of sequence pre-training and data augmentation.

Q2: Authors compare DPLM and DPLM-2 only with EvoDiff on unconditional sequence generation. Adding more baselines of different architectures (AR transformers, CARP, continuous diffusion, flow-matching, etc.) would greatly improve the work. The model is evaluated only structurally, but it is langugae model after all. Using sequence-based evaluation metrics, including sequence clustering for the diversity should benefit the soundness of the work.

A2: Thanks for your valuable suggestion on adding more baselines and sequence clustering for sequence diversity. We have added more baselines with various architectures on unconditional sequence generation, including autoregressive language models (Progen2), CNNs (CARP), continuous diffusion on ESM2 embedding space (DiMA) and flow matching (MultiFlow). We also include sequence diversity using mmseq2 clustering at sequence identy = 0.5 for good quality samples with pLDDT > 70. This quality threshold for diversity is inspired by Multiflow, which is more informative by avoiding diverse but messy sequences. We follow the experimental setting in DiMA [2], generating 2048 sequences with length sampled from the length distribution of PDB + SwissProt. The results highlight that DPLM-2 is able to generate structurally plausible and diverse sequences for protein generation. We also find that training data distillation greatly helps Multiflow's sequence quality in terms of pLDDT and diversity.

[1] Progen2: Exploring the Boundaries of Protein Language Models. 2022

[2] CARP: Convolutions are competitive with transformers for protein sequence pretraining. Cell Systems 2024

[3] DiMA: Diffusion on language model embeddings for protein sequence generation. Arxiv 2024

Q3: There are some issues with model evaluation. It is not clear, why the Authors use pdb-TM (claculate TM-score against PDB), if the model is trained also on the synthetic data. Third, the designability metric used in this work evaluates the consistensy between the generated structure and the prediction of ESMFold on the generated sequence. It does not measure protein "quality".

A3: Thank for your suggestions!

About "pdb-TM". We use "pdb-TM" to assess the novelty of proteins generated by DPLM-2 against natural proteins. As you suggested, it is indeed more reasonable to have TM-scores against swissprot and pdb structures (swissprot & pdb TM). As shown in the table, the "SwissProt & PDB TM" scores for DPLM-2 (650M) align closely with the "pdb-TM" results, consistent with the findings presented in our paper. These results from unconditional sampling indicate that DPLM-2, as a generative model of protein structure and sequence, learns the training data distribution and serves as a foundation for various potential conditional purposes. Finally, we will change to "avg. max-TM" in our revised version as used in Flowflow-2 [1] for a more accurate terminology.

About "quality". Thank you for pointing this out! We initially followed FrameDiff and Multiflow to term "quality" for self-consistency based designability for generated structures. To avoid any ambiguity, we will rephrase our wording by rephrasing "quality" to "designability".

[1] Sequence-Augmented SE(3)-Flow Matching For Conditional Protein Backbone Generation. NeurIPS 2024

Thank you so much for your constructive suggestions. We understand your concerns about the potential risks of reduced sampling diversity and mode collapse, the need for more evaluation metrics, results, and explanations for a comprehensive assessment, and the theoretical guarantee of structure-sequence consistency under the conditional independence assumption in diffusion models. To address these issues, we have made our major efforts in (1) providing a discussion on diversity metrics with updated results and an ablation study on how the training strategies influence generation diversity; (2) including sequence-based metrics for unconditional generation, conducting sampling from length distribution of natural proteins, analysis on structure tokens and discussion of sampling strategies for conditional generation (3) providing discussions in theoretical aspect on the conditional independence assumption in diffusion models, including DPLM-2 and beyond. Your insightful comments and suggestions have greatly improved our manuscript. We sincerely thank you once again and welcome any further feedback!

Q1: The DPLM-2 model is trained starting from the weights of DPLM model. As evident from e.g. Table 2 (poor diversity) and Table 5 (high AAR), DPLM suffers from mode collapse. Training DPLM-2 starting from such a checkpoint leads to severe mode collapse in DPLM-2. Generating the same result over and over again makes a generative model useless. Training both models from scratch rather than finetuning might help and provide the insight on how adding structural information to the DPLM architecture affects the generation ability. Even better would be a thourough ablation study of the training procedure. The dataset preparation included random cropping. It is not clear if it has a detrimental effect on the model behaviour.

A1: Thanks for this valuable question. We understand you are concerned about the sampling diversity and the risk of mode collapse of DPLM-2, and wondering how specific training strategies (e.g., multimodal finetuning from pre-trained DPLM) relate to the model behaviors.

Regarding sampling diversity and risk of mode collapse.

We thank the reviewer for bringing this to our attention. We notice that there are inconsistent results of diversity in Table 2 between avg inner-TM (averaged pairwise TM-score among samples) and MaxCluster (number of distinct structure clusters), which do not make sense. For example, DPLM-2 (co-generation), attains a MaxCluster=0.545 (i.e., there are 54.5 distinct fold clusters out of 100 samples determined by Foldseek with TM-score threshold = 0.5), which is better than the MaxCluster=0.500 for the official Multiflow (w/ distillation), whereas the avg inner-TM is unreasonably high (0.703 vs 0.468 for MultiFlow), which contradicts with the value of MaxCluster and the visualized observation we perceive from the predicted structures.

We have accordingly checked our code and found out that our implementation of avg inner-TM was mistakenly incorrect. Here we provide a updated results in the Table below with corrected avg inner-TM, where DPLM-2 gets reasonable avg inner-TM similar to native PDB samples (with the same length intervals) and lower than MuliFlow and RFDiffusion. These results empirically indicate that DPLM-2 does not struggle with mode collapse issue during sampling. Plus, we can further improve sampling diversity by using temperature-annealed sampling with higher initial temperature.

Conditional independence in multimodal diffusion models.

Multimodal diffusion models aim to accommodate two or more modalities using a unified models. In this case, conditional independence between modalities is generally made $p\_\theta(\mathbf{x}\_{t-1}, \mathbf{y}\_{t-1} | \mathbf{x}\_t, \mathbf{y}\_{t}) = \textstyle\prod\_i p\_\theta(x\_{t-1, [i]} | \mathbf{x}\_t, \mathbf{y}\_{t-1}) \textstyle\prod\_j p\_\theta(y\_{t-1, [j]} | \mathbf{x}\_t, \mathbf{y}\_{t})$. For instance, UniDiffuser [4] is a multimodal continuous diffusion model that handles text and image modalities independently at each timestep, conditioned on the predictions from the previous timestep. Multiflow [3], on the other hand, factorizes protein data into three modalities—translation, orientation, and amino acid type—assuming conditional independence. It establishes a multimodal diffusion/flow-based model by combining three types of stochastic processes over Euclidean, SO(3), and categorical spaces for these modalities. In DPLM-2, we adopt a unified discrete diffusion approach where structure tokens and amino acid tokens are treated as conditionally independent. While theoretical guarantees for the convergence of mixture diffusion processes are still under-explored, existing discrete diffusion theory [2,3] ensures that a well-trained DPLM-2 can converge to the tokenized structure-sequence data distribution, supporting consistency between structure and sequence tokens.

Additionally, theoretical studies on non-autoregressive Transformers (NATs) for text generation, which are akin to masked discrete diffusion, indicate that the learning difficulty of such models can be evaluated through conditional total correlation, a dataset-dependent and model-free measure captures the discrepancy between a joint data distribution and a fully factorized distribution under conditional independence [6]. These studies suggest that simplifying the original complex target data (e.g., by using instance-level knowledge distillation from other models as in NATs for text generation or in MulitFlow for sequence generation, or by using tokenized structure instead of 3D coordinates as in DPLM-2/ESM3), reduces conditional total correlation, thereby enhancing both learning and generation quality.

Given the consistency of structure tokens and amino acid can be ensured to learn in DPLM-2 by previous theoretical results [2,3,6], the overall structure and sequence consistency can be achieved with a decent structure tokenizer, such as the one proposed in this paper, which can map structure tokens to their atomic coordinates with a sufficient accuracy.

[1] Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. ICLR 2023

[2] Convergence analysis of discrete diffusion model: Exact implementation through uniformization. Arxiv 2024

[3] Convergence of Score-Based Discrete Diffusion Models: A Discrete-Time Analysis. Arixv 2024

[4] UniDiffuser: One Transformer Fits All Distributions in Multi-Modal Diffusion at Scale. ICML 2023

[5] Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Design. ICML 2024.

[6] On the Learning of Non-Autoregressive Transformers. ICML 2021

---end of Q10---

Dear Authors,

Thank you for the detailed response to my initial review. Some of my concerns have been addressed. However, several important points require further attention:

1. I appreciate the ablations. It would be valuable for the community if the analysis of the results of the ablation study end up in the final version of the manuscript.

2. DPLM-2 is a finetuned protein language model. So its evaluation should contain a great deal of sequence analysis. The authors provide only pLDDT, which is primarily a structural quality measure, and MMseqs2 clustering. MMseqs2 clustering (and also structural clustering for that matter) should be performed at different thresholds to capture different aspects of the diversity. Adding perplexity as a measure of naturalness/quality and novelty through sequence identity to nearest neighbor in the dataset to the evaluation toolkit would make analysis more sound.

3. In the presented setting, sc-TM and sc-RMSD evaluate the degree of consistency between the predictions in two modalities, sequence and structure. Both "quality" and "designability" are misnomers.

4. I appreciate the results on the novelty evaluation using TM-score on the full training set, not only PDB, which is a subset thereof. The results show that the model generates more structures similar to those in the training data, however, in this case reporting mean and std is insufficient as it may mask issues with memorization. It is advised to plot the distribution.

5. Regarding argmax vs. stochastic sampling. This important topic requires a thorough discussion, since different strategies are used for different tasks. In DPLM, which DPLM-2 is based upon, Gumbel-Max trick was used (at least) in unconditional generation to alleviate mode collapse. I have not found any mentions of this in the manuscript. Here, argmax sampling is used for inverse folding, while stochastic sampling is used for unconditional generation and scaffolding. Why different tasks use different sampling approaches? What stochastic approach is used, please provide the details? What is the reasoning behind this and how is it supported experimentally?

6. Even with added baselines, important comparisons are missing. There is a lot of active research on inverse folding, so there is no lack of good baselines and good methodology ([1-4] to name a few). At the end of the day, it is not about beating SOTA result, but about providing objective comparison for the benefit of the protein design community.

7. Thank you for adding the codebook utilization metrics. However, they are purely quantitative and don't provide insight into the semantic meaning of tokens. Since the structure tokenization is an important part of the work, it would very much benefit from a deep analysis of the used tokens. Since the codebook sizes far exceed the number of secondary structure features, it would be great to map tokens to known structural motifs beyond secondary structure and provide visualizations like in ESM-3 paper.

8. Regarding the representation learning comparison, there is a huge body of work out there for comparison on these tasks. Why not to compare against strong baselines? Again, it is not about beating SOTA, but about correctly reflecting the state of affairs in the experiments. Also, there is a more recent and more strongly performing Gear-net version (https://arxiv.org/abs/2303.06275).

9. Regarding one of the main claims of the paper, you do not answer the question. The following is outlined in the main contributions: "DPLM-2 enables unconditional co-generation of designable and diverse proteins that guarantees consistency between structure and sequence". It seems like a gross overstatement without a formal rigorous proof. If there is no guarantees of consistency, then only empirical evidence should be declared.

10. Adding a discussion of the used self-mixup training strategy would strengthen the work. Also the experiment and data in table 8 are not discussed. If the absolute numbers of clusters are reported, the number of samples that underwent clustering should be stated in the caption.

References:

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

[2] https://arxiv.org/abs/2306.16819

[3] https://arxiv.org/abs/2312.06297v2

[4] https://arxiv.org/abs/2310.11802

Q8: Regarding the representation learning comparison, there is a huge body of work out there for comparison on these tasks. Why not to compare against strong baselines? Again, it is not about beating SOTA, but about correctly reflecting the state of affairs in the experiments. Also, there is a more recent and more strongly performing Gear-net version (https://arxiv.org/abs/2303.06275).

Thanks for your valuable suggestion. We have accordingly added more recent strong baselines and updated the stronger GearNet version as you suggested. The results of newly supplemented baseline models are mainly focused on the EC and GO downstream tasks. Meanwhile, we discovered discrepancies between the new results and previous results in the EC and GO tasks. We utilized old version of SaProt codebase as the evaluation pipeline. After consulting with the authors of SaProt, it is because there was an issue in the early SaProt codebase in calculating the metrics for these tasks, which had been fixed later. As such, we have updated the EC and GO results for both SaProt and DPLM-2.

[1] Learning Complete Protein Representation by Dynamically Coupling of Sequence and Structure. NIPS 2024

[2] Continuous-Discrete Convolution for Geometry-Sequence Modeling in Proteins. ICLR 2023

[3] Structure-informed Protein Language Model. Arxiv 2024

[4] Pre-Training Protein Bi-level Representation Through Span Mask Strategy On 3D Protein Chains. ICML 2024

[5] Pre-training Sequence, Structure, and Surface Features for Comprehensive Protein Representation Learning. ICLR 2024

Q9: Regarding one of the main claims of the paper, you do not answer the question. The following is outlined in the main contributions: "DPLM-2 enables unconditional co-generation of designable and diverse proteins that guarantees consistency between structure and sequence". It seems like a gross overstatement without a formal rigorous proof. If there is no guarantees of consistency, then only empirical evidence should be declared.

We apologize for the confusion due to our inaccurate use of "guarantees". We never aimed to over-claim that the structure-sequence consistency is mathematically and rigorously guaranteed. Instead, we mainly claim our empirical findings. To make the statements more precise, we will rephrase our wording to "DPLM-2 enables unconditional protein co-generation of both structure and sequence, which demonstrates good structure-sequence consistency."

Besides, your initial question on this greatly encouraged us to delve deep into this hence we provided a discussion with literature review in our initial response. We are also truly grateful for your inspirations!

Q6: Even with added baselines, important comparisons are missing. There is a lot of active research on inverse folding, so there is no lack of good baselines and good methodology ([1-4] to name a few). At the end of the day, it is not about beating SOTA result, but about providing objective comparison for the benefit of the protein design community.

Thanks for your valuable suggestion. We have added more good inverse folding baselines as you suggested to provide objective comparison for the benefit of the protein design community. We conduct experiments on the CATH 4.2 testset, and * means results are quoted from their respective paper. We will include these results in the next version of our manuscript.

[1] Knowledge-Design: Pushing the Limit of Protein Design via Knowledge Refinement. ICLR 2024

[2] Graph Denoising Diffusion for Inverse Protein Folding. NIPS 2023

[3] Progressive Multi-Modality Learning for Inverse Protein Folding. ICME 2024

[4] De novo protein design using geometric vector field networks. ICLR 2024

[5] PiFold: Toward effective and efficient protein inverse folding. ICLR 2023

[6] Bridge-IF: Learning Inverse Protein Folding with Markov Bridges. NIPS 2024