UniIF: Unified Molecule Inverse Folding

W1 The protein inverse folding ablation seems to show the various introduced techniques do not influence model performance much. The insights of the models seem obscured from this perspective, especially considering block-based representations and architectures have already been explored in previous works (GET/EPT). The RNA design task does not include ablation, and this reviewer is not convinced by the materials-related task (see below). It is rather unclear what a reader can take away from this work in terms what techniques actually contribute to improved performance.

Reply As protein design is more mature than RNA design, we do ablation on it. However, the protein design performance has reached a bottleneck stage where any improvement is difficult. Things will be more clear on RNA dataset. Here we provide the ablation results on RNA dataset:

In addition, we have evaluated the performance of GET on protein design, which performes very poor.

W2 If "unified" is the key word of the method, perhaps the authors should include experiments where the model is trained on various data modalities and see if that improves performances on individual tasks.

Reply Such a protocol may improve performance of protein-RNA or protein-molecule interaction prediction. However, for the design task, we observe that training on single dataset lead to better results, as we do not consider the complex interactions.

W3 The materials sections are rather strange in this paper. Inverse folding for proteins/RNAs can be understood as inferring the sequence from a given structure. The materials task is not introduced in great detail. What is being predicted? How is it "inverse folding?" Q2 Can the author clarify the materials task and why it is inverse folding?

Reply This is a new task defined in Chili[1], which gives the material structure and predict the composition of atoms types that fit to this structure.

W4 To follow up on that, the baselines that the proposed model is compared against are not what's considered state-of-the-art in materials representation. For materials, there are popular benchmarks such as Matbench (https://matbench.materialsproject.org/), JARVIS (https://pages.nist.gov/jarvis_leaderboard/), or Open catalyst (https://opencatalystproject.org/index.html), where results of SOTA models are readily available. This reviewer is not convinced by the materials-related results presented as significant details regarding the dataset/goals are missing, and the baselines are not really used for materials. Q2 Can the authors compare to materials-focused SOTA methods in the space?

Reply Thanks for your suggestion. On the binary classification task of "matbench_mp_is_metal", we provide the updated results as follows:

model	mean rocauc	mean f1
UniIF	0.9617	0.9534
CGCNN	0.9520	0.9462
ALIGNN	0.9128	0.9015
coGN	0.9124	0.9012
coNGN	0.9089	0.8972
SchNet	0.8907	0.8765

where UniIF outperforms previous baseline methods.

Q1 In the experiments, what exactly does with/without ESM mean? How is ESM used? Why can't the proposed method also become "with ESM"? Why is it valuable to consider "without ESM"? Are models "with ESM" more expensive or what?

Reply If ESM is allowed, the model can fintune the pretrained ESM model to refine the designed sequence; Otherwise, the model do not allowed to access the knowledge of ESM. The proposed method can also incooparate with ESM. Without ESM, the algorithm is more efficient, and there is do risk of label leackage to avoid potential concerns.

W1 The ablation experiments in Table 1 (-GDP), which remove the geometric dot product features, show miner drawbacks in proteins with longer sequences. This makes the geometric interaction extractor mainly capture the interaction of the virtual inter-atoms. The idea for capturing long-range dependency or interaction using virtual blocks has been studied in [1,2,3,4]; the authors could consider providing a discussion on how their proposed techniques differ from these referenced works.

R1 Thanks for your recommandation. Recent works show that virtual node could help the model to learn long-term dependencies. Our work is different in the virtual frame construction and interaction: we need to construct the local frame for each virtual node and also we use geometric operators to consider the communication between virtual nodes. Other GNN-related works do not consider the 3D geometric information.

W2 In addition, the UniIF adopted both node and edge feature extractors as feature augmentation, which is model-agnostic. However, the comparison of all three tasks does not show the potential benefit of introducing such a featureizer. Q1 Could you provide some discussions or experiments on the effectiveness of the proposed featureizer?

Reply Previously, PiFold use a model-related featurizer for protein design and achieves good results. However, such a featurizer is complex and computationally expensive. Our contribution is to simplify the featurizer and mask it model-agnostic, while ensuring competitive performance. If one use PiFold's featurizer, the recovery of protein design can further improve about 0.3%. However, such a improvement is model-related and could not eazily extend to RNA data. With out simplified featurizer, RNA recovery improved much.

W3 The paper could provide further insight into the benefit of employing a unified framework for protein, RNA, and material design. For example, could the UniIF model different molecules in a unified latent space, allowing the researchers to investigate their underlying interactions? Q2 Is it possible to train UniIF with mixed molecule datasets?

Reply Yes. We have tried to train UniIF in mixed protein and RNA dataset, which achieve slightly lower performance than training on single dataset. We plan to expend the experiments to protein-RNA-molecule complexes to study the underlying interactions as you suggested.

Q3 For Equation 7, the SVD decomposition could be time-consuming when scaling the input size. Could the author explain the motivation for such a design and why a simple approach, like "mean," does not work?

Reply Simple approach, like "mean", could not obtain a orthogonal frame vectors. We use SVD to get the rotation matrix to serve as orthogonal vector basis.

W1 Some important experimental details are missing. For instance, while the paper proposes a unified inverse folding framework, it's unclear whether the model is trained on a combined dataset from different types of molecules. Specifically, the training sets for RNA design and material design are quite limited. If the model is trained with protein data, it would be beneficial to demonstrate whether incorporating data from other types improves performance. Q1 As illustrated in the previous part, do you train the model on three types of molecules?

Reply We train the model on different types of molecules. We have tried to train UniIF in mixed protein and RNA dataset, which achieve slightly lower performance than training on single dataset. We have used all the RNA data recoreded in PDB to train the model, which is actually the largest real data currently available. Maybe we can use data augmentation techiniques in the future.

W2 Additionally, if I understand correctly, the paper does not explain the decoder part of the inverse folding model. Q2 What decoder model is used?

Reply A linear layer is served as the decoder to predict 20 amini acids, 4 neuclitides, or 128 atom types.

W3 It would be helpful if the paper included a comparison of perplexity, as related works such as PiFold have used this metric.

Reply We provide the perplesxity as follows:

W4 The performance on protein design looks comparable with other models to me.

Reply As far as we know, the protein design performance has reached a bottleneck stage where any improvement is difficult.

Q5 Could the authors identify real-world applications for inverse folding, detailing the problem, how inverse folding could help, and why existing methods fall short for each specific application?

R4 If the computing benchmark is not recognized and only wetlab experiments are considered, it becomes difficult for us to explain why existing methods fail in applications due to limitations in money, time, equipment, and experience. Here are some real-world applications of inverse folding:

• Binding site design[2]: Designing binding sites for small molecules starting from previously characterized designs generated using Rosetta.
• Membrane protein design[6]: Creating soluble and functional analogues of integral membrane proteins uisng AF2+ProteinMPNN. Inverse folding methods are used for sampling new amino acid sequences for a given fold, where improving the recovery is an important goal in these research. By reducing the temperature, one can balance the recovery and diversity. In comparison, our proposed UniIF could balance recovery and diversity better than baselines.

Q6 I note that the number of sequences for CATH 4.3 differs from the referenced papers - please could the authors explain this discrepancy.

R6 Previously, the protein number of CATH4.3 follows the extual description from ESMIF[7]. However, it was recently discovered that they released a dataset with a different number of proteins than what they described. Hence, we correct the description, but actually use the same dataset as in the baseline.

Q7 For the time split protein tasks, it would be helpful to establish how similar each test sequence is to those used in training, both in terms of proximity in sequence space and in terms of 3D structure. Please then stratify performance as a function of each distance for each held out test set. It would be helpful to also compute these distances for the CATH topology-based split, to provide a watermark for comparison.

Q8 There is no way to find the CATH code for novel protein sequences. Following your suggestion, we can provide hirachical performance in terms of sequence and structure similarity using mmseq and foldseek in CASP15 dataset. We found that the sequence similarity positively correlated with the recovery more significantly.

W1: It is not clear that the unification is responsible for any performance improvements.

Reply The proposed unification method is responsible for performance improvements by using the block-level representation and modules, such as geometric featurizer, interaction operator, and virtual frame.

The recovery metric is confusing, since it is not clear that higher recovery actually corresponds to the ability to design novel molecules.

Reply One might claim that higher recovery reduces diversity without considering similar designability. However, when methods with comparable designability—measured by self-consistent TM (scTM) scores—are considered, the conclusion changes: higher recovery actually increases diversity at the same scTM score level. This is a important finding of ProteinInvBench [1]. [1] Gao, Zhangyang, et al. "Proteininvbench: Benchmarking protein inverse folding on diverse tasks, models, and metrics." NeurIPS (2024).

W2: Moreover, even if the metric were aligned with the stated goals, it is not clear to this reviewer that existing inverse folding methods are not good enough to solve any real world application. W3: The paper lacks wetlab experiments and relies on benchmark performance to demonstrate improvement. Do these benchmark improvements and metrics align with the stated goals? Q1: In the absence of wetlab experiments, the authors should provide data from published work showing that better benchmark performance in inverse folding leads to improved novel molecule design, with proper citations. Q2 Can the authors provide examples of improvements in drug and material design resulting from better performance on inverse folding benchmark tasks? If such examples are hard to find, please address this limitation. Q3: Please could the authors propose strategies to mitigate the risk that further improvements in performance on benchmark tasks, such as the improvements reported in this paper, will not result in realized impact in drug or material design.

R3: The stated goals can be decomposed into two step: computation and wetlab experiments. Most AI researchers [3,5,6,7] focus on improving the sequence recovery, which has inspired further biological research. For instance, the well-known ProteinMPNN[4] is largely based on GraphTrans[5]. AI researchers often lack the equipment and funding for wetlab experiments. If you can provide or recommend research opportunities, we would greatly appreciate it.

Improving recovery is crucial for wetlab experiments. ProteinMPNN[5] increased recovery from 39% to over 45%, achieving strong wetlab results. LigandMPNN[2] further improved recovery to over 60%, demonstrating the ability to generate small molecules and DNA-binding proteins with high affinity and specificity.

As suggested, AI researchers should collaborate with bio-labs to validate their algorithms. Additionally, evaluating proposed methods on more metrics, such as scTM and diversity, would help mitigate the risks mentioned.

[2] Dauparas, Justas, et al. "Atomic context-conditioned protein sequence design using LigandMPNN." Biorxiv (2023).

[3] Wu, Fang, and Stan Z. Li. "A hierarchical training paradigm for antibody structure-sequence co-design." NeurIPS (2024).

[4] Dauparas, Justas, et al. "Robust deep learning–based protein sequence design using ProteinMPNN." Science (2022).

[5] Ingraham, John, et al. "Generative models for graph-based protein design." NeurIPS (2019).

[6] Zheng, Zaixiang, et al. "Structure-informed language models are protein designers."ICML, 2023.

[7] Hsu, Chloe, et al. "Learning inverse folding from millions of predicted structures." ICML, 2022.

W4: While the paper focuses on maintaining the structure of the target molecules (or material), no attempt is made to ensure that the function is maintained.

R4: We appreciate the suggestion; a relevant paper has been published in Nature[6]. We also aim to extend the algorithm to protein functions, such as solubility, but have not found a well-curated dataset for these functions. If you can recommend a suitable dataset, we would be very grateful.

[6] Goverde, Casper A., et al. "Computational design of soluble and functional membrane protein analogues." Nature (2024).

Q4 Could the authors propose strategies to mitigate the risk that the metric used in this work penalizes the model for the design of novel molecules?

R4 Reducing the temperature can balance recovery and diversity: lower temperatures (temp) increase diversity while reducing recovery. At the same recovery level on CASP15 dataset, UniIF shows higher diversity compared to the baseline, aligning with findings from ProteinInvBench.

Kindly refer to official comments for more response.

Dear reviewer r3pY,

Thank you once again for your reply! We agree with you that AI researchers should consider reaching out to potential collaborators from other fields, particularly those already equipped with funding, equipment etc. We are also looking for such opportunities.

We're glad to hear that you appreciate the unified framework and that we have the chance to clarify your question regarding whether individual tasks benefit from training on datasets for other tasks.

The answer is no, based on our experiments. But our method already achieves SOTA performance on each single task.

For example, we attempted to train UniIF on a mixed dataset of proteins and RNA, which resulted in slightly lower performance compared to training on a single dataset. However, we would like to emphasize the benefits of the unified framework:

• UniIF can be trained on different molecules within a unified latent space, allowing researchers to explore underlying interactions—similar to what AF3 accomplished. We plan to extend this in the future.

• UniIF currently achieves state-of-the-art performance on individual tasks, thanks to the careful design of the model, including the interaction operator, GNN layer, and virtual frame. To evaluate the effectiveness of UniIF, we also provide the ablation results of RNA design. The performance gain obtained by virtual frame (VFrame) and geometric dot product (GDP) is more significant in RNA dataset. This is because the protein design performance has reached a bottleneck stage where any improvement is difficult. Things will be more clear on RNA dataset. Here is the RNA ablation results:

model	short	medium	long	all
w/o VFrame	44.44	49.25	37.40	45.45
w/o EAttn	47.89	49.00	36.71	48.28
w/o GDP	46.88	48.00	37.19	47.06
UniIF	48.21	49.66	37.29	48.94

• UniIF can serve as a strong baseline that can be further enhanced by incorporating domain-specific knowledge, such as adding secondary structure information for RNA design. That is, individual tasks can be enhanced by adding domain-specific knowledge to UniIF, although UniIF has already achieved SOTA performance.

We sincerely hope this could clarify your concerns, and we would be happy to answer any additional questions. If we have adequately addressed your concerns, we kindly hope that you consider raising the score to support acceptance. Your comment and suggestion are important and valuable for us. Thank you.

Best,

Authors.

Dear Reviewer r3pY,

Thanks for your review. We have tried our best to address your questions and we respectfully thank you for supporting the acceptance of our work. Also, please let us know if you have any further questions. Look forward to further discussions!

Sincerely,

The Authors