GeoDiff: a Geometric Diffusion Model for Molecular Conformation Generation

In plain words

A molecule is a set of atoms joined by bonds. Its shape in 3D decides much of what it does, but the shape is hard to compute from the bond diagram alone. Classic simulation methods get it right slowly, especially for large molecules.

GeoDiff learns to produce such shapes. It starts from atoms scattered at random and removes the randomness in many small steps, each time nudging the atoms toward a shape that suits the bond diagram. The model is trained on known shapes: it is shown the same shape with a little noise added, and it learns to predict that noise. Running the learned steps in reverse turns noise into a shape.

One detail matters a lot. Turning a molecule or moving it in space must not change how likely its shape is. The authors build the steps so that this holds by design. The method is reported to work best on larger molecules.

TL;DR

GeoDiff is a diffusion model over the 3D coordinates of atoms, conditioned on the molecular graph. Its reverse Markov kernels are built from equivariant networks, and its starting distribution has no centre-of-mass drift, so the generated density is roto-translation invariant §4.2. It is trained with a weighted variational bound that reduces to a noise-prediction loss §4.3. The authors report state-of-the-art coverage and accuracy on GEOM-QM9 and GEOM-Drugs, with the largest gains on the larger drug-like molecules §5.2. The official code differs from the paper in several places that matter for reimplementation; see Paper vs code.

Problem

Predicting stable 3D conformations of a molecule from its graph is central to cheminformatics and drug discovery §1. Molecular dynamics and Markov chain Monte Carlo are very expensive, especially for large molecules §1. Earlier learned methods mostly predict intermediate geometry such as interatomic distances, bond lengths or torsion angles, then solve for coordinates; the authors argue that this stage-wise pipeline accumulates error §2. Generating coordinates directly is harder, because the likelihood of a conformation must not change under rotation and translation §1 §3.2.

Problem setting: given a graph G = (V, E) with atom and bond types, and a conformation C in R^(n x 3), learn a conditional model p_theta(C | G) that can be sampled and that approximates the Boltzmann distribution of conformations. Training data are i.i.d. conformations per graph §3.1.

Method

Overview. GeoDiff treats each atom as a particle. The forward process adds Gaussian noise over T steps with a fixed schedule beta_1 ... beta_T eq. 2 §4.1. The reverse process is a learned Markov chain that starts from a standard Gaussian prior and denoises step by step, conditioned on G eq. 3 §4.1. The forward marginal has a closed form, so training can sample any step directly app. A. Figure 1 sketches both processes fig. 1.

Forward kernel, with alpha_t = 1 - beta_t and alpha_bar_t the running product §4.1:

q(C^t \mid C^{t-1}) = \mathcal{N}\big(C^t;\ \sqrt{1-\beta_t}\,C^{t-1},\ \beta_t I\big) \quad [\text{eq. } 2]

q(C^t \mid C^0) = \mathcal{N}\big(C^t;\ \sqrt{\bar\alpha_t}\,C^0,\ (1-\bar\alpha_t) I\big) \quad [\text{app. A}]

Equivariant reverse kernels. The reverse kernel is Gaussian with a learned mean and a user-chosen variance sigma_t^2 eq. 3 §4.2. The mean is parameterised through a noise network eps_theta, as in denoising diffusion eq. 4 §4.2:

\mu_\theta(C^t, t) = \frac{1}{\sqrt{\alpha_t}}\left(C^t - \frac{\beta_t}{\sqrt{1-\bar\alpha_t}}\,\epsilon_\theta(G, C^t, t)\right) \quad [\text{eq. } 4]

Proposition 1 §4.2: if the prior is SE(3)-invariant and every reverse kernel is SE(3)-equivariant, the marginal density over the final conformation is SE(3)-invariant. The paper's proof is in app. A.

To make the prior invariant, the paper uses a "CoM-free" standard density. A sample is drawn from an isotropic Gaussian and then shifted to zero centre of mass; the likelihood is evaluated by first shifting to zero centre of mass, then applying the Gaussian §4.2. All intermediate states are kept CoM-free in the same way §4.2.

To make eps_theta equivariant, the paper builds a graph field network (GFN). Layer l keeps node features h (invariant) and coordinates x (equivariant). For each neighbour pair the layer computes a message from invariant inputs, updates h from the summed messages, and updates x as a weighted sum of relative position vectors eq. 5 eq. 6 eq. 7 §4.2:

m_{ij} = \Phi_m\big(h_i, h_j, \|x_i - x_j\|^2, e_{ij}; \theta_m\big) \quad [\text{eq. } 5]

h_i' = \Phi_h\Big(h_i, \textstyle\sum_{j \in \mathcal{N}(i)} m_{ij}; \theta_h\Big) \quad [\text{eq. } 6]

x_i' = \textstyle\sum_{j \in \mathcal{N}(i)} \frac{1}{d_{ij}}\,(c_i - c_j)\,\Phi_x(m_{ij}; \theta_x) \quad [\text{eq. } 7]

Here d_ij is the interatomic distance and N(i) includes bonded atoms and every atom within a radius tau, which allows long-range and disconnected graphs §4.2. Proposition 2 states that stacking L such layers gives an SE(3)-equivariant noise field §4.2 app. A.

Training objective. Training maximises a variational lower bound on the log-likelihood eq. 8 §4.3. Following denoising diffusion, the bound reduces to a weighted squared error between the true noise and the predicted noise eq. 9 §4.3:

\mathcal{L}_{\mathrm{ELBO}} = \sum_{t=1}^{T} \gamma_t\, \mathbb{E}\Big[\,\|\epsilon - \epsilon_\theta(G, C^t, t)\|_2^2\,\Big] \quad [\text{eq. } 9]

with C^t = sqrt(alpha_bar_t) C^0 + sqrt(1 - alpha_bar_t) epsilon, and weights gamma_t given in the proposition text §4.3. The paper notes that the noise epsilon is not equivariant with C^t, so the loss would not be invariant as written §4.3. It proposes two fixes that replace epsilon with an equivariant target eps_hat §4.3:

• Alignment: rigidly align C^0 to C^t (centre of mass, then Kabsch rotation), and compute eps_hat from the aligned C^0 §4.3.
• Chain rule: write epsilon as a weighted score over distances, then use the chain rule through the (invariant) distances to get an equivariant target. The paper gives the form eps_hat = (dd^t/dC^t) (d^t - sqrt(alpha_bar_t) d^0) / sqrt(1 - alpha_bar_t) §4.3.

The paper reports that the chain-rule version (GeoDiff-C) is slightly better and uses it as the main model §5.2. Both are called GeoDiff-A and GeoDiff-C in the tables tab. 1.

Sampling. Algorithm 1 alg. 1 draws C^T from the prior, then for each step shifts to zero centre of mass, computes the mean with eq. 4, and samples the previous state with variance sigma_s^2 §4.4.

# Sampling, the paper's Algorithm 1  [alg. 1] [§4.4]
# Input:  molecular graph G, trained eps_model
# Output: conformation coords (N, 3)

def sample_conformation(graph, eps_model, schedule):
    # Start from an isotropic Gaussian  [§4.2] [alg. 1]
    coords = random_normal((graph.num_atoms, 3))
    for s in range(schedule.num_steps, 0, -1):
        # Move the state to zero centre of mass  [alg. 1] [§4.2]
        coords = coords - mean_over_atoms(coords, graph)
        # Predict the noise and form the reverse-kernel mean  [eq. 4]
        eps = eps_model(graph, coords, s)
        reverse_mean = (
            coords
            - schedule.beta[s] / sqrt(1 - schedule.alpha_bar[s]) * eps
        ) / sqrt(schedule.alpha[s])
        # Sample the previous state; sigma_s is a free choice  [§4.2] [alg. 1]
        noise = random_normal((graph.num_atoms, 3))
        coords = reverse_mean + schedule.sigma[s] * noise
    return coords

Training step, chain-rule target. The paper's loss, with the chain-rule target written out. The loop structure is ours, condensed from §4.3 and eq. 9; the paper does not give this as an algorithm.

# Training step, chain-rule target  [§4.3] [eq. 9]
# Input:  coords (N, 3) in zero-CoM frame; molecular graph G
# Output: scalar loss

def training_loss(coords, graph, eps_model, schedule):
    # Draw one diffusion step per molecule  [§4.3] (sampling law not stated)
    steps = sample_steps(graph.num_molecules, schedule.num_steps)
    alpha_bar = schedule.alpha_bar[steps]           # (B,)
    # Forward kernel in closed form  [app. A] [§4.1]
    noise = random_normal_like(coords)              # epsilon ~ N(0, I)
    coords_t = (
        sqrt(alpha_bar) * coords + sqrt(1 - alpha_bar) * noise
    )
    # Equivariant noise prediction from the noisy state  [§4.2] [eq. 4]
    eps_pred = eps_model(graph, coords_t, steps)
    # Chain-rule target: pull distance noise back to atoms  [§4.3]
    d_t = edge_lengths(coords_t, graph.edges)       # (E,)
    d_0 = edge_lengths(coords, graph.edges)         # (E,)
    dist_noise = (d_t - sqrt(alpha_bar) * d_0) / sqrt(1 - alpha_bar)
    eps_target = project_to_atoms(dist_noise, coords_t, graph.edges)
    # Squared error with unit weights, averaged over molecules  [eq. 9]
    return mean(sum((eps_target - eps_pred) ** 2, axis=(-2, -1)))

Hyperparameters

Values for the paper's own experiments. Where the released code differs, the code value is shown too; see Paper vs code.

Results

Author-reported. Setup: GEOM-QM9 and GEOM-Drugs, conformer ensembles of 200 test molecules per dataset, 2× samples per molecule, RMSD after alignment, no force-field refinement unless marked §5.1 §5.2. Baselines are CVGAE, GraphDG, CGCF, ConfVAE, GeoMol, ConfGF and RDKit §5.1. Each cell is a single number; the tables report no variance over seeds tab. 1 tab. 5.

GEOM-Drugs, without force field tab. 1. COV and MAT in % and Å; "-" means not reported. The table footnote says that CVGAE, GraphDG, CGCF and ConfGF COV-R/MAT-R values are taken from an earlier study, GeoMol's from another, and the rest were run by the authors tab. 1.

GEOM-Drugs, with force-field refinement tab. 2. RDKit baseline and GeoDiff + FF. The paper says this comparison follows prior work that adds an empirical force field step §5.2.

GEOM-QM9, without force field tab. 5 app. D. Bold marks the best entry per column in the source.

Ensemble properties, GEOM-QM9 subset tab. 3. MAE in eV between predicted and reference ensemble properties over 30 molecules with 50 samples each; energies computed with Psi4 §5.3. Columns: average energy (E-bar), lowest energy (E_min), average HOMO-LUMO gap (delta-epsilon-bar), minimum gap, maximum gap. Lower is better.

Fewer diffusion steps, GEOM-Drugs tab. 6 app. D. GeoDiff with T = 1000 (beta_1 = 1e-7, beta_T = 9e-3), without force field. The paper says it still beats every baseline in tab. 1 and is slightly weaker than T = 5000 app. D.

Figures: Figure 2 shows Drugs samples aligned to the reference structure for each model, qualitatively fig. 2. Figure 3 shows more drug-like samples from GeoDiff fig. 3 app. E.

Paper vs code

Code: the official repository, MinkaiXu/GeoDiff, at the commit in the frontmatter (MIT). Searched files: the training, model and sampling code under models/epsnet/, train.py, test.py, configs/, and evaluation under utils/. For "alignment" I grepped the repository for Kabsch and alignment; only RDKit RMSD code turned up.

1. Alignment approach not in code. The repository has no Kabsch or rigid-alignment step in the training path. Only the chain-rule form is implemented, through the edge-to-atom projection dualenc.py geometry.py. The paper's GeoDiff-A has no counterpart in the code.

2. Network is not the iterated GFN. The paper's GFN updates coordinates over L layers eq. 5 eq. 6 eq. 7. The code runs a SchNet encoder on the global edges and a GIN encoder on local edges, both on invariant features only. It then maps one scalar per edge to atom vectors with a single projection dualenc.py dualenc.py geometry.py.

3. Local branch is not in the paper. Bond, angle and dihedral edges (edge types above zero) form a second, local graph with its own encoder and its own loss term, weighted 5 against 2 for the global term dualenc.py dualenc.py dualenc.py. The paper uses one graph with radius neighbours §4.2.

4. No timestep input. The paper's noise network is eps_theta(G, C^t, t) eq. 4. In the code's forward pass, the time-step embedding is commented out and the diffusion branch sets the edge scale to 1, so the step index is not used dualenc.py dualenc.py. The released denoiser is therefore not explicitly step-conditioned.

5. Loss constants. The code's loss has hard-coded weights 2 (global) and 5 (local), summed over x, y and z, with global loss only on radius edges within the cutoff dualenc.py. The paper's objective uses gamma_t weights and no such constants eq. 9 §4.3.

6. Paired time steps. The code draws half the batch's steps and gives the other half T − 1 − t dualenc.py. The paper does not state how steps are drawn.

7. Sampler. The default in test.py is sampling_type 'ld', 5000 steps, with a Langevin-type update, not the ancestral DDPM step of Algorithm 1 test.py test.py dualenc.py. The DDPM-style update that matches Algorithm 1 is also present in the code dualenc.py. The paper does not say which sampler produced its numbers.

8. Global-term gate at sampling. Global-edge forces are switched off while sigma_t is at or above 0.5 by default test.py dualenc.py. Not in the paper.

9. Scaled coordinates. The code perturbs as pos + noise · sqrt((1 − alpha_bar)/alpha_bar), that is, it works in the space of C^t / sqrt(alpha_bar) dualenc.py. The sampler starts from N(0, I) scaled by sigma_T dualenc.py dualenc.py. The paper uses C^t = sqrt(alpha_bar) C^0 + sqrt(1 − alpha_bar) epsilon app. A. Both are valid but must not be mixed.

10. Off-by-one in the sampler's alpha index. Training indexes alphas[k], the product of k+1 values of (1 − beta) dualenc.py dualenc.py. The sampler's compute_alpha(t) returns the product of t values dualenc.py. As read from the code, sampler step t pairs with training index t − 1. The size of the effect is not measured here.

11. Centring. The first noise prediction uses an uncentred initial state; centring happens after each update dualenc.py dualenc.py. Algorithm 1 centres before every prediction alg. 1. The predicted noise depends only on pairwise differences geometry.py, and the first update is followed by centring, so the difference should be a translation removed at the next step. This is read from the code and not tested.

12. Training length. The paper reports 1M iterations tab. 4. The code's default is 3,000,000 drugs_default.yml. The released checkpoints and the count used for the paper's results are not stated.

13. Evaluation metric. The paper computes RMSD after Kabsch alignment §5.2. The code removes hydrogens and uses RDKit's best-alignment RMSD over symmetry-equivalent atom mappings chem.py. The paper does not say whether hydrogens were removed.

14. Force-field step. The code applies MMFF optimisation to generated molecules before scoring covmat.py. This matches the paper's "empirical force field" description §5.2, but the paper does not name the force field.

15. Unused edge encoder. The local branch calls the global edge encoder dualenc.py, while the local encoder is created and added to the local optimiser dualenc.py train.py. Its parameters receive no gradient.

Pitfalls

• Pick one coordinate convention. The paper's C^t formula app. A and the code's scaled state dualenc.py differ by sqrt(alpha_bar). Mixing them gives wrong noise targets and a wrong sampler start.
• Sigma_t is unspecified. The paper says sigma_t can be any user-defined variance §4.2 and leaves it open in Algorithm 1. The code's DDPM branch uses sigma_t^2 = beta_t dualenc.py dualenc.py. The paper does not confirm this choice.
• Missing time input in the code. The paper's noise network takes the step t eq. 4. The released diffusion model does not take it dualenc.py. A time-conditioned reimplementation will not match the released weights.
• Invariance needs the CoM-free prior and per-step centring. The proof depends on an invariant prior and equivariant kernels §4.2 app. A. If you skip the centring in Algorithm 1, the prior and the kernels no longer match the theory.
• Division by edge length. The equivariant projection divides by every edge length, with no epsilon geometry.py. Atoms that get very close give large values. The code clips the global force to a norm of 1000 by default test.py dualenc.py. On a NaN, the sampler raises and test.py retries once with a local clip of 20 dualenc.py test.py. Expect these to matter for hard molecules.
• Sampling is expensive. One conformation needs T network evaluations, 5000 at the paper's default tab. 4, each with a radius graph up to 10 Å §4.2 tab. 4. The authors report a T = 1000 variant that still beats the baselines in their table app. D, but at lower accuracy than T = 5000 tab. 6.
• Training cost. About 48 hours on one Tesla V100 app. C. Memory use is not stated.
• Margins are not uniform. On QM9, GeoDiff-A has the best COV-R mean and GeoDiff-C has the best MAT-R, and the gap to ConfGF is smaller than on Drugs tab. 5. On the property task, ConfGF is better on the maximum gap (0.1433 vs 0.1909) tab. 3. The "consistently outperform" wording in the text §5.2 does not cover every metric.
• No variance reported. The tables give one value per cell and no seeds or confidence intervals tab. 1 tab. 5.
• Force-field comparison is different. The "+FF" rows tab. 2 are a different protocol from the rows without FF tab. 1. Do not compare them directly.
• Borrowed baseline numbers. Several baseline rows come from other papers, not re-run here tab. 1 (footnote). Protocol differences are possible.
• Typos that matter. App. C says "lowest variance level beta_T" where it means beta_1 app. C. Eq. 7 writes coordinates as c_i and x_i interchangeably eq. 7.
• Ordering of figures and tables. Figure and table numbers here follow order of appearance in the LaTeX source, with the appendix lettered as in the compiled paper. Check against the PDF before quoting numbers.

Lineage

• Builds on: Denoising diffusion probabilistic models (Ho et al., arXiv:2006.11239) for the parameterisation and the simplified loss §4.1 §4.3. Equivariant message passing in the style of EGNN (Satorras et al., arXiv:2102.09844) and tensor-field networks §4.2. Equivariant flows (Köhler et al.) for the CoM-free density §4.2.
• Chain-rule target: the idea comes from the gradient-field model ConfGF (Shi et al., 2021); the bibliography gives no arXiv id §4.3. GeoDiff replaces ConfGF's learned distance score with the closed-form diffusion noise, under a Gaussian assumption on distances §4.3.
• Compared against: CVGAE, GraphDG, CGCF, ConfVAE, ConfGF, GeoMol (arXiv:2106.07802) and RDKit §5.1 tab. 1.
• Superseded by: not known from content already in this wiki.