I NDEPENDENT SE(3)-E QUIVARIANT M ODELS FOR E ND - TO -E ND R IGID P ROTEIN D OCKING
Octavian-Eugen Ganea
A BSTRACT
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
E QUI D OCK 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 signiﬁcant running time improvements and often outperform existing docking software despite not relying on heavy candidate sampling, structure reﬁnement, or templates.
I NTRODUCTION
[Figure]
In a recent breakthrough, A LPHA F OLD 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, ligand receptor receptor ligand complex [Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
<latexit sha1_ [Table] ="
[Table]
="> [Table] +3SzSbsToQS8l
[Table]
+
[Table] [Table] [Table] +
[Table]
==</latexit> Z 1 , 2 2 R 3 ⇥ ( n + m ) <latexit sha1_base64="MYvTAFkn7UBy
[Table]
[Figure]
Contributions.
We design E QUI D OCK , 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 sufﬁcient model constraints (see Section 3).
Next, we create E QUI D OCK 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-speciﬁc 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.
R ELATED 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 ﬁrst-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.
A LPHA F OLD 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: ﬁrst, 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); ﬁnally, top-ranked candidates undergo a structure reﬁnement 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, A LPHA F OLD 2 and R OSE TTAF OLD 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 C LUSPRO (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 A LPHA F OLD 2 to multiple chains during both training and inference.
Drug-Target Interaction (DTI).
DTI aims to compute drug-target binding poses and afﬁnity, playing an essential role in understanding drugs’ mechanism of action.
Prior methods (Wallach et al., 2015; Li et al., 2021) predict binding afﬁnity 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.
M ATHEMATICAL 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 ∈
R 3 × n , X ∗ 2 ∈ R 3 × 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 modiﬁed point cloud X 1 ∈ R 3 × n .
For simplicity and w.l.o.g., the receptor remains in its bound location X 2 = X ∗ 2 .
The task is to predict a rotation R ∈ SO (3) and a translation t ∈ R 3 such that RX 1 + t = X ∗ 1 , using as input the proteins and their unbound positions X 1 and X 2 .
Here, R = R ( X 1 | X 2 ) and t = t ( X 1 | 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 ﬂexible docking , but our approach offers an important step towards it.
[Table]
ï ÉËÏ ü Ï ÏÏÏÏÏ ÏïÏÏÏÏ
KÉÏ ï ËÏÊ
[Table]
ÉLÉÏ
.tt
[Table]
[Figure]
b
[Figure]
[Table]
LIÉ
LIÉ iïiiiiii c'ÏÏÏÏÏÏÏ ÏÏ Ï Ï
HÉ ï ï
[Table]
ÉLÉÏITE
[Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Figure]
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
[Figure]
we wish to guarantee that:
[Equation] X 2 + t ( X 2 | X 1 )) , (commutativity) ∀ Q 1 , Q 2 ∈ SO (3) , ∀ g 1 , g 2 ∈ R 3 , ∀ X 1 ∈ R 3 × n , X 2 ∈ R 3 × m , and Z l = Q l X l + g l , l ∈ { 1 , 2 } . (1) for any rotations Q 1 , Q 2 and translations g 1 , g 2 , 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 X 1 and receptor X 2 , i.e. we now dock Q 1 X 1 + g 1 onto Q 2 X 2 + g 2 , then the rotation matrix R ( Q 1 X 1 + g 1 | Q 2 X 2 + g 2 ) and translation vector t ( Q 1 X 1 + g 1 | Q 2 X 2 + g 2 ) can be derived from the original R ( X 1 | X 2 ) and t ( X 1 | X 2 ) assuming that we always do rotations ﬁrst.
In this case, R ( Q 1 X 1 + g 1 | Q 2 X 2 + g 2 ) can be decomposed into three rotations: i.) apply Q (cid:62) 1 to undo the rotation Q 1 applied on X 1 , ii.)
apply R ( X 1 | X 2 ) , iii.)
apply Q 2 to rotate the docked ligand together with the receptor.
This gives R ( Q 1 X 1 + g 1 | Q 2 X 2 + g 2 )
=
Q 2 R ( X 1 | X 2 ) Q (cid:62) 1 , which in turn constraints the translation vector.
We provide a formal statement and prove it in Appendix B.1:
Proposition 1.
For any Q 1 , Q 2 ∈ SO (3) , g 1 , g 2 ∈ R 3 , SE(3)-invariance of the predicted docked complex deﬁned by Eq.
(1) is guaranteed iff
[Equation]
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 X 1 , and equivariance w.r.t.
any SE(3) transformation on X 2 :
[Equation]
where Z 1 = Q 1 X 1 + g
R ( X 1 | Z 2 ) X
+ t ( X 1 | Z 2 )
=
Q 2 [ R ( X 1 | X 2 ) X 1 + t
( X 1 | X 2 )]
+ g 2 , where Z 2 = Q 2 X 2 + g∀ Q 1 , Q 2 ∈ SO (3) , ∀ g 1 , g 2 ∈ R 3 , ∀ X 1 ∈ R 3 × n , ∀ X 2 ∈ R 3 × m .
Commutativity.
Instead of docking X 1 with respect to X 2 , we can also dock X 2 with respect to
X 1 .
In this case, we require the ﬁnal complex structures to be identical after superimposition, i.e., zero
RMSD.
This property is named commutativity and it is satisﬁed as follows (proof in Appendix B.2).
Proposition 3.
Commutativity as deﬁned by Eq.
(1) is guaranteed iff
[Equation]
Point Permutation Invariance.
We also enforce residue permutation invariance.
Formally, both
R ( X 1 | X 2 ) and t ( X 1 | X 2 ) should not depend on the order or columns of X 1 and, resp., of X 2 .
E QUI D OCK M ODEL
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 deﬁne 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.
; Townshend et al.
Liu et al.
Each node i ∈ V represents one residue and has 3D coordinates x i ∈
Rcorresponding to the α -carbon atom’s location.
Edges are given by a k-nearest-neighbor (k-NN) graph using Euclidean distance of the original 3D node coordinates.
Overview of Our Approach.
Our model is depicted in Fig.
We ﬁrst build k-NN protein graphs
G 1 =
( V 1 , E 1 ) and G 2 =
( V 2 , E 2 ) .
We then design SE(3)-invariant node features F 1 ∈ R d × n , F 2 ∈
R d × m and edge features { f j → i : ∀ ( i, j ) ∈
E 1 ∪ E 2 } (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 deﬁned:
Given Z 1 , H 1 , Z 2 , H 2 = Φ( X 1 , F 1 , X 2 , F 2 ) we have Q 1 Z 1 + g 1 , H 1 , Q
[Equation]
Z 2 +
[Equation]
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 ( X 1 | X 2 ) and t ( X 1 | 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.
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.
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 G 1 =
( V 1 , E 1 ) , G 2 =
( V 2 , 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 { h ( l ) i } i ∈V 1 ∪V 2 and node coordinate embeddings { x ( l ) i } i ∈V 1 ∪V 2 as
[Equation]
where N ( i ) are the neighbors of node i ; ϕ x is a real-valued (scalar) parametric function; W is a learnable matrix; ϕ h , ϕ e are parametric functions (MLPs) outputting a vector R d ; f j → i and f i are the original edge and node features (extracted SE(3)-invariantly from the residues). Note that the parameters of W , ϕ x , ϕ h , ϕ e can be shared or different across different IEGMN layers . a j → i is an attention based coefﬁcient with trainable shallow neural networks ψ q and ψ k :
[Equation]
The output of several IEGMN layers is then denoted as:
[Equation]
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.
Keypoints for Differentiable Protein Superimposition.
Next, we use multi-head attention to obtain K points for each protein, Y 1 , Y 2 ∈ R 3 × 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 Y 1 and
Y 2 would give the corresponding ground truth superimposition of X 1 and X 2 .
Our model is :
[Equation]
where z 1 i denotes the i-th column of matrix Z 1 , and α ki = softmax i ( 1 √ d h (cid:62) [-
-]1 i W (cid:48)
k µ ( ϕ ( H 2 ))) are attention scores (similarly deﬁned for β k j ), with W (cid:48) k ∈ R
d × d a parametric matrix (different for each attention head), ϕ a linear layer plus a LeakyReLU non-linearity, and µ ( · ) is 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 Y 1 and Y 2 — see Fig.
For this, we compute a differentiable version of the Kabsch algorithm (Kabsch, 1976) as follows.
Let
A = Y 2 Y (cid:62) 1 ∈ R 3 × 3 computed using zero-mean keypoints.
The singular value decomposition (SVD) is A =
U 2 SU (cid:62) 1 , where U 2 , U 1 ∈ O (3) .
Finally, we deﬁne the differentiable functions
[Equation]
, where d = sign(det( U 2 U (cid:62) 1 ))
[Equation]
13) where µ ( · ) is the mean vector of a point cloud.
It is straightforward to prove that this model satisﬁes 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., X 2 = X ∗ 2 ), predict the SE(3) transformation using Eq.
and use it to compute the ﬁnal position of the ligand as ˜ X 1 =
R ( X 1 | X 2 ) X
1 + t ( X 1 | X 2 ) .
The mean squared error (MSE) loss is then L MSE = 1 n (cid:80)
[Equation]
Optimal Transport and Binding Pocket Keypoint Alignment.
As stated before, we desire that
Y 1 and 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 ﬁrst deﬁne 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 ∗ 1 s , x ∗ 2 s ) , s ∈ 1 , . . .
, S } , where S is variable across data pairs.
We compute midpoints of these segments, denoted as P ∗ 1 , P ∗ 2 ∈ R 3 × S , where p ∗ 1 s =
p ∗ 2 s = 0 .
5 · ( x ∗ 1 s + x ∗ 2 s ) .
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 Xand X 2 .
We denote them as P 1 , P 2 ∈ R 3 × S .
For clarity, if X 1 = QX ∗ 1 + g , then P 1 =
QP ∗ 1 + g .
We desire that Y 1 is a representative set for the 3D set P 1 (and, similarly, Y 2 for P 2 ).
However, while at training time we know that every point p 1 s corresponds to the point p 2 s (and, similarly, y 1 k aligns with y 2 k , by assumption), we unfortunately do not know the actual alignment between points in Y l and P l , for every l ∈ { 1 , 2 } .
This can be recovered using an additional optimal transport loss:
[Equation]
where C
[Equation]
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 ﬁxed during back-propagation and optimization when only the cost matrix is trained.
Note that our approach assumes that y 1 k corresponds to y 2 k , for every k ∈ { 1 , . .
, K } .
Intuitively, each attention head k will identify a speciﬁc geometric/chemical local surface feature of protein 1 by y 1 k , and match its complementary feature of protein 2 by y 2 k .
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 ﬁrst state a notion of the "interior" of a protein point cloud.
Following previous work (Sverrisson et al., 2021; Venkatraman et al., 2009), we deﬁne the surface of a protein point cloud X ∈ R 3 × n as { x ∈ R 3 : G ( x ) = γ } , where G ( x ) = − σ ln( (cid:80) ni
=1 exp( −|| x − x i || 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 ∈ R 3 : G ( x ) < γ } .
Then, the condition for non-intersecting ligand and receptor can be written as G 1 ( x 2 j ) > γ, ∀ j ∈ 1 , . . .
, m and G 2 ( x 1 i ) > γ, ∀ i ∈ 1 , . . . , n .
As a loss function, this becomes
[Equation]
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 efﬁciency 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 ﬂat.
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 ( x i ; λ
[Equation]
, where w i,
[Equation]
[Table] 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 ﬂat 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.
E XPERIMENTS
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 ﬁlter DIPS to only keep proteins with at most 10K atoms.
Datasets are then randomly partitioned in train/validation/test splits of sizes 203/25/ (DB5.5) and 39,937/974/965 (DIPS).
For DIPS, the split is based on protein family to separate similar proteins.
For the ﬁnal 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 E QUI D OCK method with popular state-of-the-art docking software C LUS P RO (P IPER ) 2 (Desta et al., 2020; Kozakov et al., 2017),A TTRACT 3 (Schindler et al., 2017; 2 https://cluspro.bu.edu/ We .gz
[Figure]
de Vries et al., 2015), P ATCH D OCK 4 (Mashiach et al., 2010; Schneidman-Duhovny et al., 2005), and
HD OCK 5 (Yan et al., 2020; 2017b;a; Huang and Zou, 2014; 2008).
These baselines provide userfriendly 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), deﬁned below.
Given the ground truth and predicted complex structures, Z ∗ ∈ R 3 × ( n + m ) and
Z ∈ R 3 × ( n + m ) , we ﬁrst superimpose them using the Kabsch algorithm (Kabsch, 1976), and then compute C-RMSD =
(cid:113) 1 n + m (cid:107) Z ∗ − Z (cid:107) 2 F .
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 ﬁrst, 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 ﬁne tune the DIPS pre-
The best DB5.5 validated model is ﬁnally 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-speciﬁc hand-crafted features, and we currently do not perform structure ﬁne-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, HD OCK 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 Efﬁciency.
We show inference times in Fig.
5 and Table 4.
We note that
E QUI D OCK 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 signiﬁcantly 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 signiﬁcantly outperforms all baselines.
C ONCLUSION
We have presented an extremely fast, end-to-end rigid protein docking approach that does not rely on candidate sampling, templates, task-speciﬁc 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 E QUI D OCK and extend it for ﬂexible 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 speciﬁc targets (e.g. for antibody generation).
Last, we hope that our architecture can inspire the design of other types of biological 3D interactions.
A CKNOWLEDGEMENTS
The authors thank Hannes Stark, Gabriele Corso, Patrick Walters, Tian Xie, Xiang Fu, Jacob Stern, Jason Yim, Lewis Martin, and Jeremy Wohlwend 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.
C ONTENTS
A Representing Proteins as Graphs
B Proofs of the Main Propositions
B.1 Proof of Proposition 1. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.2 Proof of Proposition 3. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
B.3 Proof of Proposition 4. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C Surface Features
D Datasets
E More Experimental Details and Results
A R EPRESENTING P ROTEINS AS G RAPHS
A protein is comprised of amino acid residues.
The structure of an amino acid residue is shown in
Figure Fig.
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 speciﬁc to each type of amino acid residues.
[Figure]
R
N
O
C
H
C
H
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 x i ∈ R 3 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 u i .
We denote the unit vector pointing from α -carbon atom to carbon atom of the carboxyl (-CO-) as t i .
u i and t i together deﬁne a plane, and the normal of this plane is n i =
u i × t i (cid:107) u i × t i (cid:107) .
Finally, we deﬁne v i = n i × u i .
Then n i , u i and 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 → 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 p j → i which describe relative position of j with respect to i :
[Equation]
Relative Orientation Edge Features As we mention above, each residue has orientation which carries information.
Here we introduce the edge features q j → i , k j → i and t j → i which describe relative orientation of j with respect to i :
[Equation]
Distance-Based Edge Features Distance also carries information.
Here we use radial basis function of distance as edge features:
[Equation]
Where R and scale parameters { σ r } 1 ≤ r ≤ 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 deﬁned in
Eq.
(16) using λ ∈ { 1 ., 2 ., 5 ., 10 ., 30 . } .
B P ROOFS OF THE M AIN P ROPOSITIONS B. P ROOF OF P ROPOSITION 1.
Proof.
Denote the predicted ligand position by R ( X 1 | X 2 )
X
1 + t ( X 1 | X 2 ) =
˜ X 1 .
Assume ﬁrst that SE(3)-invariance of the predicted docked complex deﬁned by Eq.
(1) is satisﬁed.
Then the transformation to dock Q 1 X 1 + g 1 with respect to Q 2 X 2 + g 2 is the same as the transformation to change Q 1 X 1 + g 1 into Q 2 ˜
X
+ g 2 .
We use the notation: R (cid:62) ( X 1 | X 2 ) =
( R ( X 1 | X 2 )) (cid:62) .
Then, we have the following derivation steps:
[Equation]
X 1 | X 2 ) t ( X 1 | X 2 ) + R (cid:62) ( X 1 | X 2 ) Q (cid:62) 2 g 2 =
R (cid:62) ( X 1 | X 2 ) Q (cid:62) 2 ( Q 2 ˜ X
1 + g 2 )
Q (cid:62) 1 ( Q 1 X 1 + g 1 − g 1 ) + R (cid:62) ( X 1 | X 2 )( t ( X 1 | X 2 ) + Q (cid:62) 2 g 2 ) =
R (cid:62) ( X 1 | X 2 ) Q (cid:62) 2 ( Q 2 ˜ X
1 + g 2 )
Q (cid:62) 1 ( Q 1 X 1 + g 1 ) − Q (cid:62) 1 g 1 + R (cid:62) ( X 1 | X 2 )( t ( X 1 | X 2 ) + Q (cid:62) 2 g 2 ) =
R (cid:62) ( X 1 | X 2 ) Q (cid:62) 2 ( Q 2 ˜ X
1 + g 2 )
R ( X 1 | X 2 ) Q (cid:62) 1 ( Q 1 X 1 + g 1 ) − R ( X 1 | X 2 ) Q (cid:62) 1 g 1 + t
( X 1 | X 2 )
+ Q (cid:62) 2 g 2 = Q (cid:62) 2 ( Q 2 ˜ X
1 + g 2 )
Q 2 R ( X 1 | X 2 ) Q (cid:62) 1 ( Q 1 X 1 + g 1 ) − Q 2 R ( X 1 | X 2 ) Q (cid:62) 1 g 1 + Q 2 t ( X 1 | X 2 ) + g 2 = Q 2 ˜ X
+ g
From the last equation above, one derives the transformation of Q 1 X 1 + g 1 into Q 2 ˜ X
1 + g 2 , which is, by deﬁnition of the functions R and t , the same as the transformation to dock Q 1 X 1 + g 1 with respect to Q 2 X
+ g 2 .
This transformation is
[Equation]
which concludes the proof.
Conversely, assuming constraints in Eq.
(2) hold, we derive that Q 1 X 1 + g 1 is transformed into
Q 2 ˜ X
1 + g 2 , which then is trivial to check that it satisﬁes SE(3)-invariance of the predicted docked complex deﬁned by Eq.
B. P ROOF OF P ROPOSITION 3.
Proof.
We use the notation R (cid:62) ( X 1 | X 2 ) :
= ( R ( X 1 | X 2 )) (cid:62) .
As in Appendix B.1, we denote
R ( X 1 | X 2 ) X 1 + t ( X 1 | X 2 ) =
˜ X 1 .
Then the transformation to dock X 2 with respect to X 1 is the same as the transformation to change ˜ X
1 back to X 1 , which is
[Equation]
From the last equation above, we derive the transformation to change ˜ X 1 back to X 1 , which is the same as the transformation to dock X 2 with respect to X 1 .
B.
P ROOF OF P ROPOSITION 4.
Proof.
Let X ( l +1), H ( l +1), X
[Equation]
l +1)
[Equation]
, H ( l +1)= IEGMN ( X ( l ) 1 , H ( l ) 1 , X ( l ) 2 , H ( l ) 2 ) be the output of an
IEGMN layer.
Then, for any matrices Q 1 , Q 2 ∈ SO (3) and any translation vectors g 1 , g 2 ∈ R 3 , we want to prove that IEGMN satisfy the pairwise independent SE(3)-equivariance property:
[Equation]
where each column of X ( l ) 1 ∈ R 3 × n , H ( l ) 1 ∈ R d × n , X ( l ) 2 ∈ R 3 × m and H ( l ) 2 ∈ R
d
× m represent an individual node’s coordinate embedding or feature embedding.
We ﬁrst note that the equations of our proposed IEGMN layer that compute messages m j → i , µ j → i , m i 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 ∈V 1 ∪V 2 , as well as the Euclidean distances on the current node coordinates { x ( l ) i } i ∈V 1 ∪V 2 .
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.
C S URFACE F EATURES
We further discuss our new surface features introduced in Eq.
We ﬁrst visualize their design intuition in Fig.
A synthetic experiment is shown in Fig.
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.
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.
D D ATASETS
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 .
[Table]
E M ORE E XPERIMENTAL D ETAILS AND R ESULTS
Baseline Failures.
On the test sets, A TTRACT 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.
[Table]
If IEGMN layers except the ﬁrst one share parameters True, False
[Table]
of coordinates
skip connection
[Table]
-5, 1e-4, 1eDetailed Running Times.
In addition to the main text, we show in Table 4 detailed running times of all methods.
Hardware speciﬁcations are as follows:
A TTRACT was run on a 6-Core Intel Core
i
[Table]
2.2 GHz CPU; HD OCK was run on a single Intel Xeon Gold 6230 2.1 GHz CPU; E QUI D OCK was run on a single Intel Core i9-9880H 2.3 GHz CPU.
C LUS P RO and P ATCH D OCK have been manually run using their respective web servers.
Plots for DB5.5.
We show the corresponding plots for DB5.5 results in Fig.
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. [Table] surface
[Table]
without intersection loss
[Table] surface [Table] ground truth ligand and receptor point clouds have been centered (0-mean), without any other modiﬁcations.
Centeredstructures
[Table] loss [Table]
complexes
[Table]