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

Introduction

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 particular, 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.

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.

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).

Here, $\mathbf{R}=\mathbf{R}(\mathbf{X}_{1}|\mathbf{X}_{2})$ and $\mathbf{t}=\mathbf{t}(\mathbf{X}_{1}|\mathbf{X}_{2})$ 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:

for any rotations $\mathbf{Q}_{1},\mathbf{Q}_{2}$ and translations $\mathbf{g}_{1},\mathbf{g}_{2}$, where $\oplus$ is concatenation along columns, and $\equiv$ 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 $\mathbf{R}(\cdot|\cdot)$ and $\mathbf{t}(\cdot|\cdot)$ ?

The rotation $\mathbf{R}$ and translation $\mathbf{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 $\mathbf{X}_{1}$ and receptor $\mathbf{X}_{2}$, i.e. we dock $\mathbf{Q}_{1}\mathbf{X}_{1}+\mathbf{g}_{1}$ onto $\mathbf{Q}_{2}\mathbf{X}_{2}+\mathbf{g}_{2}$, then the rotation matrix $\mathbf{R}(\mathbf{Q}_{1}\mathbf{X}_{1}+\mathbf{g}_{1}|\mathbf{Q}_{2}\mathbf{X}_{2}+\mathbf{g}_{2})$ and translation vector $\mathbf{t}(\mathbf{Q}_{1}\mathbf{X}_{1}+\mathbf{g}_{1}|\mathbf{Q}_{2}\mathbf{X}_{2}+\mathbf{g}_{2})$ can be derived from the original $\mathbf{R}(\mathbf{X}_{1}|\mathbf{X}_{2})$ and $\mathbf{t}(\mathbf{X}_{1}|\mathbf{X}_{2})$ assuming that we always do rotations first. In this case, $\mathbf{R}(\mathbf{Q}_{1}\mathbf{X}_{1}+\mathbf{g}_{1}|\mathbf{Q}_{2}\mathbf{X}_{2}+\mathbf{g}_{2})$ can be decomposed into three rotations: i.) apply $\mathbf{Q}_{1}^{\top}$ to undo the rotation $\mathbf{Q}_{1}$ applied on $\mathbf{X}_{1}$, ii.) apply $\mathbf{R}(\mathbf{X}_{1}|\mathbf{X}_{2})$, iii.) apply $\mathbf{Q}_{2}$ to rotate the docked ligand together with the receptor. This gives $\mathbf{R}(\mathbf{Q}_{1}\mathbf{X}_{1}+\mathbf{g}_{1}|\mathbf{Q}_{2}\mathbf{X}_{2}+\mathbf{g}_{2})=\mathbf{Q}_{2}\mathbf{R}(\mathbf{X}_{1}|\mathbf{X}_{2})\mathbf{Q}_{1}^{\top}$, which in turn constraints the translation vector. We provide a formal statement and prove it in Section B.1:

As a direct consequence of this proposition, we have the following statement.

Any model satisfying Eq. 2 guarantees invariance of the predicted complex w.r.t. any SE(3) transformation on $\mathbf{X}_{1}$, and equivariance w.r.t. any SE(3) transformation on $\mathbf{X}_{2}$:

Commutativity. Instead of docking $\mathbf{X}_{1}$ with respect to $\mathbf{X}_{2}$, we can also dock $\mathbf{X}_{2}$ with respect to $\mathbf{X}_{1}$. 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 Section B.2).

Commutativity as defined by Eq. 1 is guaranteed iff

Point Permutation Invariance. We also enforce residue permutation invariance. Formally, both $\mathbf{R}(\mathbf{X}_{1}|\mathbf{X}_{2})$ and $\mathbf{t}(\mathbf{X}_{1}|\mathbf{X}_{2})$ should not depend on the order or columns of $\mathbf{X}_{1}$ and, resp., of $\mathbf{X}_{2}$.

EquiDock Model

Protein Representation. A protein is a sequence of amino acid residues that folds in a 3D structure. Each residue has a general structure with a side-chain specifying its type, allowing us to define a local frame and derive SE(3)-invariant features for any pair of residues —see Appendix A.

Next, we apply several layers consisting of functions $\Phi$ that jointly transform node coordinates and features. Crucially, we guarantee, by design, pairwise independent SE(3)-equivariance for coordinate embeddings, and invariance for feature embeddings. This double constraint is formally defined:

We implement $\Phi$ as a novel type of message-passing neural network (MPNN). We then use the output node coordinate and feature embeddings to compute $\mathbf{R}(\mathbf{X}_{1}|\mathbf{X}_{2})$ and $\mathbf{t}(\mathbf{X}_{1}|\mathbf{X}_{2})$. These functions depend on pairwise interactions between the two proteins modeled as cross-messages, but also incorporate the 3D structure in a pairwise-independent SE(3)-equivariant way to satisfy Eq. 1, Eq. 2 and Eq. 3. We discover keypoints from each protein based on a neural attention mechanism and softly guide them to represent the respective binding pocket locations via an optimal transport based auxiliary loss. Finally, we obtain the SE(3) transformation by superimposing the two keypoint sets via a differentiable version of the Kabsch algorithm. An additional soft-constraint discourages point cloud intersections. We now detail each of these model components.

Independent E(3)-Equivariant Graph Matching Networks (IEGMNs). Our architecture for $\Phi$ satisfying Eq. 4 is called Independent E(3)-Equivariant Graph Matching Network (IEGMN) – see Fig. 3. It extends both Graph Matching Networks (GMN) (Li et al., 2019) and E(3)-Equivariant Graph Neural Networks (E(3)-GNN) (Satorras et al., 2021). IEGMNs perform node coordinate and feature embedding updates for an input pair of graphs $\mathcal{G}_{1}=(\mathcal{V}_{1},\mathcal{E}_{1})$, $\mathcal{G}_{2}=(\mathcal{V}_{2},\mathcal{E}_{2})$, and use inter- and intra- node messages, as well as E(3)-equivariant coordinate updates. The l-th layer of IEGMNs transforms node latent/feature embeddings $\{\mathbf{h}_{i}^{(l)}\}_{i\in\mathcal{V}_{1}\cup\mathcal{V}_{2}}$ and node coordinate embeddings $\{\mathbf{x}_{i}^{(l)}\}_{i\in\mathcal{V}_{1}\cup\mathcal{V}_{2}}$ as

Note that all parameters of $\mathbf{W},\varphi^{x},\varphi^{h},\varphi^{e},\psi^{q},\psi^{k}$ can be shared or different for different IEGMN layers . The output of several IEGMN layers is then denoted as:

It is then straightforward to prove the following (see Section B.3):

IEGMNs satisfy the pairwise independent SE(3)-equivariance property in Eq. 4.

where $\mu(\cdot)$ is the mean vector of a point cloud. It is straightforward to prove that this model satisfies all the equivariance properties in Eqs. 1, 2 and 3. From a practical perspective, the gradient and backpropagation through the SVD operation was analyzed by (Ionescu et al., 2015; Papadopoulo and Lourakis, 2000) and implemented in the automatic differentiation frameworks such as PyTorch.

Optimal Transport and Binding Pocket Keypoint Alignment. As stated before, we desire that $\mathbf{Y}_{1}$ and $\mathbf{Y}_{2}$ are representative points for the binding pocket location of the respective protein pair. However, this needs to be encouraged explicitly, which we achieve using an additional loss.

We desire that $\mathbf{Y}_{1}$ is a representative set for the 3D set $\mathbf{P}_{1}$ (and, similarly, $\mathbf{Y}_{2}$ for $\mathbf{P}_{2}$). However, while at training time we know that every point $\mathbf{p}_{1s}$ corresponds to the point $\mathbf{p}_{2s}$ (and, similarly, $\mathbf{y}_{1k}$ aligns with $\mathbf{y}_{2k}$, by assumption), we unfortunately do not know the actual alignment between points in $\mathbf{Y}_{l}$ and $\mathbf{P}_{l}$, for every $l\in\{1,2\}$. This can be recovered using an additional optimal transport loss:

where $\mathcal{U}(S,K)$ is the set of $S\times K$ transport plans with uniform marginals. The optimal transport plan is computed using an Earth Mover’s Distance and the POT library (Flamary et al., 2021), while being kept fixed during back-propagation and optimization when only the cost matrix is trained.

Note that our approach assumes that $\mathbf{y}_{1k}$ corresponds to $\mathbf{y}_{2k}$, for every $k\in\{1,\ldots,K\}$. Intuitively, each attention head $k$ will identify a specific geometric/chemical local surface feature of protein 1 by $\mathbf{y}_{1k}$, and match its complementary feature of protein 2 by $\mathbf{y}_{2k}$.

Surface Aware Node Features. Surface contact modeling is important for protein docking. We here design a novel surface feature type that differentiates residues closer to the surface of the protein from those in the interior. Similar to Sverrisson et al. (2021), we prioritize efficiency and avoid pre-computing meshes, but show that our new feature is a good proxy for residue’s depth (i.e. distance to the protein surface). Intuitively, residues in the core of the protein are locally surrounded in all directions by other residues. This is not true for residues on the surface, e.g., neighbors are in a half-space if the surface is locally flat. Building on this intuition, for each node (residue) $i$ in the $k$-NN protein graph, we compute the norm of the weighted average of its neighbor forces, which can be interpreted as the normalized gradient of the $G(\mathbf{x})$ surface function. This SE(3)-invariant feature is

Intuitively, as depicted in Fig. 8, residues in the interior of the protein have values close to 0 since they are surrounded by vectors from all directions that cancel out, while residues near the surface have neighbors only in a narrower cone, with aperture depending on the local curvature of the surface. We show in Appendix C that this feature correlates well with more expensive residue depth estimation methods, e.g. based on MSMS, thus offering a computationally appealing alternative. We also compute an estimation of this feature for large dense point clouds based on the local surface angle.

Experiments

Datasets. We leverage the following datasets: Docking Benchmark 5.5 (DB5.5) (Vreven et al., 2015) and Database of Interacting Protein Structures (DIPS) (Townshend et al., 2019). DB5.5 is a gold standard dataset in terms of data quality, but contains only 253 structures. DIPS is a larger protein complex structures dataset mined from the Protein Data Bank (Berman et al., 2000) and tailored for rigid body docking. Datasets information is given in Appendix D. We filter DIPS to only keep proteins with at most 10K atoms. Datasets are then randomly partitioned in train/val/test splits of sizes 203/25/25 (DB5.5) and 39,937/974/965 (DIPS). For DIPS, the split is based on protein family to separate similar proteins. For the final evaluation in Table 1, we use the full DB5.5 test set, and randomly sample 100 pairs from different protein families from the DIPS test set.

Baselines. We compare our EquiDock method with popular state-of-the-art docking software ClusPro: https://cluspro.bu.edu/, Attract: www.attract.ph.tum.de/services/ATTRACT/ATTRACT.vdi.gz, PatchDock: https://bioinfo3d.cs.tau.ac.il/PatchDock/, HDOCK: http://huanglab.phys.hust.edu.cn/software/HDOCK/ ClusPro (Piper) (Desta et al., 2020; Kozakov et al., 2017),Attract (Schindler et al., 2017; de Vries et al., 2015), PatchDock (Mashiach et al., 2010; Schneidman-Duhovny et al., 2005), and HDock (Yan et al., 2020; 2017b; 2017a; Huang and Zou, 2014; 2008). These baselines provide user-friendly local packages suitable for automatic experiments or webservers for manual submissions.

Training Details. We train our models on the train part of DIPS first, using Adam (Kingma and Ba, 2014) with learning rate 2e-4 and early stopping with patience of 30 epochs. We update the best validation model only when it achieves a score of less than 98% of the previous best validation score, where the score is the median of Ligand RMSD on the full DIPS validation set. The best DIPS validated model is then tested on the DIPS test set. For DB5.5, we fine tune the DIPS pre-trained model on the DB5.5 training set using learning rate 1e-4 and early stopping with 150 epochs patience. The best DB5.5 validated model is finally tested on DB5.5 test set. During training, we randomly assign the roles of ligand and receptor. Also, during both training and testing, we randomly rotate and translate the ligand in space (even though our model is invariant to this operation) for all baselines.

Complex Prediction Results. Results are shown in Table 1, Fig. 4 and Appendix E. We note that our method is competitive and often outperforms the baselines. However, we do not use heavy candidate sampling and re-ranking, we do not rely on task-specific hand-crafted features, and we currently do not perform structure fine-tuning, aiming to predict the SE(3) ligand transformation in a direct shot. Moreover, we note that some of the baselines might have used part of our test set in validating their models, for example to learn surface templates, thus, their reported scores might be optimistic. Notably, HDock score function was validated on DB4 which overlaps with DB5.5. A more appropriate comparison would require us to re-build these baselines without information from our test sets, a task that is currently not possible without open-source implementations.

Computational Efficiency. We show inference times in Fig. 5 and Table 4. Note that EquiDock is between 80-500 times faster than the baselines. This is especially important for intensive screening applications that aim to scan over vast search spaces, e.g. for drug discovery. In addition, it is also relevant for de novo design of binding proteins (e.g. antibodies (Jin et al., 2021)) or for use cases when protein docking models are just a component of significantly larger end-to-end architectures targeting more involved biological scenarios, e.g., representing a drug’s mechanism of action or modeling cellular processes with a single model as opposed to a multi-pipeline architecture.

Visualization. We show in Fig. 6 a successful example of a test DIPS protein pair for which our model significantly outperforms all baselines.

Conclusion

We have presented an extremely fast, end-to-end rigid protein docking approach that does not rely on candidate sampling, templates, task-specific features or pre-computed meshes. Our method smartly incorporates useful rigid protein docking priors including commutativity and pairwise independent SE(3)-equivariances, thus avoiding the computational burden of data augmentation.

We look forward to incorporating more domain knowledge in EquiDock and extend it for flexible docking and docking molecular dynamics, as well as adapt it to other related tasks such as drug binding prediction. On the long term, we envision that fast and accurate deep learning models would allow us to tackle more complex and involved biological scenarios, for example to model the mechanism of action of various drugs or to design de novo binding proteins and drugs to specific targets (e.g. for antibody generation). Last, we hope that our architecture can inspire the design of other types of biological 3D interactions.

Limitations. First, our presented model does not incorporate protein flexibility which is necessary for various protein families, e.g., antibodies. Unfortunately, both DB5 and DIPS datasets are biased towards rigid body docking . Second, we only prevent steric clashes using a soft constraint (Eq. 15) which has limitations (see Table 6). Future extensions would hard-constrain the model to prevent such artifacts.

Acknowledgements

The authors thank Hannes Stärk, Gabriele Corso, Patrick Walters, Tian Xie, Xiang Fu, Jacob Stern, Jason Yim, Lewis Martin, Jeremy Wohlwend, Jiaxiang Wu, Wei Liu, and Ding Xue for insightful and helpful discussions. OEG is funded by the Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, the Abdul Latif Jameel Clinic for Machine Learning in Health, the DTRA Discovery of Medical Countermeasures Against New and Emerging (DOMANE) threats program, and the DARPA Accelerated Molecular Discovery program. This publication was created as part of NCCR Catalysis (grant number 180544), a National Centres of Competence in Research funded by the Swiss National Science Foundation. RB and TJ also acknowledge support from NSF Expeditions grant (award 1918839): Collaborative Research: Understanding the World Through Code.

References

Appendix A Representing Proteins as Graphs

A protein is comprised of amino acid residues. The structure of an amino acid residue is shown in Figure Fig. 7. Generally, an amino acid residue contains amino (-NH-), $\alpha$-carbon atom and carboxyl (-CO-), along with a side chain (R) connected with the $\alpha$-carbon atom. The side chain (R) is specific to each type of amino acid residues.

The neighborhood of a node is the set of $k$ ($k=10$ in our experiments) nearest nodes where the distance is the Euclidean distance between 3D coordinates.

Node feature is a one dimension indicator (one-hot encoding) of the type of amino acid residue. This one dimension indicator will be passed into an embedding layer.

Similar to Ingraham et al. and Jumper et al. , we introduce a local coordinate system for each residue which denotes the orientation of a residue. Based on this, we can further design SE(3)-invariant edge features. As shown in Figure 7, for a residue $i$, we denote the unit vector pointing from $\alpha$-carbon atom to nitrogen atom as $\mathbf{u}_{i}$. We denote the unit vector pointing from $\alpha$-carbon atom to carbon atom of the carboxyl (-CO-) as $\mathbf{t}_{i}$. $\mathbf{u}_{i}$ and $\mathbf{t}_{i}$ together define a plane, and the normal of this plane is $\mathbf{n}_{i}=\frac{\mathbf{u}_{i}\times\mathbf{t}_{i}}{\lVert\mathbf{u}_{i}\times\mathbf{t}_{i}\rVert}$. Finally, we define $\mathbf{v}_{i}=\mathbf{n}_{i}\times\mathbf{u}_{i}$. Then $\mathbf{n}_{i}$, $\mathbf{u}_{i}$ and $\mathbf{v}_{i}$ together form the basis of residue $i$’s local coordinate system. They together encode the orientation of residue $i$.

Then we introduce the edge features of an edge $j\to i\in\mathcal{E}$. These features describe the relative position of $j$ with respect to $i$, the relative orientation of $j$ with respect to $i$ and the distance between $j$ and $i$.

Relative Position Edge Features

First we introduce the edge features $\mathbf{p}_{j\to i}$ which describe relative position of $j$ with respect to $i$:

Relative Orientation Edge Features

As we mention above, each residue has orientation which carries information. Here we introduce the edge features $\mathbf{q}_{j\to i}$, $\mathbf{k}_{j\to i}$ and $\mathbf{t}_{j\to i}$ which describe relative orientation of $j$ with respect to $i$:

Distance-Based Edge Features

Distance also carries information. Here we use radial basis function of distance as edge features:

Where $R$ and scale parameters $\{\sigma_{r}\}_{1\leq r\leq R}$ are hyperparameters. In experiments, the set of scale parameters we used is $\{1.5^{x}|x=0,1,2,...,14\}$. So for each edge, there are 15 distance-based edge features.

Surface Aware Node Features

We additionally compute 5 surface aware node features defined in Eq. 16 using $\lambda\in\{1.,2.,5.,10.,30.\}$.

Appendix B Proofs of the Main Propositions

B.2 Proof of Eq. 3.

B.3 Proof of Proposition 4.

We first note that the equations of our proposed IEGMN layer that compute messages $\mathbf{m}_{j\rightarrow i}$, $\mu_{j\rightarrow i}$, $\mathbf{m}_{i}$ and $\mu_{i}$ are SE(3)-invariant. Indeed, they depend on the initial features which are SE(3)-invariant by design, the current latent node embeddings $\{\mathbf{h}_{i}^{(l)}\}_{i\in\mathcal{V}_{1}\cup\mathcal{V}_{2}}$, as well as the Euclidean distances on the current node coordinates $\{\mathbf{x}_{i}^{(l)}\}_{i\in\mathcal{V}_{1}\cup\mathcal{V}_{2}}$. Thus, we also derive that the equation that computes the new latent node embeddings $\mathbf{h}_{i}^{(l+1)}$ is SE(3)-invariant. Last, the equation that updates the coordinates $\mathbf{x}_{i}^{(l+1)}$ is SE(3)-equivariant with respect to the 3D coordinates of nodes from the same graph as $i$, but SE(3)-invariant with respect to the 3D coordinates of nodes from the other graph since it only uses invariant transformations of the latter.

Appendix C Surface Features

We further discuss our new surface features introduced in Eq. 16. We first visualize their design intuition in Fig. 8. A synthetic experiment is shown in Fig. 9.

Correlation with MSMS features.

Next, we analyze how accurate are these features compared to established residue depth estimation methods, e.g. based on the MSMS software [Sanner et al., 1996]. We plot the Spearman rank-order correlation of the two methods in Fig. 10. We observe a concentrated distribution with a mean of 0.68 and a median of 0.70, suggesting a strong correlation with the MSMS depth estimation.

Closed form expression.

Finally, we prove that for points close to the protein surface and surrounded by (infinitely) many equally-distanced and equally-spaced points, one can derive a closed form expression of the surface features defined in Eq. 16. See Fig. 11. We work in 2 dimensions, but extensions to 3 dimensions are straightforward. Assume that the local surface at point $\mathbf{x}_{i}$ has angle $\alpha$. Further, assume that $\mathbf{x}_{i}$ is surrounded by $N$ equally-distanced and equally-spaced points denoted by $\mathbf{x}_{i}^{\prime}$. Then, all $w_{i,i^{\prime},\lambda}$ will be identical. Then, the summation vector in the numerator of Eq. 16 will only have non-zero components on the direction that bisects the surface angle, as the other components will cancel-out. Then, under the limit $N\rightarrow\infty$, we derive the closed form expression:

Appendix D Datasets

The overview of datasets is in Table 2. DB5.5 is obtained from https://zlab.umassmed.edu/benchmark/, while DIPS is downloaded from https://github.com/drorlab/DIPS. While DIPS contains only the bound structures, thus currently being only suitable for rigid docking, DB5.5 also includes unbound protein structures, however, mostly showing rigid structures - see Fig. 12.

Appendix E More Experimental Details and Results

On the test sets, Attract fails for ’1N2C’ in DB5.5, ’oi_4oip.pdb1_8’, ’oi_4oip.pdb1_3’ and ’p7_4p7s.pdb1_2’ in DIPS. For such failure cases, we use the unbound input structure as the prediction for metrics calculation.

Hyperparameters.

We perform hyperparameter search over the choices listed in Table 3 and select the best hyperparameters for DB5.5 and DIPS respectively based on their corresponding validation sets.

Detailed Running Times.

In addition to the main text, we show in Table 4 detailed running times of all methods. Hardware specifications are as follows: Attract was run on a 6-Core Intel Core i7 2.2 GHz CPU; HDock was run on a single Intel Xeon Gold 6230 2.1 GHz CPU; EquiDock was run on a single Intel Core i9-9880H 2.3 GHz CPU. ClusPro and PatchDock have been manually run using their respective web servers.

Plots for DB5.5.

We show the corresponding plots for DB5.5 results in Fig. 13.

Ablation Studies.

To highlight contributions of different model components, we provide ablation studies in Table 5. One can note that, as expected, removing the pocket loss results in lower interface RMSD scores compared to removing other components.

Analysis of the Intersection Loss.

We further analyze the intersection loss introduced in Eq. 15 with parameters $\gamma=10$ and $\sigma=25$ (chosen on DB5 validation set). We show in Table 6 that this loss achieves almost perfect values for the ground truth structures, being important to softly constrain non-intersecting predicted proteins.