DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking
Abstract
Predicting the binding structure of a small molecule ligand to a protein—a task known as molecular docking —is critical to drug design.
Recent deep learning methods that treat docking as a regression problem have decreased runtime compared to traditional search-based methods but have yet to offer substantial improvements in accuracy.
We instead frame molecular docking as a generative modeling problem and develop D IFF D OCK , a diffusion generative model over the non-Euclidean manifold of ligand poses. To do so, we map this manifold to the product space of the degrees of freedom (translational, rotational, and torsional) involved in docking and develop an efficient diffusion process on this space.
Empirically, D IFF D OCK obtains a 38% top-1 success rate (RMSD < 2Å) on PDBBind, significantly outperforming the previous state-of-the-art of traditional docking (23%) and deep learning (20%) methods.
Moreover, D IFF D OCK has fast inference times and provides confidence estimates with high selective accuracy.
Introduction
Biological mechanisms are mediated by proteins, whose functions may be modulated by small molecule ligands (such as drugs) that bind to them.
Thus, a crucial task in computational drug design is molecular docking —predicting the position, orientation and conformation of a ligand when bound to a target protein—from which the effect of the ligand (if any) might be inferred.
Recent works [1, 2] have developed deep learning models to predict the binding pose in one shot, framing docking as a regression problem.
However, we argue (Appendix B) that common molecular docking applications and evaluation metrics better correspond to maximizing the likelihood of the data under a generative process.
We thus frame molecular docking as a generative modeling problem —given a ligand and target protein structure, we learn a distribution over ligand poses.
We develop docking diffusion (D IFF D OCK ), a diffusion generative model (DGM) over the space of ligand poses for molecular docking.
We define a diffusion process over the degrees of freedom involved in docking: the position of the ligand relative to the protein (locating the binding pocket), its orientation in the pocket, and the m torsion angles describing its conformation.
While DGMs have been applied to other problems in molecular machine learning [3, 4, 5], existing approaches are ill-suited for molecular docking, where the space of ligand poses lies near an ( m + 6) -dimensional submanifold M ⊂ R 3 n .
To develop D IFF D OCK , we recognize that the docking degrees of freedom define M as the space of poses accessible via a set of allowed ligand pose transformations .
We use this idea to map elements in M to the product space of the groups corresponding to those transformations, where a DGM can be developed and trained more efficiently.
We then train a confidence model to rank the DGM samples and select the final prediction.
Empirically, on the standard blind docking benchmark PDBBind, D IFF D OCK achieves 37% of toppredictions with ligand root mean square distance (RMSD) below 2Å, nearly doubling performance of the previous state-of-the-art deep learning model (20%).
D IFF D OCK also significantly outperforms state-of-the-art classical search-based docking methods (23%), while still being significantly faster.
Method
A ligand pose is an assignment of atomic positions in R 3 , so in principle we can regard a pose x as an element in R 3 n , where n is the number of atoms.
However, this encompasses far more degrees of freedom than are relevant in molecular docking.
Traditional docking methods, as well as most
ML ones, take as input a seed conformation c ∈ R 3 n of the ligand in isolation, and change only the relative position and the torsion degrees of freedom in the final bound conformation.
The space of ligand poses consistent with c is, therefore, a ( m + 6)
-dimensional submanifold M c ⊂ R 3 n , where m is the number of rotatable bonds.
We follow this paradigm of taking as input a seed conformation c , and formulate molecular docking as learning a probability distribution p c ( x | y ) over the manifold
M c , conditioned on a protein structure y .
DGMs on submanifolds have been formulated by [6] in terms of projecting a diffusion in ambient space onto the submanifold.
However, the kernel p ( x t | x 0 ) of such diffusion is not available in closed form and must be sampled numerically with a geodesic random walk, making training very inefficient.
We instead define a one-to-one mapping to another, “nicer” manifold where the diffusion kernel can be sampled directly, and develop a DGM in that manifold.
To start, we restate the discussion in the last paragraph as follows:
Any ligand pose consistent with a seed conformation be reached by a combination of (1) ligand translations (2) ligand rotations and (3) changes to torsion angles.
This suggests that given a continuous family of ligand pose transformations corresponding to the m + 6 degrees of freedom, a distribution on M c can be lifted to a distribution on the product manifold of the corresponding groups.
2.1 Ligand Pose Transformations
We associate translations of ligand position with the 3D translation group T (3) , rigid rotations of the ligand with the 3D rotation group SO (3) , and changes in torsion angles at each rotatable bond with a copy of the 2D rotation group SO (2) .
More formally, we define operations of each of these groups on a ligand pose c ∈ R 3 n .
The translation A tr : T (3) × R 3 n → R 3 n is defined straightforwardly as
A tr ( r , x ) i = x
i + r using the isomorphism T (3) ∼ = R 3 and where x i ∈ R 3 is the position of the i th atom.
Similarly, the rotation A rot : SO (3) × R 3 n → R 3 n is defined by A rot ( R, x ) i =
R ( x i − ¯ x ) + ¯ x where ¯ x = 1 n (cid:80) x i , corresponding to rotations around the (unweighted) center of mass of the ligand.
Many valid definitions of a change in torsion angles are possible, as the torsion angle around any bond ( a i , b i ) can be updated by rotating the a i side, the b i side, or both.
However, we can disentangle changes in torsion angles from rotations or translations by defining them to cause a minimal perturbation (in an RMSD sense) to the structure:
Definition.
Let B k,θ k ( x ) ∈ R 3 n be any valid torsion update by θ k around the k th rotatable bond ( a k , b k ) .
We define A tor :
SO (2) m × R 3 n → R 3 n such that
[Equation]
where θ = ( θ 1 , . . .
θ m ) and RMSDAlign( x , x ′ )
=
arg min x † ∈{ g x ′ | g ∈ SE (3) } RMSD( x , x † )
This means that we apply all the m torsion updates in any manner, and then perform a global RMSD alignment with the unmodified pose.
The definition is motivated by ensuring that the infinitesimal effect of a torsion is orthogonal to any rototranslation, which has a pleasing physical interpretation as inducing no linear or angular momentum (Appendix A):
Proposition 1.
Let x ( t ) :
= A tor ( t θ , x ) for some θ and where t θ = ( tθ 1 , . . .
tθ m ) .
Then the linear and angular momentum are zero: ddt ¯ x | t =
= 0 and (cid:80) i ( x i − ¯ x ) × ddt x i | t =
= 0 where ¯ x = 1 n (cid:80) i x .
Now consider the product space P
= T 3 × SO (3) × SO (2) m and define A : P × R 3 n → R 3 n as
[Equation]
These definitions collectively provide the sought-after product space corresponding to the docking degrees of freedom.
Indeed, for a seed ligand conformation c , we can formally define the space of ligand poses M c
= { A ( g, c ) | g ∈ P } , which corresponds precisely to the intuitive notion of the space of ligand poses that can be reached by rigid-body motion plus torsion angle flexibility.
2.2 Diffusion on the Product Space
We now proceed to show how the product space can be used to learn a DGM over ligand poses in
M s .
First, we need a theoretical result (proof in Appendix A):
Proposition 2.
For a given seed conformation c , the map A ( · , c ) : P → M c is a bijection. which means that the inverse A − 1 c : M c → P given by A ( g, c ) (cid:55)→ g maps ligand poses x ∈ M c to points on the product space P .
We are now ready to develop a diffusion process on P . [6] established that the DGM framework transfers straightforwardly to Riemannian manifolds.
Thus, to implement a diffusion model on P , it suffices to develop a method for sampling from and computing the score of the diffusion kernel on P .
Since P is a product manifold, the forward diffusion proceeds independently in each manifold [7], and the tangent space is a direct sum:
T
[Equation]
SO
[Equation]
[Equation]
SO
[Equation]
[Equation]
∼ = R
[Equation]
R 3 ⊕ R m where g =
( r , R, θ ) .
Thus, it suffices to sample from the diffusion kernel and regress against its score in each group independently.
In all three groups, we define the forward SDE as d x = (cid:112) dσ 2 ( t ) /dt
d w where σ 2 = σ 2 tr , σ 2 rot , or σ 2 tor for T (3) , SO (3) , and SO (2) m respectively and where w is the corresponding Brownian motion.
Since T (3) ∼ = R 3 , the translational case is trivial and involves sampling and computing the score of a standard Gaussian with variance σ 2 ( t ) .
The diffusion kernel on SO (3) is given by the IGSO (3) distribution [8, 9], which can be sampled in the axis-angle parameterization by sampling a unit vector ˆ ω
∈ so (3) uniformly and random angle ω ∈ [0 , π ] according to
[Equation]
Further, the score of the diffusion kernel is ∇ ln p t ( R ′ | R ) =
( ddω log f ( ω ))
ˆ ω ∈ T R ′
SO (3) , where
R ′ =
R ( ω ˆ ω ) R is the result of applying Euler vector ω
ˆ ω to R .
The score computation and sampling can be accomplished efficiently by precomputing the truncated infinite series and interpolating the
CDF of p ( ω ) , respectively.
Finally, the SO (2) m group is diffeomorphic to the torus T m , on which the diffusion kernel is a wrapped normal distribution with variance σ 2 ( t ) .
This can be sampled directly and the score can be precomputed as a truncated infinite series [5].
2.3 Training, Inference and Architecture
Diffusion model.
Although we have defined the diffusion kernel and score matching objectives on P , we nevertheless develop the training and inference procedures to operate on ligand poses in 3D coordinates directly.
Providing the full 3D structure, to the score model allows it to reason about physical interactions using SE (3) equivariant models, not be dependent on arbitrary definitions of torsion angles [5] and better generalize to unseen complexes.
The training and inference procedures technically depend on the choice of seed conformation c used to define the mapping between M c and the product space.
However, providing a definite choice of c to the score model introduces an arbitrary inference-time parameter that may affect the final predicted distribution, which is undesirable.
Thus, we develop approximate training and inference procedures that remove the dependence on the c ; intuitively, these assume that updates to points in the product space P can be applied to ligand poses in M c directly.
While these are only an approximation of the theoretically correct procedures, we find that they work well in practice (Appendix C).
Confidence model.
Once the diffusion model is trained, in order to collect training data for the confidence model, we run the diffusion model to obtain a set of candidate poses for every training example and generate labels by testing whether or not each pose has RMSD below 2Å. The confidence model is then trained with cross-entropy loss to correctly predict the binary label for each pose.
During inference, the diffusion model is run to generate N poses in parallel, which are passed to the confidence model that ranks them based on its confidence that they have an RMSD below 2Å.
Model Architecture We construct the score model s ( x , y , t ) and the confidence model d ( x , y ) to take as input the current ligand pose x and protein structure y in 3D space.
The score model and confidence model have similar architectures based on SE (3) -equivariant convolutional networks.
The output of the score model must be in the tangent space T r T 3 ⊕ T R SO (3) ⊕ T θ SO (2) m .
The space T r T 3 ∼ = R 3 corresponds to translation vectors and T R SO (3) ∼ = R 3 to rotation (Euler) vectors, both of which are SE (3) -equivariant.
Finally, T θ SO (2) m corresponds to scores on SE (3) -invariant quantities.
The architectural components are detailed in Appendix D.
Experiments
[Table]
Experimental setup.
We evaluate our method on the complexes from PDBBind [14], a large collection of protein-ligand structures collected from PDB [15] [Table] which was used with time-based splits to benchmark many previous works [1, 16, 2].
We compare D IFF D OCK with state-of-the-art scoring methods SMINA [12], QuickVina-W [10], GLIDE [13] [Table] and GNINA [11] and the recent deep learning methods EquiBind and TANKBind presented above.
Extensive details about the experimental setup, data, baselines [Table] and implementation are in Appendix E. [Table]
Docking accuracy.
As the results of Table 1 show [Table] D IFF D OCK significantly outperforms all previous methods.
In particular, D IFF D OCK obtains an impressive 36.6% top-1 RMSD below 2Å,vastly surpassing the performance of state-of-the-art commercial software such as GLIDE (21.8%) and the previous state-of-the-art deep learning method TANKBind (20.4%).
Conclusion
We presented D IFF D OCK , a diffusion generative model tailored to the task of molecular docking, representing a paradigm shift from previous regression-based deep learning approaches.
To produce a fast and accurate generative model, we designed a diffusion process over the manifold describing the main degrees of freedom of the task.
Empirically, D IFF D OCK outperforms the state-of-the-art by large margins on PDBBind.
Thus, D IFF D OCK can offer great value for many existing real-world pipelines and opens up new avenues of research on similar ideas in protein-protein docking.
A Proofs A.1 Proof of Proposition
Proposition 1.
Let x ( t ) :
= A tor ( t θ , x ) for some θ and where t θ = ( tθ 1 , . . .
tθ m ) .
Then the linear and angular momentum are zero: ddt ¯ x | t =
= 0 and (cid:80) i ( x i − ¯ x ) × ddt x i | t =
= 0 where ¯ x = 1 n (cid:80) i x .
Proof.
Let x ( t ) =
R ( t )( B ( t θ , x ) − ¯ x 0 )
+ ¯ x 0 + p ( t ) where B ( t θ , · ) =
B 1 ,tθ 1 ◦ · · · B m,tθ m and
R ( t ) , p ( t ) are the rotation (around ¯ x 0 :=
¯ x (0) ) and translation associated with the optimal alignment between B ( t θ , x ) and x .
By definition of RMSD , we have
[Equation]
which, in the limit of t → 0 , becomes
[Equation]
The derivative in the LHS of Equation 5 at t = 0 is
[Equation] [Equation]
and represents the instantaneous velocity of the points
x i at t = 0 .
Denoting r i
= x i − ¯ x 0 , our objective is to minimize (cid:80) i
|| r
i
|| 2 .
Then, if we describe R ′ ( t )( x − ¯ x 0 )
=
R ′ ( t ) r by an angular velocity ω and abbreviate B ′ ( t θ , x (0))
i :=
b i and p ′
= v , we have
[Equation]
where we have used the fact that (cid:80) i r i = 0 and where I ( r )
=
( (cid:80) i r i · r i )
I − (cid:80) i r i r
Ti is the 3 ×inertia tensor .
Then setting gradients with respect to v , ω gives
[Equation]
Now with r ′ i
=
b i + ω
× r i + v we evaluate the linear momentum
[Equation]
Similarly, we evaluate the angular momentum
[Equation]
Thus, the linear and angular momentum are zero at t = 0 for arbitrary x (0) .
Note that since we did not use the particular form of B ( t θ , x ) in the above proof, we have shown that RMSD alignment can be used to disentangle rotations and translations from the infinitesimal action of any arbitrary function.
A.2 Proof of Proposition
Proposition 2.
For a given seed conformation c , the map A ( · , c ) : P → M c is a bijection.
Proof.
Since we defined M c = { A ( g, c ) | g ∈ P } , A ( · , c ) is automatically surjective.
We now show that it is injective.
Assume for the sake of contradiction that A ( · , c ) is not injective, so that there exist elements of the product space g 1 , g 2 ∈ P with g 1 ̸
= g 2 but with A ( g 1 , c )
=
A ( g 2 , c ) =
c ′ .
That is,
[Equation]
which we abbreviate as c (1)
=
c (2) .
Since only A tr changes the center of mass (cid:80)
i c i
/n , we have (cid:80) i c (1) i /n
= (cid:80)
i c i /n
+ r 1 and (cid:80) i c (2)
i /n
=
(cid:80)
i c
i /n
+ r 2 .
However, since c (1) =
c (2) , this implies r 1 = r 2 .
Next, consider the torsion angles τ 1 = ( τ (1)1 , . . .
τ (1) m ) of c (1) corresponding to some choice of dihedral angles at each rotatable bond.
Because A tr and A rot are rigid-body motions, only A tor changes the dihedral angles; in particular, by definition we have τ (1) i ∼ =
τ i + θ (1) i mod 2 π and τ (2) i ∼ =
τ i + θ (2)
i mod 2 π for all i = 1 , . . .
m .
However, because τ (1) i = τ (2) i , this means θ (1) i ∼ = θ (2) i for all i and therefore θ 1 = θ 2 (as elements of SO (2) m ).
Now denote c
⋆ =
A tor ( θ 1 , c ) =
A tor ( θ 2 , c ) and apply A tr ( − r 1 , · ) =
A tr ( − r 2 , · ) to both sides of Equation 11.
We then have
[Equation]
(12) which further leads to
[Equation]
In general, this does not imply that R 1 = R 2 .
However, R 1 ̸
= R 2 is possible only if c ⋆ is degenerate , in the sense that all points are collinear along the shared axis of rotation of R 1 , R 2 .
However, in practice conformers never consist of a collinear set of points, so we can safely assume R
= R 2 .
We now have ( r 1 , R 1 , θ 1 ) =
( r 2 , R 2 , θ 2 ) , or g 1 = g 2 , contradicting our initial assumption.
We thus conclude that A ( · , c ) is injective, completing the proof.
B Docking as Generative Modeling
Although EquiBind and other ML methods have provided strong runtime improvements by avoiding an expensive optimization process over ligand poses, their performance has not reached yet that of traditional methods.
As our analysis below argues, this may be caused by the models’ uncertainty and the optimization of an objective function that does not correspond with how molecular docking is used and evaluated in practice.
Molecular docking objective.
Molecular docking plays a critical role in drug discovery because the prediction of the 3D structure of a bound protein-ligand complex enables further computational and human expert analyses on the strength and properties of the binding interaction.
Therefore, a docked prediction is only useful if its deviation from the true structure does not significantly affect the output of such analysis.
Concretely, a prediction is considered acceptable when the distance between the structures (measured in terms of ligand RMSD) is below some small tolerance on the order of the length scale of atomic interactions (a few angstroms).
Consequently, the standard evaluation metric used in the field has been the percentage of predictions with a ligand RMSD (relative to the ground truth) within some value ϵ .
However, the objective of maximizing the proportion of predictions with RMSD within some tolerance ϵ is not differentiable and therefore cannot be used for the training of models with stochastic gradient descent.
Instead, we observe that maximizing the expected proportion of predictions with RMSD < ϵ corresponds to maximizing the likelihood of the true structure under the model’s output distribution, in the limit as ϵ goes to 0.
Generative model.
This observation motivates training a generative model to minimize an upper bound on the negative log-likelihood of the observed structures under the model’s distribution.
Thus, we view molecular docking as the problem of learning a distribution over ligand poses conditioned on the protein structure and develop a diffusion generative model over this space.
Confidence model.
With a trained diffusion model, it is possible to sample an arbitrary number of ligand poses from the posterior distribution of poses according to the model.
However, researchers are often interested in seeing only one or a small number of predicted poses for downstream analysis and an associated confidence measure 2 .
Thus, we train a confidence model over the poses sampled by the diffusion model and ranks them based on its confidence they are within the error tolerance.
The top-ranked ligand pose and the associated confidence are then taken as D IFF D OCK ’s top-1 prediction and the confidence score.
Problem with regression methods.
The difficulty with the development of deep learning models for molecular docking lies in the intrinsic aleatoric uncertainty on the pose often present and the [Figure] complexity of the task compared with the limited model capacity and data available (epistemic uncertainty).
Therefore, given the available co-variate information (only protein structure and ligand identity), any method will be unsure how to favor the correct among many viable alternative poses.
Any regression-style method that is forced to select a single configuration that minimizes the expected square error would learn to predict the (weighted) mean of such alternatives.
In contrast, a generative model with the same co-variate information would instead aim to capture the distribution over the alternatives, populating all/most of the significant modes even if similarly unable to distinguish the correct target.
This behavior, illustrated in Figure 2, causes the regression-based models to produce significantly more physically implausible poses than our method.
In particular, we observe frequent steric clashes (e.g., 26% of EquiBind’s predictions) and self-intersections in EquiBind’s and
TANKBind’s predictions (Figures 3 and 6).
We found no intersections in D IFF D OCK ’s predictions.
Visualizations and quantitative evidence of these phenomena are in Appendix F.1.
C Training and Inference
In this section we present the training and inference procedures of the diffusion generative model.
First, however, there are a few subtleties of the generative approach to molecular docking that are worth mentioning.
Unlike the standard generative modeling setting where the dataset consists of many samples drawn from the data distribution, each training example ( x ⋆ , y ) of protein structure y and ground-truth ligand pose x
⋆ is the only sample from the corresponding conditional distribution p x ⋆ ( · | y ) defined over M x ⋆ .
Thus, the innermost training loop iterates over distinct conditional distributions p x ⋆ ( · | y ) , along with a single sample from that distribution, rather than over samples from a common data distribution p data ( x ) .
As discussed in Section 2, during inference c is the ligand structure generated with a method such as RDKit.
However, during training we require M c = M x ⋆ in order to define a bijection between c ∈ M x ⋆ and P .
If we take c ∈ M x ⋆ , there will be a distribution shift between the manifolds M c considered at training time, and those considered at inference time.
To circumvent this issue, at training time we predict c with RDKit and replace x ⋆ with arg min x † ∈M c RMSD( x ⋆ , x † ) using the conformer matching procedure described in [5].
The above paragraph may be rephrased more intuitively as follows: during inference, the generative model docks a ligand structure generated by RDKit, keeping its non-torsional degrees of freedom (e.g., local structures) fixed.
At training time, however, if we train the score model with the local structures of the ground truth pose, this will not correspond to the local structures seen at inference time.
Thus, at training time we replace the ground truth pose by generating a ligand structure with
RDKit and aligning it to the ground truth pose while keeping the local structures fixed.
With these preliminaries, we now continue to the full procedures.
The training and inference procedures of a score-based diffusion generative model on a Riemannian manifold consist of (1) sampling and regressing against the score of the diffusion kernel during training; and (2) sampling a geodesic random walk with the score as a drift term during inference [6].
Because we have developed the diffusion process on P but continue to provide the score model with elements in M c ⊂ R 3 n , the full training and inference procedures involve repeatedly interconverting between the two spaces using the bijection given by the seed conformation c .
Input: Training pairs { ( x ⋆ , y ) } , RDKit predictions { c } foreach c , x ⋆ , y do
Let x 0 ← arg min x † ∈M c RMSD( x ⋆ , x † ) ;
Compute ( r 0 , R 0 , θ 0 ) ← A − 1 c ( x 0 ) ;
Sample t ∼ Uni([0 , 1]) ;
Sample ∆ r , ∆ R, ∆ θ from diffusion kernels p tr t ( · | 0) , p rot t ( · | 0) , p tor t ( · | 0) ;
Set r t ← r 0 + ∆ r ;
Set R t ← (∆ R ) R 0 ;
Set θ t ← θ 0 + ∆ θ mod 2 π ;
Compute x t ← A (( r t , R t , θ t ) , c ) ;
Predict scores α ∈ R 3 , β ← R 3 , γ ∈ R m = s ( x t , c , y , t ) ;
Take optimization step on loss
L = ||
α − ∇ p tr t (∆ r | 0) ||
+ || β − ∇ p rot t (∆ R | 0)
||
2 + || γ − ∇ p tor t (∆ θ | 0)
||
Input: RDKit prediction c , protein structure y (both centered at origin)
Output: Sampled ligand pose x
Sample θ N ∼ Uni( SO (2) m ) , R N ∼ Uni( SO (3)) , r N ∼ N (0 , σ 2 tor ( T )) ;
Let x N =
A ( ; for n ← N to 1 do
Let t = n/N and ∆ σ 2 tr = σ 2 tr ( n/N ) − σ 2 tr (( n − 1) /N ) and similarly for ∆ σ 2 rot , ∆ σ 2 tor ;
Predict scores
However, as noted in the main text, the dependence of these procedures on the exact choice of c is potentially problematic, as it suggests that at inference time, the model distribution may be different depending on the orientation and torsion angles of c .
Simply removing the dependence of the score model on c is not sufficient, since the update steps themselves still occur on P and require a choice of c to be mapped to M c .
However, notice that the update steps—in both training and inference—consist of (1) sampling the diffusion kernels at the origin; (2) applying these updates to the point on P ; and (3) transferring the point on P to M c via A ( · , c ) .
Might it instead be possible to apply the updates to 3D ligand poses x ∈ M c directly ?
It turns out that the notion of applying these steps to ligand poses “directly” corresponds to the formal notion of group action .
The operations A tr , A rot , A tor that we have already defined are formally group actions if they satisfy A ( · ) ( g 1 g 2 , x ) =
A ( g 1 , A
( g 2 , x )) .
While true for A tr , A rot , this is not generally true for A tor if we take SO (2) m to be the direct product group; however, the approximation is increasingly good as the magnitude of the torsion angle updates decreases.
If we then define P to be the direct product group of its constituent groups, A is a group action of P on M c , as the operations of A tr , A rot , A tor commute and are (under the approximation) individually group actions.
The implication of A being a group action can be seen as follows.
Let δ =
g b g − 1 a be the update which brings g a ∈ P to g b ∈ P via left multiplication, and let x a , x b be the corresponding ligand poses A ( g a , c ) , A ( g b , c ) .
Then
[Equation]
which means that the updates δ can be applied directly to x a using the operation A .
The training and inference procedures then become Algorithm 3 and 4 below.
The initial conformer c is no longer used, except in the the initial steps to define the manifold—to find the closest point to x ⋆ in training, and to sample x N from the prior over M c in inference.
Conceptually speaking, this procedure corresponds to “forgetting” the location of the origin element on M c , which is permissible because a change of the origin to some equivalent seed c ′ ∈
M c merely translates—via right multiplication by A − 1 c ( c ′ ) —the original and diffused data distributions on P , but does not cause any changes on M c itself.
The training and inference routines involve updates—formally left multiplications—to group elements, but as left multiplication on the group corresponds to group actions on M c , the updates can act on M c directly, without referencing the origin c .
We find that the approximation of A as a group action works quite well in practice, and use Algorithms 3 and 4 for all training and experiments discussed in the paper.
Of course, disentangling the torsion updates from rotations in a way that makes A tor exactly a group action would justify the procedure further, and we regard this as a possible direction of future work.
Input: Training pairs { ( x ⋆ , y ) } , RDKit predictions { c } foreach c , x ⋆ , y do
Let x 0 ← arg min x † ∈M c RMSD( x ⋆ , x † ) ;
Sample t ∼ Uni([0 ,
Sample ∆ r , ∆ R, ∆ θ from diffusion kernels p tr t ( · | 0) , p rot t ( · | 0) , p tor t ( · |
Compute x t ← A ((∆ r , ∆ R, ∆ θ ) , x 0 ) ;
Predict scores α ∈ R 3 , β ← R 3 , γ ∈ R m = s ( x t , y , t ) ;
Take optimization step on loss
[Equation]
Input: RDKit prediction c , protein structure y (both centered at origin)
Output: Sampled ligand pose x
Sample θ N ∼ Uni( SO (2) m ) , R N ∼ Uni( SO (3)) , r N ∼ N (0 , σ 2 tor ( T )) ;
Let x N =
A (( r N , R N , θ N ) , c ) ; for n ← N to 1 do
Let t = n/N and ∆ σ 2 tr = σ 2 tr ( n/N ) − σ 2 tr (( n − 1) /N ) and similarly for ∆ σ 2 rot , ∆ σ 2 tor ;
Predict scores α ∈ R 3 , β ∈ R 3 , γ ∈ R m ← s ( x n , y , t ) ;
Sample z tr , z rot , z tor from N (0 , ∆ σ 2 tr ) , N (0 , ∆ σ 2 rot ) , N (0 , ∆ σ 2 tor ) respectively;
Set ∆ r ← r 0 + ∆ σ 2 tr α + z tr ;
Set ∆ R ← R (∆ σ 2 rot β + z rot ) ;
Set ∆ θ ← ∆ σ 2 tor γ + z tor ;
Compute x n − 1 ← A ((∆ r , ∆ R, ∆ θ ) , x n ) ; Return x 0 ; D Architecture Details
We use convolutional networks based on tensor products of irreducible representations (irreps) of SO(3) [20] as architecture for both the score and confidence models. In particular, these are implemented using the e3nn library [21]. Below, ⊗ w refers to the spherical tensor product of irreps with path weights w and ⊕ refers to normal vector addition (with possibly padded inputs).
Features have multiple channels for each irrep. Both the architectures can be decomposed in three main parts: embedding layer, interaction layers and output layer. We outline each of them below. D.1 Embedding layer
Geometric heterogeneous graph. Structures are represented as heterogeneous geometric graphs with nodes representing ligand (heavy) atoms, receptor residues (located in the position of the α -carbon atom) and receptor (heavy) atoms (only for the confidence model).
Because of the high number of nodes involved it is necessary for the graph to be sparsely connected for runtime and memory constraints.
Moreover, sparsity can act as a useful inductive bias for the model, however, it is critical, for the model to find the right pose, that nodes that might have a strong interaction in the final pose to be connected during the diffusion process.
Therefore, to build the radius graph, we connect nodes using cutoffs that are dependent on the types of nodes they are connecting:
Ligand atoms-ligand atoms, receptor atoms-receptor atoms, and ligand atoms-receptor atoms interactions all use a cutoff of 5Å, standard practice for atomic interactions.
For the ligand atoms-ligand atoms interactions we also preserve the covalent bonds as separate edges with some initial embedding representing the bond type (single, double, triple and aromatic).
For receptor atoms-receptor atoms interactions, we limit at 8 the maximum number of neighbors of each atom.
Note that the ligand atoms-receptor atoms only appear in the confidence model where the final structure is already set.
Receptor residues-receptor residues use a cutoff of 15 Å with 24 as the maximum number of neighbors for each residue.
Receptor residues-ligand atoms use a cutoff of 20 + 3 ∗ σ tr Å where σ tr represents the current standard deviation of the diffusion translational noise present in each dimension (zero for the confidence model).
Intuitively this guarantees that with high probability any of the ligand and receptors that will be interacting in the final pose the diffusion model converges to are connected in the message passing at every step.
Finally, receptor residues are connected to the receptor atoms that form the corresponding amino-acid.
Node and edge featurization.
For the receptor residues , we use the residue type as a feature as well as a language model embedding obtained from ESM2 [22].
The ligand atoms have the following features: atomic number; chirality; degree; formal charge; implicit valence; the number of connected hydrogens; the number of radical electrons; hybridization type; whether or not it is in an aromatic ring; in how many rings it is; and finally, 6 features for whether or not it is in a ring of size 3, 4, 5, 6, 7, or 8.
This are concatenated with sinusoidal embeddings of the diffusion time [23] and, in the case of edges, radial basis embeddings of edge length [24].
These scalar features of each node and edge are then transformed with learnable two-layer MLPs (different for each node and edge type) into a set of scalar features that are used as initial representations by the interaction layers.
Notation Let ( V , E ) represent the heterogeneous graph, with V =
( V ℓ , V r ) respectively ligand atoms and receptor residues (receptor atoms V a , present in the confidence model, are for simplicity not included here), and similarly E =
( E ℓℓ , E ℓr , E rℓ , E rr ) .
Let h a be the node embeddings (initially only scalar channels) of node a , e ab the edge embeddings of ( a, b ) , and µ ( r ab ) radial basis embeddings of the edge length.
Let σ 2 tr , σ 2 rot , and σ 2 tor represent the variance of the diffusion kernel in each of the three components: translational, rotational and torsional.
D.2 Interaction layers
At each layer, for every pair of nodes in the graph, we construct messages using tensor products of the current node features with the spherical harmonic representations of the edge vector.
The weights of this tensor product are computed based on the edge embeddings and the scalar features—denoted h 0 a —of the outgoing and incoming nodes.
The messages are then aggregated at each node and used to update the current node features.
For every node a of type t a :
[Equation]
(cid:33) with ψ ab =
Ψ ( t a ,t ) ( e ab , h 0 a , h 0 b ) (15)
Here, t indicates an arbitrary node type, N ( t ) a = { b | ( a
, b ) ∈
E t a t } the neighbors of a of type t , Y are the spherical harmonics up to ℓ = 2 , and BN the (equivariant) batch normalisation.
The orders of the output are restricted to a maximum of ℓ = 1 .
All learnable weights are contained in Ψ , a dictionary of MLPs, which uses different sets of weights for different edge types (as an ordered pair so four types for the score model and nine for the confidence) and different rotational orders.
D.3 Output layer
The ligand atom representations after the final interaction layer are used in the output layer to produce the required outputs.
This is where the score and confidence architecture differ significantly.
On one hand, the score model’s output is in the tangent space T r T 3 ⊕ T R SO (3) ⊕ T θ SO (2) m .
This corresponds to having two SE (3) -equivariant output vectors representing the translational and rotational score predictions and m SE (3) -invariant output scalars representing the torsional score.
For each of these we design final tensor-product convolutions inspired by classical mechanics.
On the other hand, the confidence model outputs a single SE (3) -invariant scalar representing the confidence score.
Below we detail how each of these outputs is generated.
Translational and rotational scores.
The translational and rotational score intuitively represent, respectively, the linear acceleration of the center of mass of the ligand and the angular acceleration of the rest of the molecule around the center.
Considering the ligand as a rigid object and given a set of forces and masses at each ligand a tensor product convolution between the atoms and the center of mass would be capable of computing the desired quantities.
Therefore, for each of the two outputs, we perform a convolution of each of the ligand atoms with the (unweighted) center of mass c . [Equation] with ψ ca = Ψ
[Equation]
r [Equation] We restrict the output of v to a single odd and a single even vectors (for each of the two scores).
Since we are using coarse-grained representations of the protein, the score will neither be even nor odd, therefore, we sum the even and odd vector representations of v .
Finally, the magnitude (but not direction) of these vectors is adjusted with an MLP taking as input the current magnitude and the sinusouidal embeddings of the diffusion time.
Finally, we (revert the normalization) by multiplying the outputs by 1 /σ tr for the translational score and by the expected magnitude of a score in SO (3) with diffusion parameter σ rot (precomputed numerically).
Torsional score.
To predict the m SE (3) -invariant scalar describing the torsional score, we use a pseudotorque layer similar to that of [5].
This predicts a scalar score δτ for each rotatable bond from the per-node outputs of the atomic convolution layers.
For rotatable bond g = ( g 0 , g 1 ) and b ∈ V ℓ , let r gb and ˆ r gb be the magnitude and direction of the vector connecting the center of bond g and b .
We construct a convolutional filter T g for each bond g from the tensor product of the spherical harmonics with a ℓ = 2 representation of the bond axis ˆ r g :[-
-]
[Equation]
(17) ⊗ is the full (i.e., unweighted) tensor product as described in [25], and the second term contains the spherical harmonics up to ℓ = 2 (as usual).
This filter (which contains orders up to ℓ = 3 ) is then bond direction. used to convolve with the representations of every neighbor on a radius graph: [Equation] with γ gb =
{+Γ( e gb , h 0 b , h 0 g 
+}
+ h 0 g 1 )
[Equation]
Here, N g = { b | ( g, b ) ∈ E τ } and Υ ( τ ) and Γ are MLPs with learnable parameters.
Since unlike [5], we use coarse-grained representations the parity also here is neither even nor odd, the irreps in the output are restricted to arrays both even h ′ g and odd h ′′ g scalars.
Finally, we produce a single scalar prediction for each bond:
[Equation]
where Π is a two-layer MLP with tanh nonlinearity and no biases.
This is also "denormalised" multiplying by by the expected magnitude of a score in SO (2) with diffusion parameter σ tor .
Confidence output.
The single SE (3) -invariant scalar representing the confidence score output is instead obtained by concatenating the even and odd final scalar representation of each ligand atom, averaging these feature vectors among the different atoms and finally applying a three layers MLP (with batch normalization).
E Experimental Details E.1 Experimental Setup
Data.
We use the molecular complexes in PDBBind [14] that were extracted from the Protein Data
Bank (PDB) [15].
We employ the time-split of PDBBind proposed by [1] with 17k complexes from 2018 or earlier for training/validation and 363 test structures from 2019.
This is motivated by the further adoption of the same split [2] and the critical assessment of PDBBind splits by [16] who favor temporal splits over artificial splits based on molecular scaffolds or protein sequence/structure similarity.
For completeness we also report the results on protein sequence similarity splits in
Appendix F.
Metrics.
To evaluate the generated complexes we compute the heavy-atom RMSD between the predicted and the crystal ligand atoms when the protein structures are aligned.
To account for permutation symmetries in the ligand, we use the symmetry-corrected RMSD of sPyRMSD [26].
For these RMSD values, we report the percentage of predictions that have an RMSD that is less than 2Å. We choose 2Å since much prior work considers poses with an RMSD less that 2Å as "good" or successful [27, 10, 11].
This is a chemically relevant metric unlike the mean RMSD as detailed in
Appendix B since for further downstream analyses such as determining function changes, a prediction is only useful below a certain RMSD error threshold.
Less relevant metrics such as the mean RMSD are provided in Appendix F.
We report the metrics for the highest ranked prediction as the top-1; top-5 refers to selecting the best performing pose out of the 5 highest ranked predictions, which can be important for many applications where multiple predictions are used for downstream tasks.
Implementation.
We use Adam [28] as optimizer for the diffusion and the confidence model.
The diffusion model with which we run inference uses the exponential moving average of the weights during training and we update the moving average after every optimization step with a decay factor of 0.999.
The batchsize is 16.
We run inference with 20 denoising steps on 500 validation complexes every 5 epochs and use the set of weights with the highest percentage of RMSDs less than 2Å as the final diffusion model.
We train on four 48GB RTX A6000 GPUs for 400 epochs (around 9 days).
For inference only a single GPU is required.
Scaling up the model size seems to improve performance and future work could explore whether this trend continues further.
For the confidence model uses the validation cross-entropy loss is used for early stopping and training only takes 70 epochs.
E.2 Baselines: implementation, used scripts, and runtime details
For obtaining the runtimes of the different methods we always used 16 CPUs except for GLIDE as explained below.
The runtimes do not include any preprocessing time for any of the methods.
For instance, the time that it takes to run P2Rank is not included for TANKBind and P2Rank +
SMINA/GNINA since this receptor preparation only needs to be ran once when docking many ligands to the same protein.
In applications where different receptor are processed (such as reverse screening), the experienced runtimes for TANKBind and P2Rank + SMINA/GNINA will thus be higher.
We note that for all these baselines we have used the default hyperparameters unless specified differently below.
Modifying some of these hyperparameters (for example the scoring method’s exhaustiveness) will change the runtime and performance tradeoffs (e.g. if the searching routine is left running for longer then better poses are likely to be found), however, we leave these analysis to future work.
SMINA [12] improves Autodock Vina with a new scoring-function and user friendliness.
The default parameters were used with the exception of setting –num_modes 10 .
To define the search box, we use the automatic box creation option around the receptor with the default buffer of 4Å on all 6 sides.
GNINA [11] builds on SMINA by additionally using a learned 3D CNN for scoring.
The default parameters were used with the exception of setting –num_modes 10 .
To define the search box, we use the automatic box creation option around the receptor with the default buffer of 4Å on all 6 sides.
QuickVina-W [10] extends the speed-optimized QuickVina 2 [27] for blind docking.
We reuse the numbers from [1] which had used the default parameters except for increasing the exhaustiveness to 64.
GLIDE [13] is a strong heavily used commercial docking tool.
These methods all use biophysics based scoring-functions.
We reuse the numbers from [1] since we do not have a license.
As explained by [1], the very high runtime of GLIDE with 1405 seconds per complex is partially explained by the fact that GLIDE only uses a single thread when processing a complex.
This fact and the parallelization options of GLIDE are explained here https://www.schrodinger.com/kb/1165 .
With GLIDE, it is possible to start data-parallel processes that compute the docking results for a different complex in parallel.
However, in the mentioned link, it is explained that each process also requires a separate software license.
EquiBind [1], we reuse the numbers reported in their paper and generate the predictions that we visualize with their code at https://github.com/HannesStark/EquiBind .
TANKBind [2], we use the code associated with the paper at https://github.com/luwei0917/
TankBind .
The runtimes do not include the runtime of P2Rank or any preprocessing steps.
In Tablewe report two runtimes for GPU (0.72/2.5 sec) and CPU (1.95/3.7 sec).
The first is the runtime when making only the top-1 prediction and the second is for producing the top-5 predictions.
Producing only the top-1 predictions is faster since TANKBind produces distance predictions that need to be converted to coordinates with a gradient descent algorithm and this step only needs to be run once for the top-1 prediction while it needs to be run 5 times for producing 5 outputs.
To obtain our runtimes we run the forward pass of TANKBind on an RTX A6000 GPU or on an Intel Xeon Gold[-
-]
CPU with the default batchsize of 5 that is used in their GitHub repository.
To compute the time the distances-to-coordinates conversion step takes, we parallelizes the computation across 16 processes which we also run on an Intel Xeon Gold 6230 CPU.
This way we obtain 0.44 seconds runtime for the conversion step of the top-1 prediction (averaged over the 363 complexes of the testset).
P2Rank [29], is a tool that predicts multiple binding pockets and ranks them.
We use it for running
TANKBind and P2Rank + SMINA/GNINA.
We download the program from https://github. com/rdk/p2rank and run it with its default parameters.
EquiBind + SMINA/GNINA [1], the bounding box in which GNINA/SMINA searches for binding poses is constructed around the prediction of EquiBind with the –autobox_ligand option of
GNINA/SMINA.
EquiBind is thus used to find the binding pocket and SMINA/GNINA to find the exact final binding pose.
We use –autobox_add 10 to add an additional 10Å on all 6 sides of the bounding box following [1].
P2Rank + SMINA/GNINA.
The bounding box in which GNINA/SMINA searches for binding poses is constructed around the pocket center that P2Rank predicts as the most likely binding pocket.
P2Rank is thus used to find the binding pocket and SMINA/GNINA to find the exact final binding pose.
The diameter of the search box is the diameter of a ligand conformer generated by RDKit with an additional 10Å on all 6 sides of the bounding box.
F Additional Results F.1 Physically plausible predictions
[Table]
Due to the averaging phenomenon of regression based methods such as TANKBind and EquiBind, they make predictions at the mean of the distribution.
If aleatoric uncertainty is present, such as in case of symmetric complexes, this leads to predicting the ligand to be at an un-physical state in the middle of the possible binding pockets as visualized in Figure 7.
The Figure also illustrates how
D IFF D OCK does not suffer from this issue and is able to accurately sample from the modes.
In the scenario when epistemic uncertainty about the correct ligand conformation is present, this often results in "squashed-up" predictions of the regression based methods as visualized in Figure 3.
If there is uncertainty about the correct conformer, the square error minimizing option is to put all atoms close to the mean.
These averaging phenomena in the presence of either aleatoric or epistemic uncertainty cause the regression based methods to often generate steric clashes and self intersections.
To investigate this quantitatively, we determine the fraction of test complexes for which the methods exhibit steric clashes.
We define a ligand as exhibiting a steric clash, if one of its heavy atoms is within 0.4Å of a heavy receptor atom.
This cutoff is used by protein quality assessment tools and in previous literature [30].
Table 2 shows that D IFF D OCK , as a generative model, produces fewer steric clashes than the regression based baselines.
We generally observe no unphysical predictions from D IFF D OCK unlike the self intersections that, e.g., TANKBind produces (Figure 3) or its incorrect local structures (Figure 4).
F.2 Further results and selective accuracy
Figure 5- left shows the proportion of RMSDs below an arbitrary threshold ϵ with D IFF D OCK exceeding previous methods for almost every possible ϵ .
With the exception of very small ϵ < 1Å where GLIDE performs better than the other methods because it is the only one directly including a final energy relaxation of the local structures inside the method, but this final relaxation could be added to any of the other methods.
As the top-1 results show, D IFF D OCK ’s confidence model is very accurate in ranking the sampled poses for a given complex and picking the best one.
We also investigate the selective accuracy of this score across different complexes, by evaluating how D IFF D OCK ’s accuracy increases if it only makes predictions when the confidence is above a certain threshold, known as selective prediction .
In Figure 5 we plot the "success" rate as we decrease the percentage of complexes for which we make predictions ( selective accuracy ), corresponding to increasing confidence threshold.
When only making predictions for the top one-third of complexes in terms of model confidence, the (RMSD < 2Å)-rate improves from 37% to 85%.
Additionally, there is a high Spearman correlation of 0. between D IFF D OCK ’s confidence and the RMSD.
Evidently, the confidence score is informative for judging the quality of D IFF D OCK ’s top-ranked sampled pose and provides a highly valuable confidence measure for practitioners, a critical factor for many downstream applications.
[Figure]
selective accuracy, the percentage of predictions with RMSD below 2Å when only making predictions for the percentage of the dataset on the x-axis for which D IFF D OCK is the most confident. F.3 Visualizations
see that, since it is a generative model, D IFF D is able to produce multiple correct to sample around them.
Meanwhile, as a regression based model, EquiBind is only able to predict a structure at the mean of the modes.
The complexes are unseen during training.
The PDB IDs in reading order: 6agt, 6gdy, 6ckl, 6dz