DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking
AffiliationAddress
Abstract
Predicting the binding structure of a small molecule ligand to a protein—a taskknown as molecular docking —is critical to drug design.
Recent deep learning methods that treat docking as a regression problem have decreased runtime compared totraditional search-based methods but have yet to offer substantial improvements inaccuracy.
We instead frame molecular docking as a generative modeling problemand develop D IFF D OCK , a diffusion generative model over the non-Euclideanmanifold of ligand poses.
To do so, we map this manifold to the product spaceof the degrees of freedom (translational, rotational, and torsional) involved indocking 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%) anddeep learning (20%) methods.
Moreover, D IFF D OCK has fast inference times andprovides confidence estimates with high selective accuracy.
Introduction
The biological functions of proteins can be modulated by small molecule ligands (such as drugs)binding to them.
Thus, a crucial task in computational drug design is molecular docking —predictingthe position, orientation and conformation of a ligand when bound to a target protein—from whichthe effect of the ligand (if any) might be inferred.
Recent works [1, 2] have developed deep learningmodels to predict the binding pose in one shot, treating docking as a regression problem.
Whilethese methods are much faster than traditional search-based methods, they have yet to demonstratesignificant improvements in accuracy.
We argue that this may be because the regression-basedparadigm corresponds imperfectly with the objectives of molecular docking, which is reflected inthe fact that standard accuracy metrics resemble the likelihood of the data under the predictivemodel rather than a regression loss.
We thus frame molecular docking as a generative modelingproblem —given a ligand and target protein structure, we learn a distribution over ligand poses.
We therefore develop D IFF D OCK , a diffusion generative model (DGM) over the space of ligandposes for molecular docking.
We define a diffusion process over the degrees of freedom involved indocking: the position of the ligand relative to the protein (locating the binding pocket), its orientationin the pocket, and the torsion angles describing its conformation.
D IFF D OCK samples poses byrunning the learned (reverse) diffusion process, which iteratively transforms an uninformed, noisyprior distribution over ligand poses into the learned model distribution (Figure 1).
Intuitively, thisprocess can be viewed as the progressive refinement of random poses via updates of their translations,rotations, and torsion angles.
While DGMs have been applied to other problems in molecular machine learning [3, 4, 5], existingapproaches are ill-suited for molecular docking, where the space of ligand poses is an ( m + 6) -dimensional submanifold M ⊂ R 3 n , where n and m are, respectively, the number of atoms and t=T t= ligand & protein reverse [Figure]
over
[Figure] and torsions ranked poses & confidence score
DiffDock torsion angles.
To develop D IFF D OCK , we recognize that the docking degrees of freedom define Mas the space of poses accessible via a set of allowed ligand pose transformations .
We use this ideato map elements in M to the product space of the groups corresponding to those transformations,where a DGM can be developed and trained efficiently.
As applications of docking models often require only a fixed number of predictions and a confidencescore over these, we train a confidence model to provide confidence estimates for the poses sampledfrom the DGM and to pick out the most likely sample.
This two-step process can be viewed as anintermediate approach between brute-force search and one-shot prediction: we retain the ability toconsider and compare multiple poses without incurring the difficulties of high-dimensional search.
Empirically, on the standard blind docking benchmark PDBBind, D IFF D OCK achieves 38% oftop-1 predictions with ligand root mean square distance (RMSD) below 2Å, nearly doubling theperformance of the previous state-of-the-art deep learning model (20%).
D IFF D OCK significantlyoutperforms even state-of-the-art search-based methods (23%), while still being 3 to 12 times fasteron GPU.
Moreover, it provides an accurate confidence score of its predictions, obtaining 83%
RMSD < 2Å on its most confident third of the previously unseen complexes.
To summarize, the main contributions of this work are:
We frame the molecular docking task as a generative problem and highlight the issues with previous deep learning approaches.
We formulate a novel diffusion process over ligand poses corresponding to the degrees of freedom involved in molecular docking.
We achieve a new state-of-the-art 38% top-1 prediction with RMSD < 2Å on PDBBind blinddocking benchmark, considerably surpassing the previous best search-based (23%) and deeplearning methods (20%).
Background and Related Work
Molecular docking.
The molecular docking task is usually divided between known-pocket and blinddocking.
Known-pocket docking algorithms receive as input the position on the protein where themolecule will bind (the binding pocket ) and only have to find the correct orientation and conformation.
Blind docking instead does not assume any prior knowledge about the binding pocket; in this work,we will focus on this general setting.
Due to the relative rigidity of the protein in the majority ofthe cases, docking methods typically assume the knowledge of the bound protein structure, andwe will follow this assumption [6].
Methods are normally evaluated by the percentage of hits, orapproximately correct predictions, commonly considered to be those where the ligand RMSD error isbelow 2Å [7, 8, 9]. [Figure]
Search-based docking methods.
Traditional docking methods [10, 11, 12] consist of a parameterizedphysics-based scoring function and a search algorithm.
The scoring-function takes in 3D structuresand returns an estimate of the quality/likelihood of the given pose, while the search stochasticallymodifies the ligand pose (position, orientation, and torsion angles) with the goal of finding the globaloptimum of the scoring function.
Recently, machine learning has been applied to parameterize thescoring-function [9, 13].
These search-based methods have offered relative improvements whendocking to a known pocket but are typically very computationally expensive to run and must stillgrapple with the very large search space that characterizes blind docking.
Machine learning for blind docking.
Recently, EquiBind [1] has tried to tackle the blind dockingtask by directly predicting pocket keypoints on both ligand and protein and aligning them.
TANKBind[2] improved over this by independently predicting a docking pose (in the form of an interatomicdistance matrix) for each possible pocket and then ranking them.
Although these one-shot or few-shotregression-based prediction methods are orders of magnitude faster, their performance has not yetreached that of traditional search-based methods.
Docking as Generative Modeling
Molecular docking objective.
Molecular docking plays a critical role in drug discovery because theprediction of the 3D structure of a bound protein-ligand complex enables further computational andhuman expert analyses on the strength and properties of the binding interaction.
Therefore, a dockedprediction is only useful if its deviation from the true structure does not significantly affect the outputof such analyses.
Concretely, a prediction is considered acceptable when the distance between thestructures (measured in terms of ligand RMSD) is below some small tolerance on the order of thelength scale of atomic interactions (a few Ångström).
Consequently, the standard evaluation metricused in the field has been the percentage of predictions with a ligand RMSD (to the crystal ligandpose) below some value ϵ .
However, the objective of maximizing the proportion of predictions with RMSD within some toleranceϵ is not differentiable and cannot be used for training with stochastic gradient descent.
Instead,maximizing the expected proportion of predictions with RMSD <
ϵ corresponds to maximizingthe likelihood of the true structure under the model’s output distribution, in the limit as ϵ goesto 0.
This observation motivates training a generative model to minimize an upper bound on thenegative log-likelihood of the observed structures under the model’s distribution.
Thus, we viewmolecular docking as the problem of learning a distribution over ligand poses conditioned on theprotein structure and develop a diffusion generative model over this space (Section 4).
Confidence model.
With a trained diffusion model, it is possible to sample an arbitrary number ofligand poses from the posterior distribution according to the model.
However, researchers are often interested in seeing only one or a small number of predicted poses and an associated confidencemeasure 1 for downstream analysis.
Thus, we train a confidence model over the poses sampled by thediffusion model and rank them based on its confidence that they are within the error tolerance.
Thetop-ranked ligand pose and the associated confidence are then taken as D IFF D OCK ’s top-1 predictionand confidence score.
Problem with regression-based methods.
The difficulty with the development of deep learningmodels for molecular docking lies in the data inherent (aleatoric) uncertainty on the pose (multipleposes could be correct) and the complexity of the task compared with the limited model capacity anddata available (epistemic uncertainty).
Therefore, given the available co-variate information (onlyprotein structure and ligand identity), any method will exhibit uncertainty about the correct bindingpose among many viable alternatives.
Any regression-style method that is forced to select a singleconfiguration that minimizes the expected square error would learn to predict the (weighted) meanof such alternatives.
In contrast, a generative model with the same co-variate information wouldinstead aim to capture the distribution over the alternatives, populating all/most of the significantmodes 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 thanour 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 5 and 9).
We foundno intersections in D IFF D OCK ’s predictions.
Visualizations and quantitative evidence of thesephenomena are in Appendix E.1.
Method
4.1 Overview
A ligand pose is an assignment of atomic positions in R 3 , so in principle, we can regard a pose x asan element in R 3 n , where n is the number of atoms.
However, this encompasses far more degrees offreedom than are relevant in molecular docking.
In particular, bond lengths, angles, and small ringsin the ligand are essentially rigid, such that the ligand flexibility lies almost entirely in the torsionangles at rotatable bonds.
Traditional docking methods, as well as most ML ones, take as input a seedconformation c ∈ R 3 n of the ligand in isolation and change only the relative position and the torsiondegrees of freedom in the final bound conformation.
The space of ligand poses consistent with cis, therefore, an ( m + 6) -dimensional submanifold M c ⊂ R 3 n , where m is the number of rotatablebonds, and the six additional degrees of freedom come from rototranslations relative to the fixedprotein.
We follow this paradigm of taking as input a seed conformation c , and formulate moleculardocking as learning a probability distribution p c ( x | y ) over the manifold M c , conditioned on aprotein structure y .
DGMs on submanifolds have been formulated by [17] in terms of projecting a diffusion in ambientspace onto the submanifold.
However, the kernel p ( x t | x 0 ) of such a diffusion is not availablein closed form and must be sampled numerically with a geodesic random walk, making trainingvery inefficient.
We instead define a one-to-one mapping to another, “nicer” manifold where thediffusion kernel can be sampled directly and develop a DGM in that manifold.
To start, we restate thediscussion in the last paragraph as follows:
Any ligand pose consistent with a seed conformation can be reached by a combination of (1) ligand translations, (2) ligand rotations, and (3) changes to torsion angles.
This can be viewed as an informal definition of the manifold M c .
Simultaneously, it suggests thatgiven a continuous family of ligand pose transformations corresponding to the m + 6 degrees offreedom, a distribution on M c can be lifted to a distribution on the product space of the corresponding groups—which is itself a manifold.
We will then show how to sample the diffusion kernel on thisproduct space and train a DGM over it.
4.2 Ligand Pose Transformations
We associate translations of ligand position with the 3D translation group T (3) , rigid rotations of theligand with the 3D rotation group SO (3) , and changes in torsion angles at each rotatable bond with acopy of the 2D rotation group SO (2) .
More formally, we define operations of each of these groupson 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 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 anybond ( a i , b i ) can be updated by rotating the a i side, the b i side, or both.
However, we can specifychanges of torsion angles to be disentangled from rotations or translations.
One way of doing so isto identify a central motif in the molecule, such as a ring, and change torsion angles in a way thatkeeps the motif fixed.
However, this special treatment of the central motif introduces an arbitraryasymmetry into the problem and could be difficult for a score model to reason about.
Thus, weinstead define the operation of elements of SO (2) m such that it causes 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
[Equation]
This means that we apply all the m torsion updates in any order and then perform a global RMSDalignment with the unmodified pose.
The definition is motivated by ensuring that the infinitesimaleffect of a torsion is orthogonal to any rototranslation, i.e., it induces no linear or angular momentum .
These properties can be stated more formally as follows (proof in Appendix A):
Proposition 1.
Let x ( t ) :
= A tor ( t θ , x ) for some θ and where t θ = ( tθ 1 , . . .
tθ m ) .
Then the linearand 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 4 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 dockingdegrees of freedom.
Indeed, for a seed ligand conformation c , we can formally define the space ofligand poses M c =
{ A ( g, c ) | g ∈ P } .
This corresponds precisely to the intuitive notion of thespace of ligand poses that can be reached by rigid-body motion plus torsion angle flexibility.
4.3 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 c .
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 topoints on the product space P .
We are now ready to develop a diffusion process on P .[17] established that the DGM framework transfers straightforwardly to Riemannian manifolds withthe score and score model as elements of the tangent space and with the geodesic random walk as thereverse SDE solver.
Further, the score model can be trained in the standard manner with denoisingscore matching [18].
Thus, to implement a diffusion model on P , it suffices to develop a methodfor sampling from and computing the score of the diffusion kernel on P .
Furthermore, since P is a product manifold, the forward diffusion proceeds independently in each manifold [19], and thetangent space is a direct sum: T g P = T r T 3 ⊕ T R SO (3) ⊕ T θ SO (2) m ∼ = R 3 ⊕ R 3 ⊕
R m whereg =
( r , R, θ ) .
Thus, it suffices to sample from the diffusion kernel and regress against its score ineach 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 ofa standard Gaussian with variance σ 2 ( t ) .
The diffusion kernel on SO (3) is given by the IGSO (3)distribution [20, 21], which can be sampled in the axis-angle parameterization by sampling a unitvector ˆ ω ∈ so (3) uniformly 5 and random angle ω ∈ [0 , π ] according top
[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 samplingcan 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 whichthe diffusion kernel is a wrapped normal distribution with variance σ 2 ( t ) .
This can be sampleddirectly, and the score can be precomputed as a truncated infinite series [4].
4.4 Training and Inference
Diffusion model.
Although we have defined the diffusion kernel and score matching objectives on 207 P , we nevertheless develop the training and inference procedures to operate on ligand poses in 3Dcoordinates directly.
Providing the full 3D structure, rather than abstract elements of the productspace, to the score model allows it to reason about physical interactions using SE (3) equivariantmodels, not be dependent on arbitrary definitions of torsion angles [4], and better generalize to unseencomplexes.
The training and inference procedures technically depend on the choice of seed conformation c usedto define the mapping between M c and the product space.
However, providing a definite choiceof c to the score model introduces an arbitrary inference-time parameter that may affect the finalpredicted distribution, which is undesirable.
In other words, while c defines the manifold of ligandposes, the precise location of c within that manifold should not affect the predicted distribution.
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 posesin M c directly, without referencing the origin conformer c .
While these are only an approximationof the theoretically correct procedures, we find that they work well in practice.
In Appendix B, wepresent the training and inference procedures in more detail and further discussion on this point.
Confidence model.
In order to collect training data for the confidence model d ( x , y ) , we run thetrained diffusion model to obtain a set of candidate poses for every training example and generatelabels by testing whether or not each pose has RMSD below 2Å. The confidence model is then trainedwith cross-entropy loss to correctly predict the binary label for each pose.
During inference, thediffusion model is run to generate N poses in parallel, which are passed to the confidence model thatranks them based on its confidence that they have RMSD below 2Å.
Architecture 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 3Dcoordinates directly to leverage SE (3) equivariant networks.
The score model and confidence modelhave similar architectures except for the space of their output.
The output of the score model must bein the tangent space T r T 3 ⊕ T R SO (3) ⊕ T θ SO (2) m , which corresponds to equivariant translationand rotation (Euler) vectors, and scores on invariant torsion angles.
The confidence model outputs asingle invariant scalar.
The architectural components are detailed in Appendix C.
Experiments
Experimental setup.
We evaluate our method on the complexes from PDBBind [22], a largecollection of protein-ligand structures collected from PDB [23], which was used with time-basedsplits to benchmark many previous works [1, 24, 2].
We compare D IFF D OCK with state-of-the-artsearch-based methods SMINA [25], QuickVina-W [8], GLIDE [11], and GNINA [9] and the recentdeep learning methods EquiBind and TANKBind presented above.
Extensive details about theexperimental setup, data, baselines, and implementation are in Appendix D.3 and all code is availableat https://anonymous.4open.science/r/DiffDock .
The repository also contains videos ofthe reverse diffusion process (images of the same are in Figure 11).
As we are evaluating blind docking, the methods receive two inputs: the ligand with a predictedseed conformation (e.g., from RDKit) and the crystal structure of the protein.
Since search-basedmethods work best when given a starting binding pocket to restrict the search space, we also testthe combination of using an ML-based method, such as P2Rank [26] (also used by TANKBind) or
EquiBind to find an initial binding pocket, followed by a search-based method to predict the exactpose in the pocket.
[Table]
Docking accuracy.
D IFF D OCK significantly outperforms all previous methods (Table 1).
In particular, D IFF D OCK obtains an impressive 38.2% top-1 success rate (i
[Table]
e., percentage of predictions with
RMSD < 2Å 6 ) when sampling 40 poses and 35.0% when sampling just 10.
This performance vastlysurpasses that of state-of-the-art commercial software such as GLIDE (21.8%) and the previousstate-of-the-art deep learning method TANKBind (20.4%).
The use of ML-based pocket prediction incombination with search-based docking methods improves over the baseline performances [Table] but eventhe best of these (P2Rank+GNINA) reaches a success rate of only 28.8%.
proportion of RMSDs below an arbitrary threshold ϵ with D IFF D OCK exceeding previous methods for performance
Inference runtime.
D IFF D OCK holds its superior accuracy while being (on GPU) 3 to 12 times fasterthan the best search-based method, GNINA (Table 1).
This high speed is critical for applicationssuch as high throughput virtual screening for drug candidates or reverse screening for protein targets,where one often searches over a vast number of complexes.
As a diffusion model, D IFF D OCK isinevitably slower than the one-shot deep learning method E QUI B IND , but as shown in Figure 3- rightand Appendix E.3, it can be significantly sped up without significant loss of accuracy.
[Figure] [Figure]
only making predictions for the portion of the dataset where D IFF D OCK is most confident.
Selective accuracy of confidence score.
As thetop-1 results show, D IFF D OCK ’s confidence modelis very accurate in ranking the sampled poses for agiven complex and picking the best one.
We alsoinvestigate the selective accuracy of the confidencemodel across different complexes by evaluating how D IFF D OCK ’s accuracy increases if it only makespredictions when the confidence is above a certainthreshold, known as selective prediction .
In Figure 4, we plot the success rate as we decrease thepercentage of complexes for which we make predictions, i.e., increase the confidence threshold.
Whenonly making predictions for the top one-third ofcomplexes in terms of model confidence, the success rate improves from 38% to 83%.
Additionally,there is a high Spearman correlation of 0.68 between D IFF D OCK ’s confidence and the negative RMSD.
Thus, the confidence score is a good indicator of the quality of D IFF D OCK ’s top-ranked sampled pose and provides a highly valuable confidencemeasure for downstream applications.
Conclusion
We presented D IFF D OCK , a diffusion generative model tailored to the task of molecular docking.
This represents a paradigm shift from previous deep learning approaches, which use regression-basedframeworks, to a generative modeling approach that is better aligned with the objective of moleculardocking.
To produce a fast and accurate generative model, we designed a diffusion process over themanifold describing the main degrees of freedom of the task via ligand pose transformations spanningthe manifold.
Empirically, D IFF D OCK outperforms the state-of-the-art by very large margins on PDBBind, has fastinference times, and provides confidence estimates with high selective accuracy.
Thus, D IFF D OCKcan offer great value for many existing real-world pipelines and opens up new avenues of research onhow to best integrate downstream tasks, such as affinity prediction, into the framework and applysimilar ideas to protein-protein and protein-nucleic acid docking.
A Proofs
A.1 Proof of Proposition 1
Proposition 1.
Let x ( t ) :
= A tor ( t θ , x ) for some θ and where t θ = ( tθ 1 , . . .
tθ m ) .
Then the linearand 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 alignmentbetween B ( t θ , x ) and x .
By definition of RMSD , we have|| x
( t ) − x (0)
||
= min R, p || R ( t )( B ( t θ , x (0)) − ¯ x 0 )
+ ¯ x 0 + p ( t ) − x (0) ||
(4) which, in the limit of t → 0 , becomes(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
d dt [Equation]
t ) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) [Equation] = min [Equation]
p (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) d
[Equation]
t
[Equation]
−
[Equation]
[Equation]
)) (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[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 , ourobjective is to minimize (cid:80) i ||
r ′
i
|| 2 .
Then, if we describe R ′ ( t )( x − ¯ x 0 )
=
R ′ ( t ) r by an angularvelocity ω 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 × 3inertia tensor .
Then setting gradients with respect to v , ω gives
[Equation]
i
[Equation]
and
[Equation]
Now with r ′ i
=
b i + ω
× r i + v we evaluate the linear momentum
[Equation]
which is zero by direct substitution of v .
Similarly, we evaluate the angular momentum
[Equation]
which is zero by direct substitution of ω .
Thus, the linear and angular momentum are zero at t = 0for arbitrary x (0) .
Note that since we did not use the particular form of B ( t θ , x ) in the above proof, we have shownthat RMSD alignment can be used to disentangle rotations and translations from the infinitesimalaction of any arbitrary function.
A.2 Proof of Proposition 2
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 showthat it is injective.
Assume for the sake of contradiction that A ( · , c ) is not injective, so that there existelements 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]
[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 impliesr 1 = r 2 .
Next, consider the torsion angles τ 1 = ( τ (1)1 , . . .
τ (1) m ) of c (1) corresponding to some choiceof dihedral angles at each rotatable bond.
Because A tr and A rot are rigid-body motions, only A torchanges 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) ifor all i and therefore θ = θ 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]
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, inpractice, 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 thusconclude that A ( · , c ) is injective, completing the proof.
B 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 areworth mentioning.
Unlike the standard generative modeling setting where the dataset consists ofmany samples drawn from the data distribution, each training example ( x ⋆ , y ) of protein structure yand ground-truth ligand pose x
⋆ is the only sample from the corresponding conditional distributionp x ⋆ ( · | y ) defined over M x ⋆ .
Thus, the innermost training loop iterates over distinct conditionaldistributions p x ⋆ ( · | y ) , along with a single sample from that distribution, rather than over samplesfrom a common data distribution p data ( x ) .
As discussed in Section 4, during inference, c is the ligand structure generated with a method suchas RDKit.
However, during training we require M c = M x ⋆ in order to define a bijection betweenc ∈ 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, attraining time we predict c with RDKit and replace x ⋆ with arg min x † ∈M c RMSD( x ⋆ , x † ) using theconformer matching procedure described in [4].
The above paragraph may be rephrased more intuitively as follows: during inference, the generativemodel 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 localstructures of the ground truth pose, this will not correspond to the local structures seen at inferencetime.
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 (Algorithms 1 and 2).
The trainingand inference procedures of a score-based diffusion generative model on a Riemannian manifoldconsist 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 [17].
Becausewe have developed the diffusion process on P but continue to provide the score model with elementsin M c ⊂ R 3 n , the full training and inference procedures involve repeatedly interconverting betweenthe 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
[Equation]
Algorithm 2:
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 , c , 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 n − 1 ← r 0 + ∆ σ 2 tr α + z tr ; Set R n − 1 ← R (∆ σ 2 rot β + z rot ) R n ) ; Set θ n − 1 ← θ n + (∆ σ 2 tor γ + z tor ) mod 2 π ; Compute x n − 1 ← A (( r n − 1 , R n − 1 , θ n − 1 ) , c ) ;
However, as noted in the main text, the dependence of these procedures on the exact choice of c ispotentially problematic, as it suggests that at inference time, the model distribution may be differentdepending on the orientation and torsion angles of c .
Simply removing the dependence of the scoremodel on c is not sufficient since the update steps themselves still occur on P and require a choice of cto be mapped to M c .
However, notice that the update steps—in both training and inference—consistof (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 to3D ligand poses x ∈ M c directly ?
It turns out that the notion of applying these steps to ligand poses “directly” corresponds to theformal notion of group action .
The operations A tr , A rot , A tor that we have already defined are formallygroup 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 notgenerally true for A tor if we take SO (2) m to be the direct product group; however, the approximationis increasingly good as the magnitude of the torsion angle updates decreases.
If we then define P to bethe direct product group of its constituent groups, A is a group action of P on M c , as the operationsof 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 updatewhich brings g a ∈ P to g b ∈ P via left multiplication, and let x a , x b be the corresponding ligandposes A ( g a , c ) , A ( g b , c ) .
Then
[Equation]
14) which means that the updates δ can be applied directly to x a using the operation A .
The training andinference procedures then become Algorithm 3 and 4 below.
The initial conformer c is no longer used, except in the initial steps to define the manifold—to find the closest point to x ⋆ in training, andto sample x N from the prior over M c in inference.
Conceptually speaking, this procedure corresponds to “forgetting” the location of the origin elementon M c , which is permissible because a change of the origin to some equivalent seed c ′ ∈
M cmerely translates—via right multiplication by A − 1 c ( c ′ ) —the original and diffused data distributionson P , but does not cause any changes on M c itself.
The training and inference routines involveupdates—formally left multiplications—to group elements, but as left multiplication on the groupcorresponds to group actions on M c , the updates can act on M c directly, without referencing theorigin c .
We find that the approximation of A as a group action works quite well in practice and use Algorithms3 and 4 for all training and experiments discussed in the paper.
Of course, disentangling the torsionupdates from rotations in a way that makes A tor exactly a group action would justify the procedurefurther, and we regard this as a possible direction for future work.
Algorithm
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 , 1]) ;
Sample ∆ r , ∆ R, ∆ θ from diffusion kernels p tr t ( · | 0) , p rot t ( · | 0) , p tor t ( · | 0) ;
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
L = ||
[Equation]
− ∇ p tr t (
[Equation]
r |
[Equation]
{+|| 
+}
+ ||
[Equation]
− ∇ p rot t (
[Equation]
R
[Equation]
[Equation]
{+|| 
+}
+ ||
[Equation]
∇ p tor t (
[Equation]
θ |
[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 ;
C Architecture Details
We use convolutional networks based on tensor products of irreducible representations (irreps)of SO(3) [27] as architecture for both the score and confidence models.
In particular, these areimplemented using the e3nn library [28].
Below, ⊗ w refers to the spherical tensor product of irrepswith path weights w , and ⊕ refers to normal vector addition (with possibly padded inputs).
Featureshave multiple channels for each irrep.
Both the architectures can be decomposed into three mainparts: embedding layer, interaction layers, and output layer.
We outline each of them below.
C.1 Embedding layer
Geometric heterogeneous graph.
Structures are represented as heterogeneous geometric graphs withnodes representing ligand (heavy) atoms, receptor residues (located in the position of the α -carbonatom), and receptor (heavy) atoms (only for the confidence model).
Because of the high numberof nodes involved, it is necessary for the graph to be sparsely connected for runtime and memoryconstraints.
Moreover, sparsity can act as a useful inductive bias for the model, however, it is criticalfor the model to find the right pose that nodes that might have a strong interaction in the final pose tobe connected during the diffusion process.
Therefore, to build the radius graph, we connect nodesusing 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 withsome initial embedding representing the bond type (single, double, triple and aromatic)
Forreceptor atoms receptor atoms interactions, we limit at 8 the maximum number of neighborsof each atom.
Note that the ligand atoms-receptor atoms only appear in the confidencemodel 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 ligands 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 wellas a language model embedding obtained from ESM2 [29].
The ligand atoms have the followingfeatures: atomic number; chirality; degree; formal charge; implicit valence; the number of connectedhydrogens; the number of radical electrons; hybridization type; whether or not it is in an aromaticring; 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.
These are concatenated with sinusoidal embeddings of the diffusion time [30] and, in the caseof edges, radial basis embeddings of edge length [31].
These scalar features of each node and edgeare then transformed with learnable two-layer MLPs (different for each node and edge type) into aset 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 atomsand receptor residues (receptor atoms V a , present in the confidence model, are for simplicity notincluded here), and similarly E =
( E ℓℓ , E ℓr , E rℓ , E rr ) .
Let h a be the node embeddings (initially onlyscalar channels) of node a , e ab the edge embeddings of ( a, b ) , and µ ( r ab ) radial basis embeddings ofthe edge length.
Let σ 2 tr , σ 2 rot , and σ 2 tor represent the variance of the diffusion kernel in each of thethree components: translational, rotational and torsional.
C.2 Interaction layers
At each layer, for every pair of nodes in the graph, we construct messages using tensor products ofthe current node features with the spherical harmonic representations of the edge vector.
The weightsof this tensor product are computed based on the edge embeddings and the scalar features—denotedh 0 a —of the outgoing and incoming nodes.
The messages are then aggregated at each node and usedto update the current node features.
For every node a of type t a :
[Equation]
with ψ ab = Ψ ( t a ,t ) ( e ab , h 0 a , h 0 b )
[Equation]
[Equation]
Here, t indicates an arbitrary node type, N ( t ) a = { b | ( a
, b ) ∈
E t a t } the neighbors of a of type t , Yare 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 Ψ , adictionary of MLPs, which uses different sets of weights for different edge types (as an ordered pairso four types for the score model and nine for the confidence) and different rotational orders.
C.3 Output layer
The ligand atom representations after the final interaction layer are used in the output layer toproduce 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 androtational 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 theother hand, the confidence model outputs a single SE (3) -invariant scalar representing the confidencescore.
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 ofthe rest of the molecule around the center.
Considering the ligand as a rigid object and given a set offorces and masses at each ligand, a tensor product convolution between the atoms and the center ofmass 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]
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 norodd; therefore, we sum the even and odd vector representations of v .
Finally, the magnitude (but notdirection) of these vectors is adjusted with an MLP taking as input the current magnitude and thesinusoidal embeddings of the diffusion time.
Finally, we (revert the normalization) by multiplyingthe 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 apseudotorque layer similar to that of [4].
This predicts a scalar score δτ for each rotatable bond fromthe per-node outputs of the atomic convolution layers.
For rotatable bond g = ( g 0 , g 1 ) and b ∈ V ℓ , letr gb and ˆ r gb be the magnitude and direction of the vector connecting the center of bond g and b .
Weconstruct a convolutional filter T g for each bond g from the tensor product of the spherical harmonicswith a ℓ = 2 representation of the bond axis ˆ r g : {+
+}
[Equation]
(17) ⊗ is the full (i.e., unweighted) tensor product as described in [32], and the second term contains thespherical harmonics up to ℓ = 2 (as usual).
This filter (which contains orders up to ℓ = 3 ) is thenused to convolve with the representations of every neighbor on a radius graph:
[Equation]
γ gb =
[Equation]
Here, N g = { b | ( g, b ) ∈ E τ } and Υ ( τ ) and Γ are MLPs with learnable parameters.
Since unlike [4],we use coarse-grained representations the parity also here is neither even nor odd, the irreps in theoutput are restricted to arrays both even h ′ g and odd h ′′ g scalars.
Finally, we produce a single scalarprediction for each bond:
[Equation]
where Π is a two-layer MLP with tanh nonlinearity and no biases.
This is also “denormalized" bymultiplying 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 isinstead 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).
D Experimental Details
In general, all our code is available at https://anonymous.4open.science/r/DiffDock .
Thisincludes running the baselines, runtime calculations, training and inference scripts for D IFF D OCK ,the PDB files of D IFF D OCK ’s predictions for all 363 complexes of the test set, and visualizationvideos of the reverse diffusion.
D.1 Experimental Setup
Data.
We use the molecular complexes in PDBBind [22] that were extracted from the Protein Data
Bank (PDB) [23].
We employ the time-split of PDBBind proposed by [1] with 17k complexes from2018 or earlier for training/validation and 363 test structures from 2019 with no ligand overlapwith the training complexes.
This is motivated by the further adoption of the same split [2] and thecritical assessment of PDBBind splits by [24] who favor temporal splits over artificial splits based onmolecular scaffolds or protein sequence/structure similarity.
For completeness, we also report theresults on protein sequence similarity splits in Appendix E.
Metrics.
To evaluate the generated complexes, we compute the heavy-atom RMSD between thepredicted and the crystal ligand atoms when the protein structures are aligned.
To account forpermutation symmetries in the ligand, we use the symmetry-corrected RMSD of sPyRMSD [33].
Forthese RMSD values, we report the percentage of predictions that have an RMSD that is less than2Å. We choose 2Å since much prior work considers poses with an RMSD less that 2Å as “good"or successful [7, 8, 9].
This is a chemically relevant metric, unlike the mean RMSD as detailed in Section 3 since for further downstream analyses such as determining function changes, a predictionis only useful below a certain RMSD error threshold.
Less relevant metrics such as the mean RMSDare provided in Appendix E.
D.2 Implementation details: hyperparameters, training, and runtime measurement
Training Details.
We use Adam [34] as optimizer for the diffusion and the confidence model.
Thediffusion model with which we run inference uses the exponential moving average of the weightsduring training, and we update the moving average after every optimization step with a decay factorof 0.999.
The batch size is 16.
We run inference with 20 denoising steps on 500 validation complexesevery 5 epochs and use the set of weights with the highest percentage of RMSDs less than 2Å as thefinal diffusion model.
We trained our final score model on four 48GB RTX A6000 GPUs for 850epochs (around 18 days).
The confidence model is trained on a single 48GB GPU.
For inference,only a single GPU is required.
Scaling up the model size seems to improve performance and futurework could explore whether this trend continues further.
For the confidence model uses the validationcross-entropy loss is used for early stopping and training only takes 75 epochs.
Code to reproduceall results including running the baselines or to perform docking calculations for new complexes isavailable at https://anonymous.4open.science/r/DiffDock .
Hyperparameters.
For determining the hyperparameters of D IFF D OCK ’s score model, we trainedsmaller models (3.97 million parameters) that fit into 48GB of GPU RAM before scaling it up to thefinal model (20.24 million parameters) that was trained on four 48GB GPUs.
The smaller modelswere only trained for 250 or 300 epochs, and we used the fraction of predictions with an RMSD below2Å on the validation set to choose the hyperparameters.
Table 2 shows the main hyperparameterswe tested and the final parameters of the large model we use to obtain our results.
We only did littletuning for the minimum and maximum noise levels of the three components of the diffusion.
For thetranslation, the maximum standard deviation is 19Å. We also experimented with second-order featuresfor the Tensorfield Network but did not find them to help.
The complete set of hyperparameters nextto the main ones we describe here can be found in our repository.
The confidence model has 4.77 million parameters and the parameters we tried are in Table 3.
Wegenerate 28 different training poses for the confidence model (for which it predicts whether or notthey have an RMSD below 2Å) with a score model that has a maximum translation standard deviationof 34Å. The score model used to generate the training samples for the confidence model is thusdifferent from the large score model we used during inference.
[Table] [Table]
[Table] [Table]
Runtime.
Similar to all the baselines, the preprocessing times are not included in the reportedruntimes.
For D IFF D OCK the preprocessing time is negligible compared to the rest of the inferencetime where multiple reverse diffusion steps are performed.
Preprocessing mainly consists of a forwardpass of ESM2 to generate the protein language model embeddings, RDKit’s conformer generation,and the conversion of the protein into a radius graph.
We measured the inference time when runningon an RTX A100 40GB GPU when generating 10 samples.
The runtimes we report for generating40 samples and ranking them are extrapolations where we multiply the runtime for 10 samples by 4.
In practice, this only gives an upper bound on the runtime with 40 samples, and the actual runtimeshould be faster.
D.3 Baselines: implementation, used scripts, and runtime details
Our scripts to run the baselines are available at https://anonymous.4open.science/r/
DiffDock .
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 ofthe 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 run once when dockingmany ligands to the same protein.
In applications where different receptors are processed (such asreverse screening), the experienced runtimes for TANKBind and P2Rank + SMINA/GNINA will thusbe higher.
We note that for all these baselines we have used the default hyperparameters unless specifieddifferently below.
Modifying some of these hyperparameters (for example the scoring method’sexhaustiveness) will change the runtime and performance tradeoffs (e.g., if the searching routine isleft running for longer then better poses are likely to be found), however, we leave these analyses tofuture work.
SMINA [25] improves Autodock Vina with a new scoring-function and user friendliness.
The defaultparameters were used with the exception of setting –num_modes 10 .
To define the search box, weuse the automatic box creation option around the receptor with the default buffer of 4Å on all 6 sides.
GNINA [9] builds on SMINA by additionally using a learned 3D CNN for scoring.
The defaultparameters were used with the exception of setting –num_modes 10 .
To define the search box, weuse the automatic box creation option around the receptor with the default buffer of 4Å on all 6 sides.
QuickVina-W [8] extends the speed-optimized QuickVina [-
 [Alhossary et al., 2015]-]2 [7] for blind docking.
We reuse thenumbers from [1] which had used the default parameters except for increasing the exhaustiveness to64.
GLIDE [11] is a strong heavily used commercial docking tool.
These methods all use biophysicsbased scoring-functions.
We reuse the numbers from [1] since we do not have a license.
As explainedby [1], the very high runtime of GLIDE with 1405 seconds per complex is partially explained by thefact that GLIDE only uses a single thread when processing a complex.
This fact and the parallelizationoptions 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 complexin parallel.
However, each process also requires a separate software license.
EquiBind [1], we reuse the numbers reported in their paper and generate the predictions that wevisualize 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 Table 1 we report two runtimes (0.72/2.5 sec).
The first is the runtime when making only the top-1prediction and the second is for producing the top-5 predictions.
Producing only the top-1 predictionsis faster since TANKBind produces distance predictions that need to be converted to coordinates witha gradient descent algorithm and this step only needs to be run once for the top-1 prediction, whileit needs to be run 5 times for producing 5 outputs.
To obtain our runtimes we run the forward passof TANKBind on GPU (0.28 seconds) with the default batch size of 5 that is used in their GitHubrepository.
To compute the time the distances-to-coordinates conversion step takes, we run thefile baseline_run_tankbind_parallel.sh in our repository, which parallelizes the computationacross 16 processes which we also run on an Intel Xeon Gold 6230 CPU.
This way, we obtain 0.44seconds runtime for the conversion step of the top-1 prediction (averaged over the 363 complexes ofthe testset).
P2Rank [26], 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 bindingposes 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 theexact final binding pose.
We use –autobox_add 10 to add an additional 10Å on all 6 sides of thebounding box following [1].
P2Rank + SMINA/GNINA.
The bounding box in which GNINA/SMINA searches for bindingposes 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 bindingpose.
The diameter of the search box is the diameter of a ligand conformer generated by RDKit withan additional 10Å on all 6 sides of the bounding box.
E Additional Results
E.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 incase of symmetric complexes, this leads to predicting the ligand to be at an un-physical state in themiddle of the possible binding pockets as visualized in Figure 10.
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, thisoften results in “squashed-up" predictions of the regression-based methods as visualized in Figure 5.
If there is uncertainty about the correct conformer, the square error minimizing option is to put allatoms close to the mean.
These averaging phenomena in the presence of either aleatoric or epistemic uncertainty cause theregression-based methods to often generate steric clashes and self intersections.
To investigate thisquantitatively, we determine the fraction of test complexes for which the methods exhibit stericclashes.
We define a ligand as exhibiting a steric clash if one of its heavy atoms is within 0.4Å of aheavy receptor atom.
This cutoff is used by protein quality assessment tools and in previous literature[35].
Table 4 shows that D IFF D OCK , as a generative model, produces fewer steric clashes thanthe regression-based baselines.
We generally observe no unphysical predictions from D IFF D OCKunlike the self intersections that, e.g., TANKBind produces (Figure 5) or its incorrect local structures(Figure 6).
This is also visible in the randomly chosen examples of Figure 9 and can be examined inour repository, where we provide all predictions of D IFF D OCK for the test set.
E.2 Further Results and Metrics
In this section, we present further evaluation metrics on the results presented in Table 1.
In particular,for both top-1 (Table 5) and top-5 (Table 6) we report: 25th, 50th and 75th percentiles, the proportionbelow 2Å and below 5Å of both ligand RMSD and centroid distance.
Moreover, while [24] advocatedagainst artificial protein set splits and for time-based splits, for completeness, in Table 7 and Figure 7,we report the performances of the different methods when evaluated exclusively on the portion of thetest set where the UniProt IDs of the proteins are not contained in the data that is seen by D IFF D
OCKin its training and validation.
E.3 Ablation studies
Below we report the performance of our method over different hyperparameter settings.
In particular,we highlight the different ways in which it is possible to control the tradeoff between runtime and accuracy in our method.
These mainly are: (1) model size, (2) diffusion time, and (3) diffusionsamples.
Model size.
The final D IFF D OCK score model has 20.24 million parameters from its 6 convolutionlayers with 48 scalar and 10 vector features.
In Table 8 we show the results for a smaller score modelwith 5 convolutions, 24 scalar, and 6 vector features resulting in 3.97 million parameters that can betrained on a single 48GB GPU.
The confidence model used is the same for both score models.
Wefind that scaling up the model size helped improve performance which we did as far as possible usingfour 48GB GPUs for training.
Scaling the model size further is a promising avenue for future work.
Diffusion steps.
Another hyperparameter determining the runtime of the method during inferenceis the number of steps we take during the reverse diffusion.
Since these are applied sequentially
D IFF D OCK ’s runtime scales approximately linearly with the number of diffusion steps.
In the rest ofthe paper, we always use 20 steps, but in Figure 8 we show how the performance of the model varieswith the number of steps.
We note that the model reaches nearly the full performance even with just10 steps, suggesting that the model can be sped up 2x with a small drop in accuracy.
Diffusion samples.
Given a score-based model and a number of steps for the diffusion model, itremains to be determined how many independent samples N to query from the diffusion model andthen feed to the confidence model.
As expected the more samples the confidence model receivesthe more likely it is that it will find a pose that it is confident about and, therefore, the higher theperformance.
The runtime of D IFF D OCK on GPU scales sublinearly until the different samples fit in
[Table]
parallel in the model (depends on the protein size and the GPU memory) and approximately linearlyfor larger sample sizes (however it can be easily parallelized across different GPUs).
In Figure 3we show how the success rate for the top-1, top- [Table] and top-10 prediction change as a function of N .
For example, for the top-1 prediction, the proportion of the prediction with RMSD below 2Åvaries between 22% of a random sample of the diffusion model ( N = 1 ) to 38% when the confidencemodel is allowed to choose between 40 samples [Table] [Table] [Table]
[Table] ANK [Table] [Table]
[Table]
[Table]
[Table]
[Table]
[Table]
DiﬀDock
GLIDE
GNINA
SMINA
QVinaW
TANKBind
EquiBind
RMSD (Å)
F r a c t i o n w i t
h l o w e r
R M
S D
[Table]
E.4 Affinity prediction
To validate the quality of the predicted poses [Table] we also do some experiments in predicting the bindingaffinity labels already present in PDBBind.
In this section we report some preliminary results on thistask that show that a simple approach can already achieve results competitive with the state-of-the-art.
We leave a more thorough and sophisticated analysis on how to best use the D IFF D OCK frameworkfor binding affinity to future work.
Affinity prediction framework.
We train the binding affinity predictor by generating a fixednumber of poses with the diffusion model and then feeding them to an affinity prediction model with [Table] Top-1 performance Top-5 performance
Top-10 performance
Number of reverse diﬀusion steps
F r a c t i o n w i t h R M S D < 2 Å architecture almost analogous to the confidence model. This affinity prediction model takes in theposes as a single heterogeneous graph with a single receptor but multiple sets of ligand nodes, whichhave edges to the same receptor but not among themselves.
After the final interaction layer, the scalarrepresentations of nodes in each ligand are aggregated with a mean pooling and passed through a setof dense layers (as it is done for the confidence prediction).
Then, the representations of the differentligands are aggregated using multiple permutation invariant aggregators (mean, maximum, minimum,and standard deviation) as in [36], and transformed with another set of dense layers producing asingle output, the predicted affinity.
Dataset, baselines, and training.
To train we use PDBBind with the same splits used to trainthe diffusion and confidence models.
This provides for each of the complexes an affinity measurethat consists of inhibiting concentration ( IC 50 ), inhibition constant ( K i ), or dissociation constant( K d ) and its conversion to the − log K d /K i metric.
As baselines, we use a series of state-of-the-artsequence-based and structure-based methods: TransformerCPI [37], MONN [38], IGN [39], PIGNet[40], HOLOPTOT [41], STAMPDPI [42] and TANKBind [2]. We take the baselines’ performancesfrom [2]. Prediction of − log K d /K i on PDBBind.
The baseline numbers are from [2]. No hyperparameter tuning was performed for D IFF D OCK ’s performance. Methods [Table]
Results.
The results presented in Table 9 highlight how even preliminary results with a straightforwardway of using D IFF D OCK ’s predictions for affinity prediction achieve a performance that is on par withthe state-of-the-art.
We hope this can motivate future work on better integrating affinity prediction inthe method and scaling to larger amounts of data
[Table]
see
, since is a generative model, D 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