END-TO-END RIGID PROTEIN DOCKING
Year: 2021

Abstract:
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.

text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 MATHEMATICAL CONSTRAINTS FOR RIGID BODY DOCKING
text: 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.

heading: 4 EQUIDOCK MODEL
text: 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.
We represent a protein as a graph G = (V, E), similar to Fout et al. (2017); Townshend et al. (2019); Liu et al. (2020). Each node i ∈ V represents one residue and has 3D coordinates xi ∈ R3 corresponding to the α-carbon atom’s location. Edges are given by a k-nearest-neighbor (k-NN) graph using Euclidean distance of the 3D node coordinates.
Overview of Our Approach. Our model is depicted in Fig. 3. We first build k-NN protein graphs G1 = (V1, E1) and G2 = (V2, E2). We then design SE(3)-invariant node features F1 ∈ Rd×n,F2 ∈ Rd×m and edge features {fj→i : ∀(i, j) ∈ E1 ∪ E2} (see Appendix A). Next, we apply several layers consisting of functions Φ 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:
Given Z1,H1,Z2,H2 = Φ(X1,F1,X2,F2) we have Q1Z1 + g1,H1,Q2Z2 + g2,H2 = Φ(Q1X1 + g1,F1,Q2X2 + g2,F2),
∀Q1,Q2 ∈ SO(3),∀g1,g2 ∈ R3. (4)
We implement Φ as a novel type of message-passing neural network (MPNN). We then use the output node coordinate and feature embeddings to compute R(X1|X2) and t(X1|X2). 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), Proposition 1 and Proposition 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 Φ 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 G1 = (V1, E1), G2 = (V2, E2), and use inter- and intranode messages, as well as E(3)-equivariant coordinate updates. The l-th layer of IEGMNs transforms node latent/feature embeddings {h(l)i }i∈V1∪V2 and node coordinate embeddings {x (l) i }i∈V1∪V2 as
mj→i = ϕ e(h (l) i ,h (l) j , exp(−‖x (l) i − x (l) j ‖2/σ), fj→i),∀ej→i ∈ E1 ∪ E2 (5)
µj→i = aj→iWh (l) j ,∀i ∈ V1, j ∈ V2 or i ∈ V2, j ∈ V1 (6)
mi = 1 |N (i)| ∑
j∈N (i) mj→i,∀i ∈ V1 ∪ V2 (7)
µi = ∑
j∈V2 µj→i,∀i ∈ V1, and µi =
∑ j∈V1 µj→i,∀i ∈ V2 (8)
x (l+1) i = ηx (0) i + (1− η)x (l) i +
∑
j∈N (i) (x
(l) i − x (l) j )ϕ x(mj→i),∀i ∈ V1 ∪ V2 (9)
h (l+1) i = (1− β) · h (l) i + β · ϕh(h (l) i ,mi, µi, fi),∀i ∈ V1 ∪ V2, (10)
whereN (i) are the neighbors of node i; ϕx is a real-valued (scalar) function; W is a learnable matrix; ϕh, ϕe are functions outputting a vector Rd; fj→i and fi are the original edge and node features (extracted SE(3)-invariantly from the residues); aj→i is an attention based coefficient with trainable shallow neural networks ψq and ψk:
aj→i = exp(〈ψq(h(l)i ), ψk(h (l) j )〉)∑
j′ exp(〈ψq(h (l) i ), ψ k(h (l) j′ )〉)
, (11)
The output of several IEGMN layers is then denoted as:
Z1 ∈ R3×n,H1 ∈ Rd×n,Z2 ∈ R3×m,H2 ∈ Rd×m = IEGMN(X1,F1,X2,F2). (12)
It is then straightforward to prove the following (see Appendix B.3):
Proposition 4. IEGMNs satisfy the pairwise independent SE(3)-equivariance property in Eq. (4).
Keypoints for Differentiable Protein Superimposition. Next, we use multi-head attention to obtain K points for each protein, Y1,Y2 ∈ R3×K , which we name keypoints. We train them to become representative points for the binding pocket of the respective protein pair (softly-enforced by an additional loss described later). If this would holds perfectly, then the superimposition of Y1 and Y2 would give the corresponding ground truth superimposition of X1 and X2. Our model is :
y1k :=
n∑
i=1
αki z1i; y2k :=
m∑
j=1
βkj z2j ,
where z1i denotes the i-th column of matrix Z1, and αki = softmaxi( 1√ d h>1iWµ(ϕ(H2))) are attention scores (similarly defined for βkj ), with W ∈ Rd×d a parametric matrix, ϕ a linear layer plus a LeakyReLU non-linearity, and µ(·) being the mean vector.
Differentiable Kabsch Model. We design the rotation and translation that docks protein 1 into protein 2 to be the same transformation used to superimpose Y1 and Y2 — see Fig. 3. For this, we compute a differentiable version of the Kabsch algorithm (Kabsch, 1976) as follows. Let A = Y2Y > 1 ∈ R3×3 computed using zero-mean keypoints. The singular value decomposition (SVD) is A = U2SU>1 , where U2,U1 ∈ O(3). Finally, we define the differentiable functions
R(X1|X2; θ) = U2 (
1 0 0 0 1 0 0 0 d
) U>1 , where d = sign(det(U2U > 1 ))
t(X1|X2; θ) = µ(Y2)−R(X1|X2; θ)µ(Y1), (13)
where µ(·) is the mean vector of a point cloud. It is straightforward to prove that this model satisfies all the equivariance properties in Eqs. (1) to (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.
MSE Loss. During training, we randomly decide which protein is the receptor (say protein 2), keep it in the docked position (i.e., X2 = X∗2), predict the SE(3) transformation using Eq. (13) and use it to compute the final position of the ligand as X̃1 = R(X1|X2)X1 + t(X1|X2). The mean squared error (MSE) loss is then LMSE = 1n ∑n i=1 ‖x∗i − x̃i‖2.
Optimal Transport and Binding Pocket Keypoint Alignment. As stated before, we desire that Y1 and Y2 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 first define the binding pocket point sets, inspiring from previous PPI work (Section 2). Given the residues’ α-carbon locations of the bound (docked) structures, X∗1 and X ∗ 2, we select all pairs of residues at less than τ Euclidean distance (τ = 8Å in our experiments). We assume these are all interacting residues. Denote these pairs as {(x∗1s,x∗2s), s ∈ 1, . . . , S}, where S is variable across data pairs. We compute midpoints of these segments, denoted as P∗1,P ∗ 2 ∈ R3×S , where p∗1s = p ∗ 2s = 0.5 · (x∗1s + x∗2s). We view P∗1 and P∗2 as binding pocket points. In the unbound state, these sets are randomly moved in space together with the respective protein residue coordinates X1 and X2. We denote them as P1,P2 ∈ R3×S . For clarity, if X1 = QX∗1 + g, then P1 = QP∗1 + g.
We desire that Y1 is a representative set for the 3D set P1 (and, similarly, Y2 for P2). However, while at training time we know that every point p1s corresponds to the point p2s (and, similarly, y1k aligns with y2k, by assumption), we unfortunately do not know the actual alignment between points in Yl and Pl, for every l ∈ {1, 2}. This can be recovered using an additional optimal transport loss:
LOT = min T∈U(S,K) 〈T,C〉, where Cs,k = ‖y1k − p1s‖2 + ‖y2k − p2s‖2, (14)
where U(S,K) is the set of S ×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 y1k corresponds to y2k, for every k ∈ {1, . . . ,K}. Intuitively, each attention head k will identify a specific geometric/chemical local surface feature of protein 1 by y1k, and match its complementary feature of protein 2 by y2k.
Avoiding Point Cloud Intersection. In practice, our model does not enforce a useful inductive bias, namely that proteins forming complexes are never "intersecting" with each other. To address this issue, we first state a notion of the "interior" of a protein point cloud. Following previous work (Sverrisson et al., 2021; Venkatraman et al., 2009), we define the surface of a protein point cloud X ∈ R3×n as {x ∈ R3 : G(x) = γ}, where G(x) = −σ ln(∑ni=1 exp(−||x − xi||2/σ)). The parameters σ and γ are chosen such that there exist no "holes" inside a protein (we found γ = 10, σ = 25 to work well, see Appendix E). As a consequence, the interior of the protein is given by {x ∈ R3 : G(x) < γ}. Then, the condition for non-intersecting ligand and receptor can be written as G1(x2j) > γ, ∀j ∈ 1, . . . ,m and G2(x1i) > γ, ∀i ∈ 1, . . . , n. As a loss function, this becomes
LNI = 1
n
n∑
i=1
max(0, γ −G2(x1i)) + 1
m
m∑
j=1
max(0, γ −G1(x2j)). (15)
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(x) surface function. This SE(3)-invariant feature is
ρi(xi;λ) = ‖∑i′∈Ni wi,i′,λ(xi − xi′)‖∑
i′∈Ni wi,i′,λ‖xi − xi′‖ , where wi,i′,λ = exp(−||xi − xi′ ||2/λ)∑ j∈Ni exp(−||xi − xj ||2/λ) . (16)
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. For example, in 2D, 2π is the ρ value for flat surface points with uniformly spread neighbors. 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.

heading: 5 EXPERIMENTS
text: 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/validation/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 (PIPER)1 (Desta et al., 2020; Kozakov et al., 2017),ATTRACT2 (Schindler et al., 2017; de Vries et al., 2015), PATCHDOCK3 (Mashiach et al., 2010; SchneidmanDuhovny et al., 2005), and HDOCK4 (Yan et al., 2020; 2017b;a; Huang and Zou, 2014; 2008). These baselines provide user-friendly local packages suitable for automatic experiments or webservers for manual submissions.
Evaluation Metrics. For measuring the quality of the predicted complex structures, we report Complex Root Mean Square Deviation (CRMSD) and Interface Root Mean Square Deviation (IRMSD), defined below. Given the ground truth and predicted complex structures, Z∗ ∈
R3×(n+m) and Z ∈ R3×(n+m), we first superimpose them using the Kabsch algorithm (Kabsch, 1976), and then compute C-RMSD = √ 1
n+m‖Z∗ − Z‖2F . We compute I-RMSD similarly, but using only the coordinates of the interface residues with distance less than 8Å to the other protein’s residues. For a fair comparison among baselines, we use only the α-carbon coordinates to compute both metrics.
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
1https://cluspro.bu.edu/ 2We use the VirtualBox: www.attract.ph.tum.de/services/ATTRACT/ATTRACT.vdi.gz 3https://bioinfo3d.cs.tau.ac.il/PatchDock/ 4http://huanglab.phys.hust.edu.cn/software/HDOCK/
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. 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. We 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, for example representing a drug’s mechanism of action or modeling cellular processes with a single machine learning 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.

heading: 6 CONCLUSION
text: 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.
ETHICS STATEMENT
The protein-protein docking problem is a key step for understanding cell mechanisms since protein complexes are essential for virtually all cellular processes. Technologies that target this problem have the potential to drastically change the current paradigm in computational methods for impactful areas related to drug discovery, but could also potentially be mistakenly misused. However, we note that our model, as well as existing protein docking software, have yet to achieve a more realistic qualitative level and proved to be more interpretable and reliable when transferred to different types of proteins.

heading: A Representing Proteins as Graphs 15

heading: B Proofs of the Main Propositions 16
text: B.1 Proof of Proposition 1. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16
B.2 Proof of Proposition 3. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 17
B.3 Proof of Proposition 4. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 17

heading: C Surface Features 17

heading: D Datasets 19

heading: E More Experimental Details and Results 19

heading: A REPRESENTING PROTEINS AS GRAPHS
text: 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-), α-carbon atom and carboxyl (-CO-), along with a side chain (R) connected with the α-carbon atom. The side chain (R) is specific to each type of amino acid residues.
We work on residue level (our approaches can be extended to atom level as well). A protein is represented by a set of nodes where each node is an amino acid residue in the protein. Each node i has a 3D coordinate xi ∈ R3 which is the 3D coordinate of α-carbon atom of the residue. 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.
Local Coordinate System. Similar to Ingraham et al. (2019) and Jumper et al. (2021), 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 α-carbon atom to nitrogen atom as ui. We denote the unit vector pointing from α-carbon atom to carbon atom of the carboxyl (-CO-) as ti. ui and ti together define a plane, and the normal of this plane is ni = ui×ti‖ui×ti‖ . Finally, we define vi = ni × ui. Then ni, ui and vi 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 → i ∈ 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 pj→i which describe relative position of j with respect to i:
pj→i =   n>i u>i v>i   [xj − xi]
Relative Orientation Edge Features As we mention above, each residue has orientation which carries information. Here we introduce the edge features qj→i, kj→i and tj→i which describe relative orientation of j with respect to i:
qj→i =   n>i u>i v>i   [nj ] , kj→i =   n>i u>i v>i   [uj ] , tj→i =   n>i u>i v>i   [vj ]
Distance-Based Edge Features Distance also carries information. Here we use radial basis function of distance as edge features:
fj→i,r = e − (‖xj−xi‖)
2
2σ2r , r = 1, 2, ..., R
Where R and scale parameters {σr}1≤r≤R are hyperparameters. In experiments, the set of scale parameters we used is {1.5x|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 λ ∈ {1., 2., 5., 10., 30.}.

heading: B PROOFS OF THE MAIN PROPOSITIONS
text: B.1 PROOF OF PROPOSITION 1.
Proof. Denote the predicted ligand position by R(X1|X2)X1 + t(X1|X2) = X̃1.
Assume first that SE(3)-invariance of the predicted docked complex defined by Eq. (1) is satisfied. Then the transformation to dock Q1X1 + g1 with respect to Q2X2 + g2 is the same as the transformation to change Q1X1 +g1 into Q2X̃1 +g2. We use the notation: R>(X1|X2) = (R(X1|X2))>. Then, we have the following derivation steps: R(X1|X2)X1 + t(X1|X2) = X̃1 X1 + R
>(X1|X2)t(X1|X2) = R>(X1|X2)X̃1 X1 + R
>(X1|X2)t(X1|X2) = R>(X1|X2)Q>2 (Q2X̃1 + g2 − g2) X1 + R
>(X1|X2)t(X1|X2) = R>(X1|X2)Q>2 (Q2X̃1 + g2)−R>(X1|X2)Q>2 g2 X1 + R
>(X1|X2)t(X1|X2) + R>(X1|X2)Q>2 g2 = R>(X1|X2)Q>2 (Q2X̃1 + g2) Q>1 (Q1X1 + g1 − g1) + R>(X1|X2)(t(X1|X2) + Q>2 g2) = R>(X1|X2)Q>2 (Q2X̃1 + g2) Q>1 (Q1X1 + g1)−Q>1 g1 + R>(X1|X2)(t(X1|X2) + Q>2 g2) = R>(X1|X2)Q>2 (Q2X̃1 + g2) R(X1|X2)Q>1 (Q1X1 + g1)−R(X1|X2)Q>1 g1 + t(X1|X2) + Q>2 g2 = Q>2 (Q2X̃1 + g2) Q2R(X1|X2)Q>1 (Q1X1 + g1)−Q2R(X1|X2)Q>1 g1 + Q2t(X1|X2) + g2 = Q2X̃1 + g2
From the last equation above, one derives the transformation of Q1X1 + g1 into Q2X̃1 + g2, which is, by definition of the functions R and t, the same as the transformation to dock Q1X1 + g1 with respect to Q2X2 + g2. This transformation is
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
which concludes the proof.
Conversely, assuming constraints in Eq. (2) hold, we derive that Q1X1 + g1 is transformed into Q2X̃1 + g2, which then is trivial to check that it satisfies SE(3)-invariance of the predicted docked complex defined by Eq. (1).
B.2 PROOF OF PROPOSITION 3.
Proof. We use the notation R>(X1|X2) := (R(X1|X2))>. As in Appendix B.1, we denote R(X1|X2)X1 + t(X1|X2) = X̃1. Then the transformation to dock X2 with respect to X1 is the same as the transformation to change X̃1 back to X1, which is
R(X1|X2)X1 + t(X1|X2) = X̃1 X1 + R
>(X1|X2)t(X1|X2) = R>(X1|X2)X̃1 X1 = R >(X1|X2)X̃1 −R>(X1|X2)t(X1|X2)
From the last equation above, we derive the transformation to change X̃1 back to X1, which is the same as the transformation to dock X2 with respect to X1.
B.3 PROOF OF PROPOSITION 4.
Proof. Let X(l+1)1 ,H (l+1) 1 ,X (l+1) 2 ,H (l+1) 2 = IEGMN(X (l) 1 ,H (l) 1 ,X (l) 2 ,H (l) 2 ) be the output of an IEGMN layer. Then, for any matrices Q1,Q2 ∈ SO(3) and any translation vectors g1,g2 ∈ R3, we want to prove that IEGMN satisfy the pairwise independent SE(3)-equivariance property:
Q1X (l+1) 1 +g1,H (l+1) 1 ,Q2X (l+1) 2 +g2,H (l+1) 2 = IEGMN(Q1X (l) 1 +g1,H (l) 1 ,Q2X (l) 2 +g2,H (l) 2 )
where each column of X(l)1 ∈ R3×n,H (l) 1 ∈ Rd×n,X (l) 2 ∈ R3×m and H (l) 2 ∈ Rd×m represent an individual node’s coordinate embedding or feature embedding.
We first note that the equations of our proposed IEGMN layer that compute messages mj→i, µj→i, mi and µi are SE(3)-invariant. Indeed, they depend on the initial features which are SE(3)-invariant by design, the current latent node embeddings {h(l)i }i∈V1∪V2 , as well as the Euclidean distances on the current node coordinates {x(l)i }i∈V1∪V2 . Thus, we also derive that the equation that computes the new latent node embeddings h(l+1)i is SE(3)-invariant. Last, the equation that updates the coordinates x (l+1) i 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.

heading: C SURFACE FEATURES
text: 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.
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.

heading: D DATASETS
text: 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.

heading: E MORE EXPERIMENTAL DETAILS AND RESULTS
text: Baseline Failures. 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. 11.
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 γ = 10 and σ = 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.

Review summaries:
In this paper, the authors propose a method for the “rigid body” docking of protein-protein complexes, i.e., a complex in which conformational changes in protein structures are not allowed. This method, called independent E(3)-equivariant graph matching networks (IEGMNs), finds the optimal rotation and translation to place proteins in a manner that the distances between residues in the binding site is minimized. A core feature of this method is SE(3) invariance, which means the optimal solution is invariant to the rotations and translations of the two proteins. This is achieved in an elaborate manner employing several ideas and technics:  - A k-NN graph is built for each protein using its CA coordinates and a set of additional features are also extracted - Message-passing neural network (MPNN) is used for graph-matching - Binding pocket residues or “keypoints” are identified - A differentiable Kabsch model is used for superimposing is the “keypoints” - Optimal transport is used for the alignment of residues in the binding pocket   This method, as reported by the authors, runs very fast and can identify the protein complex efficiently.  

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

The paper proposes a SE(3) equivariant graph matching network for end-to-end rigid protein docking. They propose a novel optimal transport loss to approximate binding pocket and a differentiable Kabsch algorithm to predict the docking pose. They achieve significant running time improvements over existing protein docking software with competitive results, and do not rely on heavy sampling, structure refinement, or templates