Torsional Diffusion for Molecular Conformer Generation

Bowen Jing; Gabriele Corso; Jeffrey Chang; Regina Barzilay; Tommi Jaakkola

Source-linked AI summary

Torsional Diffusion for Molecular Conformer Generation

Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, Tommi Jaakkola

arXiv:2206.01729v2Jun 1, 2022physics.chem-phcs.LGq-bio.BM

TL;DR

Conformer generation seeks accurate and efficient ensembles of low-energy molecular structures, but prior machine-learning methods had not surpassed leading cheminformatics approaches. The paper introduces torsional diffusion on a hypertorus with an extrinsic-to-intrinsic score model, achieving state-of-the-art benchmark performance and enabling a generalizable Boltzmann generator.

Problem
Conformer generation must balance the accuracy of slow metadynamics methods with the speed of less-accurate cheminformatics methods, while prior diffusion models had not surpassed the best cheminformatics methods.
Method
Torsional diffusion applies diffusion only to torsion angles on a hypertorus and predicts molecule-specific intrinsic torsional scores from extrinsic 3D conformer coordinates.
Results
Torsional diffusion achieves state-of-the-art GEOM-DRUGS results, reduces average minimum recall RMSD by 30% and precision RMSD by 16% versus the previous state of the art, and consistently outperforms OMEGA.
Takeaways & Limitations
Exact likelihoods support the first system-agnostic Boltzmann generator described by the paper, enabling approximate conditional Boltzmann sampling for unseen molecules.
Takeaways & Limitations
The method assumes RDKit can provide local structures; the resulting distributional shift between training and test local structures significantly hurts performance.

Abstract

from arXiv · show

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 first generalizable Boltzmann generator. Code is available at https://github.com/gcorso/torsional-diffusion.

1 Introduction

Conformer generation requires balancing accuracy and speed, while prior diffusion models have not surpassed leading cheminformatics methods. Torsional diffusion addresses this by modeling torsional degrees of freedom on a molecule-specific hypertorus with an extrinsic-to-intrinsic score model.

Motivation: Conformer generation targets low-energy 3D molecular structures, but traditional methods trade accuracy against speed.Metadynamics-based methods are accurate but slow, whereas cheminformatics methods are fast but less accurate.
Motivation: Prior diffusion models use Euclidean noise on distances or 3D coordinates, require many denoising steps, and have not outperformed leading cheminformatics methods.
Approach: Torsional diffusion restricts diffusion to torsion angles, reducing the modeled dimensionality from 3n Cartesian coordinates to m torsions.For drug-like molecules, the cited averages are n = 44 atoms and m = 7.9 rotatable-bond torsion angles.
Approach: The model represents torsion angles on a molecule-specific hypertorus and predicts intrinsic torsional scores from extrinsic 3D conformer coordinates.Its score model is designed to handle varying torsional spaces and required molecular symmetries.
Contributions: Exact likelihoods enable torsional Boltzmann generators that approximately sample conditional Boltzmann distributions for unseen molecules.The paper presents this as a generalizable alternative to Boltzmann generators specific to one chemical system.
Contributions: Torsional diffusion achieves state-of-the-art GEOM-DRUGS results, consistently outperforming OMEGA while using two orders of magnitude fewer denoising steps than GeoDiff.

2 Background

Conformer generation spans accurate but expensive energy-surface exploration, faster cheminformatics approaches, and newer machine-learning models. Diffusion models traditionally operate in Euclidean spaces, while this work relies on the manifold formulation of diffusion for torsional coordinates.

Diffusion generative models: Diffusion models learn scores of progressively noised data and generate samples by approximately solving the reverse diffusion process.The score also defines a probability-flow ODE that deterministically maps the prior to the data distribution.
Diffusion generative models: The diffusion framework extends with relatively few modifications from Euclidean spaces to compact Riemannian manifolds, including the hypertorus used for torsional diffusion.
Molecular conformer generation: Conformers are energetically favorable 3D structures corresponding to local minima of a molecule’s potential energy surface.
Molecular conformer generation: Metadynamics methods explore the potential energy surface accurately but require an average of 90 core-hours per drug-like molecule.
Boltzmann generators: Boltzmann generators use exact-likelihood generative models to approximate target distributions proportional to e^(-E(x)/kT).

3 Torsional Diffusion

The method models conditional torsion distributions on a hypertorus while predicting intrinsic updates from extrinsic 3D conformers. It avoids arbitrary torsion definitions through coordinate-space relative rotations and enforces the required geometric symmetries.

Coordinate representation: A conformer is represented by local structures L and torsion angles τ around freely rotatable bonds, with ring torsions included in L.A bond is freely rotatable when severing it creates two components containing at least two atoms each.
Coordinate representation: The model samples approximate local structures with RDKit and learns pG(τ | L) over torsion angles conditioned on the molecular graph and local structure.
Diffusion on the hypertorus: Because each torsion lies in [0, 2π), the torsional coordinates form an m-dimensional hypertorus Tm, where diffusion is defined using a manifold score-based framework.
Diffusion on the hypertorus: The forward process uses rescaled Brownian motion with exponential noise, and its compact-manifold prior is uniform rather than Gaussian.The cited noise range is σmin = 0.01π to σmax = π.
Extrinsic-to-intrinsic score: Instead of learning directly on a variable, ambiguously parameterized torus, the score model maps 3D conformer point clouds to one intrinsic scalar output per rotatable bond.This construction yields an SE(3)-invariant model over Cartesian conformers with outputs in the torsional tangent space.
Extrinsic-to-intrinsic score: Relative rotations of atoms across a bond implement torsion updates independently of the arbitrary neighbor choices used to define torsion angles.All-bond updates are applied sequentially to 3D point clouds during training and inference.
Symmetries: The score model is SE(3)-invariant and parity-equivariant, changing sign under inversion to respect the required density symmetry.
Score network architecture: The final bond-convolution layer multiplies spherical-harmonic filters by bond representations, forming a pseudotorque layer that predicts torsional pseudoscalars.

4 Experiments

Experiments evaluate torsional diffusion on GEOM datasets using ensemble RMSD, coverage, chemical properties, runtime, and Boltzmann sampling. The method outperforms prior approaches while addressing local-structure distribution shift through conformer matching.

4.1 Conformer matching: Conformer matching addresses test-time distribution shift caused by training on ground-truth local structures while sampling approximate local structures with RDKit.The procedure replaces each training conformer with a similar synthetic conformer whose local structures are sampled from the RDKit distribution, while preserving useful dependence between local structures and torsions.
4.2 Experimental setup: GEOM-DRUGS contains 304k drug-like molecules averaging 44 atoms, with gold-standard ensembles generated by metadynamics in CREST.GEOM-XL tests extrapolation to molecules larger than 100 atoms, while GEOM-QM9 contains smaller molecules averaging 11 atoms.
4.2 Experimental setup: Evaluation compares generated and ground-truth ensembles using Average Minimum RMSD and Coverage for both Recall and Precision, generating 2K conformers for K reference conformers.The reported train/validation/test splits follow Ganea et al. (2021).
4.3 Ensemble RMSD: 30% lower average minimum recall RMSD and 16% lower precision RMSD are achieved relative to the previous state-of-the-art method on GEOM-DRUGS.Torsional diffusion also consistently generates better ensembles than OMEGA, the established commercial baseline.
4.3 Ensemble RMSD: 20 denoising steps suffice for torsional diffusion versus 5000 for GeoDiff, and the model outperforms GeoDiff with as few as 5 steps.The reduced number of denoising steps produces enormous runtime improvements.
4.4 Ensemble properties: Torsional diffusion produces the most chemically accurate ensembles, finding conformers only 0.13 kcal/mol above the global minimum on median.Chemical evaluation uses relaxed conformers and compares energy, dipole moment, HOMO-LUMO gap, and minimum energy.
4.5 Torsional Boltzmann generator: The torsional Boltzmann generator significantly outperforms AIS on effective sample size and improves with increased step size, despite training with a 5-step resampler.Evaluation uses 32 samples per unseen test molecule and targets the conditional Boltzmann distribution for molecules with 3–7 rotatable bonds.

5 Conclusion

Torsional diffusion generates conformers by diffusing only over torsional degrees of freedom while keeping local structures fixed, and extends this framework to likelihood-based Boltzmann generation. The authors report faster and more accurate generation than prior approaches, while identifying full intrinsic-coordinate diffusion and larger systems as future directions.

Conclusion: Torsional diffusion restricts diffusion to the molecule’s most flexible torsional degrees of freedom while fixing local structures.The method treats bond lengths and angles as fixed and modifies only selected torsion angles.
Conclusion: The model significantly outperforms standard cheminformatics methods and is orders of magnitude faster than previous Euclidean diffusion models.The conclusion characterizes these as the method’s principal performance and efficiency advantages.
Conclusion: Exact likelihoods enable training with ground-truth energies and motivate torsional Boltzmann generators for approximately sampling conditional Boltzmann distributions on unseen molecules.This extends torsional diffusion beyond conformer generation to energy-based generative modeling.
Future work: A remaining limitation is the rigid local-structure assumption, with future work targeting diffusion over the full intrinsic-coordinate space.The authors also identify larger molecular systems as a potential application area.
Method: The framework represents conformers with molecule-specific torsional coordinates on a hypertorus and predicts intrinsic torsional scores from extrinsic 3D point clouds.The score model maps Euclidean conformer representations to torsional-space scores while addressing symmetry requirements.

C Training and inference procedures

Training perturbs matched conformers with wrapped-normal torsional noise and learns denoising scores, while inference starts from randomized torsions and reverses the diffusion process over a fixed number of steps.

Implementation: Training uses at most 30 CREST conformers per molecule, typically those with the largest Boltzmann weight.Molecules are batched and optimized with Adam plus a learning-rate scheduler.
Training: During training, the procedure samples a time and conformer, adds wrapped-normal torsional noise, and updates the score model toward the noise score.The noise scale is σ = σ1−t, and optimization uses the squared error between predicted and target torsional scores.
Inference: During inference, RDKit supplies local structures, torsions are randomized uniformly on [0, 2π]^m, and reverse diffusion iteratively updates each conformer.Each reverse step combines the predicted score with wrapped-normal noise.

D Score network architecture

The score network combines chemical and geometric embeddings, E(3)-equivariant interaction layers, and a pseudotorque layer that predicts pseudoscalar torsional scores for rotatable bonds.

Architecture: The architecture consists of an embedding layer, K interaction layers, and a final pseudotorque layer producing one pseudoscalar torsional score per rotatable bond.The pseudotorque layer converts per-node interaction outputs into bond-specific scores.
Embedding layer: The embedding layer builds a radius graph with chemical features, time embeddings, and radial basis representations of interatomic distances.The distance cutoff is rmax = 5 Å.
Interaction layers: Interaction layers construct E(3)-equivariant messages from node representations and spherical harmonics, then aggregate them with Clebsch-Gordan coefficients.Learnable channel weights depend on edge embeddings and scalar features of neighboring nodes.
Pseudotorque layer: For each rotatable bond, the pseudotorque layer builds a bond-centered tensor filter from spherical harmonics and an l = 2 bond-axis representation.Odd-function dense layers extract pseudoscalar products to produce the final prediction.
Pseudotorque layer: The bond-axis representation avoids choosing a bond direction because the l = 2 spherical harmonic has even parity.This design supports the required torsional symmetry treatment.

E Conformer matching

Conformer matching aligns ground-truth conformers with RDKit local-structure estimates before training, using torsional optimization and a cost-minimizing assignment.

Matching procedure: The procedure generates K random RDKit local-structure estimates and computes a K × K cost matrix of optimized RMSDs to the ground-truth conformers.Each cost is the best RMSD obtained by modifying the estimate’s torsions.
Matching procedure: An optimal linear-sum assignment pairs true conformers with local-structure estimates, after which differential evolution finds the optimal torsions for each matched pair.The assignment is performed over the approximate cost matrix before final torsional optimization.
Results: 0.324 Å is the average RMSD achieved by conformer matching on GEOM-DRUGS.This value provides an approximate lower bound for methods that retain RDKit local structures.

F.1 RDKit local structures

RDKit provides accurate local molecular structures, so torsional sampling can focus on conformational variability while preserving stereochemical constraints. The section supports this through narrow local-geometry errors and matched-conformer RMSD analysis.

RDKit bond lengths and angles differ narrowly from CREST structures, with RMSEs of 0.03 Å and 4.1° on GEOM-DRUGS.
0.324 Å is the average RMSDmin when RDKit local structures are optimally matched to GEOM-DRUGS ground-truth conformers.RMSDmin minimizes RMSD over torsion-angle assignments while holding the local structure fixed.
The 0.324 Å RDKit matching error is only slightly above the 0.284 Å variability obtained by matching one ground-truth local structure to another.
These observations support improving conformer generation through better conditional sampling of torsion angles rather than replacing accurate local structures.
Torsional updates can be defined as relative rotations of the two molecular sides, avoiding dependence on a chosen torsion-angle reference.
Torsional diffusion preserves chirality because RDKit supplies stereoisomer-specific local structures that the diffusion process does not modify.
The framework does not currently model E/Z isomerism because double-bond torsions are treated as freely rotatable.

F.4 Limitations of torsional diffusion

Torsional diffusion improves conformer-generation accuracy and runtime, but its scope is bounded by local-structure quality, ring modeling, likelihood coverage, and molecular size.

Conformer generation: Conformer-generation error is lower bounded by the selected cheminformatics method’s local-structure quality.With RDKit local structures on DRUGS, the mean matched-conformer RMSD is 0.324 Å.
Conformer generation: 0.324 Å is the mean RMSD lower bound for RDKit local structures on DRUGS.
Rings: Ring conformations are treated as local structure, leaving puckered, fused, larger, and macrocyclic rings dependent on the local-structure sampler.
Boltzmann generation: The Boltzmann generator samples pG(C | L), not the full pG(C), because RDKit lacks exact likelihoods for local structures.
Proteins: Direct application to proteins is not considered promising because long-range effects make torsional-score influences nonlocal.

G.1 Dataset details

The experiments use GEOM-DRUGS, GEOM-QM9, and GEOM-XL datasets with substantially different molecular sizes, including a DRUGS-to-XL domain shift. Boltzmann generators are evaluated on a restricted DRUGS subset.

Splits: GEOM-DRUGS and GEOM-QM9 use random train/validation/test splits of 243473/30433/1000 and 106586/13323/1000 molecules, respectively.
Splits: GEOM-XL contains a test-only set of 102 molecules with at least 100 atoms.
Dataset statistics: The datasets differ substantially in atom and rotatable-bond counts, with a particularly significant domain shift between DRUGS and XL.
Boltzmann generator: The Boltzmann-generator experiment uses GEOM-DRUGS molecules with 3–7 rotatable bonds.It trains on 10000 molecules, validates on 400, and tests on 453.
Boltzmann generator: Annealed training starts from a smoother temperature-dependent target and gradually approaches the desired Boltzmann distribution.

G.3 Evaluation details

Evaluation measures conformer-ensemble quality with RMSD and coverage metrics, runtime on a controlled CPU setup, and Boltzmann sampling quality with effective sample size.

Ensemble RMSD: Average Minimum RMSD and Coverage measure precision and recall when generating twice as many conformers as provided by CREST.
Ensemble RMSD: XL-dataset RMSDs are upper bounds because molecular symmetries are not exhaustively tested, though they are close to permutation-aware RMSDs in practice.
Runtime evaluation: Runtime is benchmarked on an Intel i9-9920X CPU using 10 random DRUGS test molecules and 8 generated conformers per molecule.Script-loading and model-loading times are excluded.
Boltzmann generator: Boltzmann sampling quality is reported using median effective sample size from importance weights over 32 samples per molecule.
Boltzmann generator: Effective sample size approximates the number of independent target-distribution samples needed to achieve the same estimator variance.

H Additional results

Additional results show that torsional diffusion degrades more slowly as molecular flexibility increases, remains competitive across molecule sizes, and achieves most diffusion gains with few reverse steps. Its performance is sensitive to local-structure quality, while relaxation is necessary for chemically useful ensemble properties.

Performance vs size: As rotatable-bond count increases, torsional diffusion deteriorates more slowly than other conformer-generation methods.For molecules with very few rotatable bonds, its error is similar to RDKit because local-structure quality dominates.
Performance vs size: 30% recall AMR and 12% precision AMR reductions versus RDKit were obtained on GEOM-XL molecules averaging 136 atoms and 32 rotatable bonds.These molecules are larger, more flexible, and out of distribution for models trained on GEOM-DRUGS; larger-molecule training may improve results.
Small molecules: 0.178Å mean recall AMR was achieved on GEOM-QM9 using RDKit local structures, near the 0.17Å lower bound for those structures.Torsional diffusion outperformed other machine-learning methods but was only on par with or slightly worse than OMEGA on these small molecules.
Ablation experiments: About 5% higher average error accompanied a 41% runtime speed-up when node irreducible representations were truncated at order ℓ=1 instead of ℓ=2.The ablation evaluates a speed–accuracy trade-off within the torsional diffusion model.
Reverse diffusion steps: Most of torsional diffusion’s performance gain over prior diffusion-based methods is attained with only 10 reverse-diffusion steps.Increasing the default from 20 to 50 steps produces only minor additional gains.
Ensemble properties: Without relaxation, ensemble-property errors are too large for chemical use; realistic settings require relaxation of local structures.The paper uses 0.59 kcal/mol, the room-temperature thermal energy, as a reference scale for chemical usefulness.