Equivariant Energy-Guided SDE for Inverse Molecular Design

TL;DR

This work introduces EEGSDE, a framework that guides equivariant diffusion-based 3D molecule generation with an energy function to achieve targeted quantum properties and structures. By enforcing orthogonal invariance through an EGNN-backed noise predictor and an energy term invariant to rotations, EEGSDE outperforms prior conditional diffusion models on QM9 and enables multi-property and structure-guided design via linear energy combinations. The approach demonstrates notable reductions in mean absolute error for several properties and improved structural similarity to targets, including challenging GEOM-Drug cases. The method provides a flexible, scalable pathway for multi-objective inverse molecular design with strong theoretical guarantees on symmetry preservation. Overall, EEGSDE offers a principled, controllable route to accelerate drug and material discovery by integrating energy-based guidance into equivariant diffusion in 3D molecular space.

Abstract

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.

Figures (7)

• Figure 1: Overview of our EEGSDE. EEGSDE iteratively generates molecules with desired properties (represented by the condition $c$) by adopting the guidance of energy functions in each step. As the energy function is invariant to rotational transformation ${\bm{R}}$, its gradient (i.e., the energy guidance) is equivariant to ${\bm{R}}$, and therefore the distribution of generated samples is invariant to ${\bm{R}}$.
• Figure 1: How generated molecules align with the target quantum property. The L-bound hoogeboom2022equivariant represents the loss of $\phi_p$ on $D_b$ and can be viewed as a lower bound of the MAE metric. The conditional EDM results are reproduced, and are consistent with hoogeboom2022equivariant (see Appendix \ref{['sec:reproduce']}). "#Atoms" uses public results from hoogeboom2022equivariant.
• Figure 2: Generated molecules on QM9 targeted to specific structures (unseen during training). The molecular structures of EEGSDE align better with target structures then conditional EDM.
• Figure 3: Generate molecules on QM9 targeted to both the quantum property $\alpha$ and the molecular structure. As the scaling factor $s_2$ grows, the substructure of generated molecule gradually change from the symmetric ring to a less isometrically shaped structure. Meanwhile the generated molecule aligns better with the target structure as the scaling factor $s_1$ grows.
• Figure 4: Visualization of the effect of the scaling factor on QM9. As the scaling factor grows, the generated structures align better with the target structure. $S=0$ corresponds to the conditional EDM.

Theorems & Definitions (18)

• theorem 1
• theorem 2
• proposition 1
• proof
• proposition 2
• proposition 3
• definition 1: Gaussian distributions in the zero CoM subspace
• proposition 4: Transition kernels of a SDE in the zero CoM subspace
• proof
• proposition 5: Transition kernels of the SDE in the product space