PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivariant Energy-Guided SDE for Inverse Molecular Design

Fan Bao, Min Zhao, Zhongkai Hao, Peiyao Li, Chongxuan Li, Jun Zhu

arXiv:2209.15408v3physics.chem-phcs.LGq-bio.BM

TL;DR

Inverse molecular design seeks molecules with desired properties, but brute-force search is challenging and conditional generation alone may not yield accurate target molecules. EEGSDE adds energy guidance to equivariant 3D diffusion generation. It improves on the conditional EDM baseline for quantum-property and molecular-structure targets, and combines energy functions to target multiple properties.

  • Problem

    Brute-force search across the overwhelming molecular space is challenging, and the paper argues that a single conditional generative model may be insufficient for accurate molecules with desired properties.

    Source: §1 · p. 1

  • Method

    EEGSDE guides 3D diffusion generation with energy functions and exploits geometric symmetry when the energy is invariant to orthogonal transformations.

    Source: §1 · p. 1

  • Results

    EEGSDE improves on conditional EDM for quantum-property and molecular-structure targets, with dipole-moment mean absolute error reduced by more than 30% and target-structure similarity improved by more than 10%.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    Combining energy functions linearly lets EEGSDE generate molecules targeted to multiple properties.

    Source: §7 · p. 9

    The rotational-invariance guarantee depends on an equivariant noise predictor and an energy invariant to orthogonal transformations; conformations are represented in the zero-CoM subspace.

    Source: §3 · p. 3, §4.2 · p. 5

Abstract

from arXiv · show

Inverse molecular design is critical in material science and drug discovery, where the generated molecules should satisfy certain desirable properties. In this paper, we propose equivariant energy-guided stochastic differential equations (EEGSDE), a flexible framework for controllable 3D molecule generation under the guidance of an energy function in diffusion models. Formally, we show that EEGSDE naturally exploits the geometric symmetry in 3D molecular conformation, as long as the energy function is invariant to orthogonal transformations. Empirically, under the guidance of designed energy functions, EEGSDE significantly improves the baseline on QM9, in inverse molecular design targeted to quantum properties and molecular structures. Furthermore, EEGSDE is able to generate molecules with multiple target properties by combining the corresponding energy functions linearly.

1 INTRODUCTION

Inverse molecular design offers a way to search for molecules with desired properties, but conditional generation alone may not yield sufficiently accurate molecules. EEGSDE adds energy guidance to 3D diffusion generation, exploiting geometric symmetry and improving results for quantum properties and molecular structures.

1 INTRODUCTION

  • Inverse molecular design predicts promising molecules with desired properties, offering an alternative to brute-force search across an overwhelming molecular space.The paper notes that discovering such molecules matters in drug and material design.Source: §1 · p. 1
  • EEGSDE guides iterative 3D molecule generation with energy functions and exploits geometric symmetry when the energy is invariant to orthogonal transformations.The gradient of an invariant energy is equivariant, so the generated-sample distribution is invariant to rotations.Source: §1 · p. 1, Figure 1 · p. 2
  • EEGSDE improves on EDM for quantum-property and structure targets, reducing dipole-moment mean absolute error by more than 30% and improving target-structure similarity by more than 10%.The framework also targets multiple properties by combining their energy functions linearly.Source: §1 · pp. 1–2

2 RELATED WORK

Prior work spans diffusion models, guidance methods, and molecular generation, but existing approaches leave a gap in guidance that accounts for 3D molecular symmetry. EEGSDE proposes an invariant-energy guidance framework for 3D molecules.

2 RELATED WORK

3 BACKGROUND

The paper represents molecules through 3D coordinates and atom features, using zero-center-of-mass coordinates to represent conformations. Its background also defines the symmetry conditions and diffusion-model components used by the proposed method.

3 BACKGROUND

  • A molecule is represented by atom coordinates and atom features, with the coordinate collection determining its conformation.The paper groups these into the tuple z = (x, h).Source: §3 · p. 3
  • In the paper's definitions, invariance leaves a function or distribution unchanged under a transformation, while equivariance transforms the coordinate output correspondingly.For equivariance, atom-feature outputs remain unchanged as coordinate outputs transform with R.Source: §3 · p. 3
  • The paper uses zero-center-of-mass representatives because translations do not change molecular identity, while translational invariance is impossible for a distribution over the full coordinate space.Coordinates in the same translation-equivalence class represent the same conformation.Source: §3 · p. 3
  • EGNN layers are equivariant to orthogonal transformations, and their atom-feature output is invariant to those transformations.The paper notes that rotational transformations are a special case of orthogonal transformations.Source: §3 · p. 3
  • EDM perturbs molecular coordinates and features through a forward diffusion process and generates molecules by reversing that process.The forward and reverse processes use a noise schedule, and the reverse model is trained with a mean-squared-error objective.Source: §3 · p. 4
  • The cited EDM result states that equivariant noise prediction and a rotationally invariant prior yield a rotationally invariant generated-sample distribution (Hoogeboom et al. (2022)).The paper extends this proposition to the SDE formulation of molecular diffusion modeling.Source: §3 · p. 4

4 EQUIVARIANT ENERGY-GUIDED SDE

EEGSDE formulates molecular diffusion as an SDE on zero-center-of-mass coordinates and atom features, then adds energy gradients to guide generation. Its rotational-invariance guarantee requires equivariant noise prediction and an energy invariant to orthogonal transformations.

4 EQUIVARIANT ENERGY-GUIDED SDE

  • EEGSDE guides molecule generation with an energy function in a continuous-time SDE framework, with derivations supplied in the appendix.The method is introduced through the product-space SDE, its equivariance properties, and energy-function design.Source: §4 · p. 4

4.1 SDE IN THE PRODUCT SPACE

  • The forward SDE adds noise to molecular coordinates in the zero-CoM subspace and to atom features, whose transition kernels are Gaussian.The reverse process generates molecules by time reversal; its coordinate score has its center of mass subtracted to keep coordinates in the subspace.Source: §4.1 · pp. 4–5

4.2 EQUIVARIANT SDE

  • If the noise predictor is equivariant to orthogonal transformations and the prior is invariant, the generated-sample distribution is rotationally invariant.The paper parameterizes the noise predictor with an EGNN, which satisfies the equivariance constraint.Source: §4.2 · p. 5

4.3 EQUIVARIANT ENERGY-GUIDED SDE

  • EEGSDE adds the energy gradient to the reverse SDE; subtracting its coordinate-gradient center of mass keeps the process in the product space.If the energy is invariant to orthogonal transformations and the assumptions of Theorem 1 hold, the conditioned generated distribution is rotationally invariant.Source: §4.3 · pp. 5–6

4.4 HOW TO DESIGN THE ENERGY FUNCTION

  • Different energy functions let EEGSDE target different properties; as generation proceeds, their gradients encourage lower energy and thus better consistency with the target.The paper describes energy as a general scalar function capturing dependencies between input variables, and notes that it can be an MSE loss for property alignment.Source: §4.4 · p. 6

5 GENERATING MOLECULES WITH DESIRED QUANTUM PROPERTIES

EEGSDE guides 3D molecule generation with task-specific energy functions; for quantum-property targets, it improves alignment with desired properties over conditional EDM, including when targeting multiple properties.

5 GENERATING MOLECULES WITH DESIRED QUANTUM PROPERTIES

  • For quantum-property guidance, the energy is the scaled squared error between the predicted property and its target, with the scaling factor controlling guidance strength.The property predictor can use equivariant models to ensure energy invariance; this paper uses EGNN, and multiple target energies can be combined linearly.Source: §5 · p. 6

5.1 SETUP

  • Evaluation uses QM9’s six quantum properties and MAE, with generated samples assessed by a separate predictor trained on a non-overlapping subset from the energy model’s training data.The evaluation also computes properties for 100 generated molecules with Gaussian software; conditional EDM is the primary baseline.Source: §5.1 · pp. 6–7, §5.2 · p. 7

5.2 RESULTS

  • EEGSDE achieves lower MAE than conditional EDM across six quantum properties, with reductions exceeding 25% for ∆ε and εLUMO and 30% for µ.The results also report better MAE with Gaussian-software evaluation and for multiple quantum-property targets.Source: Table 1 · p. 7, Table 2 · p. 7, Table 3 · p. 7
  • The paper concludes that EEGSDE aligns generated molecules better with desired properties than conditional EDM for both single-property and multiple-property targets.The authors suggest guided chemical-space exploration for downstream applications such as virtual screening.Source: §5.2 · p. 8

6 GENERATING MOLECULES WITH TARGET STRUCTURES

This section defines fingerprint-based energy guidance for target structures, evaluates structural similarity on QM9 and GEOM-Drug, and shows that linearly combined energies can target structure and quantum properties together.

6 GENERATING MOLECULES WITH TARGET STRUCTURES

  • For target structures, EEGSDE uses a squared-error energy between a time-dependent fingerprint classifier and the target fingerprint.The guidance scale controls energy strength; the authors report that using binary cross-entropy as the energy destabilized generation.Source: §6 · p. 8

6.1 SETUP

  • The structure-alignment evaluation uses Tanimoto similarity between molecular fingerprints, with conditional EDM and cG-SchNet as baselines.The setup also evaluates the method on QM9 and GEOM-Drug.Source: §6.1 · p. 8

6.2 RESULTS

  • On QM9, EEGSDE improves target-structure similarity by more than 10% versus conditional EDM, and it also improves similarity over conditional EDM on GEOM-Drug.The section reports significant QM9 gains over conditional EDM and cG-SchNet; the GEOM-Drug comparison uses the baseline's default setting.Source: §6.2 · pp. 8–9
  • When targeting both structure and polarizability α, increasing s1 improves target alignment while increasing s2 removes a ring substructure and yields a less isometrically shaped molecule.The example uses α = 100 Bohr3; the conclusion says the corresponding energy functions are combined linearly.Source: §6.2 · p. 9, Figure 3 · p. 9

7 CONCLUSION

EEGSDE guides diffusion-based 3D molecule generation with energy functions and exploits geometric symmetry when those functions are orthogonally invariant. The paper reports improvements over conditional EDM for quantum-property and structure targets, and linear energy combinations for multiple targets.

7 CONCLUSION

  • EEGSDE is a flexible framework for controllable 3D molecule generation guided by energy functions in diffusion models.Source: §7 · p. 9
  • The method naturally exploits geometric symmetry in 3D molecular conformation when its energy function is invariant to orthogonal transformations.Source: §7 · p. 9
  • The paper reports improved conditional-EDM performance for quantum-property and molecular-structure design, and multi-property targeting by linearly combining energy functions.Source: §7 · p. 9

ETHICS STATEMENT

The ethics statement frames inverse molecular design as relevant to material science and drug discovery, while noting that negative consequences are not currently obvious.

ETHICS STATEMENT

REPRODUCIBILITY STATEMENT

The reproducibility statement locates code and experiment details in the supplementary material and paper sections, with training and model details and theoretical proofs in appendices.

REPRODUCIBILITY STATEMENT

A DERIVATIONS

The derivations formulate molecular coordinates in a zero-center-of-mass subspace and extend the SDE to a product space. They establish rotational invariance under equivariant noise prediction and orthogonally invariant energy, then give the sampling and model parameterization procedures.

A.1 ZERO COM SUBSPACE

  • The zero-CoM coordinate space is represented isometrically in a lower-dimensional space, with its projection removing the center-of-mass component.The derivation also establishes equivalence of distributions under this representation.Source: §A.1 · p. 13

A.2 SDE IN THE PRODUCT SPACE

  • The product-space SDE combines zero-CoM coordinates and atom features, with factorized transition kernels and a reverse-time form expressed through scores or predicted noise.The coordinate marginal has a differentiable extension to the full coordinate space, allowing gradients to be taken there.Source: §A.2 · pp. 14–15

A.3 EQUIVARIANCE

  • If the noise predictor is equivariant and the energy is orthogonally invariant, the guided generation distribution is rotationally invariant.The result applies to the conditional distribution in Theorem 2; the unguided result also follows under the stated invariant-prior assumption.Source: §A.3 · pp. 15–16

B SAMPLING

  • Algorithm 1 samples from EEGSDE by stepping backward from a zero-CoM Gaussian prior and incorporating the energy gradient at each step.The algorithm subtracts the center of mass from the energy gradient and noise.Source: B SAMPLING · pp. 16–17

C CONDITIONAL NOISE PREDICTION NETWORKS

D PARAMETERIZATION OF NOISE PREDICTION NETWORKS

  • The noise prediction network uses an EGNN, concatenating time and atom features, and the conditional version additionally concatenates the condition.The coordinate output has its center-of-mass component subtracted.Source: D PARAMETERIZATION OF NOISE PREDICTION… · p. 17

E.1 PARAMETERIZATION OF TIME-DEPENDENT MODELS

  • Time-dependent property and multi-label prediction models use EGNN-based parameterizations that keep the energy invariant to orthogonal transformations.The property model uses a decoder; the multi-label classifier outputs a vector for fingerprint prediction.Source: §E.1 · pp. 17–18

E.2 TRAINING OBJECTIVES OF ENERGY FUNCTIONS

The appendix specifies training losses and energy construction for quantum-property and structure guidance, alongside evaluation procedures and baselines. Its ablation reports that conditioning and energy guidance improve property MAE in sequence, with little change in stability metrics.

E.2 TRAINING OBJECTIVES OF ENERGY FUNCTIONS

  • Quantum-property guidance trains a time-dependent property predictor with an ℓ1 loss, while fingerprint guidance trains a time-dependent multi-label classifier with binary cross-entropy.The property target is scalar; the fingerprint is treated as a bitmap.Source: §E.2 · p. 18

F.1 HOW TO GENERATE MOLECULES TARGETED TO MULTIPLE QUANTUM PROPERTIES

  • For multiple quantum targets, the method sums the single-property energy functions linearly and guides the reverse SDE with the resulting energy gradient.Each component uses its own scaling factor and time-dependent predictor.Source: §F.1 · p. 18

F.2 THE “U-BOUND” AND “#ATOMS” BASELINES

  • The U-bound baseline shuffles training labels to estimate an upper bound on MAE, while the #Atoms baseline predicts a property from molecule size alone..Source: §F.2 · p. 18

F.3 GENERATING MOLECULES WITH DESIRED QUANTUM PROPERTIES

  • Quantum-property evaluation samples molecule sizes and target properties from their training-data distributions before generating molecules.The noise model uses EDM settings, and the energy model is trained separately; the passage specifies the training procedures.Source: §F.3 · p. 18

F.4 GENERATING MOLECULES WITH TARGET STRUCTURES

  • Structure guidance evaluates target similarity with Tanimoto overlap of fingerprint bits; the GEOM-Drug evaluation uses a random 10K test-molecule subset.The training and energy models are trained on the specified GEOM-Drug training split.Source: §F.4 · p. 19

G.1 ABLATION STUDY ON NOISE PREDICTION NETWORKS AND ENERGY GUIDANCE

  • Energy guidance improves MAE over a sole conditional model, while conditioning and guidance do not affect atom or molecule stability much.The ablation compares conditional and unconditional noise prediction networks and energy guidance; the tables cover single- and multiple-property settings.Source: §G.1 · p. 19, Table 5 · p. 19, Table 6 · p. 20

G.2 RESULTS ON NOVELTY, ATOM STABILITY AND MOLECULE STABILITY

The appendix reports novelty and stability trends, illustrates the effect of guidance scaling, and checks reproduced EDM results. It also shows generated GEOM-Drug structures and reports that generated and training-set property distributions match well.

G.2 RESULTS ON NOVELTY, ATOM STABILITY AND MOLECULE STABILITY

  • Larger scaling factors generally improve novelty, while conditional EDM and EEGSDE with small scaling factors have slightly better stability.The authors note a novelty–stability tradeoff, so stability may decrease slightly at large scaling factors; novelty is evaluated against Db and is inflated relative to evaluation on the whole dataset.Source: §G.2 · p. 20, Table 8 · p. 21, Table 9 · p. 21

G.3 VISUALIZATION OF THE EFFECT OF THE SCALING FACTOR

  • Generated structures align better with the target as the scaling factor grows.S = 0 corresponds to conditional EDM.Source: §G.3 · p. 21, Figure 4 · p. 21

G.4 REPRODUCE

G.5 GENERATING MOLECULES WITH TARGET STRUCTURES ON GEOM-DRUG

  • On GEOM-Drug, EEGSDE-generated molecules often match target atom types better than conditional EDM in the plotted examples.The figure shows generated molecules on GEOM-Drug.Source: §G.5 · p. 22, Figure 5 · p. 22

G.6 THE DISTRIBUTIONS OF PROPERTIES

  • The distributions of generated-molecule properties, calculated with Gaussian software, match the training-set distributions well.The comparison is visualized in Figure 6.Source: §G.6 · p. 22, Figure 6 · p. 23
Loading 2209.15408v3…