Protein-Nucleic Acid Complex Modeling with Frame Averaging Transformer

Authors: tinglin  huang
zhenqiao  song
rex  ying
wengong  jin
aptamer  screening
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 07:52:21

Abstract: Nucleic acid-based drugs like aptamers have recently demonstrated great therapeutic potential. However, experimental platforms for aptamer screening are costly, and the scarcity of labeled data presents a challenge for supervised methods to learn protein-aptamer binding. To this end, we develop an unsupervised learning approach based on the predicted pairwise contact map between a protein and a nucleic acid and demonstrate its effectiveness in protein-aptamer binding prediction. Our model is based on FAFormer, a novel equivariant transformer architecture that seamlessly integrates frame averaging (FA) within each transformer block. This integration allows our model to infuse geometric information into node features while preserving the spatial semantics of coordinates, leading to greater expressive power than standard FA models. Our results show that FAFormer outperforms existing equivariant models in contact map prediction across three protein complex datasets, with over 10% relative improvement. Moreover, we curate five real-world protein-aptamer interaction datasets and show that the contact map predicted by FAFormer serves as a strong binding indicator for aptamer screening.

Main paper: 
 introduction: 
 Nucleic acids have recently shown significant potential in drug discovery, as shown by the success of mRNA vaccines [26, 61, 60] and aptamers [15, 32, 14, 41] . Aptamers are single-stranded nucleic acids capable of binding to a wide range of molecules, including previously undruggable targets [20, 11] . Currently, aptamer discovery is driven by high-throughput screening, which is time-consuming and labor-intensive. While machine learning can potentially accelerate this process, the limited availability of labeled data presents a significant challenge in ML-guided aptamer discovery [57, 44, 16] . Given this challenge, our goal is to build an unsupervised protein-nucleic acid interaction predictor for large-scale aptamer screening.Motivated by previous work on unsupervised protein-protein interaction prediction [34] , we focus on predicting the contact map between proteins and nucleic acids at the residue/nucleotide level. The main idea is that a predicted contact map offers insights into the likelihood of a protein forming a complex with an aptamer, thereby encoding the binding affinity between them. Concretely, as shown in Figure 1 (a), our model is trained to identify specific contact pairs between residues and nucleotides when forming a complex. The maximum contact probability across all pairs is then interpreted as the binding affinity, which is subsequently used for aptamer screening. One key factor contributing to the accuracy of contact map prediction models is their capacity to learn equivariant transformations for symmetry groups [67, 39, 66, 74, 52] . A novel line of research focuses on adapting Transformer [83] to equivariant frameworks, leveraging its great expressive power. However, these studies have encountered issues with either (1) high computational overhead with spherical harmonics-based models [47, 27] , which complicate encoding by introducing irreducible representations; or (2) limited expressive capability with frame averaging (FA) [63] , which diminishes geo-information exploitation by simply concatenating coordinates with node representations.In light of this, we propose FAFormer, an equivariant Transformer architecture that integrates FA as a geometric module within each layer. FA as a geometric component offers flexibility to effectively integrate geometric information into node representations while preserving the spatial semantics of coordinates, eliminating the need for complex geometric feature extraction. FAFormer consists of a Local Frame Edge Module that embeds local pairwise interactions between each node and its neighbors; a Biased MLP Attention Module that integrates relational bias from edge representation within MLP attention and equivariantly updates coordinates; and a Global Frame FFN that integrates geometric features into node representations within the global context.To validate the advantage of our model architecture, we evaluate FAFormer on two tasks: (1) proteinnucleic acid contact prediction and (2) unsupervised aptamer virtual screening. In the first task, our model consistently surpasses state-of-the-art equivariant models with over 10% relative improvement across three protein complex datasets. For the second task, we collected five real-world proteinaptamer interaction datasets with experimental binding labels. Our results show that FAFormer, trained on the contact map prediction task, is an effective binding indicator for aptamer screening. Compared to RoseTTAFoldNA, a large pretrained model for complex structure prediction, FAFormer achieves comparable performance on contact map prediction and better results on aptamer screening, while offering 20-30x speedup. 
related work: 
 Aptamer screening Aptamers are single-stranded RNA or DNA oligonucleotides that can bind various molecules with high affinity and specificity [14, 15, 32, 41] . SELEX (Systematic Evolution of Ligands by EXponential Enrichment) is a conventional technique used for high-throughput screening of aptamers [23, 72, 23] , which iteratively selects and amplifies target-bound sequences. Despite its effectiveness, SELEX needs to take substantial time to identify a small number of aptamers [73, 51] . There are some recent studies applying machine learning techniques to aptamer research [18] , including generating aptamer structures [36] , optimizing SELEX protocols [6] , and predicting and modeling protein-aptamer interactions [24, 46, 68, 57] . However, these methods require the user to provide labeled data from time-consuming SELEX assays, thus do not apply to new targets without any SELEX data. Our work focuses on predicting the contact map between protein and nucleic acids using 3D structures and conducting unsupervised screening based on the predicted contact maps, which has not yet been thoroughly explored.The prediction and understanding of interactions between proteins and molecules play a crucial role in biomedicine. For example, some prior studies focus on developing a geometric learning method to predict the conformation of a small molecule when it binds to a protein target [71, 19, 53, 50] . As for the protein-protein complex, [25, 29] explore the application of machine learning in predicting the structure of protein multimer. Some studies [77, 56] investigate the protein-protein interface prediction in the physical space. Jin et al. [38] studies the proteinprotein affinity prediction in an unsupervised manner. For protein-nucleic acid, some prior works explore the identification of the nucleic-acid-binding residues on protein [65, 86, 89, 37, 85] or predict the binding probability of RNA on proteins [79, 82, 87, 54] . Some previous studies focus on modeling protein-nucleic acid complex by computational method [81, 80, 28] . AlphaFold3 [1] and RoseTTAFoldNA [5] are the recent progresses in this field, which both are pretrained language models for complex structure prediction.Recently, geometric deep learning achieved great success in chemistry, biology, and physics domains [13, 90, 40, 55, 52, 64, 10, 70] . The previous methods roughly fall into four categories: 1) Invariant methods extract invariant geometric features from the molecules, such as pairwise distance and torsion angles, to exhibit invariant transformations [67, 31, 30] ; 2) Spherical harmonics-based models leverage the functions derived from spherical harmonics and irreducible representations to transform data equivariantly [27, 47, 75] ; 3) Some methods encode the coordinates and node features in separate branches, interacting these features through the norm of coordinates [66, 39] ; 4) Frame averaging (FA) [63, 21] framework proposes to model the coordinates in eight different frames extracted by PCA, achieving equivariance by averaging the encoded representations.The proposed FAFormer combines the strengths of FA and the third category of methods by encoding and interacting coordinates with node features using FA-based components. Besides, FAFormer can also be viewed as a combination of GNN and Transformer architectures, as the edge representation calculation functions similarly to message passing in GNN. This integration is made possible by the flexibility provided by FA as an integrated component. 
frame averaging transformer: 
 In this section, we present our proposed FAFormer, a frame averaging (FA)-based transformer architecture. We first introduce the FA framework in Section 3.1 and elaborate on the proposed FAFormer in Section 3.2. Discussion on the equivariance is provided in Section 3.3, and the computational complexity analysis can be found in Appendix B.Frame averaging (FA) [63] is an encoder-agnostic framework that can make a given encoder equivariant to the Euclidean symmetry group. FA applies the principle components derived via Principal Component Analysis (PCA) to construct the frame capable of achieving E(3) equivariance. Specifically, the frame function F(•) maps a given set of coordinates X to eight transformations:EQUATIONwhere u 1 , u 2 , u 3 are the three principal components of X, U ∈ R 3×3 denotes the rotation matrix based on the principal components, and c ∈ R 3 is the centroid of X. The main idea of FA is to encode the coordinates as projected by the transformations, followed by averaging these representations. We introduce f F (•) to represent the projections of given coordinates via F(•):f F (X) := {(X -c)U | (U , c) ∈ F(X)} := {X (g) } F (2)where X (g) denotes the coordinates transformed by g-th transformations. We can apply any encoder Φ(•) to the projected coordinates and achieve equivariance by averaging, which can be formulated as an inverse mapping f F -1 (•): where U -1 g is the inverse matrix of g-th transformations and the result exhibit E(3) equivariance. The outcome is invariant when simply averaging the representations without inverse matrix.EQUATIONInstead of serving FA as an external encoder wrapper, we propose instantiating FA as an integral geometric component within Transformer. This integration preserves equivariance and enables the model to encode coordinates effectively in the latent space, ensuring compatibility with Transformer architecture. The overall architecture is illustrated in Figure 2 .Graph Construction Each molecule can be naturally represented as a graph [47, 35, 43] where the residues/nucleic acids are the nodes and the interactions between them represent the edges. To efficiently model the macro-molecules (i.e., protein and nucleic acid), we restrict the attention for each node i to its K-nearest neighbors N top-K (i), within a predetermined distance cutoff c:EQUATIONwhere we use N (i) to denote the valid neighbors of node i, and d ij denotes the distance between node i and j. The long-range context information can be captured by iterative attention within the local neighbors for each node i.Overall Architecture As shown in Figure 2 (a), the input of FAFormer comprises the node features Z ∈ R N ×D , coordinates X ∈ R N ×3 , and edge representations E ∈ R N ×K×D derived by our proposed edge module, where N is the number nodes and D is the hidden size. FAFormer processes and updates the input features at each layer:EQUATIONwhere f (l) (•) represents l-th layer of FAFormer. In each model layer, we first update coordinates and node features using the Biased MLP Attention Module. Then, these features are fed into Local Frame Edge Module to refine the edge representations. Finally, the node representations undergo further updates through Global Frame FFN.FA Linear Module Based on FA, we generalize the vanilla linear module to encode coordinates in the latent space invariantly:EQUATIONwhere {X (g) } F is obtained using Equ.(2), W g ∈ R 3×D is a learnable matrix for g-th transformations, and Norm(X) := X/ 1 ν ||X|| 2 2 is the normalization which scales the coordinates such that their root-mean-square norm is one [39] . Linear F (•) is an invariant transformation and will serve as a building block within each layer of the model.We explicitly embed the interactions between each node and its neighbors as the edge representation E ∈ R N ×K×D , where K is the number of the neighbors. It encodes the relational information and represents the bond interaction between nodes, which is critical in understanding the conformation of molecules [40, 4, 35] .As shown in Figure 2 (f), unlike the vanilla FA which globally encodes the geometric context of the entire molecule, the edge module builds frame locally around each node's neighbors. Specifically, given a node i and its neighbor j ∈ N (i), the geometric context is encoded within the local neighborhood:EQUATIONwhere {X i -X j } N (i) denotes the direction vectors from center node i to its neighbors, andX ′ i→j ∈ R d is the encoded representation.With the local frame, the spatial information sent from one source node depends on the target node, which is compatible with the attention mechanism. Then the node features are engaged with geometric features, and the edge representation is finalized through the residual connection with the gate mechanism:EQUATIONwhere (•||•) is the concatenation operation, and the calculated gate g ij = σ(Linear(m ij )) provides flexibility in regulating the impact of updated edge representation.The encoded edge representation in FAFormer plays a crucial role in modeling the pairwise relationships between nodes, especially for nucleic acid due to their specific base pairing rules [58, 45] .The incorporation of FA facilitates the encoding of the pairwise relationships in a geometric context, resulting in an expressive representation.Biased MLP Attention Module As shown in Figure 2 (b), the attention module of FAFormer first transforms the node features Z into query, key, and value representations:EQUATIONwhere W Q , W K , W V ∈ R D×D are the learnable projections. We adopt MLP attention [12] to derive the attention weight between node pairs, which can effectively capture any attention pattern. The relational information from the edge representation is integrated as an additional bias term:EQUATIONwhere b ij = Linear(LN(E ij )) represents the scalar bias term based on the edge representation, a ij denotes the attention score between i-th and j-th nodes, Z * ,i is i-th representation of the matrix Z * , Softmax i (•) is the softmax function operated on the attention scores of node i's neighbors, and LN(•) is layernorm function [3] .Besides the value embeddings, the edge representation will also be aggregated to serve as the context for the update of node feature in FAFormer:EQUATIONEQUATIONwhere Z ′ i is the update representation of node i. The above attention can be extended to a multi-head fashion by performing multiple parallel attention functions. For the coordinates, we employ an equivariant aggregation function that supports multi-head attention:EQUATIONwhere {X (g) } F = f F (X), H is the number of attention heads, [•] is the tensor stack operation and W ∈ R H×1 is a linear transformation for aggregating coordinates in different heads. An additional gate function that uses node representations as input to modulate the aggregation is applied:EQUATIONwhere g attn = σ(Linear(Z)) is the vector-wise gate designed to modulate the integration between the aggregated and the original coordinates. This introduced gate mechanism further encourages the communication between node features and geometric features.Global Frame FFN To further exploit the interaction between node features and coordinates, we extend the conventional FFN to Global Frame FFN which integrates spatial locations with node features through FA, which is illustrated in Figure 2(d) :EQUATIONwhere FFN(•) denotes a two-layer fully connected feed-forward network. This integration of spatial information X v into the feature vectors enables the self-attention mechanism to operate in a geometricaware manner. Unlike the edge module that focuses on each node's local neighbors, global frame FFN encodes the coordinates of all nodes, thereby capturing the long-range correlation among nodes.The function Linear F (•) exhibits invariance since results are simply averaged across different transformations. In light of this, the node representation generated by the edge module and FFN are also invariant. The biased attention is based on the scalar features so the output is always invariant.The update of coordinates within FAFormer leverages a multi-head attention aggregation with a gate function. Both functions are E(3)-equivariant: attention aggregation is based on frame averaging, while gate function is linear and also exhibits E(3)-equivariance, with a formal proof in Appendix C.In conclusion, FAFormer is symmetry-aware which exhibits invariance for node representations and E(3)-equivariance for coordinates. 
experiments: 
 In this section, we present three protein complex datasets and five aptamer datasets to explore protein complex interactions and evaluate the effectiveness of FAFormer. More details regarding the experiments and datasets can be found in Appendix A and D. Additional experiments, including the ablation studies, and the comparison with AlphaFold3, can be found in Appendix E. All the datasets used in this study are included in our anonymous repository.Protein Complexes We cleaned up and constructed three 3D structure datasets of protein complexes from multiple sources [8, 7, 2, 77] . A residue-nucleotide pair is determined to be in contact if any of their atoms are within 6Å from each other [77, 78] . We conduct dataset splitting based on the protein sequence identity, using a threshold of 50% for protein-RNA/DNA complexes3 and 30% for protein-protein complexes. The details of all datasets are shown in Table 1 .The protein's and nucleic acid's structures in the validation/test sets are generated by ESMFold [48] (proteins) or RoseTTAFoldNA (nucleic acids). This offers a more realistic scenario, given that the crystal structures are often unavailable. Aptamers Our aptamer datasets come from the previous studies [76, 46, 73] , including five protein targets and their corresponding aptamer candidates. The affinity of each candidate to the target is experimentally determined. Each dataset is equally split into validation and test sets.We construct the protein-RNA training set by excluding complexes from our collected dataset with over 30% protein sequence identity to these protein targets, resulting in 1,238 training cases. The 3D structures of proteins are obtained from AlphaFold Database [5] , and the structures of RNAs are generated by RoseTTAFoldNA without using MSAs. The statistics are presented in Table 2 .Feature The coordinates of the C α atoms from residue and the C 3 atoms from nucleotide are used as coordinate features. For node feature generation, we employ ESM2 [49] for proteins and RNA-FM [17] for RNA. The one-hot embedding is utilized as DNA's node feature.As shown in Figure 1 (a), this task aims to predict the exact contact pairs between protein {S i } N and nucleic acid {S ′ j } N ′ which conducts binary classification over all pairs:EQUATIONBaselines We compare FAFormer with four classes of methods: 1) Vanilla Transformer [83] which doesn't utilize 3D structure; 2) Spherical harmonics-based models Equiformer [47] and SE(3)Transformer [27] ; 3) GNN-based models EGNN [66] and GVP-GNN [39] ; 4) Transformer with FA [63] . The protein and nucleic acid will be separately encoded with two encoders to avoid label leakage. The representations of residues and nucleotides are concatenated from all pairs and fed into a MLP classifier to conduct prediction.To comprehensively evaluate the performance of label-imbalanced datasets, we apply F1 and PRAUC as the evaluation metrics. The comparison results are presented in Table 3 from which FAFormer reaches the best performance over all the baselines with a relative improvement over 10%.Additionally, some geometric methods fail to outperform the vanilla Transformer in certain cases.We attribute this to overfitting on crystal structures during training, which hinders their ability to generalize well to unbound structures during evaluation. Compared with serving FA as an external equivariant framework on Transformer, the performance gain on FAFormer verifies the effectiveness of embedding FA as a geometric component within Transformer. In addition to contact map prediction, we examine the nucleic acid binding site prediction task, which is a node-level task, to comprehensively evaluate our model. This task solely takes a protein {S i } N as input and aims to identify the nucleic-acid-binding residues on the protein:EQUATIONPredicting the nucleic acid binding site offers promising therapeutic potential for undruggable targets by conventional small molecule drug [33, 9, 88] , expanding the range of potential therapeutic targets.Baselines We compare FAFormer with two state-of-the-art geometric deep learning models: Graph-Bind [86] and GraphSite [89] in this task. The protein structure is embedded with the geometric encoder and the residue's representations are fed into a classifier for prediction.Results F1 and PRAUC are applied as the evaluation metrics and we report the average results over three different seeds in Table 4 . We can observe that FAFormer achieves the best performance over all the baselines, demonstrating the effectiveness of FAFormer on nucleic acid-related tasks and modeling the geometric 3D structure. This task aims to screen the positive aptamers from a large number of candidates for a given protein target. We quantify the binding affinities between RNA and the protein target as the highest contact probability among the residue-nucleotide pairs. The main idea is that two molecules with a high probability of contact are very likely to form a complex [34] . The models are first trained on the protein-RNA complexes training set using the contact map prediction, then the aptamer candidates are ranked based on the calculated highest contact probabilities.Results Top-10 precision, Top-50 precision, and PRAUC are used as the metrics. As shown in Table 5 , the geometric encoders outperform sequence-based Transformer in most cases, and FAFormer generally reaches the best performance. This demonstrates the great potential of an accurate interaction predictor in determining unsupervisedly promising aptamers. In this section, we investigate the performance of RoseTTAFoldNA [5] which is a pretrained protein complex structure prediction model and compare it with FAFormer. The performance of FAFormer is evaluated on the individual predicted protein and nucleic acid structures by ESMFold and RoseTTAFoldNA. We additionally test the performance of AlphaFold3 [1] on a subset of the screening tasks due to AlphaFold3 server submission limits (Appendix E).Dataset For the contact map prediction task, we select the test cases used in RoseTTAFoldNA to create the test set, yielding 86 protein-DNA and 16 protein-RNA cases. Furthermore, the complexes from our collected dataset that have more than 30% protein sequence identity to these test examples are removed. This leads to 1,962 training cases for protein-DNA and 1,094 for protein-RNA, which are used for training FAFormer. The MSAs of proteins and RNAs are retrieved for RoseTTAFoldNA.For the aptamer screening task, we construct a smaller candidate set for each protein target by randomly sampling 10% candidates, given that the inference of RoseTTAFoldNA with MSA searching is time-consuming. The datasets will be equally split into validation and test sets. More details of these datasets can be found in Appendix D. Table 6 : Comparison results with RoseTTAFoldNA using the sampled datasets, which accounts for the performance differences of FAFormer as shown in Table 5 . The comparison results of contact map prediction are presented in Figure 3 , where FAFormer can achieve comparable performance to RoseTTAFoldNA using unbounded structures. Specifically, our method has higher F1 scores for protein-DNA complexes (0.103 vs. 0.087) and performs comparably for protein-RNA complexes (0.108 vs. 0.12). Besides, Table 6 shows that RoseTTAFoldNA fails to receive positive aptamers for some targets, e.g., NELF and UBLCP1, while FAFormer consistently outperforms RoseTTAFoldNA for all the targets.Aligning Figure 3 with Table 6 , we observe that while FAFormer does not surpass RoseTTAFoldNA in contact map prediction, it significantly excels in aptamer screening. We attribute this to two reasons: 1) Similar to AlphaFold3 (Table 10 ), RoseTTAFoldNA as a foundational structure prediction model is optimized for general complex structures, 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.2) The protein-RNA test set used by RoseTTAFoldNA for contact map prediction is limited, which may not comprehensively evaluate FAFormer.Two examples of protein-DNA (PDB id: 7DVV) and protein-RNA (PDB id: 7KX9) complexes are provided in Figure 4 , which shows a visual comparison between the actual (left) and the predicted (right) contact maps. Note that only the residues involved in the actual contact map are presented for a clear demonstration 4 . The complete contact map can be found in Appendix E.1. We can find that despite the sparsity of contact pairs, the predicted contact maps show a high degree of accuracy when compared to the ground truth.  
conclusion: 
 This research focuses on predicting contact maps for protein complexes in the physical space and reformulates the task of large-scale unsupervised aptamer screening as a contact map prediction task.To this end, we propose FAFormer, a frame averaging-based Transformer, with the main idea of incorporating frame averaging within each layer of Transformer. Our empirical results demonstrate the superior performance of FAFormer on contact map prediction and unsupervised aptamer screening tasks, which outperforms eight baseline methods on all the tasks.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: A code link is included in the paper, containing the training pipelines for all tasks. Additionally, we have detailed the dataset preprocessing methods and provided the corresponding scripts in the code link.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: As shown in Appendix A and Appendix D, we have introduced the dataset splitting strategies based on protein sequence identity, the search range of each hyperparameter, and the optimal hyperparameters for each model.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer Justification: This paper presents work whose goal is to advance the field of geometric deep learning on protein complex modeling. There are some potential societal consequences of our work, none of which we feel must be specifically highlighted here.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The datasets and models used in this study don't have such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: All the involved models and datasets, including pretrained protein/RNA sequence models, pretrained structure prediction models (RoseTTAFoldNA and AlphaFold3), complex datasets, and aptamer datasets, are properly credited and cited. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: The paper does not introduce new assets; it only uses existing models and datasets that are properly credited and cited. Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: 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. Besides, FAFormer introduces a novel approach to equivariant model design by leveraging the flexibility of FA. This idea opens up numerous possibilities for future research, including exploring different ways to integrate FA with various neural network architectures.Limitation In this study, the geometric features utilized in the encoders are limited to the coordinates of the C α and C 3 atoms. Features extracted from the backbone or sidechains have not been used, which may limit the performance of the geometric encoders.Running environment. The experiments are conducted on a single Linux server with The AMD EPYC 7513-32 Core Processor, 1024G RAM, and 4 Tesla A40-48GB. Our method is implemented on PyTorch 1.13.1 and Python 3.9.6.Training details. For all the baseline models and FAFormer, we fix the batch size as 8, the number of layers as 3, the dimension of node representation as 64, and the optimizer as Adam [42] . Binary cross-entropy loss is used for contact map identification tasks with a positive weight of 4. The gradient norm is clipped to 1.0 in each training step to ensure learning stability. We report the model's performance on the test set using the best-performing model selected based on its performance on the validation set. All the results are reported based on three different random seeds.The learning rate is tuned within {1e-3, 5e-4, 1e-4} and is set to 1e-3 by default, as it generally yields the best performance. For each model, we search the hyperparameters in the following ranges: dropout rate in [0, 0.5], the number of nearest neighbors for the GNN-based methods in {10, 20, 30}, and the number of attention heads in {1, 2, 4, 8}. The hyperparameters used in each method are shown below:• FAFormer: The number of attention heads, dropout rate, and attention dropout rate are 4, 0.2, and 0.2 respectively. We initialize the weight of the gate function with zero weights, and bias with a constant value of 1, ensuring a mostly-opened gate. SiLU [22] is used as the activation function.The distance threshold c is set as 1e5Å and the number of neighbors is 30.• GVP-GNN5 : The dimensions of node/edge scalar features and node/edge vector features are set as 64 and 16 respectively. The dropout rate is fixed at 0.2. For a fair comparison, we only extract the geometric feature based on C α , i.e., the forward and reverse unit vectors oriented in the direction of C α between neighbor residues.• EGNN6 : The number of neighbors is set as 30. Besides, we apply gate attention to each edge update module and residue connection to the node update module. SiLU [22] is used as the activation function.• Equiformer 7 and SE(3)Transformer8 : The number of attention heads, and the hidden size of each attention head are set as 4 and 16. We exclude the neighbor nodes with a distance greater than 100Å and set the number of neighbors as 30. Based on our experiments, we set the degree of spherical harmonics to 1, as higher degrees tend to lead to performance collapse according to our experiments.• Transformer and FA9 : The Transformer is applied as FA's backbone encoder. The dropout and attention dropout rates are 0.2 and 0.2. The number of attention heads is set as 4.• GraphBind10 : The dropout ratio and the number of neighbors are 0.5 and 30. We apply addition aggregation to the node and edge update module, following the suggested setting presented in the paper.• GraphSite11 : The number of neighbors and dropout ratio are 30 and 0.2. The number of attention layers and attention heads are 2 and 4 respectively. Besides, we additionally use the DSSP features as the node features, as suggested in the paper.• RoseTTAFoldNA12 : We employ the released pretrained weight of RoseTTAFoldNA, and set up all the required databases, including UniRef, BFD, structure templates, Rfam, and RNAcentral following the instructions.As a core component in the model, frame averaging's time complexity mainly comes from the calculation of PCA among the input coordinates. This operation is practically efficient due to the low dimensionality of the input (only 3 for coordinates). Besides, the calculation of eigenvalue decomposition can be significantly accelerated by some libraries, such as PyTorch [62] and SciPy [84] .We ignore the complexity used for calculating projected coordinates in the following analysis.As for the local frame edge module in FAFormer, linear transformations are employed to compute the pairwise representation (Equ. (7) and Equ. ( 8)), and an additional gate is applied to regulate these messages (Equ. ( 8)), resulting in a computational complexity of O(N KD + N KD 2 ). Considering the residue connection, the overall complexity of the local frame edge module is O(N KD+N KD 2 + N D).As for the self-attention module, linear transformations are performed on token embeddings and edge representations (Equ. ( 9)), and the multi-head MLP attention computation is limited to nearest neighbors (Equ. ( 10)), leading to the complexity of O(N HD 2 + N KHD) where H is the number of attention heads. Further operations, including aggregation, linear projection, and residual connections (Equ. (11) and Equ. ( 12)), add O(N KHD + N KHD 2 + N D). Moreover, applying a gate function to the combination of aggregated coordinates and original coordinates (Equ. ( 13) and Equ. ( 14)) adds O(N KH + N D + N D 2 ). As a result, the total complexity for the self-attention module isO(N KHD + N KHD 2 + N KH + N D + N D 2 ).Regarding the FFN, most operations are linear transformations that have the complexity of O(N D 2 ).The gate function and linear combination of coordinates contribute O(N D 2 + N D). So the total computational complexity of FFN is O(N D 2 + N D).We conduct a comparison of wall-clock time performance between FAFormer and other geometric baseline models under the same computational environment. Specifically, we measure the average training time for one epoch of each model, with the results illustrated in Figure 5 . It can be observed that FAFormer demonstrates greater efficiency compared to spherical harmonics-based models and achieves performance comparable to the GNN-based method GVP-GNN. Such efficiency is attributed to the utilization of top-K neighbor graphs and FA-based modules in FAFormer, which enable efficient modeling of coordinates through linear transformations. In this section, we investigate the equivariance of the gate function (Figure 2 (e)), which can be formulated as:Gate(X i , f ∆ (X i ), Z i ) := g i ∆X i + (1 -g i )X i (18) := X ′ i (19)where X i represents the coordinates of i-th node, g i = σ(Linear(Z i )) is the gate score based on i-th node's feature, and ∆X i = f ∆ (X i ) denotes the coordinate update function, i.e., the multi-head aggregation function Equ. (13) which is based on the frame averaging and thus is E(3) equivariant:EQUATIONwhere t ∈ R 3 is a translation vector and Q ∈ R 3×3 is an orthogonal matrix.We aim to prove that the gate function is E(3) equivariant, meaning it is translation equivariant for any translation vector t ∈ R 3 and rotation/reflection equivariant for any orthogonal matrix Q ∈ R 3×3 . Specifically, we want to show:EQUATIONDerivation.EQUATIONEQUATIONEQUATIONTherefore, we have proven that applying rotation and translation to X i results in the identical rotation and translation being applied to X ′ i .Protein Complex Datasets We collect the complexes from PDB [8] , NDB [7] , RNASolo [2] and DIPS [77] databases. Complexes are excluded if they have protein sequences shorter than 5 or longer than 800 residues, or nucleic acid sequences shorter than 5 or longer than 500 nucleotides. The redundant proteins with over 90% sequence similarity to other sequences within the datasets are removed. The protein and the binder structures will be separated and decentered.Aptamer Datasets Detailed information on each protein target and the aptamer candidates is presented below. The threshold for categorizing sequences as positive or negative aptamers is determined by referencing previous studies or identifying natural cutoffs in the affinity score distributions.• GFP 13 : Green fluorescent protein. The aptamer candidates are mutants of GFPapt [69] , with K d values ranging from 0nM to 125nM as affinity measures. Candidates with K d values lower than 10nM are considered positive cases.• NELF14 : Negative elongation factor E. The aptamer candidates are mutants of NELFapt [59] , with K d values ranging from 0nM to 183nM as affinity measures. Candidates with K d values lower than 5nM are considered positive cases.• HNRNPC 15 : Heterogeneous nuclear ribonucleoproteins C1/C2. The aptamer candidates are the randomly generated RNA 7mers and we apply the affinity values provided by the previous studies [46] . Candidates with affinity scores lower than -0.5 are positive.• CHK2 16 : Serine/threonine-protein kinase Chk2. Szeto et al [73] applied SELEX (Systematic Evolution of Ligands by EXponential Enrichment) [23] to screen aptamers from a large library of random nucleic acid sequences through multiple rounds. During each round, the bound sequences were amplified and isolated. We use the final round of sequences as the candidates, with sequences having multiplicities over 100 considered positive aptamers.• UBLCP117 : Ubiquitin-like domain-containing CTD phosphatase 1. Similar to CHK2, the final round of sequences are considered candidates, with sequences having multiplicities over 200 considered positive aptamers. The statistics are shown in Table 8 . We equally split each dataset into a validation set and a test set. The performance of FAFormer on the test set is reported with the best performance on the validation set. RoseTTAFoldNA is directly evaluated on the test set.Test datasets of RoseTTAFoldNA The test datasets used to evaluate RoseTTAFoldNA are available at this accessible link. We downloaded the dataset and filtered out non-dimer complexes, resulting in 86 protein-DNA and 16 protein-RNA complexes.Figure 6 presents the complete groundtruth and predicted contact maps of the cases used in Figure 4 , where the model accurately captures the sparse pattern. In this section, we conduct an ablation study to investigate the impact of FAFormer's core modules. Specifically, we individually disable the edge module and attention mechanism, and replace the proposed FFN with the conventional FFN in FAFormer. The results are presented in Table 9 . Removing any core module results in a significant performance decline. Notably, the model degrades to either an attention-only or message-passing-only architecture when the edge or attention module is removed, both of which lead to significant declines in performance. This demonstrates the advantage of combining these two architectures with FA.AlphaFold3 [1] represents the latest advancement in biomolecular complex structure prediction and is accessible through an online server 18 . Due to its limited quota (20 jobs per day), we evaluated it on two of the smallest sampled aptamer datasets, GFP and HNRNPC. The statistics for these datasets are presented in Table 8 . We used the contact probability produced by the server and the maximum probability as the estimated affinity. The results are shown in Table 10 . Although AlphaFold3 performs well in predicting the complex structure, it fails to identify promising aptamers in our cases. We attribute this to fine-grained optimization and overfitting on molecule interaction patterns in the structure prediction task, which have biased its screening performance on specific protein targets. (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility.In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.Note that we don't use 30% as the threshold since it results in a very limited validation and test set.We sort the residue IDs alongside the row ID so that the contact map appears diagnostic.https://github.com/drorlab/gvp-pytorchhttps://github.com/vgsatorras/egnnhttps://github.com/atomicarchitects/Equ.iformerhttps://github.com/FabianFuchsML/se3-transformer-publichttps://github.com/omri1348/Frame-Averaginghttp://www.csbio.sjtu.edu.cn/bioinf/GraphBind/sourcecode.htmlhttps://github.com/biomed-AI/GraphSitehttps://github.com/uw-ipd/RoseTTAFold2NAhttps://www.uniprot.org/uniprotkb/P42212/entryhttps://www.uniprot.org/uniprotkb/P92204/entryhttps://www.uniprot.org/uniprotkb/P07910/entryhttps://www.uniprot.org/uniprotkb/O96017/entryhttps://www.uniprot.org/uniprotkb/Q8WVY7/entryhttps://golgi.sandbox.google.com/ 


Decision: This work introduces FAFormer, a novel equivariant transformer architecture that seamlessly integrates frame averaging within each transformer block. By incorporating FA as a core component, the architecture allows for the design of expressive and equivariant modules. Results demonstrate that FAFormer outperforms existing equivariant models in contact map prediction across three protein complex datasets. The reviewers found this framework generally interesting and innovative, with valuable yet underexplored application scenarios. Despite some minor concerns, this is a strong paper that I recommend for acceptance to NeurIPS.

Reviews: review: paper_summary: paper_summary: The authors propose a novel equivariant model, FAFormer, which leverages the frame averaging operation as an integral geometric component within the Transformer. The authors prove the invariance and equivariance of the architecture. They further conduct experiments showing that FAFormer performs well in contact map prediction and could serve as a strong binding indicator for aptamer screening.
strength_weakness: strength_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. The model potentially performs worse due to separate modeling of nucleic acid and protein, the linear prediction head, and the difficulty of the pairwise prediction task. I suggest involving other node-level and graph-level prediction tasks for comparing FAFormer with other equivariant neural networks.
    
- The main idea of the model is confusing. 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?
questions: questions: - The module structures with formulas are a bit difficult to follow. Also, could you provide more intuition on the design of these modules?
    
- Since we are constructing a KNN graph and always leveraging the coordinate X and edge E information, why should the model be called a “former” instead of a GNN, which could be misleading?
    
- 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?
    
- I am curious about the efficiency of FAFormer compared with other equivariant NNs. Could you provide some analyses?
    
- During experiments, do other methods include the same featurization? And are you training all models from scratch?
limitations: limitations: The work provides an innovative framework, FAFormer, and performs well on contact map prediction and aptamer screening tasks. However, the experimental setting is not convincing enough and requires refinement.
paper_summary: paper_summary: This paper mainly focuses on protein-nucleic acid contact prediction and unsupervised aptamer virtual screening. The latter is based on an unsupervised learning approach by predicting the contact maps.  An equivariant architecture that integrates frame averaging and transformer blocks is proposed for this task. Experiments on the two tasks demonstrate the superiority of the proposed architecture compared with other geometric deep learning models.
strength_weakness: strength_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.
- 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.
questions: questions: See the above weaknesses.
limitations: limitations: The authors have pointed out the limitations due to only modeling certain atoms in the related tasks.
paper_summary: paper_summary: The authors propose am method for predicting contacts between proteins and aptamers based on frame averaging transformers. They showcase their methods on contact prediction and unsupervised aptamer screening, showing improvements over some baselines. The authors also compare to RoseTTAFoldNA, which seems to outperform their method.
strength_weakness: strength_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.
questions: questions: - Could the authors expand section 4.4 and discuss the combination of Figure 3 and Table 5?
- 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?
- Could the authors expand a bit on possible future research (e.g. applying the same method to other screening tasks)?
limitations: limitations: Authors discuss some scientific limitations. I do not foresee any negative societal impact.

Comments: comment: Thanks for providing the experimental results! I have two related questions:

- Why not use the standard protein complex dataset? In the dataset you proposed in the paper, there are only about 4k protein-protein complexes, while in the widely used protein-protein complex dataset, i.e., the Database of Interacting Protein Structures (DIPS), there are about 40k samples. The experiments on such a small dataset may not be very convincing.

- The results of complex RMSD and interface RMSD are too high. In this case, the numerical results (~18) may indicate invalid docking results. Please refer to the paper of EquiDock for more details. Their results are much lower than the reported in the above response.
Thank you once again for your valuable comments on our paper! As the deadline is approaching, we may not be able to respond after that. Please let us know if you have any further questions or concerns. We are committed to addressing your feedback.
We greatly appreciate your valuable comments on our paper! All the suggestions will be added to our updated manuscript.
```
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. 
```
**Ans**: We believe the key contribution and novelty of our architecture lie in demonstrating the potential of designing and applying equivariant modules based on FA, which, to the best of our knowledge, has not been explored previously. Moreover, the primary motivation behind our architectural design is to address the inherent issues of the prior Transformers, as mentioned in Lines 35-38, which is non-trivial and can effectively enhance performance.

```
Comments: And some important baselines are missing, e.g., [1].
```
**Ans**: Thank you for bringing this to our attention! We have applied GATr’s official implementation and evaluated its performance on the contact map prediction task, as shown below:

|  |  | EGNN | GATr | FAFormer |
| --- | --- | --- | --- | --- |
| Protein-RNA | F1 | 0.1093 | 0.1107 | 0.1284 |
|  | PRAUC | 0.0964 | 0.0976 | 0.1113 |
| Protein-DNA | F1 | 0.1208 | 0.1077 | 0.1457 |
|  | PRAUC | 0.1139 | 0.1075 | 0.1279 |

The input features for GATr are consistent with those used in our method and other baselines, including coordinates and pretrained residue/nucleotide embeddings. Following GATr’s setup, the input coordinates are first mapped to tri-vectors. However, GATr’s original implementation is designed for general point clouds and uses full attention among points, which is memory-inefficient for large biomolecules. To address this, we limited the attention to K-nearest neighbors. The hyperparameters were tuned using the same process outlined in Appendix A and are shown below:
- Hidden dim: 64
- Hidden dim for multivector: 32
- Number of neighbors: 30
- Number of attention blocks: 5
- Dropout ratio: 0.2

```
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.
```
**Ans**: Existing docking benchmarks [1,3] primarily compare different docking models, such as diffusion model-based approaches [4] and alignment-based models [2]. FAFormer, as a geometric backbone, is not directly comparable to these docking methods but is **compatible** with them by extending or replacing the geometric encoder they use.

To demonstrate this compatibility, we present the results of integrating FAFormer with Equidock [2] on the protein-protein docking task using our protein complex dataset:

|  | complex RMSD | interface RMSD |
| --- | --- | --- |
| Equidock | 20.266 | 22.769 |
| Equidock w/ FAFormer | 18.572 | 18.089 |

In this table, "complex RMSD" represents the mean distance between each pair of residues, while "interface RMSD" focuses specifically on the binding interface. These preliminary results demonstrate that incorporating FAFormer can effectively enhance the performance of existing docking methods. 

We believe that extending FAFormer to different fields, such as structure prediction and docking, is a promising research direction that we plan to explore further. A discussion will be added to our updated manuscript.

[1] Deep Learning for Protein-Ligand Docking: Are We There Yet?

[2] Independent SE(3)-Equivariant Models for End-to-End Rigid Protein Docking

[3] Accurate structure prediction of biomolecular interactions with AlphaFold 3

[4] DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking
Thank you for the responses. I believe all of my concerns have been addressed, and I look forward to seeing additional results/discussion in the updated manuscript.

I have raised my score from 5 to 7.
Thanks for your response!

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. And some important baselines are missing, e.g., [1].

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.

My score currently remains the same.

References:

[1] Brehmer, Johann, et al. "Geometric algebra transformer." NeurIPS 2023
Thank you again for providing valuable comments on our paper! All the suggestions will be added to our updated manuscript.
Thank you again for your hard work in reviewing our paper!

We hope you've had a chance to review our responses to your comments. Please let us know if you have any further questions or concerns. We greatly appreciate your feedback and are committed to clarifying the model innovation and the evaluation of the additional benchmarks.
Thank you again for your hard work in reviewing our paper!

We hope you've had a chance to review our responses to your comments. Please let us know if you have any further questions or concerns. We greatly appreciate your feedback and are committed to addressing any potential issues.
I thank the authors for addressing my concerns. My score (accept) remains the same.