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

Protein complex formation is a central problem in biology , being involved in most of the cell ’ s processes , and essential for applications , e.g . drug design or protein engineering . We tackle rigid body protein-protein docking , i.e. , computationally predicting the 3D structure of a protein-protein complex from the individual unbound structures , assuming no conformational change within the proteins happens during binding . We design a novel pairwise-independent SE ( 3 ) -equivariant graph matching network to predict the rotation and translation to place one of the proteins at the right docked position relative to the second protein . We mathematically guarantee a basic principle : the predicted complex is always identical regardless of the initial locations and orientations of the two structures . Our model , named EQUIDOCK , approximates the binding pockets and predicts the docking poses using keypoint matching and alignment , achieved through optimal transport and a differentiable Kabsch algorithm . Empirically , we achieve significant running time improvements and often outperform existing docking software despite not relying on heavy candidate sampling , structure refinement , or templates . 1 INTRODUCTION protein Z protein Z protein Z-dependent inhibitor protein Z-dependent inhibitor a. b. PDB ID : 3F1S Figure 1 : Different views of the 3D structure of a protein complex . a . Surface and b. cartoon view of protein Z and its inhibitor . In a recent breakthrough , ALPHAFOLD 2 ( Jumper et al. , 2021 ; Senior et al. , 2020 ) provides a solution to a grand challenge in biology—inferring a protein ’ s three-dimensional structure from its amino acid sequence ( Baek et al. , 2021 ) , following the dogma sequence determines structure . Besides their complex three-dimensional nature , proteins dynamically alter their function and structure in response to cellular signals , changes in the environment , or upon molecular docking . In par- ticular , protein interactions are involved in various biological processes including signal transduction , protein synthesis , DNA replication and repair . Molecular docking is key to understanding protein interactions ’ mechanisms and effects , and , subsequently , to developing therapeutic interventions . We here address the problem of rigid body protein-protein docking which refers to computationally predicting the 3D structure of a protein-protein complex given the 3D structures of the two proteins in unbound state . Rigid body means no deformations occur within any protein during binding , which is a realistic assumption in many biological settings . Popular docking software ( Chen et al. , 2003 ; Venkatraman et al. , 2009 ; De Vries et al. , 2010 ; Torchala et al. , 2013 ; Schindler et al. , 2017 ; Sunny and Jayaraj , 2021 ) are typically computationally expensive , taking between minutes and hours to solve a single example pair , while not being guaranteed to find accurate complex structures . These methods largely follow the steps : i . ) randomly sample a large number ( e.g. , millions ) of candidate initial complex structures , ii . ) employ a scoring function to rank the candidates , iii . ) adjust and refine the top complex structures based on an energy model ( e.g. , force field ) . We here take a first step towards tackling these issues by using deep learning models for direct prediction of protein complex structures . Contributions . We design EQUIDOCK , a fast , end-to-end method for rigid body docking that directly predicts the SE ( 3 ) transformation to place one of the proteins ( ligand ) at the right location and orientation with respect to the second protein ( receptor ) . Our method is based on the principle that the exact same complex structure should be predicted irrespectively of the initial 3D placements and roles of both constituents ( see Fig . 2 ) . We achieve this desideratum by incorporating the inductive biases of pairwise SE ( 3 ) –equivariance and commutativity , and deriving novel theoretical results for necessary and sufficient model constraints ( see Section 3 ) . Next , we create EQUIDOCK to satisfy these properties by design , being a combination of : i ) a novel type of pairwise independent SE ( 3 ) -equivariant graph matching networks , ii ) an attention-based keypoint selection algorithm that discovers representative points and aligns them with the binding pocket residues using optimal transport , and iii ) a differentiable superimposition model to recover the optimal global rigid transformation . Unlike prior work , our method does not use heavy candidate sampling or ranking , templates , task-specific geometric or chemical hand-crafted features , or pre-computed meshes . This enables us to achieve plausible structures with a speed-up of 80-500x compared to popular docking software , offering a promising competitive alternative to current solutions for this problem . 2 RELATED WORK . Geometric Deep Learning . Graph Neural Networks ( GNNs ) are becoming the de facto choice for learning with graph data ( Bruna et al. , 2013 ; Defferrard et al. , 2016 ; Kipf and Welling , 2016 ; Gilmer et al. , 2017 ; Xu et al. , 2018 ; Li et al. , 2019 ) . Motivated by symmetries naturally occurring in different data types , architectures are tailored to explicitly incorporate such properties ( Cohen and Welling , 2016a ; b ; Thomas et al. , 2018 ; Fuchs et al. , 2020 ; Finzi et al. , 2020 ; Eismann et al. , 2020 ; Satorras et al. , 2021 ) . GNNs are validated in a variety of tasks such as particle system dynamics or conformation-based energy estimation ( Weiler and Cesa , 2019 ; Rezende et al. , 2019 ) . Euclidean Neural Networks ( E ( 3 ) -NNs ) . However , plain GNNs and other deep learning methods do not understand data naturally lying in the 3D Euclidean space . For example , how should the output deterministically change with the input , e.g . when it is rotated ? The recent Euclidean neural networks address this problem , being designed from geometric first-principles . They make use of SE ( 3 ) - equivariant and invariant neural layers , thus avoiding expensive data augmentation strategies . Such constrained models ease optimization and have shown important improvements in biology or chemistry – e.g . for molecular structures ( Fuchs et al. , 2020 ; Hutchinson et al. , 2020 ; Wu et al. , 2021 ; Jumper et al. , 2021 ; Ganea et al. , 2021 ) and different types of 3D point clouds ( Thomas et al. , 2018 ) . Different from prior work , we here derive constraints for pairs of 3D objects via pairwise independent SE ( 3 ) -equivariances , and design a principled approach for modeling rigid body docking . Protein Folding . Deep neural networks have been used to predict inter-residue contacts , distance and/or orientations ( Adhikari and Cheng , 2018 ; Yang et al. , 2020 ; Senior et al. , 2020 ; Ju et al. , 2021 ) , that are subsequently transformed into additional constraints or differentiable energy terms for protein structure optimization . ALPHAFOLD 2 ( Jumper et al. , 2021 ) and Rosetta Fold ( Baek et al. , 2021 ) are state-of-the-art approaches , and directly predict protein structures from co-evolution information embedded in homologous sequences , using geometric deep learning and E ( 3 ) -NNs . Protein-Protein Docking and Interaction . Experimentally determining structures of protein complexes is often expensive and time-consuming , rendering a premium on computational methods ( Vakser , 2014 ) . Protein docking methods ( Chen et al. , 2003 ; Venkatraman et al. , 2009 ; De Vries et al. , 2010 ; Biesiada et al. , 2011 ; Torchala et al. , 2013 ; Schindler et al. , 2017 ; Weng et al. , 2019 ; Sunny and Jayaraj , 2021 ; Christoffer et al. , 2021 ; Yan et al. , 2020 ) typically run several steps : first , they sample thousands or millions of complex candidates ; second , they use a scoring function for ranking ( Moal et al. , 2013 ; Basu and Wallner , 2016 ; Launay et al. , 2020 ; Eismann et al. , 2020 ) ; finally , top-ranked candidates undergo a structure refinement process using energy or geometric models ( Verburgt and Kihara , 2021 ) . Relevant to protein-protein interaction ( PPI ) is the task of protein interface prediction where GNNs have showed promise ( Fout et al. , 2017 ; Townshend et al. , 2019 ; Liu et al. , 2020 ; Xie and Xu , 2021 ; Dai and Bailey-Kellogg , 2021 ) . Recently , ALPHAFOLD 2 and ROSETTAFOLD have been utilized as subroutines to improve PPIs from different aspects ( Humphreys et al. , 2021 ; Pei et al. , 2021 ; Jovine ) , e.g. , combining physics-based docking method CLUSPRO ( Kozakov et al. , 2017 ; Ghani et al. , 2021 ) , or using extended multiple-sequence alignments to predict the structure of heterodimeric protein complexes from the sequence information ( Bryant et al. , 2021 ) . Concurrently to our work , Evans et al . ( 2021 ) extend ALPHAFOLD 2 to multiple chains during both training and inference . Drug-Target Interaction ( DTI ) . DTI aims to compute drug-target binding poses and affinity , playing an essential role in understanding drugs ’ mechanism of action . Prior methods ( Wallach et al. , 2015 ; Li et al. , 2021 ) predict binding affinity from protein-ligand co-crystal structures , but such data is expensive to obtain experimentally . These models are typically based on heavy candidate sampling and ranking ( Trott and Olson , 2010 ; Koes et al. , 2013 ; McNutt et al. , 2021 ; Bao et al. , 2021 ) , being tailored for small drug-like ligands and often assuming known binding pocket . Thus , they are not immediately applicable to our use case . In contrast , our rigid docking approach is generic and could be extended to accelerate DTI research as part of future work . 3 MATHEMATICAL CONSTRAINTS FOR RIGID BODY DOCKING . We start by introducing the rigid body docking problem and derive the geometric constraints for enforcing same output complex prediction regardless of the initial unbound positions or roles ( Fig . 2 ) . Rigid Protein-Protein Docking – Problem Setup . We are given as input a pair of proteins forming a complex . They are ( arbitrarily ) denoted as the ligand and receptor , consisting of n and m residues , respectively . These proteins are represented in their bound ( docked ) state as 3D point clouds X∗1 ∈ R3×n , X∗2 ∈ R3×m , where each residue ’ s location is given by the coordinates of its corresponding α-carbon atom . In the unbound state , the docked ligand is randomly rotated and translated in space , resulting in a modified point cloud X1 ∈ R3×n . For simplicity and w.l.o.g. , the receptor remains in its bound location X2 = X∗2 . The task is to predict a rotation R ∈ SO ( 3 ) and a translation t ∈ R3 such that RX1 + t = X∗1 , using as input the proteins and their unbound positions X1 and X2 . Here , R = R ( X1|X2 ) and t = t ( X1|X2 ) are functions of the two proteins , where we omit residue identity or other protein information in this notation , for brevity . Note that we assume rigid backbone and side-chains for both proteins . We therefore do not tackle the more challenging problem of flexible docking , but our approach offers an important step towards it . We desire that the predicted complex structure is independent of the initial locations and orientations of the two proteins , as well as of their roles – see Fig . 2 . Formally , we wish to guarantee that : ( R ( Z1|Z2 ) Z1 + t ( Z1|Z2 ) ) ⊕ Z2 ≡ ( R ( X1|X2 ) X1 + t ( X1|X2 ) ) ⊕X2 , ( SE ( 3 ) -invariance ) ( R ( X1|X2 ) X1 + t ( X1|X2 ) ) ⊕X2 ≡ X1 ⊕ ( R ( X2|X1 ) X2 + t ( X2|X1 ) ) , ( commutativity ) ∀Q1 , Q2 ∈ SO ( 3 ) , ∀g1 , g2 ∈ R3 , ∀X1 ∈ R3×n , X2 ∈ R3×m , and Zl = QlXl + gl , l ∈ { 1 , 2 } . ( 1 ) KÉÏ ïËÏÊ PIRtbÉLÉÏ r.tt Ï HÉ ï ï PFRibÉLÉÏITE for any rotations Q1 , Q2 and translations g1 , g2 , where ⊕ is concatenation along columns , and ≡ denotes identity after superimposition , i.e . zero Root-Mean-Square Deviation ( RMSD ) between the two 3D point sets after applying the Kabsch algorithm ( Kabsch , 1976 ) . An immediate question arises : How do the constraints in Eq . ( 1 ) translate into constraints for R ( ·|· ) and t ( ·|· ) ? The rotation R and translation t change in a systematic way when we apply SE ( 3 ) transformations or swap proteins ’ roles . These properties restrict our class of functions as derived below . SE ( 3 ) -equivariance Constraints . If we apply any distinct SE ( 3 ) transformations on the unbound ligand X1 and receptor X2 , i.e . we now dock Q1X1 + g1 onto Q2X2 + g2 , then the rotation matrix R ( Q1X1 +g1|Q2X2 +g2 ) and translation vector t ( Q1X1 +g1|Q2X2 +g2 ) can be derived from the original R ( X1|X2 ) and t ( X1|X2 ) assuming that we always do rotations first . In this case , R ( Q1X1 + g1|Q2X2 + g2 ) can be decomposed into three rotations : i . ) apply Q > 1 to undo the rotation Q1 applied on X1 , ii . ) apply R ( X1|X2 ) , iii . ) apply Q2 to rotate the docked ligand together with the receptor . This gives R ( Q1X1 + g1|Q2X2 + g2 ) = Q2R ( X1|X2 ) Q > 1 , which in turn constraints the translation vector . We provide a formal statement and prove it in Appendix B.1 : Proposition 1 . For any Q1 , Q2 ∈ SO ( 3 ) , g1 , g2 ∈ R3 , SE ( 3 ) -invariance of the predicted docked complex defined by Eq . ( 1 ) is guaranteed iff R ( Q1X1 + g1|Q2X2 + g2 ) = Q2R ( X1|X2 ) Q > 1 t ( Q1X1 + g1|Q2X2 + g2 ) = Q2t ( X1|X2 ) −Q2R ( X1|X2 ) Q > 1 g1 + g2 . ( 2 ) As a direct consequence of this proposition , we have the following statement . Proposition 2 . Any model satisfying Proposition 1 guarantees invariance of the predicted complex w.r.t . any SE ( 3 ) transformation on X1 , and equivariance w.r.t . any SE ( 3 ) transformation on X2 : R ( Z1|X2 ) Z1 + t ( Z1|X2 ) = R ( X1|X2 ) X1 + t ( X1|X2 ) , where Z1 = Q1X1 + g1 R ( X1|Z2 ) X1 + t ( X1|Z2 ) = Q2 [ R ( X1|X2 ) X1 + t ( X1|X2 ) ] + g2 , where Z2 = Q2X2 + g2 ∀Q1 , Q2 ∈ SO ( 3 ) , ∀g1 , g2 ∈ R3 , ∀X1 ∈ R3×n , ∀X2 ∈ R3×m . Commutativity . Instead of docking X1 with respect to X2 , we can also dock X2 with respect to X1 . In this case , we require the final complex structures to be identical after superimposition , i.e. , zero RMSD . This property is named commutativity and it is satisfied as follows ( proof in Appendix B.2 ) . Proposition 3 . Commutativity as defined by Eq . ( 1 ) is guaranteed iff R ( X2|X1 ) = R > ( X1|X2 ) ; t ( X2|X1 ) = −R > ( X1|X2 ) t ( X1|X2 ) , ( 3 ) Point Permutation Invariance . We also enforce residue permutation invariance . Formally , both R ( X1|X2 ) and t ( X1|X2 ) should not depend on the order or columns of X1 and , resp. , of X2 .

Summary: This paper proposes an end-to-end deep learning architecture to model the rigid body protein docking problem. By incorporating the inductive biases of SE(3)-invariance of the final docking position and commutativity, the proposed method avoids the millions of sampling and achieves a competitive performance with much faster speed. In addition, they discover and align keypoints by the attention-based selection algorithm and use optimal transport to predict the binding pocket location based on those selected points. The main contribution of this paper is the combination of the novel graph matching networks and keypoint selection algorithm to predict the 3D position of the docking model. Results are shown on the task of 1) Protein docking complex prediction 2) Runtime