Torsional Diffusion for Molecular Conformer Generation
Abstract
Molecular conformer generation is a fundamental task in computational chemistry.
Several machine learning approaches have been developed, but none have outperformed state-of-the-art cheminformatics methods.
We propose torsional diffusion , a novel diffusion framework that operates on the space of torsion angles via a diffusion process on the hypertorus and an extrinsic-to-intrinsic score model.
On a standard benchmark of drug-like molecules, torsional diffusion generates superior conformer ensembles compared to machine learning and cheminformatics methods in terms of both RMSD and chemical properties, and is orders of magnitude faster than previous diffusion-based models.
Moreover, our model provides exact likelihoods, which we employ to build the ﬁrst generalizable Boltzmann generator.
Code is available at https://github.com/gcorso/torsional-diffusion .
Introduction
Many properties of a molecule are determined by the set of low-energy structures, called conformers , that it adopts in 3D space.
Conformer generation is therefore a fundamental problem in computational chemistry [Hawkins, 2017] and an area of increasing attention in machine learning.
Traditional approaches to conformer generation consist of metadynamics-based methods, which are accurate but slow [Pracht et al., 2020]; and cheminformatics-based methods, which are fast but less accurate [Hawkins et al., 2010, Riniker and Landrum, 2015].
Thus, there is growing interest in developing deep generative models to combine high accuracy with fast sampling.
Diffusion or score-based generative models [Ho et al., 2020, Song et al., 2021]—a promising class of generative models—have been applied to conformer generation under several different formulations.
These have so far considered diffusion processes in Euclidean space in which Gaussian noise is injected independently into every data coordinate—either pairwise distances in a distance matrix [Shi et al., 2021, Luo et al., 2021] or atomic coordinates in 3D [Xu et al., 2022].
However, these models require a large number of denoising steps and have so far failed to outperform the best cheminformatics methods.
We instead propose torsional diffusion , in which the diffusion process over conformers acts only on the torsion angles and leaves the other degrees of freedom ﬁxed.
This is possible and effective because the ﬂexibility of a molecule, and thus the difﬁculty of conformer generation, lies largely in torsional degrees of freedom [Axelrod and Gómez-Bombarelli, 2022]; in particular, bond lengths and angles can already be determined quickly and accurately by standard cheminformatics methods.
Leveraging this insight signiﬁcantly reduces the dimensionality of the sample space; drug-like molecules 2 have, on average, n = 44 atoms, corresponding to a 3 n -dimensional Euclidean space, but only m =
9 torsion angles of rotatable bonds.
As measured from the standard dataset GEOM-DRUGS [Axelrod and Gómez-Bombarelli, 2022] 36th
Torsion angle coordinates deﬁne not a Euclidean space, but rather an m -dimensional torus T m (Figure 1, left ).
However, the dimensionality and distribution over the torus vary between molecules and even between different ways of deﬁning the torsional space for the same molecule.
To resolve these difﬁculties, we develop an extrinsic-to-intrinsic score model (Figure 1, right ) that takes as input a 3D point cloud representation of the conformer in Euclidean space (extrinsic coordinates), and predicts as output a score on a torsional space speciﬁc to that molecule (intrinsic coordinates).
To do so, we consider a torsional score for a bond as a geometric property of a 3D point cloud, and use
SE (3) -equivariant networks to predict them directly for each bond.
Unlike prior work, our model provides exact likelihoods of generated conformers, enabling training with the ground-truth energy function rather than samples alone.
This connects with the literature on Boltzmann generators —generative models which aim to sample the Boltzmann distribution of physical systems without expensive molecular dynamics or MCMC simulations [Noé et al., 2019,
Köhler et al., 2021].
Thus, as a variation on the torsional diffusion framework, we develop torsional
Boltzmann generators that can approximately sample the conditional Boltzmann distribution for unseen molecules.
This starkly contrasts with existing Boltzmann generators, which are speciﬁc for the chemical system on which they are trained.
Our main contributions are:
• We formulate conformer generation in terms of diffusion modeling on the hypertorus— the ﬁrst demonstration of non-Euclidean diffusion on complex datasets—and develop an extrinsic-to-intrinsic score model that satisﬁes the required symmetries: SE (3) invariance, torsion deﬁnition invariance, and parity equivariance.
• We obtain state-of-the-art results on the GEOM-DRUGS dataset [Axelrod and GómezBombarelli, 2022] and are the ﬁrst method to consistently outperform the established commercial software OMEGA [Hawkins, 2017].
We do so using two orders of magnitude fewer denoising steps than GeoDiff [Xu et al., 2022], the best Euclidean diffusion approach.
• We propose torsional Boltzmann generators—the ﬁrst Boltzmann generator based on diffusion models rather than normalizing ﬂows and the ﬁrst to be useful for a class of molecules rather than a speciﬁc system.
Background
Diffusion generative models Consider the data distribution as the starting distribution p 0 ( x ) of a forward diffusion process described by an Ito stochastic differential equation (SDE):
[Equation]
where w is the Wiener process and f ( x , t ) , g ( t ) are chosen functions.
With sufﬁciently large T , the distribution p T ( x ) —the prior —approaches a simple Gaussian.
Sampling from the prior and solving the reverse diffusion
[Equation]
yields samples from the data distribution p 0 ( x )
[Anderson, 1982, Song et al., 2021].
Diffusion, or score-based, generative models [Ho et al., 2020, Song et al., 2021] learn the score ∇ x log p t ( x ) of the diffused data with a neural network and generate data by approximately solving the reverse diffusion.
The score of the diffused data also deﬁnes a probability ﬂow ODE —a continuous normalizing ﬂow that deterministically transforms the prior into the data distribution [Song et al., 2021].
We leverage the insight that, in many cases, this ﬂow makes it possible to use diffusion models in place of normalizing ﬂows and highlight one such case with the torsional Boltzmann generator.
Diffusion generative models have traditionally been used to model data on Euclidean spaces (such as images)
; however, De Bortoli et al.
[2022] recently showed that the theoretical framework holds with relatively few modiﬁcations for data distributions on compact Riemannian manifolds.
The hypertorus T m , which we use to deﬁne torsional diffusion, is a speciﬁc case of such a manifold.
Several methods [Salimans and Ho, 2022, Vahdat et al., 2021, Nichol and Dhariwal, 2021] have been proposed to improve and accelerate diffusion models in the domain of image generation.
Among these, the most relevant to this work is subspace diffusion
[Jing et al., 2022], in which the diffusion is progressively restricted to linear subspaces.
Torsional diffusion can be viewed in a similar spirit, as it effectively restricts Euclidean diffusion to a nonlinear manifold given by ﬁxing the non-torsional degrees of freedom.
Molecular
conformer
generation The conformers of a molecule are the set of its energetically favorable 3D structures, corresponding to local minima of the potential energy surface.
The gold standards for conformer generation are metadynamics-based methods such as CREST [Pracht et al., 2020], which explore the potential energy surface while ﬁlling in local minima [Hawkins, 2017].
However, these require an average of 90 core-hours per drug-like molecule [Axelrod and GómezBombarelli, 2022] and are not considered suitable for high-throughput applications.
Cheminformatics methods instead leverage approximations from chemical heuristics, rules, and databases for significantly faster generation [Lagorce et al., 2009, Cole et al., 2018, Miteva et al., 2010, Bolton et al., 2011, Li et al., 2007]; while these can readily model highly constrained degrees of freedom, they fail to capture the full energy landscape.
The most well-regarded of such methods include the commercial software OMEGA [Hawkins et al., 2010] and the open-source RDKit ETKDG [Landrum et al., 2013,
Riniker and Landrum, 2015].
A number of machine learning methods for conformer generation has been developed [Xu et al., 2021a,b, Shi et al., 2021, Luo et al., 2021], the most recent and advanced of which are GeoMol [Ganea et al., 2021] and GeoDiff [Xu et al., 2022].
GeoDiff is a Euclidean diffusion model that treats conformers as point clouds x ∈ R 3 n and learns an SE (3) equivariant score.
On the other hand,
GeoMol employs a graph neural network that, in a single forward pass, predicts neighboring atomic coordinates and torsion angles from a stochastic seed.
Boltzmann generators
An important problem in physics and chemistry is that of generating independent samples from a Boltzmann distribution p ( x ) ∝ e − E ( x ) /kT
with known but unnormalized density.
4 Generative models with exact likelihoods, such as normalizing ﬂows, can be trained to match such densities [Noé et al., 2019] and thus provide independent samples from an approximation of the target distribution.
Such Boltzmann generators have shown high ﬁdelity on small organic molecules [Köhler et al., 2021] and utility on systems as large as proteins [Noé et al., 2019].
However, a separate model has to be trained for every molecule, as the normalizing ﬂows operate on intrinsic coordinates whose deﬁnitions are speciﬁc to that molecule.
This limits the utility of existing Boltzmann generators for molecular screening applications.
Torsional Diffusion
Consider a molecule as a graph G =
( V , E ) with atoms v ∈ V and bonds e ∈ E , 5 and denote the space of its possible conformers C G .
A conformer C ∈ C G can be speciﬁed in terms of its intrinsic (or internal) coordinates: local structures L consisting of bond lengths, bond angles, and cycle conformations; and torsion angles τ consisting of dihedral angles around freely rotatable bonds (precise deﬁnitions in Appendix A).
We consider a bond freely rotatable if severing the bond creates two connected components of G , each of which has at least two atoms.
Thus, torsion angles in cycles (or rings), which cannot be rotated independently, are considered part of the local structure L .
Conformer generation consists of learning probability distributions p G ( L, τ ) .
However, the set of possible stable local structures L for a particular molecule is very constrained and can be accurately predicted by fast cheminformatics methods, such as RDKit ETKDG [Riniker and Landrum, 2015]
(see Appendix F.1 for veriﬁcation).
Thus, we use RDKit to provide approximate samples from p G ( L ) , and develop a diffusion-based generative model to learn distributions p G ( τ | L ) over torsion angles—conditioned on a given graph and local structure.
Our method is illustrated in Figure 1 and detailed as follows.
Section 3.1 formulates diffusion modeling on the torus deﬁned by torsion angles.
Section 3.2 describes the torsional score framework,
Section 3.3 the required symmetries, and Section 3.4 our score model architecture.
Section 3. discusses likelihoods, and Section 3.6 how likelihoods can be used for energy-based training.
3.1 Diffusion modeling on T m
Since each torsion angle coordinate lies in [0 , 2 π ) , the m torsion angles of a conformer deﬁne a hypertorus T m .
To learn a generative model over this space, we apply the continuous score-based framework of Song et al.
[2021], which holds with minor modiﬁcations for data distributions on compact Riemannian manifolds (such as T m )
[De Bortoli et al., 2022].
Speciﬁcally, for Riemannian manifold M let x ∈ M , let w be the Brownian motion on the manifold, and let the drift f ( x , t ) , score ∇ x log p t ( x ) , and score model output s ( x , t ) be elements of the tangent space T x M .
Then equation 2 remains valid—that is, discretizing and solving the reverse SDE on the manifold as a geodesic random walk starting with samples from p T ( x ) approximately recovers the original data distribution p 0 ( x ) [De Bortoli et al., 2022].
For the forward diffusion we use rescaled Brownian motion given by f ( x , t ) = 0 , g ( t ) =
(cid:113) ddt σ 2 ( t ) where σ ( t ) is the noise scale.
Speciﬁcally, we use an exponential diffusion σ ( t ) =
σ 1 − t min σ t max as in
Song and Ermon [2019], with σ min = 0 .
01 π , σ max =
π, t ∈ (0 , 1) .
Due to the compactness of the manifold, however, the prior p T ( x ) is no longer a Gaussian, but a uniform distribution over M .
Training the score model with denoising score matching requires a procedure to sample from the perturbation kernel p t | 0 ( x (cid:48) | x ) of the forward diffusion and compute its score.
We view the torus T m ∼ =
[0 , 2 π ) m as the quotient space R m /
2 π Z m with equivalence relations ( τ 1 , . . .
τ m ) ∼ ( τ 1 + 2 π, . . .
, τ m ) . . .
∼ ( τ 1 , . . .
τ m + 2 π ) .
Hence, the perturbation kernel for rescaled Brownian motion on T m is the wrapped normal distribution on R m ; that is, for any τ , τ (cid:48) ∈ [0 , 2 π ) m , we have
[Equation]
[Equation]
where σ ( t ) is the noise scale of the perturbation kernel p t | 0 .
We thus sample from the perturbation kernel by sampling from the corresponding unwrapped isotropic normal and taking elementwise mod 2 π .
The scores of the kernel are pre-computed using a numerical approximation.
During training, we sample times t at uniform and minimize the denoising score matching loss
[Equation]
where the weight factors λ ( t )
= 1 / E τ ∼ p t | 0 ( ·| 0) (cid:2)
||∇ τ log p t | 0 ( τ | 0 )
|| 2 (cid:3)
are also precomputed.
As the tangent space T τ T m is just R m , all the operations in the loss computation are the familiar ones.
For inference, we ﬁrst sample from a uniform prior over the torus.
We then discretize and solve the reverse diffusion with a geodesic random walk; however, since the exponential map on the torus (viewed as a quotient space) is just exp τ ( δ )
=
τ + δ mod 2 π , the geodesic random walk is equivalent to the wrapping of the random walk on R m .
3.2 Torsional score framework
While we have deﬁned the diffusion process over intrinsic coordinates, learning a score model s ( τ , t ) directly over intrinsic coordinates is potentially problematic for several reasons.
First, the dimensionality m of the torsional space depends on the molecular graph G .
Second, the mapping from torsional space to physically distinct conformers depends on G and local structures L , but it is unclear how to best provide these to a model over T m .
Third, there is no canonical choice of independent intrinsic coordinates ( L, τ ) ; in particular, the torsion angle at a rotatable bond can be deﬁned as any of the dihedral angles at that bond, depending on an arbitrary choice of reference neighbors (Figure 2 and Appendix A).
Thus, even with ﬁxed G and L , the mapping from T m to conformers is ill-deﬁned.
This posed a signiﬁcant challenge to prior works using intrinsic coordinates [Ganea et al., 2021].
To circumvent these difﬁculties, we instead consider a conformer
C ∈ C G in terms of its extrinsic (or Cartesian) coordinates—that is, as a point cloud in 3D space, deﬁned up to global roto-translation:
C G ∼ =
R 3 n /SE (3) .
Then, we construct the score model
s G ( C, t ) as a function over C G rather than T m .
The outputs remain in the tangent space of T m , which is just R m .
Such a score model is simply an SE (3) - invariant model over point clouds in 3D space s G : R 3 n × [0 , T ] (cid:55)→ R m conditioned on G .
Thus, we have reduced the problem of learning a score on the torus, conditioned on the molecular graph and local structure, to the much more familiar problem of predicting SE (3) -invariant scalar quantities—one for each bond—from a 3D conformer.
It may appear that we still need to choose a deﬁnition of each torsion angle τ i so that we can sample from p t | 0 ( · | τ ) during training and solve the reverse SDE over τ during inference.
However, we leverage the following insight: given ﬁxed local structures , the action on C of changing a single torsion angle τ i by some ∆ τ i can be applied without choosing a deﬁnition (Figure 2).
Geometrically, this action is a (signed) relative rotation of the atoms on opposite sides of the bond and can be applied directly to the atomic coordinates in 3D. The geometric intuition can be stated as follows (proven in Appendix B and discussed further in
Appendix F.2).
Proposition 1.
Let ( b i , c i ) be a rotatable bond, let x V ( b i ) be the positions of atoms on the b i side of the molecule, and let R ( θ , x c i ) ∈ SE (3) be the rotation by Euler vector θ about x c i .
Then for
C, C (cid:48) ∈ C G , if τ i is any deﬁnition of the torsion angle around bond ( b i , c i ) ,
[Equation]
where ˆr b i c i
=
( x c i − x b i ) /
||
x c i − x b i || .
To apply a torsion update ∆ τ = (∆ τ 1 , . . .
∆ τ m ) involving all bonds, we apply ∆ τ i sequentially in any order.
Then, since training and inference only make use of torsion updates ∆ τ , we work solely in terms of 3D point clouds and updates applied to them.
To draw local structures L from RDKit, we draw full 3D conformers C ∈ C G and then randomize all torsion angles to sample uniformly over T m .
To solve the reverse SDE, we repeatedly predict torsion updates directly from, and apply them directly to, the 3D point cloud.
Therefore, since our method never requires a choice of reference neighbors for any τ i , it is manifestly invariant to such a choice.
These procedures are detailed in
Appendix C.
3.3 Parity equivariance
The torsional score framework presented thus far requires an SE (3) -invariant model.
However, an additional symmetry requirement arises from the fact that the underlying physical energy is invariant, or extremely nearly so, under parity inversion [Quack, 2002].
Thus our learned density should respect p ( C ) =
p ( − C ) where − C = {− x | x ∈ C } .
In terms of the conditional distribution over torsion angles, we require p ( τ ( C )
| L ( C ))
=
p ( τ ( − C )
| L ( − C )) .
Then,
Proposition 2.
If p ( τ ( C ) | L ( C ))
=
p ( τ ( − C )
| L ( − C )) , then for all diffusion times t ,
[Equation]
Because the score model seeks to learn s G ( C, t ) = ∇ τ log p t ( τ ( C )
| L ( C )) , we must have s G ( C, t ) =
− s G ( − C, t ) .
Thus, the score model must be invariant under SE (3) but equivariant (change sign) under parity inversion of the input point cloud— i.e. it must output a set of pseudoscalars in R m .
3.4 Score network architecture
Based on sections 3.2 and 3.3, the desiderata for the score model are:
Predict a pseudoscalar δτ i := ∂ log p/∂τ i ∈ R that is SE (3) -invariant and parity equivariant for every rotatable bond in a 3D point cloud representation of a conformer.
While there exist several GNN architectures which are SE (3) -equivariant
[Jing et al., 2021, Satorras et al., 2021], their SE (3) -invariant outputs are also parity invariant and, therefore, cannot satisfy the desired symmetry.
Instead, we leverage the ability of equivariant networks based on tensor products [Thomas et al., 2018, Geiger et al., 2022] to produce pseudoscalar outputs.
Our architecture, detailed in Appendix D, consists of an embedding layer, a series of atomic convolution layers, and a ﬁnal bond convolution layer.
The ﬁrst two closely follow the architecture of Tensor
Field Networks [Thomas et al., 2018], and produce learned feature vectors for each atom.
The ﬁnal bond convolution layer constructs tensor product ﬁlters spatially centered on every rotatable bond and aggregates messages from neighboring atom features.
We extract the pseudoscalar outputs of this ﬁlter to produce a single real-valued pseudoscalar prediction δτ i for each rotatable bond.
Naively, the bond convolution layer could be constructed the same way as the atomic convolution layers, i.e., with spherical harmonic ﬁlters.
However, to supply information about the orientation of the bond about which the torsion occurs, we construct a ﬁlter from the product of the spherical harmonics with a representation of the bond (Figure 1D).
Because the convolution conceptually resembles computing the torque, we call this ﬁnal layer the pseudotorque layer.
3.5 Likelihood
By using the probability ﬂow ODE, we can compute the likelihood of any sample τ as follows [Song et al., 2021, De Bortoli et al., 2022]:
[Equation]
In Song et al.
[2021], the divergence term is approximated via Hutchinson’s method [Hutchinson, 1989], which gives an unbiased estimate of log p 0 ( τ ) .
However, this gives a biased estimate of p 0 ( τ ) , which is unsuitable for our applications.
Thus, we compute the divergence term directly, which is feasible here (unlike in Euclidean diffusion) due to the reduced dimensionality of the torsional space.
The above likelihood is in torsional space p G ( τ | L ) , τ ∈ T m , but to enable compatibility with the
Boltzmann measure e − E ( x ) /kT , it is desirable to interconvert this with a likelihood in Euclidean space p ( x | L ) , x ∈ R 3 n .
A factor is necessary to convert between the volume element in torsional space and in Euclidean space (full derivation in Appendix B):
Proposition 3.
Let x ∈ C
( τ , L ) be a centered 7 conformer in Euclidean space.
Then,
[Equation]
where the indices α, β are integers between 1 and m + 3 .
For 1 ≤ α ≤ m , J ( k ) α is deﬁned as
[Equation]
and for α ∈ { m + 1 , m + 2 , m + 3 } as
[Equation]
where ( b i , c i ) is the freely rotatable bond for torsion angle
i , V ( b i ) is the set of all nodes on the same side of the bond as b i , and ˆ x, ˆ y, ˆ z are the unit vectors in the respective directions.
3.6 Energy-based training
By computing likelihoods, we can train torsional diffusion models to match the Boltzmann distribution over torsion angles using the energy function.
At a high level, we minimize the usual score matching loss, but with simulated samples from the Boltzmann distribution rather than data samples.
The procedure therefore consists of two stages: resampling and score matching, which are tightly coupled during training (Algorithm 1).
In the resampling stage, we use the model as an importance sampler for the Boltzmann distribution, where Proposition 3 is used to compute the (unnormalized) torsional
Boltzmann density ˜ p G ( τ | L ) .
In the score-matching stage, the importance weights are used to approximate the denoising score-matching loss with expectations taken over ˜ p G ( τ | L ) .
As the model learns the score, it improves as an importance sampler.
[Table]
This training procedure differs substantially from that of existing Boltzmann generators, which are trained as ﬂows with a loss that directly depends on the model density.
In contrast, we train the model as a score-based model, but use it as a ﬂow—both during training and inference—to generate samples.
The model density is needed only to reweight the samples to approximate the target density.
Since in principle the model used for resampling does not need to be the same as the model being trained, 8 we can use very few steps (a shallow ﬂow) during resampling to accelerate training, and then increase the number of steps (a deeper ﬂow) for better approximations during inference—an option unavailable to existing Boltzmann generators.
Experiments
We evaluate torsional diffusion by comparing the generated and ground-truth conformers in terms of ensemble RMSD (Section 4.3) and properties (Section 4.4).
Section 4.1 ﬁrst discusses a preprocessing procedure required to train a conditional model p G ( τ | L ) .
Section 4.5 concludes with torsional
Boltzmann generators.
See Appendix H for additional results, including ablation experiments.
4.1 Conformer matching
In focusing on p G ( τ | L ) , we have assumed that we can sample local structures L ∼ p G ( L ) with
RDKit.
While this assumption is very good in terms of RMSD, the RDKit marginal ˆ p G ( L ) is only an approximation of the ground truth p G ( L ) .
Thus, if we train on the denoising score-matching loss with ground truth conformers—i.e., conditioned on ground truth local structures—there will be a distributional shift at test time, where only approximate local structures from ˆ p G ( L ) are available.
We found that this shift signiﬁcantly hurts performance.
We thus introduce a preprocessing procedure called conformer matching .
In brief, for the training split only, we substitute each ground truth conformer C with a synthetic conformer ˆ C with local structures ˆ L ∼ ˆ p G ( L ) and made as similar as possible to C .
That is, we use RDKit to generate
[Table]
ˆ L and change torsion angles ˆ τ to minimize RMSD( C, ˆ C ) .
Naively, we could sample ˆ L ∼ ˆ p G ( L ) independently for each conformer, but this eliminates any possible dependence between L and τ that could serve as training signal.
Instead, we view the distributional shift as a domain adaptation problem that can be solved by optimally aligning p G ( L ) and ˆ p G ( L ) .
See Appendix E for details.
4.2 Experimental setup
Dataset We evaluate on the GEOM dataset [Axelrod and Gómez-Bombarelli, 2022], which provides gold-standard conformer ensembles generated with metadynamics in CREST [Pracht et al., 2020].
We focus on GEOM-DRUGS—the largest and most pharmaceutically relevant part of the dataset— consisting of 304k drug-like molecules (average 44 atoms).
To test the capacity to extrapolate to the largest molecules, we also collect from GEOM-MoleculeNet all species with more thanatoms into a dataset we call GEOM-XL and use it to evaluate models trained on DRUGS.
Finally, we train and evaluate models on GEOM-QM9, a more established dataset but with signiﬁcantly smaller molecules (average 11 atoms).
Results for GEOM-XL and GEOM-QM9 are in Appendix H.
Evaluation We use the train/val/test splits from Ganea et al.
[2021] and use the same metrics to compare the generated and ground truth conformer ensembles: Average Minimum RMSD (AMR) and
Coverage.
These metrics are reported both for Recall (R)—which measures how well the generated ensemble covers the ground-truth ensemble—and Precision (P)—which measures the accuracy of the generated conformers.
See Appendix G for exact deﬁnitions and further details.
Following the literature, we generate 2 K conformers for a molecule with K ground truth conformers.
Baselines We compare with the strongest existing methods from Section 2.
Among cheminformatics methods, we evaluate RDKit ETKDG [Riniker and Landrum, 2015], the most established open-source package, and OMEGA [Hawkins et al., 2010, Hawkins and Nicholls, 2012], a commercial software in continuous development.
Among machine learning methods, we evaluate GeoMol [Ganea et al., 2021] and GeoDiff [Xu et al., 2022], which have outperformed all previous models on the evaluation metrics.
Note that GeoDiff originally used a small subset of the DRUGS dataset, so we retrained it using the splits from Ganea et al.
4.3 Ensemble RMSD
Torsional diffusion signiﬁcantly outperforms all previous methods on GEOM-DRUGS (Table 1 and
Figure 3), reducing by 30% the average minimum recall RMSD and by 16% the precision RMSD relative to the previous state-of-the-art method.
Torsional diffusion is also the ﬁrst ML method to consistently generate better ensembles than OMEGA.
As OMEGA is a well-established product used in industry, this represents an essential step towards establishing the utility of conformer generation with machine learning.
Torsional diffusion offers speciﬁc advantages over both GeoDiff and GeoMol, the most advanced prior machine learning methods.
GeoDiff, a Euclidean diffusion model, requires 5000 denoising steps to obtain the results shown, whereas our model—thanks to the reduced degrees of freedom—requires only 20 steps.
In fact, our model outperforms GeoDiff with as few as 5 denoising steps.
As seen in
Table 2, this translates to enormous runtime improvements.
[Figure]
Method
Steps
[Table] [Table]
Compared to torsional diffusion, GeoMol similarly makes use of intrinsic coordinates.
However, since GeoMol can only access the molecular graph, it is less suited for reasoning about relationships that emerge only in a spatial embedding, especially between regions of the molecule that are distant on the graph.
Our extrinsic-to-intrinsic score framework—which gives direct access to spatial relationships—addresses precisely this issue.
The empirical advantages are most evident for the large molecules in GEOM-XL, on which GeoMol fails to improve consistently over RDKit (Appendix H).
On the other hand, because GeoMol requires only a single-forward pass, it retains the advantage of faster runtime compared to diffusion-based methods.
4.4 Ensemble properties
While RMSD gives a geometric way to evaluate ensemble quality, we also consider the chemical similarity between generated and ground truth ensembles.
For a random 100-molecule subset of
DRUGS, we generate min(2 K, 32) conformers per molecule, relax the conformers with GFN2-xTB
[Bannwarth et al., 2019], 9 and compare the Boltzmann-weighted properties of the generated and ground truth ensembles.
Speciﬁcally, the following properties are computed with xTB [Bannwarth et al., 2019]: energy E , dipole moment µ , HOMO-LUMO gap ∆ (cid:15) , and the minimum energy E min .
The median errors for torsional diffusion and the baselines are shown in Table 4.
Our method produces the most chemically accurate ensembles, especially in terms of energy.
In particular, we signiﬁcantly improve over GeoMol and GeoDiff in ﬁnding the lowest-energy conformers that are only (on median) 0.13 kcal/mol higher in energy than the global minimum.
4.5 Torsional Boltzmann generator
[Table]
Finally, we evaluate how well a torsional Boltzmann generator trained with MMFF [Halgren, 1996] energies can sample the corresponding Boltzmann density over torsion angles.
We train and test on GEOMDRUGS molecules with 3–7 rotatable bonds and use the local structures of the ﬁrst ground-truth conformers.
For the baselines, we implement annealed importance samplers (AIS)
[Neal, 2001] with MetropolisHastings steps over the torsional space and tune the variance of the transition kernels.
Table 4 shows the quality of the samplers in terms of the effective sample size (ESS) given by the weights of 32 samples for each test molecule, which measures the α -divergence (with α = 2 ) between the model and Boltzmann distributions [Midgley et al., 2021].
Our method signiﬁcantly outperforms the AIS baseline, and improves with increased step size despite being trained with only a 5-step resampler.
Note that, since these evaluations are done on unseen molecules, they are beyond the capabilities of existing Boltzmann generators.
Conclusion
We presented torsional diffusion , a method for generating molecular conformers based on a diffusion process restricted to the most ﬂexible degrees of freedom.
Torsional diffusion is the ﬁrst machine learning model to signiﬁcantly outperform standard cheminformatics methods and is orders of magnitude faster than previous Euclidean diffusion models.
Using the exact likelihoods provided by our model, we also train the ﬁrst system-agnostic Boltzmann generator.
There are several exciting avenues for future work.
A natural extension is to relax the rigid local structure assumption by developing an efﬁcient diffusion-based model over the full space of intrinsic coordinates while still incorporating chemical constraints.
Moreover, torsional diffusion—or similar ideas—could be applicable to larger molecular systems, for which fast, parsimonious models of structural ﬂexibility could beneﬁt applications such as drug discovery and protein design.
Acknowledgments
We pay tribute to Octavian-Eugen Ganea (1987-2022), dear colleague, mentor, and friend without whom this work would have never been possible.
We thank Hannes Stärk, Wenxian Shi, Xiang Fu, Felix Faltings, Jason Yim, Adam Fisch, Alex Wu,
Jeremy Wohlwend, Peter Mikhael, and Saro Passaro for helpful feedback and discussions.
We thank
Lagnajit Pattanaik, Minkai Xu, and Simon Axelrod for their advice and support when working with, respectively, GeoMol, GeoDiff and the GEOM dataset.
This work was supported by the Machine
Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium, the Abdul Latif Jameel
Clinic for Machine Learning in Health, the DTRA Discovery of Medical Countermeasures Against
New and Emerging (DOMANE) threats program, the DARPA Accelerated Molecular Discovery program and the Sanoﬁ Computational Antibody Design grant.
We acknowledge support from the Department of Energy Computational Science Graduate Fellowship (BJ), the Robert Shillman
Fellowship (GC), and the NSF Graduate Research Fellowship (JC).
A Deﬁnitions
Consider a molecular graph G =
( V , E ) and its space of possible conformers C G .
A conformer is an assignment V (cid:55)→ R 3 of each atom to a point in 3D-space, deﬁned up to global rototranslation.
For notational convenience, we suppose there is an ordering of nodes such that we can regard a mapping as a vector in R 3 n where n = |V| .
Then a conformer C ∈ C G is a set of SE (3) -equivalent vectors in R 3 n —that is, C G ∼
=
R 3 n /SE (3) .
This deﬁnes the space of conformers in terms of extrinsic (or
Cartesian) coordinates.
An intrinsic (or internal) coordinate is a function over C G —i.e., it is an SE (3) -invariant function over R 3 n .
There are four types of such coordinates typically considered:
Bond lengths .
For ( a, b ) ∈ E , the bond length l ab ∈ [0 , ∞ ) is deﬁned as | x a − x b | .
Bond angles .
For a, b, c ∈ V such that a, c ∈ N ( b ) , the bond angle α abc ∈ [0 , π ] is deﬁned by
[Equation]
Chirality .
For a ∈ V with 4 neighbors b, c, d, e ∈ N
[Equation]
a ) , the chirality z abcd ∈ {− 1 , 1 } is deﬁned as z abcde := sign det
[Equation]
Similar quantities are deﬁned for atoms with other numbers of neighbors.
Chirality is often considered part of the speciﬁcation of the molecule, rather than the conformer.
See Appendix F.3 for additional discussion on this point.
Torsion angles .
For ( b, c ) ∈ E , with a choice of reference neighbors a ∈ N ( b ) \{ c } , d ∈ N ( c ) \{ b } , the torsion angle τ abcd ∈ [0 , 2 π ) is deﬁned as the dihedral angle between planes abc and bcd :
[Equation]
where u ab = x
b − x a and n abc is the normal vector u ab × u bc .
Note that τ abcd =
τ dcba —i.e., the dihedral angle is the same for four consecutively bonded atoms regardless of the direction in which they are considered.
A complete set of intrinsic coordinates of the molecule is a set of such functions ( f 1 , f 2 , . . . )
such that F ( C ) =
( f 1 ( C ) , f 2 ( C ) , . . . ) is a bijection.
In other words, they fully specify a unique element of C G without overparameterizing the space.
In general there exist many possible such sets for a given molecular graph.
We will not discuss further how to ﬁnd such sets, as our work focuses on manipulating molecules in a way that holds ﬁxed all l, α, z and only modiﬁes (a subset of) torsion angles τ .
As presently stated, the torsion angle about a bond ( b, c ) ∈
E is ill-deﬁned, as it could be any τ abcd with a ∈ N ( b ) \ { c } , d ∈ N ( c ) \ { b } .
However, any complete set of intrinsic coordinates needs to only have at most one such τ abcd for each bond ( b, c )
[Ganea et al., 2021].
Thus, we often refer to the torsion angle about a bond ( b i , c i ) as τ i when reference neighbors a i , b i are not explicitly stated.
B Propositions B.1 Torsion update
Given a freely rotatable bond ( b i , c i ) , by deﬁnition removing ( b i , c i ) creates two connected components V ( b i ) , V ( c i ) .
Then, consider torsion angle τ j at a different bond ( b j , c j ) with neighbor choices a j ∈ N ( b j ) , d j ∈ N ( c j ) , a j (cid:54) =
c j , d j (cid:54) =
b j .
Without loss of generality, there are two cases
• Case 1: a j , b j , c j , d j ∈ V ( b i ) • Case 2: d j ∈ V ( c i ) and a j , b j , c j ∈ V ( b i )
Note that in Case 2, c j = b i and d j = c i must hold because there is only one edge between
V ( b i ) , V ( c i ) .
With these preliminaries we now restate the proposition:
Proposition 1.
Let ( b i , c i ) be a rotatable bond, let x V ( b i ) be the positions of atoms on the b i side of the molecule, and let R ( θ , x c i ) ∈ SE (3) be the rotation by Euler vector θ about x c i .
Then for
C, C (cid:48) ∈ C G , if τ i is any deﬁnition of the torsion angle around bond ( b i , c i ) ,
[Equation]
where ˆr b i c i
=
( x c i − x b i ) /
||
x c i − x b i || .
Proof.
First we show τ i ( C (cid:48) )
=
τ i ( C ) + θ , for which it sufﬁces to show τ i ( x (cid:48) )
=
τ i ( x ) + θ .
Because a i , b i ∈ V ( b i ) , x (cid:48)
a i = x a i and x (cid:48)
b i = x b i .
Since the rotation of x V ( c i ) is centered at x c i , we have x (cid:48)
c i = x
c i as well.
Now we consider d i and u (cid:48) cd =
x
(cid:48)
d i − x c i .
By the Rodrigues rotation formula,
[Equation] [Equation]
Then we have
[Equation] [Equation]
To obtain | n (cid:48) bcd | , note that since n bcd ⊥ u bc ,
[Equation]
17) which gives | n (cid:48) bcd | = | n bcd | .
Thus,
[Equation]
Similarly, sin
[Equation]
=
[Equation]
Therefore, τ (cid:48)
i =
τ i + θ
Now we show τ (cid:48)
j = τ j for all j (cid:54) =
i .
Consider any such j .
For Case 1, x (cid:48) a j = x a j , x (cid:48)
b j = x b j
[Equation]
x
[Equation]
c j = x c j , x (cid:48)
d j = x d j so clearly τ (cid:48)
j = τ j .
For Case 2, x (cid:48) a j = x a j , x (cid:48)
b j = x b j , x (cid:48) c
j = x c j immediately.
But because d j = c i , we also have x (cid:48)
d j = x d j .
Thus, τ (cid:48)
j = τ j .
B.2 Parity equivariance
Proposition 2.
If p ( τ ( C ) | L ( C ))
=
p ( τ ( − C )
| L ( − C )) , then for all diffusion times t ,
[Equation]
Proof.
From Equation 13 we see that for any torsion
τ i , we have τ i ( − C )
= −
τ i ( C ) ; therefore τ i ( − C )
= −
τ i ( C ) , which we denote τ − .
Also denote τ := τ ( C ) , p t ( τ ) :
= p t ( τ | L ( C )) and p (cid:48) t ( τ − ) :
= p t ( τ − | L ( − C )) .
We claim p t ( τ ) = p (cid:48) t ( τ − ) for all t .
Since the perturbation kernel (equation 3) is parity invariant,
[Equation]
Next, we have
[Equation]
which concludes the proof.
B.3 Likelihood conversion
Proposition 3.
Let x ∈ C
( τ , L ) be a centered conformer in Euclidean space.
Then,
[Equation]
where the indices α, β are integers between 1 and m + 3 .
For 1 ≤ α ≤ m , J ( k ) α is deﬁned as
[Equation]
and for α ∈ { m + 1 , m + 2 , m + 3 } as
[Equation]
where ( b i , c i ) is the freely rotatable bond for torsion angle
i , V ( b i ) is the set of all nodes on the same side of the bond as b i , and ˆ x, ˆ y, ˆ z are the unit vectors in the respective directions.
Proof.
Let M be ( m + 3) -dimensional manifold embedded in 3 n -dimensional Euclidean space formed by the set of all centered conformers with ﬁxed local structures but arbitrary torsion angles and orientation.
A natural set of coordinates for M is q α = { τ 1 , τ 2 , . .
, τ m , ω x , ω y , ω z } , where τ i is the torsion angle at bond i and ω x , ω y ,
ω z deﬁne the global rotation about the center of mass:
[Equation]
˜ x (cid:96) where ˜ x (cid:96) =
e Λ( ω ) x (cid:48) k ,
[Equation] [Equation]
Here x (cid:48) k is the position of atom k as determined by the torsion angles, without centering or global rotations, and ω x , ω y , ω z are rotation about the x , y , and z axis respectively.
Consider the set of covariant basis vectors
[Equation]
[Equation]
and corresponding the covariant components of the metric tensor,
[Equation]
[Equation]
The conversion factor between torsional likelihood and Euclidean likelihood is given by
[Equation]
where √ det g d m +3 q is the invariant volume element on M [Carroll, 2019], and the integration over ω marginalizes over the uniform distribution over global rotations.
The calculation of Eq.
29 proceeds as follows.
Let the position of the k ’th atom be x k , and let the three corresponding components of J α be J ( k ) α .
For 1 ≤ i ≤ m , J ( k ) i is given by
[Equation]
where ˜ J ( k ) i :
= ∂
˜ x
k /∂τ
i is the displacement of atom k upon an inﬁnitesmal change in the torsion angle τ i , without considering the change in the center of mass.
Clearly ˜ J ( b i ) i =
˜ J
( c i ) i = 0 because neither b i nor c i itself is displaced; furthermore, all atoms on the b side of torsioning bond are not displaced, so J ( k ) i
=
for all k ∈ N ( b i ) .
The remaining atoms, in N ( c i ) , are rotated about the axis of the ( b i , c i ) bond.
The displacement per inﬁnitesimal ∂τ i is given by the cross product of the unit normal along the rotation axis, (˜ x c i − ˜ x b i ) / ||
˜ x
c i − ˜ x
b i || , with the displacement from rotation axis, ˜ x k − ˜ x
b i .
This cross product yields J ( k ) α in Eq.
24, where the tildes are dropped as relative positions do not depend on center of mass.
For α ∈ { m + 1 , m + 2 , m + 3 } , a similar consideration of the cross product with the rotation axis yields Eq.
Finally, since none of the components of the metric tensor depend explicitly on ω , the integration over ω in Eq.
29 is trivial and yields the volume over
SO (3) of 8 π 2 [Chirikjian, 2011], proving the proposition.
C Training and inference procedures
Algorithms 2 and 3 summarize, respectively, the training and inference procedures used for torsional diffusion.
In practice, during training, we limit K G to 30 i.e. we only consider the ﬁrst 30 conformers found by CREST (typically those with the largest Boltzmann weight).
Moreover, molecules are batched and an Adam optimizer with a learning rate scheduler is used for optimization.
For inference, to fairly compare with other methods from the literature, we follow Ganea et al.
and set K to be twice the number of conformers returned by CREST.
Input: molecules [ G 0 , ..., G N ] each with true conformers [ C G, 1 , ...
C G,K G ] , learning rate α
Output: trained score model s θ conformer matching process for each G to get [ ˆ C G, 1 , ...
ˆ C G,K G ] ; for epoch ← 1 to epoch max do for G in [ G 0 , ..., G N ] do sample t ∈ [0 , 1] and ˆ C ∈
[ ˆ C G, 1 , ...
ˆ C G,K G ] ; sample ∆ τ from wrapped normal p t | 0 ( · | 0 ) with σ = σ 1 − t min σ t max ; apply ∆ τ to ˆ C ; predict δ τ =
s θ,G ( ˆ C, t ) ; update θ ← θ − α ∇ θ (cid:107)
δ τ − ∇ ∆ τ p t | 0 (∆ τ | 0 ) (cid:107) 2 ;
Algorithm
Input: molecular graph G , number conformers K , number steps N
Output: predicted conformers [ C 1 , ...
C K ] generate local structures by obtaining conformers [ C 1 , ...
C K ] from RDKit; for C in [ C 1 , ...
C K ] do sample ∆ τ ∼ U [0 , 2 π ] m and apply to C to randomize torsion angles; for n ← N to 1 do let t = n/N, g ( t ) =
σ 1 − t min σ t max (cid:112) 2 ln( σ max /σ min ) ; predict δ τ = s
θ,G ( ˆ C, t ) ; draw z from wrapped normal with σ 2 = 1 /N ; set ∆ τ = ( g 2 ( t ) /N )
δ τ + g
( t ) z ; apply ∆ τ to C ;
D Score network architecture
Overview To perform the torsion score prediction under these symmetry constraints we design an architecture formed by three components: an embedding layer, a series of K interaction layers and a pseudotorque layer.
The pseudotorque layer produces pseudoscalar torsion scores δτ := ∂ log p/∂τ for every rotatable bond.
Following the notation from Thomas et al.
[2018] and Batzner et al.
[2022], we represent the node representations as V ( k,l,p ) acm a dictionary with keys the layer k , rotation order l and parity p that contains tensors with shapes [ |V| , n l , 2 l +1] corresponding to the indices of the node, channel and representation respectively.
We use the e3nn library [Geiger et al., 2022] to implement our architecture.
Embedding layer In the embedding layer, we build a radius graph ( V , E r max ) around each atom on top of the original molecular graph and generate initial scalar embeddings for nodes V (0 , 0 , 1) a and edges e ab combining chemical properties, sinusoidal embeddings of time φ ( t )
[Vaswani et al., 2017]
[Figure]
[Equation]
Δτ for each bond
[Figure]
[Equation]
for each atom with 𝑝 ∈ [−1, 1]
[Figure]
[Equation]
and, for the edges, a radial basis function representation of their length µ ( r ab )
[Schütt et al., 2017]:
[Equation]
where Υ ( e ) and Υ ( v ) are learnable two-layers MLPs, r ab is the Euclidean distance between atoms a and b , r max = 5 Å is the distance cutoff, f a are the chemical features of atom a , f ab are the chemical features of bond ( a, b ) if it was part of E and 0 otherwise.
The node and edge chemical features f a and f ab are constructed as in Ganea et al.
Brieﬂy, the node features include atom identity, atomic number, aromaticity, degree, hybridization, implicit valence, formal charge, ring membership, and ring size, constituting a 74-dimensional vector for
GEOM-DRUGS and 44-dimensional for QM9 (due to fewer atom types).
The edge features are adimensional one-hot encoding of the bond type.
Interaction layers The interaction layers are based on E(3)NN [Geiger et al., 2022] convolutional layers.
At each layer, for every pair of nodes in the graph, we construct messages using tensor products of the current irreducible representation of each node with the spherical harmonic representations of the normalized edge vector.
These messages are themselves irreducible representations, which are weighted channel-wise by a scalar function of the current scalar representations of the two nodes and the edge and aggregated with Clebsch-Gordan coefﬁcients.
At every layer k , for every node a , rotation order l o , and output channel c (cid:48) :
[Equation]
k −
[Equation]
l
[Equation]
p
[Equation]
with ψ ( k,l o ,l f ,l i ,p i ) abc
[Equation]
l
[Equation]
l
[Equation]
l i
[Equation]
p
[Equation]
) where the outer sum is over values of l f , l i , p i such that | l i − l f | ≤ l o ≤ l i + l f and ( − 1) l f p i
= p
o ,
C indicates the Clebsch-Gordan coefﬁcients [Thomas et al., 2018], N a = { b | ( a, b )
∈ E max } the neighborhood of a and Y the spherical harmonics.
The rotational order of the nodes representations l o and l i and of the spherical harmonics of the edges ( l f ) are restricted to be at most 2.
All the learnable weights are contained in Ψ , a dictionary of MLPs that compute per-channel weights based on the edge embeddings and scalar features of the outgoing and incoming node.
Pseudotorque layer
The ﬁnal part of our architecture is a pseudotorque layer that predicts a pseudoscalar score δτ for each rotatable bond from the per-node outputs of the interaction layers.
For every rotatable bond, we construct a tensor-valued ﬁlter, centered on the bond, from the tensor product of the spherical harmonics with a l = 2 representation of the bond axis .
Since the parity of the l = 2 spherical harmonic is even, this representation does not require a choice of bond direction.
The ﬁlter is then used to convolve with the representations of every neighbor on a radius graph, and the products which produce pseudoscalars are passed through odd-function (i.e., with tanh nonlinearity and no bias) dense layers (not shown in equation 33) to produce a single prediction.
For all rotatable bonds g = ( g 0 , g 1 ) ∈
E rot 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 .
[Equation]
with γ ( l,p i ) gcb =
[Equation]
where Υ ( τ ) and Γ are MLPs with learnable parameters and N g = { b | ( g, b ) ∈
E τ } .
E Conformer matching
The conformer matching procedure, summarised in Algorithm 4, proceeds as follows.
For a molecule with K conformers, we ﬁrst generate K random local structure estimates ˆ L from RDKit.
To match with the ground truth local structures, we compute the cost of matching each true conformer C with each estimate ˆ L (i.e. a K × K cost matrix), where the cost is the best RMSD that can be achieved by modifying the torsions of the RDKit conformer with local structure ˆ L to match the ground truth conformer C .
Note that in practice, we compute an upper bound to this optimal RMSD using the fast von Mises torsion matching procedure proposed by Stärk et al.
We then ﬁnd an optimal matching of true conformers C to local structure estimates ˆ L by solving the linear sum assignment problem over the approximate cost matrix [Crouse, 2016].
Finally, for each matched pair, we ﬁnd the true optimal ˆ C by running a differential evolution optimization procedure over the torsion angles [Méndez-Lucio et al., 2021].
The complete assignment resulting from the linear sum solution guarantees that there is no distributional shift in the local structures seen during training and inference.
Algorithm 4: Conformer matching
Input: true conformers of G [ C 1 , ...
C K ]
Output: approximate conformers for training [ ˆ C 1 , ...
ˆ C K ] generate local structures [ˆ L 1 , ...
ˆ L K ] with RDKit; for ( i, j ) in [1 , K ] × [1 , K ] do
C temp =
von_Mises_matching( C i , ˆ L j ); cost[i,j] = RMSD( C i , C temp ) ; assignment = linear_sum_assignment(cost); for i ← 1 to K do j = assignment[i]; ˆ C i =
differential_evolution( C i , ˆ L j , RMSD);
Table 5 shows the average RMSD between a ground truth conformer C i and its matched conformer ˆ
C i .
The average RMSD of 0.324 Å obtained via conformer matching provides an approximate lower bound on the achievable AMR performance for methods that do not change the local structure and take those from RDKit (further discussion in Appendix F.1).
[Table]
F Additional discussion F.1 RDKit local structures
In this section
[Figure]
we provide empirical justiﬁcation for the claim that cheminformatics methods like
RDKit already provide accurate local structures.
It is well known in chemistry that bond lengths and angles take on a very narrow range of values due to strong energetic constraints.
However, it is not trivial to empirically evaluate the claim due to the difﬁculty in deﬁning a distance measure between a pair of local structures.
In this section, we will employ two sets of observations: marginal error distributions and matched conformer RMSD.
[Figure]
Marginal error distributions We examine the distribution of errors of the bond lengths and angles in a random RDKit conformer relative to the corresponding lengths and angles in a random CREST conformer (Figure 5).
The distributions are narrow and uni-modal distributions around zero, with a
RMSE of 0.03 Å for bond lengths and 4.1° for bond angles on GEOM-DRUGS.
Comparing DRUGS and QM9, the error distribution does not depend on the size of the molecule.
Although it is difﬁcult to determine how these variations will compound or compensate for each other in the global conformer structure, the analysis demonstrates that bond lengths and angles have little ﬂexibility (i.e., no strong variability among conformers) and are accurately predicted by RDKit.
Matched conformer RMSD We can more rigorously analyze the quality of a local structure ˆ L with respect to a given reference conformer C by computing the minimum RMSD that can be obtained by combining ˆ L with optimal torsion angles.
That is, we consider the RMSD distance of C to the closest point on the manifold of possible conformers with local structure ˆ L : RMSD min ( C, ˆ L ) :
=
min τ RMSD( C, ˆ C ) where ˆ C
=
(ˆ L, τ ) .
Conveniently, ˆ C is precisely the output of the differential evolution in Appendix E. Thus, the average
RSMD reported in the last row of Table 5 is the expected RMSD min of an optimal assignment of
RDKit local structures to ground-truth conformers.
This distance—0.324 Å on GEOM-DRUGS—is signiﬁcantly smaller than the error of the current state-of-the-art conformer generation methods.
Further, it is only slightly larger than the average RMSD min of 0.284 Å resulting from matching a ground truth conformer to the local structure of another randomly chosen ground truth conformer, which provides a measure of the variability among ground truth local structures.
These observations support the claim that the accuracy of existing approaches on drug-like molecules can be signiﬁcantly improved via better conditional sampling of torsion angles.
F.2 Torsion updates
In the main text, we viewed updates ∆ τ as changes to a torsion angle τ , and asserted that the same update applied to any torsion angle at a given bond (i.e., with any choice of reference neighbors) results in the same conformer.
Given this, a potentially more intuitive presentation is to deﬁne ∆ τ for a bond as a relative rotation around that bond, without reference to any torsion angle.
Consider a rotatable bond ( b, c ) and the connected components V ( b ) , V ( c ) formed by removing the bond.
Let ˆr bc =
( x c − x b ) / | x c − x b | and similarly ˆr cb = − ˆr bc .
Because the bond is freely rotatable, consider rotations of each side of the molecule around the bond axis given by ˆr bc .
Speciﬁcally, let
V ( b ) be rotated by some Euler vector θ b := θ b ˆr bc around x b , and V ( c ) by θ c :
= θ c ˆr bc around x c .
Then the rotations induce a torsion update ∆ τ if θ c − θ b = ∆ τ ; or equivalently
[Equation]
bc (34)
The expression remains unchanged if we swap the indices b, c ; thus there is no sign ambiguity.
Some less formal but possibly more intuitive restatements of the sign convention are:
• Looking down a bond, a positive update is given by a CCW rotation of the nearer side; or a CW rotation of the further side • For a viewer positioned in the middle of the bond, a positive update is given by the CW rotation of any one side • A positive update is given by Euler vectors that point outwards from the bond These are illustrated in Figure
Since the Euler vector θ is a pseudovector that remains unchanged under parity inversion, while ˆr is a normal vector, it is apparent that ∆ τ —and any model predicting ∆ τ —must be a pseudoscalar.
Because the update is determined by a relative rotation, it is not necessary to specify which side to rotate.
That is, the same torsion update can be accomplished by rotating only one side, both sides in opposite directions, or both sides in the same direction.
In practical implementation, we rotate the side of the molecule with fewer atoms, and keep the other side ﬁxed.
F.3 Chemical isomerism
We have deﬁned a molecule in terms of its bond connectivity, i.e., as a graph G =
( V , E ) with atoms
V and E .
In chemistry, however, it is universal to consider molecules with the same connectivity, but whose conformers cannot interconvert, as different molecules called stereoisomers .
In our formalism, stereoisomers correspond to subsets of the space of conformers C G for some molecular graph G .
Many types of stereoisomerism exist, but the two most important are:
• Chirality .
Conformers with distinct values of chirality tags { z i } —one for each chiral atom meeting certain criteria—are considered different molecules.
• E/Z isomerism , also called cis/trans isomerism.
For each double bond meeting certain criteria, the space [0 , 2 π ) is partitioned into two halves, such that conformers are considered different molecules depending on the value of the torsion angle.
These are not meant to be formal deﬁnitions, and we refer to standard chemistry texts for a more detailed treatment.
For our purposes, the key implication is that conformer generation requires generating conformers consistent with a given stereoisomer .
For a molecular graph G with k relevant chiral centers and l relevant double bonds, there are 2 k + l possible stereoisomers, corresponding to the partition of C G into 2 k + l disjoint subsets—only one of which corresponds to the molecule under consideration.
Torsional diffusion automatically handles chirality.
Because we have considered chirality to be part of the local structure, it is drawn from the cheminformatics package RDKit, which is given the full identiﬁcation of the stereoisomer along with the molecular graph, and is not modiﬁed by the torsional diffusion.
Hence, our method always generates conformers with the correct chirality at each chiral center.
On the other hand, GeoDiff does not consider chirality at all, while GeoMol generates molecules without any chirality constraints, and merely inverts the chiral centers that were generated incorrectly.
E/Z isomerism is signiﬁcantly trickier, as it places a constraint on the torsion angles at double bonds, which are considered freely rotatable in our framework.
Presently, torsional diffusion does not attempt to capture E/Z isomerism.
One possible way of doing so is to augment the molecular graph with edges of a special type, and we leave such augmentation to future work.
GeoDiff and GeoMol also do not attempt to treat E/Z isomerism.
More generally, while the abstract view of molecules as graphs has enabled rapid advances in molecular machine learning, stereoisomerism shows that it is clearly a simpliﬁcation.
As stereoisomers can have signiﬁcantly different chemical properties and bioactivities, a more complete view of molecular space will be essential for further advances in molecular machine learning.
F.4 Limitations of torsional diffusion
As demonstrated in Section 4, torsional diffusion signiﬁcantly improves the accuracy and reduces the denoising runtime for conformer generation.
However, torsional diffusion also has a number of limitations that we will discuss in this section.
Conformer generation
The ﬁrst clear limitation is that the error that torsional diffusion can achieve is lower bounded by the quality of the local structure from the selected cheminformatics method.
As discussed in Appendix F.1, this corresponds to the mean RMSD obtained after conformer matching, which is 0.324 Å with RDKit local structures on DRUGS.
Moreover, due to the the local structure distributional shift discussed in Section 4.1, conformer matching (or another method bridging the shift) is required to generate the training set.
However, the resulting conformers are not the minima of the (unconditional or even conditional) potential energy function.
Thus, the learning task becomes less physically interpretable and potentially more difﬁcult; empirically we observe this clearly in the training and validation score-matching losses.
We leave to future work the exploration of relaxations of the rigid local structures assumption in a way that would still leverage the predominance of torsional ﬂexibility in molecular structures, while at the same time allowing some ﬂexibility in the independent components.
Rings
The largest source of ﬂexibility in molecular conformations that is not directly accounted for by torsional diffusion is the variability in ring conformations .
Since the torsion angles at bonds inside cycles cannot be independently varied, our framework treats them as part of the local structure.
Therefore, torsional diffusion relies on the local structure sampler p G ( L ) to accurately model cycle conformations.
Although this is true for a large number of relatively small rings (especially aromatic ones) present in many drug-like molecules, it is less true for puckered rings, fused rings, and larger cycles.
In particular, torsional diffusion does not address the longstanding difﬁculty that existing cheminformatics methods have with macrocycles—rings with 12 or more atoms that have found several applications in drug discovery [Driggers et al., 2008].
We hope, however, that the idea of restricting diffusion processes to the main sources of ﬂexibility will motivate future work to deﬁne diffusion processes over cycles conformations combined with free torsion angles.
Boltzmann generation
With Boltzmann generators we are typically interested in sampling the
Boltzmann distribution over the entire (Euclidean) conformational space p G ( C ) .
However, the procedure detailed in Section 3.6 generates (importance-weighted) samples from the Boltzmann distribution conditioned on a given local structure p G ( C | L ) .
To importance sample from the full Boltzmann distribution p G ( C ) , one would need a model p G ( L ) over local structures that also provides exact likelihoods.
This is not the case with RDKit or, to the best of our knowledge, other existing models, and therefore an interesting avenue for future work.
Proteins As protein conformations are often described with backbone dihedral (i.e., torsion angles), it is natural to consider whether torsional diffusion may be useful for modeling protein ﬂexibility.
However, we do not believe that the direct application of the framework to proteins or other macromolecules is very promising.
Small changes in torsional coordinates cause large displacements in distant regions of the molecule, so the inﬂuence on a torsional score is not limited to the local neighborhood of the bond.
For small molecules—even the ones in GEOM-XL—this is not a problem because of their limited spatial and graph theoretic diameters.
In proteins, however, the graph diameter is 3 times the sequence length and can easily reach over 1000; and interactions between distant residues are extremely important in determining the structure and constraining ﬂexibility.
Although torsional diffusion may not be the right framework for modeling proteins, we believe that similar ideas (i.e., well-chosen diffusions over the ﬂexible degrees of freedom) could be useful for generative models of protein structure and is a promising avenue of work.
G Experimental details G.1 Dataset details
Splits We follow the data processing and splits from Ganea et al.
The splits are random with train/validation/test of 243473/30433/1000 for GEOM-DRUGS and 106586/13323/1000 for
GEOM-QM9.
GEOM-XL consists of only a test split (since we do not train on it), which consists of all 102 molecules in the MoleculeNet dataset with at least 100 atoms.
For all splits, the molecules whose CREST conformers all have a canonical SMILES different from the SMILES of the molecule (meaning a reacted conformer), or that cannot be handled by RDKit, are ﬁltered out.
Dataset statistics As can be seen in Figure 7, the datasets differ signiﬁcantly in molecule size as measured by number of atoms or rotatable bonds.
Particularly signiﬁcant is the domain shift between DRUGS and XL, which we leverage in our experiments by testing how well models trained on DRUGS generalize to XL.
Boltzmann
generator The torsional Boltzmann generator described in Section 4.5 is trained and tested on molecules from GEOM-DRUGS with 3–7 rotatable bonds.
The training (validation) set consists of 10000 (400) such randomly selected molecules from the DRUGS training (validation) set.
The test set consists of all the 453 molecules present in the DRUGS test set with 3–7 rotatable bonds.
[Figure]
l a t i v e f r e q u e n c y
[Figure]
R e l a t i v e f r e q u
e n c y
Rotatable bonds count
[Figure]
QM
XL
G.2 Training and tuning details
Conformer generation For conformer ensemble generation on GEOM-DRUGS, the torsional diffusion models were trained on NVIDIA RTX A6000 GPUs for 250 epochs with the Adam optimizer (taking from 4 to 11 days on a single GPU).
The hyperparameters tuned on the validation set were (in bold the value that was chosen): initial learning rate (0.0003, 0.001 , 0.003), learning rate scheduler patience (5, 20 ), number of layers (2, 4 , 6), maximum representation order (1st, 2nd ), r max ( 5Å , 7Å, 10Å) and batch norm ( True , False).
All the other default hyperparameters used can be found in the attached code.
For GEOM-XL the same trained model was used; for GEOM-QM9 a new model with the same hyperparameters was trained.
Torsional Boltzmann generators We start from a torsional diffusion model pretrained on GEOMDRUGS, and train for 250 epochs (6-9 days on a single GPU).
A separate model is trained for every temperature.
The resampling procedure with 5 steps is run for every molecule every max(5 , ESS ) epochs, where ESS is computed for the current set of 32 samples.
The only hyperparameter tuned (at temperature 300K) is σ min , the noise level at which to stop the reverse diffusion process.
We further improve the training procedure of torsional Boltzmann generators by implementing annealed training .
The Boltzmann generator for some temperature T is trained at epoch k by using the Boltzmann distribution at temperature T (cid:48)
=
T + (3000 − T ) /k
as the target distribution for that epoch.
Intuitively, this trains the model at the start with a smoother distribution that is easier to learn, which gradually transforms into the desired distribution.
G.3 Evaluation details
Ensemble RMSD
As evaluation metrics for conformer generation, Ganea et al.
[2021] and following works have used the so-called Average Minimum RMSD (AMR) and Coverage (COV) for
Precision (P) and Recall (R) measured when generating twice as many conformers as provided by
CREST.
For K = 2 L let { C ∗ l } l ∈ [1 ,L ] and { C k } k ∈ [1 ,K ] be respectively the sets of ground truth and generated conformers:
[Equation]
where δ is the coverage threshold.
The precision metrics are obtained by swapping ground truth and generated conformers.
In the XL dataset, due to the size of the molecules, we compute the RMSDs without testing all possible symmetries of the molecules, therefore the obtained RMSDs are an upper bound, which we ﬁnd to be very close in practice to the permutation-aware RSMDs.
Runtime evaluation We benchmark the methods on CPU (Intel i9-9920X) to enable comparison with RDKit.
The number of threads for RDKit, numpy , and torch is set to 8.
We select 10 molecules at random from the GEOM-DRUGS test set and generate 8 conformers per molecule using each method.
Script loading and model loading times are not included in the reported values.
Boltzmann generator To evaluate how well the torsional Boltzmann generator and the AIS baselines sample from the conditional Boltzmann distribution, we report their median effective sample size (ESS)
[Kish, 1965] given the importance sampling weights w i of 32 samples for each molecule:
[Equation]
This approximates the number of independent samples that would be needed from the target Boltzmann distribution to obtain an estimate with the same variance as the one obtained with the importanceweighted samples.
For the baseline annealed importance samplers, the transition kernel is a single Metropolis-Hastings step with the wrapped normal distributions on T m as the proposal.
We run with a range of kernel variances: 0 . 25 , 0 . 5 , 0 . 3 , 0 . 5 , 0 .
75 , 1 , 1 . 5 . 2 ; and report the best result.
We use an exponential annealing schedule; i.e., p n ∝ p 1 − n/Np n/NN where p 0 is the uniform distribution and p N is the target
Boltzmann density.
H Additional results
Performance vs size Figure 8 shows the performance of different models as a function of the number of rotatable bonds.
Molecules with more rotatable bonds are more ﬂexible and are generally larger; it is therefore expected that the RMSD error will increase with the number of bonds.
With very few rotatable bonds, the error of torsional diffusion depends mostly on the quality of the local structures it was given, and therefore it has a similar error as RDKit.
However, as the number of torsion angles increases, torsional diffusion deteriorates more slowly than other methods.
The trend continues with the very large molecules in GEOM-XL (average 136 atoms and 32 rotatable bonds).
These not only are larger and more ﬂexible, but—for machine learning models trained on
GEOM-DRUGS—are also out of distribution.
As shown in Table 6, on GEOM-XL GeoMol only performs marginally better than RDKit, while torsional diffusion reduces RDKit AMR by 30% on recall and 12% on precision.
These results can very likely be improved by training and tuning the torsional diffusion model on larger molecules.
[Figure]
[Table] [Table]
Small molecules We also train and evaluate our model on the small molecules from GEOM-QM and report the performance in Table 7.
For these smaller molecules, cheminformatics methods already do very well and, given the very little ﬂexibility and few rotatable bonds present, the accuracy of local structure signiﬁcantly impacts the performance of torsional diffusion.
RDKit achieves a mean recall
AMR just over 0.23Å, while torsional diffusion based on RDKit local structures results in a mean recall AMR of 0.178Å.
This is already very close lower bound of 0.17Å that can be achieved with
RDKit local structures (as approximately calculated by conformer matching).
Torsional diffusion does signiﬁcantly better than other ML methods, but is only on par with or slightly worse than
OMEGA, which, evidently, has a better local structures for these small molecules.
Ablation experiments In Table 8 we present a set of ablation studies to evaluate the importance of different components of the proposed torsional diffusion method:
Baseline refers to the model described and tested throughout the paper.
Probability ﬂow ODE refers to using the ODE formulation of the reverse diffusion process (not an ablation, strictly speaking).
As expected, it obtains similar results to the baseline SDE formulation.
Only D.E. matching refers to a model trained on conformers obtained by a random assignment of RDKit local structures to ground truth conformers (without ﬁrst doing an optimal assignment as in Appendix E); this performs only marginally worse than full conformer matching.
First order irreps refers to the same model but with node irreducible representations kept only until order (cid:96)
= 1 instead of (cid:96)
= 2 ; this worsens the average error by about 5%, but results in a 41% runtime speed-up. 5.
Train on ground truth L refers to a model trained directly on the ground truth conformers without conformer matching but tested (as always) on RDKit local structures; although the training and validation score matching loss of this model is signiﬁcantly lower, its inference performance reﬂects the detrimental effect of the local structure distributional shift.
No parity equivariance refers to a model whose outputs are parity invariant instead of parity equivariant; the model cannot distinguish a molecule from its mirror image and fails to learn, resulting in performance on par with a random baseline.
Random τ refers to a random baseline using RDKit local structures and uniformly random torsion angles.
[Table] [Table]
[Table]
Reverse diffusion steps In Table 9 we vary the number of steps used in the reverse diffusion process and evaluate the ensemble RMSD results on GEOM-DRUGS.
We ﬁnd that torsional diffusion is remarkably parsimonious in terms of number of steps required: the majority of gain in performance over prior diffusion-based methods is attained with only 10 steps.
We conﬁrm that increasing the number of steps from the default of 20 to 50 only results in minor performance gains.
Ensemble properties In Table 10, we report the median absolute errors of the Boltzmann-weighted properties of the generated vs CREST ensembles, with and without GFN2-xTB relaxation.
For all methods, the errors without relaxation are far too large for the computed properties to be chemically useful—for reference, the thermal energy at room temperature is 0.59 kcal/mol.
In realistic settings, relaxation of local structures is necessary for any method, after which errors from global ﬂexibility become important.
After relaxation, torsional diffusion obtains property approximations on par or better than all previous methods.
Torsional Boltzmann generator Figure 9 shows the histograms of ESSs at 500K for the torsional
Boltzmann generator and the AIS baseline.
While AIS fails to generate more than one effective sample for most molecules (tall leftmost column), torsional Boltzmann generators are much more efﬁcient, with more than ﬁve effective samples for a signiﬁcant fraction of molecules.
[Figure]