Protein-Nucleic Acid Complex Modeling with Frame Averaging Transformer

Weakness: Comparing to a pre-trained structure prediction method is clearly something to be explored well, and the authors do that in section 4.4 (with additional results with AF3 in the appendix). At the same time, this part seems to me not very well developed in the paper. I am especially confused by Figure 3, which the authors comment with "[..] where FAFormer can achieve comparable performance to RoseTTAFoldNA using unbounded structures." Maybe I am not understanding something but the right hand side of the plot seems to clearly indicate that at least for RNAs, RoseTTAFoldNA strongly outperforms the authors' method.

Ans: We apologize for any confusion and will make this clearer in our updated manuscript. Our claim is based on the F1 scores, which show better performance for protein-DNA complexes (0.103 vs 0.087) and comparable performance for protein-RNA complexes (0.108 vs 0.12). It should be noted that the number of test protein-RNA complexes from RoseTTAFoldNA is limited (16 complexes), which might not be able to comprehensively assess our method's performance.

Question: Could the authors expand section 4.4 and discuss the combination of Figure 3 and Table 5?

Ans: Aligning Figure 3 with Table 5, we observe that while FAFormer does not achieve higher performance than RoseTTAFoldNA in terms of contact map prediction, it significantly outperforms RoseTTAFoldNA in aptamer screening tasks. We attribute this to two main reasons:

• RoseTTAFoldNA as a foundational structure prediction model is optimized for generally predicting the protein-RNA complex structure, which might bias its performance on some specific protein targets. For example, it can achieve good performance on proteins GFP and HNRNPC but fails for NELF.
• For the contact map prediction comparison, the protein-RNA test set used by RoseTTAFoldNA is limited (16 complexes), which may not comprehensively evaluate the performance of FAFormer.

All these discussions will be added to our revised manuscript.

Question: Do the times in Table 6 include MSA search for Rosetta? If the protein is unchanged, could you make it more efficient by doing the MSA search only once?

Ans: The times reported in Table 6 include the time required for searching MSAs. For the screening task, RoseTTAFoldNA only needs to search MSAs for protein targets once. However, it still needs to search MSAs for each RNA sequence individually. Besides, the forward process of RoseTTAFoldNA is computationally expensive because the inclusion of MSAs results in a large input sequence matrix. We will add this discussion to our revised manuscript.

Question: Could the authors expand a bit on possible future research (e.g. applying the same method to other screening tasks)?

Ans: For future applications, the proposed paradigm for aptamer screening can be extended to other modalities, such as protein-small molecules and antibody-antigen. Moreover, the strong correlation between contact prediction and affinity estimation demonstrated in our paper can guide future model development. Specifically, this correlation suggests promising directions for designing new objectives and collecting datasets that better capture the nuances of molecular interactions.

Regarding the architecture, FAFormer introduces a novel approach to equivariant model design by leveraging the flexibility of Frame Averaging (FA). This idea opens up numerous possibilities for future research, including exploring different ways to integrate FA with various neural network architectures.

Weakness: The novelty of the proposed architecture is limited. The proposed architecture simply combines frame averaging and transformer blocks. The idea that builds the attention module based on SE(3)-invariant features is proposed in Equiformer. The dataset and benchmark track may be more suitable for this paper.

Ans: We respectfully disagree with the statements and would like to emphasize the contributions of our paper:

First, this is the first work designing a new equivariant Transformer based on the Frame Averaging (FA). Instead of simply combining them, each module is dedicatedly integrated with FA, which has effectively enhanced model performance, as demonstrated in our ablation study:

• The local frame edge module focuses on local spatial context by constructing the frames on the point cloud centered around each node;
• The biased MLP attention module applies FA to enable equivariant multi-head attention on the geometric features;
• Global frame FFN extends the FFN by incorporating geometric information within node representations using FA, allowing the attention to be conducted in a geometry-aware manner.

Second, the way FAFormer achieves equivariance is completely different from Equiformer which relies on spherical harmonics operations. As discussed in the Introduction (Lines 32-38), models like Equiformer compromise efficiency due to the complexity of irreducible representations, as shown in our efficiency comparison (Appendix B.2). We believe that incorporating FA within the model opens new possibilities for designing equivariant architectures in this domain.

Third, besides the architectural contributions, we propose a new paradigm for unsupervised aptamer screening. This paradigm connects contact map prediction and affinity estimation between two molecules, going beyond simply collecting datasets and constructing benchmarks.

Weakness: The expressiveness of the proposed architecture is demonstrated in the contact prediction tasks. More experiments on more general tasks are needed to further show the effectiveness of the proposed architecture, such as protein-protein docking and others.

Ans: The contact map prediction result between proteins is presented in our paper (Table 3). Besides, we have provided results on binding site prediction in Appendix E.1. This task solely takes a protein as input and aims to predict the NA-binding residues on the protein, which is a node-level prediction task.

To further demonstrate the performance of FAFormer, we here present its performance on two protein understanding tasks, including:

• Fold prediction: This task aims to predict the fold class (a total of 1195 classes) for each protein, with three different test sets (Fold, Family, and Super-Family). Accuracy is used as the evaluation metric.
• Reaction prediction: This task aims to predict the class of reactions for a given protein catalyzes (a total of 384 classes). Accuracy is used as the evaluation metric.

We follow the experimental settings in ProteinWorkshop[1], including dataset splits and input features (Ca-only). The statistics of the datasets are shown below:

The comparison results are presented below from which we can observe the best performance of FAFormer across most of the tasks:

Due to the limited time and computational resources, we can only show the results of these two tasks. More benchmarking results will be completed and added to the updated manuscript.

[1] ICLR2024-Evaluating representation learning on the protein structure universe

Comments: It is true that the proposed transformer which is based on FA is more efficient than Equiformer which is based on higher-order tensors, which is an advantage. However, to my knowledge, the idea of such architectural design is not so novel. 

Comments: And some important baselines are missing, e.g., [1].

Comments: I wonder the performance of the model on SE(3)-equivariant property prediction tasks. To be more precise, contact map prediction, binding site prediction, fold prediction and reaction prediction are SE(3)-invariant property prediction tasks, while protein-protein docking or protein-ligand docking tasks that I mentioned in the previous review can be formulated as SE(3)-equivariant prediction tasks with well-established datasets and benchmarks. I highly recommend the authors to test the proposed architecture on these tasks.

Weakness: The F1 score for contact map prediction tasks is about 0.1, which is very low. I question whether this score can be used to judge which model is better. I suggest involving other node-level and graph-level prediction tasks for comparing FAFormer with other equivariant neural networks.

Ans: We would like to clarify that no matter the contact map prediction performance, the main focus of this paper is unsupervised aptamer screening. Contact map prediction is proposed as a prerequisite task due to its strong correlation with the screening task. While contact map prediction is challenging due to sparse labels, all comparisons are fair, with the same input features and learning pipelines. For benchmarking:

• We have provided results on binding site prediction in Appendix E.1. This task predicts the NA-binding residues for a given protein, which is a node-level task. The metrics exceed 0.4 since this task is much easier than the contact map prediction.
• The aptamer screening task is a graph-level task aimed at retrieving the positive RNAs for a protein.

Moreover, we present FAFormer’s performance on two protein-related tasks:

• Fold prediction: Predict the fold class (1195 classes) for protein, with three test sets (Fold, Family, and Superfamily). Accuracy is the evaluation metric.
• Reaction prediction: Predict the reaction class catalyzed by a given protein (384 classes). Accuracy is the evaluation metric.

Following the experimental settings of ProteinWorkshop[1], including dataset splits and Ca-only input features, the comparison results are shown below. FAFormer demonstrates the best performance across most tasks:

Due to limited time and computational resources, we only show these two tasks' results. More results will be completed and added to the updated manuscript.

[1] ICLR2024-Evaluating representation learning on the protein structure universe

Weakness: If we already know the coordinates of the protein and nucleic acid, why do we need further encoding and prediction for contact prediction, instead of directly counting the contacts?

Ans: We use the individually predicted structures of proteins/nucleic acids during the evaluation (Lines 207-209), and all the input structures are decentralized (Line 597), meaning the coordinates are moved to the zero centers to avoid label leakage. Consequently, the distances between two input structures cannot be used to directly indicate the complex's contacts.

Question: Could you provide more intuition on the design of these modules?

Ans: FA is a general framework that endows a given encoder with equivariance, allowing for the flexible design of equivariant modules. Using FA as a module rather than a model wrapper also avoids an 8x increase in computation. We are glad to elaborate more on each module’s intuition:

• Local Frame Edge Module: In biomolecules, atoms interact through chemical bonds and electronic forces. Embedding these pairwise relationships allows the model to represent these intricate dependencies explicitly. We differentiate the spatial context around each node by considering the local point cloud centered on each target node. FA embeds the directed vectors between source and target nodes, capturing the direction of interactions.
• Biased MLP Attention Module: The attention mechanism is extended to include edge representations, biasing the attention map to prioritize/deprioritize certain weights based on interactions. Node coordinates are also updated to simulate conformational changes during molecular interactions, which is crucial for optimal binding.
• Global Frame FFN: We extend the FFN by incorporating FA to update representations in a geometry-aware context. The attention between two atom representations, combined with coordinate information, functions similarly to distance calculations.

We will add these discussions to the revised manuscript.

Question: Why should the model be called a “former” instead of a GNN, which could be misleading?

Ans: The core module of FAFormer is the biased attention module, which is the primary reason for considering our model a Transformer. We limit the attention to k-nearest neighbors to reduce the computation, which is commonly used in previous geometric Transformers. In the paper (Lines 90-94), we have summarized that FAFormer is a hybrid of GNN and Transformer, as the edge module functions similarly to GNN aggregation.

Question: As far as I know, RosettaFoldNA is used mainly for protein and NA structure prediction. However, FAFormer uses coordinate information as input. Is that a fair comparison?

Ans: Besides the input sequence, RosettaFoldNA utilizes searched MSAs and protein structure templates as input. The protein structure templates are the coordinates of the MSAs. In contrast, the input structures for FAFormer are predicted solely based on input sequences.

Question: I am curious about the efficiency of FAFormer compared with other equivariant NNs.

Ans: We have provided an analysis in Appendix B, which includes computational complexity and wall-clock time comparisons. In summary, our proposed FAFormer demonstrates greater efficiency than spherical harmonics and FA-based Transformers.

Question: Do other methods include the same featurization? And are you training all models from scratch?

Ans: Yes, all the baseline models use the same features and are all trained from scratch. The hyperparameter and training details are provided in Appendix A.