Geometric-Facilitated Denoising Diffusion Model for 3D Molecule Generation

TL;DR

GFMDiff introduces a geometry-aware diffusion framework for de novo 3D molecule generation by integrating a Dual-Track Transformer Network (DTN) to capture multi-body interatomic relationships and a Geometric-Facilitated Loss (GFLoss) to steer bond formation during training. By operating on 3D coordinates with pairwise distances and triplet angles under an $E(n)$-equivariant denoising kernel, GFMDiff achieves more accurate conformations and chemically valid structures. Empirical results on QM9 and GEOM-Drugs show superior stability, validity, and uniqueness compared to state-of-the-art baselines, including under conditional generation tasks. The approach advances diffusion-based molecular design by explicitly coupling geometry-aware representation learning with bond-formation dynamics, enabling scalable and reliable 3D molecule generation.

Abstract

Denoising diffusion models have shown great potential in multiple research areas. Existing diffusion-based generative methods on de novo 3D molecule generation face two major challenges. Since majority heavy atoms in molecules allow connections to multiple atoms through single bonds, solely using pair-wise distance to model molecule geometries is insufficient. Therefore, the first one involves proposing an effective neural network as the denoising kernel that is capable to capture complex multi-body interatomic relationships and learn high-quality features. Due to the discrete nature of graphs, mainstream diffusion-based methods for molecules heavily rely on predefined rules and generate edges in an indirect manner. The second challenge involves accommodating molecule generation to diffusion and accurately predicting the existence of bonds. In our research, we view the iterative way of updating molecule conformations in diffusion process is consistent with molecular dynamics and introduce a novel molecule generation method named Geometric-Facilitated Molecular Diffusion (GFMDiff). For the first challenge, we introduce a Dual-Track Transformer Network (DTN) to fully excevate global spatial relationships and learn high quality representations which contribute to accurate predictions of features and geometries. As for the second challenge, we design Geometric-Facilitated Loss (GFLoss) which intervenes the formation of bonds during the training period, instead of directly embedding edges into the latent space. Comprehensive experiments on current benchmarks demonstrate the superiority of GFMDiff.

Figures (6)

• Figure 1: An overview of GFMDiff. For each noise sample at arbitrary time step, the denoising kernel predicts atom types, valencies, and coordinates through the DTN. A loss term named GFLoss that intervenes in the formation of bonds is added to the training objective as well.
• Figure 2: The illustration of Dual-Track Transformer Network (DTN). The atom-pair track and pair-triplet with multi-head attention modules update atom and pair-wise features, respectively. The pair-wise and triplet-wise features are further fused with the latest position information.
• Figure 3: The illustration of GFLoss. We leverages chemical rules to predict the existence of bonds and then calculate potential valencies based on probabilities of atom types and bond predictions. The loss minimizes the difference between valencies predicted by DTN and valencies calculated from molecule geometries.
• Figure 4: Molecule samples generated by GFMDiff for GEOM-QM9
• Figure 5: Generated samples of GFMDiff on QM9 conditioned with increasing values of $\alpha$