Geometry-Complete Diffusion for 3D Molecule Generation

Abstract

Denoising diffusion probabilistic models (DDPMs) (Ho et al. (2020)) have recently taken the field of generative modeling by storm, pioneering new state-of-the-art results in disciplines such as computer vision and computational biology for diverse tasks ranging from text-guided image generation (Ramesh et al. (2022); Saharia et al. (2022); Rombach et al. (2022)) to structure-guided protein design (Ingraham et al. (2022); Watson et al. (2022)). Along this latter line of research, methods such as those of Hoogeboom et al. (2022) have been proposed for unconditionally generating 3D molecules using equivariant graph neural networks (GNNs) within a DDPM framework. Toward this end, we propose GCDM, a geometry-complete diffusion model that achieves new state-of-the-art results for 3D molecule diffusion generation by leveraging the representation learning strengths offered by GNNs that perform geometry-complete message-passing. Our results with GCDM also offer preliminary insights into how physical inductive biases impact the generative dynamics of molecular DDPMs. The source code, data, and instructions to train new models or reproduce our results are freely available at https://github.com/BioinfoMachineLearning/Bio-Diffusion.

1. Introduction

Generative modeling has recently been experiencing a renaissance in modeling efforts driven largely by denoising diffusion probabilistic models (DDPMs). At a high level, DDPMs are trained by learning how to denoise a noisy version of an input example. For example, in the context of computer vision, Gaussian noise may be successively added to an input image, and, with the goals of a DDPM in mind, we would desire for a generative model of images to learn how to successfully distinguish between the original input image’s feature signal and the noise signal added to the image thereafter. If a model can achieve such outcomes, we can use the model to generate novel images by first sampling multivariate Gaussian noise and then iteratively removing from the current state of the image the noise predicted by our model. This classic formulation of DDPMs has achieved significant results in the space of image generation (Rombach et al. (2022)), audio synthesis (Kong et al. (2020)), and even meta-learning by learning how to conditionally generate neural network checkpoints (Peebles et al. (2022)). Furthermore, such an approach to generative modeling has expanded its reach to encompass scientific disciplines such as computational biology (Anand & Achim (2022)), computational chemistry (Xu et al. (2022)), and even computational physics (Mudur & Finkbeiner (2022)).

Concurrently, the field of geometric deep learning (GDL) (Bronstein et al. (2021)) has seen a sizeable increase in research interest lately, driven largely by theoretical advances within the discipline (Joshi et al. (2023)) as well as by novel applications of such methodology (Stärk et al. (2022)). Notably, such applications even include what is considered by many researchers to be a solution to the problem of predicting 3D protein structures from their corresponding amino acid sequences (Jumper et al. (2021)). Such an outcome arose, in part, from recent advances in sequence-based language modeling efforts (Vaswani et al. (2017)) as well as from innovations in equivariant neural network modeling (Thomas et al. (2018)).

With such diverse, successful use cases of DDPMs and GDL in mind, in this work, we explore the intersection of geometric graph representation learning and DDPMs to answer the following questions.

• What is the impact of geometric representation learning on DDPMs designed to generate molecular data?

• What are the limitations of current equivariant graph neural networks empowering contemporary molecular DDPMs?

• What role do physical inductive biases play within the generative denoising of molecular DDPMs?

2. Related Work

Generative Modeling.

The field of deep generative modeling (Ruthotto & Haber (2021)) has pioneered a variety of techniques by which to train deep neural networks to create new content similar to that of an existing data repository (e.g., a text dataset of English sentences). Language models such as GPT-3 and ChatGPT (Brown et al. (2020); Schulman et al. (2022)) have become known as hallmark examples of successful generative modeling of text data. In the domains of computer vision and computational biology, techniques such as latent diffusion (Rombach et al. (2022)) and equivariant graph diffusion (Luo et al. (2022)) have established some of the latest state-of-the-art results in generative modeling of images and biomolecules such as proteins, respectively.

Geometric Deep Learning.

Data residing in a geometric or physical space (e.g., ${ℝ}^3$) can be processed by machine learning algorithms in a plethora of ways. However, in recent years, the field of geometric deep learning has become known for its proficiency in introducing powerful new deep learning methods designed specifically to process geometric data efficiently (Cao et al. (2020)). Examples of popular GDL algorithms include convolutional neural networks designed for working with image data (LeCun et al. (1995)), recurrent neural networks for processing sequence-based data (Medsker & Jain (1999)), and graph neural networks for handling graph-structured model inputs (Zhou et al. (2020)).

Equivariant Neural Networks.

To process geometric data efficiently, however, recent GDL research (Cohen & Welling (2016); Bronstein et al. (2021); Bulusu et al. (2021)) has specifically shown that designing one’s machine learning algorithm to be equivariant to the symmetry groups the input data points naturally respect (e.g., 3D rotation symmetries) often helps such an algorithm generalize to datasets beyond those used for its cross-validation (e.g., training and testing datasets). As a particularly relevant example of a neural network that is equivariant to several important and common symmetry groups of geometric data, equivariant graph neural networks (Fuchs et al. (2020); Satorras et al. (2021b); Kofinas et al. (2021); Morehead & Cheng (2022)) that are translation and rotation equivariant to inputs residing in ${ℝ}^3$ have become known as hallmark examples of geometric deep learning algorithms that generalize remarkably well to new inputs and require notably fewer training iterations to converge.

Representation Learning of Scientific Data.

Scientific data, in particular, requires careful consideration in the context of representation learning. As much scientific data contains within it a notion of geometry or latent structure, equivariance has become a key algorithmic component for processing such inputs as well (Han et al. (2022)). Moreover, equivariant graph representation learning algorithms have recently become a de facto methodology for processing scientific data of many shapes and origins Musaelian et al. (2022); Batzner et al. (2022).

Contributions.

In this work, we connect ideas at the forefront of GDL and generative modeling to advance the state-of-the-art (SOTA) for 3D molecule generation. In detail, we provide the following contributions.

• We introduce the Geometry-Complete Diffusion Model (GCDM) which establishes new SOTA results for unconditional 3D molecule generation.

• We investigate the impact of geometric message-passing on the behavior and performance of DDPMs trained to generate 3D molecular data.

• Our experiments demonstrate the importance of incorporating physical inductive biases within DDPM denoising neural networks when training them on data from physical domains.

3. Preliminaries

3.1. Diffusion Models

Key to understanding our contributions in this work are denoising diffusion probabilistic models. As alluded to previously, once trained, DDPMs can generate new data of arbitrary shapes, sizes, formats, and geometries by learning to reverse a noising process acting on each model input. More precisely, for a given data point x, a diffusion process adds noise to x for time step t = 0, 1, …, T to yield z_t, a noisy representation of the input x at time step t. Such a process is defined by a multivariate Gaussian distribution:

$$q\left({\mathbf{z}}_t\mid x\right)=𝒩\left({\mathbf{z}}_t\mid {\alpha }_t{\mathbf{x}}_t,{\sigma }_t^2\mathbf{I}\right),$$ (1)

where ${\alpha }_t\in {ℝ}^{+}$ regulates how much feature signal is retained and ${\sigma }_t^2$ modulates how much feature noise is added to input x. Note that we typically model α as a function defined with smooth transitions from α₀ = 1 to α_T = 0, where a special case of such a noising process, the variance preserving process (Sohl-Dickstein et al. (2015); Ho et al. (2020)), is defined by ${\alpha }_t=\sqrt{1-{\sigma }_t^2}$. To simplify notation, in this work, we define the feature signal-to-noise ratio as $\mathrm{SNR}(t)={\alpha }_t^2/{\sigma }_t^2$. Also interesting to note is that this diffusion process is Markovian in nature, indicating that we have transition distributions as follows:

$$q\left({\mathbf{z}}_t\mid {\mathbf{z}}_s\right)=𝒩({\mathbf{z}}_t\mid {\alpha }_{t\mid s}{\mathbf{z}}_s,{\sigma }_{t\mid s}^2\mathbf{I}),$$ (2)

for all t > s with α_t|s = α_t/α_s and ${\sigma }_{t\mid s}^2={\sigma }_t^2-{\alpha }_{t\mid s}^2{\sigma }_s^2$. In total, then, we can write the noising process as:

$$q\left({\mathbf{z}}_0,{\mathbf{z}}_1,\dots ,{\mathbf{z}}_T\mid \mathbf{x}\right)=q\left({\mathbf{z}}_0\mid x\right)\prod _{t=1}^{T}q\left({\mathbf{z}}_t\mid {\mathbf{z}}_{t-1}\right).$$ (3)

If we then define μ_t→s(x, z_t) and σ_t→s as

$${\mathit{\mu }}_{t\to s}\left(\mathbf{x},{\mathbf{z}}_t\right)=\frac{{\alpha }_{t\mid s}{\sigma }_s^2}{{\sigma }_t^2}{\mathbf{z}}_t+\frac{{\alpha }_s{\sigma }_{t\mid s}^2}{{\sigma }_t^2}\mathbf{x}\text{and}{\sigma }_{t\to s}=\frac{{\sigma }_{t\mid s}{\sigma }_s}{{\sigma }_t},$$

we have that the inverse of the noising process, the true denoising process, is given by the posterior of the transitions conditioned on x, a process that is also Gaussian:

$$q\left({\mathbf{z}}_s\mid \mathbf{x},{\mathbf{z}}_t\right)=𝒩\left({\mathbf{z}}_s\mid {\mathit{\mu }}_{t\to s}\left(\mathbf{x},{\mathbf{z}}_t\right),{\sigma }_{t\to s}\mathbf{I}\right).$$ (4)

The Generative Denoising Process.

In diffusion models, we define the generative process according to the true denoising process. However, for such a denoising process, we do not know the value of x a priori, so we typically approximate it as $\widehat{\mathbf{x}}=\varphi \left({\mathbf{z}}_t,t\right)$ using a neural network ϕ. Doing so then lets us express the generative transition distribution p(z_s|z_t) as $q\left({\mathbf{z}}_s\mid \widehat{\mathbf{x}}\left({\mathbf{z}}_t,t\right),{\mathbf{z}}_t\right)$. As a practical alternative to Eq. 4, we can represent this expression using our approximation for $\widehat{\mathbf{x}}$:

$$p\left({\mathbf{z}}_s\mid {\mathbf{z}}_t\right)=𝒩\left({\mathbf{z}}_s\mid {\mathit{\mu }}_{t\to s}\left(\widehat{\mathbf{x}},{\mathbf{z}}_t\right),{\sigma }_{t\to s}^2\mathbf{I}\right).$$ (5)

If we choose to define s as s = t − 1, then we can derive the variational lower bound on the log-likelihood of x given our generative model as:

$$\log p(\mathbf{x})\ge {ℒ}_0+{ℒ}_{base}+\sum _{t=1}^T{ℒ}_t,$$ (6)

where we note that ${ℒ}_0=\log p\left(\mathbf{x}\mid {\mathbf{z}}_0\right)$ models the likelihood of the data given its noisy representation z₀, ${ℒ}_{base}=-\mathrm{KL}\left(q\left({\mathbf{z}}_T\mid \mathbf{x}\right)\mid p\left({\mathbf{z}}_T\right)\right)$ models the difference between a standard normal distribution and the final latent variable q(z_T|x), and

$${ℒ}_t=-\mathrm{KL}\left(q\left({\mathbf{z}}_s\mid \mathbf{x},{\mathbf{z}}_t\right)\mid p\left({\mathbf{z}}_s\mid {\mathbf{z}}_t\right)\right)\text{for}t=1,2,\dots ,T.$$

Note that, in this formation of diffusion models, our neural network ϕ directly predicts $\widehat{\mathbf{x}}$. However, Ho et al. (2020) and others have found optimization of ϕ to be made much easier when instead predicting the Gaussian noise added to x to create $\widehat{\mathbf{x}}$. An intuition for how this changes the neural network’s learning dynamics is that, when predicting back the noise added to the model’s input, the network is being trained to more directly differentiate which part of z_t corresponds to the input’s feature signal (i.e., the underlying data point x) and which part corresponds to added feature noise. In doing so, if we let z_t = α_tx + σ_tϵ, our neural network can then predict $\widehat{\mathit{ϵ}}=\varphi \left({\mathbf{z}}_t,t\right)$ such that:

$$\widehat{\mathbf{x}}=\left(1/{\alpha }_t\right){\mathbf{z}}_t-\left({\sigma }_t/{\alpha }_t\right)\widehat{\mathit{ϵ}}.$$ (7)

Kingma et al. (2021) and others have since shown that, when parametrizing our denoising neural network in this way, the loss term ${ℒ}_t$ reduces to:

$${ℒ}_t={\mathbb{E}}_{\mathit{ϵ}~𝒩(\mathbf{0},\mathbf{I})}\left[\frac{1}{2}(1-\mathrm{SNR}(t-1)/\mathrm{SNR}(t))\Vert \mathit{ϵ}-\widehat{\mathit{ϵ}}{\Vert }^2\right]$$ (8)

Note that, in practice, the loss term ${ℒ}_{base}$ should be close to zero when using a noising schedule defined such that α_T ≈ 0. Moreover, if and when α₀ ≈ 1 and x is a discrete value, we will find ${ℒ}_0$ to be close to zero as well.

3.2. SE(3) Equivariance

In this work, we will consider designing denoising neural networks, here denoted as f, that are equivariant to the action of the special Euclidean group (i.e., SE(3)). We say that a function f is equivariant to the action of a group G if, for all g ∈ G, we have T_g(f(x)) = f(S_g(x)), where T_g and S_g are linear representations associated with the group element g (Serre et al. (1977)). Given that we are considering the SE(3) group, which is generated by 3D translations and rotations, we can represent T_g and S_g by a translation t and an orthogonal matrix R that rotates coordinates. Then we consider f to be equivariant to a rotation R if transforming its input x yields an equivalent transformation of its output, that is, we have Rf(x) = f(Rx).

Diffusion Models and Equivariant Distributions.

In this work, we desire for the marginal distribution p(x) of our denoising neural network to be an invariant distribution. We begin by observing that a conditional distribution p(y|x) is equivariant to the action of 3D rotations by meeting the criterion:

$$p(y\mid x)=p(\mathbf{R}y\mid \mathbf{R}x)\text{for}\text{all}\text{orthogonal}\mathbf{R}\text{.}$$ (9)

Moreover, a distribution is invariant to rotation transformations R when

$$p(y)=p(\mathbf{R}y)\text{for}\text{all}\text{orthogonal}\mathbf{R}\text{.}$$ (10)

As Köhler et al. (2020) and Xu et al. (2022) have collectively demonstrated, we know that if p(z_T) is invariant and the neural network we use to parametrize p(z_t−1|z_t) is equivariant, we have, as desired, that the marginal distribution p(x) of our denoising model is an invariant distribution.

SE(3)-Equivariant Points and Features.

In the context of our work, we represent a point cloud as a fully-connected 3D graph $𝒢=(𝒱,ℰ)$ with $\mathbf{X}=\left({\mathbf{x}}_1,{\mathbf{x}}_2,\dots ,{\mathbf{x}}_N\right)\in {ℝ}^{N\times 3}$ as the respective Cartesian coordinates for each node, where $N=|𝒱|$ and $E=\left|ℰ\right|$. Each node is described by scalar features $\mathbf{H}\in {ℝ}^{N\times h}$ and m vector-valued features $\chi \in {ℝ}^{N\times (m\times 3)}$. Likewise, each edge is described by scalar features $\mathbf{E}\in {ℝ}^{E\times e}$ and x vector-valued features $\mathbf{\xi }\in {ℝ}^{E\times (x\times 3)}$. Important to note is that the features H and E are invariant to rotations, reflections, and translations, whereas the features χ and ξ are equivariant to rotations and reflections. In particular, we describe a function Φ as SE(3)-equivariant if it satisfies the following constraint:

Definition 3.1.

(SE(3) Equivariance).

Given (H′, E′, X′, χ′, ξ′) = Φ(H, E, X, χ, ξ),

we have $({\mathbf{H}}^{\prime },{\mathbf{E}}^{\prime },\mathbf{Q}{\mathbf{X}}^{{\prime }^{\text{T}}}+\mathbf{g},\mathbf{Q}{\mathbf{\chi }}^{\prime \text{T}},\mathbf{Q}{\mathbf{\xi }}^{{\prime }^{\text{T}}})=\mathbf{\Phi }\left(\mathbf{H},\mathbf{E},\mathbf{Q}{\mathbf{X}}^{\text{T}}+\mathbf{g},\mathbf{Q}{\mathbf{\chi }}^{\text{T}},\mathbf{Q}{\mathbf{\xi }}^{\text{T}}\right)$, $\forall \mathbf{Q}\in SO(3),\forall \mathbf{g}\in {ℝ}^{3\times 1}$.

Geometry-Complete Perceptron Networks.

GCPNets are a type of geometric Graph Neural Network that satisfies our SE(3) equivariance constraint (3.1). In this setting, with (h_i ∈ H, ${\chi }_i\in \chi$, e_ij ∈ E, ξ_ij ∈ ξ), GCPNet consists of a composition of Geometry-Complete Graph Convolution (GCPConv) layers $\left(h_i^l,{\chi }_i^l\right),x_i^l=\mathbf{G}\mathbf{C}\mathbf{P}\mathbf{C}\mathbf{o}\mathbf{n}\mathbf{v}[\left(h_i^{l-1},{\chi }_i^{l-1}\right),(e_{ij}^0,{\xi }_{ij}^0),x_i^{l-1},{ℱ}_{ij}]$ which are defined as:

$$n_i^l={\varphi }^l(n_i^{l-1},{𝒜}_{\forall j\in 𝒩(i)}{\mathrm{\Omega }}_{\omega }^l(n_i^{l-1},n_j^{l-1},e_{ij},{ℱ}_{ij})),$$ (11)

where $n_i^l=(h_i^l,{\chi }_i^l);e_{ij}=(e_{ij}^0,{\xi }_{ij}^0)$; ϕ is a trainable function; l signifies the representation depth of the network; $𝒜$ is a permutation-invariant aggregation function; Ω_ω represents a message-passing function corresponding to the ω-th GCP message-passing layer; and ${𝓕}_{ij}^t=(a_{ij}^t,b_{ij}^t,c_{ij}^t)$, with $a_{ij}^t=\frac{x_i^t-x_j^t}{\Vert x_i^t-x_j^t\Vert }$, $b_{ij}^t=\frac{x_i^t\times x_j^t}{x_i^t\times x_j^t}$ and $c_{ij}^t=a_{ij}^t\times b_{ij}^t$, respectively.

Lastly, if one desires to update the positions of each node in $𝒢$, as we do in the context of 3D molecule generation, GCPConv provides a simple, SE(3)-equivariant method to do so using a dedicated GCP module as follows:

$$\left(h_{p_i}^l,{\chi }_{p_i}^l\right)=\mathbf{G}\mathbf{C}{\mathbf{P}}_p^l\left(n_i^l,{ℱ}_{ij}\right)$$ (12)

$$x_i^l=x_i^{l-1}+{\chi }_{p_i}^l,\text{where}{\chi }_{p_i}^l\in {ℝ}^{1\times 3},$$ (13)

where $\mathbf{G}\mathbf{C}{\mathbf{P}}_{.}^l\left(\cdot ,{ℱ}_{ij}\right)$ is defined as in (Morehead & Cheng (2022)) to provide rotation and translation invariant updates to h_i and rotation equivariant updates to χ_i following centralization of the input point cloud’s coordinates X (Du et al. (2022)). The effect of using feature updates to χ_i to update x_i is, after decentralizing X following the final GCPConv layer, that updates to x_i then become SE(3)-equivariant. As such, all the transformations described above collectively satisfy the required equivariance constraint in Def. 3.1. Important to note is that, (1) GCPNet performs message passing directly using vector-valued features corresponding to nodes and edges instead of performing insufficient approximations of such geometric quantities using only scalar features, and (2) GCPNet incorporates a biophysical inductive bias concerning reflection symmetries (e.g., molecular chirality) into the network architecture by encoding into the network’s updates to scalar and vector-valued features geometric frames that are not reflection equivariant. For a more detailed description of the subcomponents within GCPNet, we refer interested readers to Morehead & Cheng (2022).

4. GCDM: A Geometry-Complete Diffusion Model

In this section, we describe GCDM, a new Geometry-Complete SE(3)-Equivariant Diffusion Model, which is illustrated in Figure 1. In particular, we describe how GCDM defines a noising process jointly on equivariant node positions X and invariant node features H and learns a generative denoising process using GCPNet.

Figure 1:

A framework overview for our proposed Geometry-Complete Diffusion Model (GCDM). Our framework consists of (i.) a graph (topology) definition process, (ii.) a GCPNet-based graph neural network for 3D graph representation learning, (iii.) denoising of 3D input graphs using GCPNet, and (iv.) application of a trained GCPNet denoising network for 3D molecule generation. Zoom in for the best viewing experience.

4.1. The Diffusion Process

In this work, we define an equivariant diffusion process for equivariant node coordinates x_i and invariant node features h_i, one that adds random noise to such input data. Recall that the graph inputs to our model, $𝒢$, associate with each node a coordinate representation ${\mathbf{x}}_i\in {ℝ}^3$ and a feature vector $\mathbf{h}\in {ℝ}^h$. Subsequently, let [·, ·] denote the concatenation of two variables. We then define our equivariant noising process on latent variables ${\mathbf{z}}_t=[{\mathbf{z}}_t^{(x)},{\mathbf{z}}_t^{(h)}]$ as:

$$q\left({\mathbf{z}}_t\mid \mathbf{x},\mathbf{h}\right)={𝒩}_{xh}\left({\mathbf{z}}_t\mid {\alpha }_t[\mathbf{x},\mathbf{h}],{\sigma }_t^2\mathbf{I}\right)\text{for}t=1,2,\dots ,T,$$ (14)

where we use ${𝒩}_{xh}$ as concise notation to denote the product of two distributions. In this context, the former distribution represents the noised node coordinates ${𝒩}_x$, and the latter distribution represents the noised node features ${𝒩}_h$. Being as such, ${𝒩}_{xh}$ is given by:

$${𝒩}_x({\mathbf{z}}_t^{(x)}\mid {\alpha }_t\mathbf{x},{\sigma }_t^2\mathbf{I})\cdot {𝒩}_h({\mathbf{z}}_t^{(h)}\mid {\alpha }_t\mathbf{h},{\sigma }_t^2\mathbf{I}).$$ (15)

Note that, with the context of a standard diffusion model in mind, these two equations, Eqs. 14 and 15, correspond to Eq. 1. To address the translation invariance issue raised by Satorras et al. (2021a) in the context of handling a distribution over 3D coordinates, we adopt the zero center of gravity trick proposed by Xu et al. (2022) to define ${𝒩}_x$ as a normal distribution on the subspace defined by ${\sum }_i{\mathbf{x}}_i=\mathbf{0}$. However, to handle node features h_i that are rotation, reflection, and translation-invariant, we can instead use a conventional normal distribution $𝒩$.

GCDM Generative Denoising Process.

Recall that we need to address the noise posteriors to define a generative process for GCDM. In a similar manner as in Eq. 4, we can directly use the noise posteriors q(z_s|x, h, z_t) of Eq. 14. To do so, we must replace the input variables x and h with the approximations $\widehat{\mathbf{x}}$ and $\widehat{\mathbf{h}}$ predicted by our denoising neural network:

$$p({\mathbf{z}}_s\mid {\mathbf{z}}_t)={𝒩}_{xh}({\mathbf{z}}_s\mid {\mathit{\mu }}_{t\to s}([\widehat{\mathbf{x}},\widehat{\mathbf{h}}],{\mathbf{z}}_t),{\sigma }_{t\to s}^2\mathbf{I}),$$ (16)

where our values for $\widehat{\mathbf{x}}$ and $\widehat{\mathbf{h}}$ depend on z_t, t, and our denoising neural network ϕ. As mentioned previously, it is often easier to optimize a diffusion model using a noise parametrization to predict the noise $\widehat{ϵ}$ In this work, we use such a parametrization to predict $\widehat{ϵ}=[{\widehat{ϵ}}^{(x)},{\widehat{ϵ}}^{(h)}]$, which represents the noise individually added to $\widehat{\mathbf{x}}$ and $\widehat{\mathbf{h}}$. We can then use the predicted $\widehat{ϵ}$ to derive:

$$[\widehat{\mathbf{x}},\widehat{\mathbf{h}}]={\mathbf{z}}_t/{\alpha }_t-{\widehat{\mathit{ϵ}}}_t\cdot {\sigma }_t/{\alpha }_t.$$ (17)

Observe that rotating z_t with R yields $\mathbf{R}\widehat{\mathit{ϵ}}=\left(\varphi \left(\mathbf{R}{\mathbf{z}}_t,t\right)\right)$. Moreover, since the mean of the denoising distribution (16), one that uses isotropic noise, rotates as $\mathbf{R}\widehat{\mathbf{x}}=\mathbf{R}{\mathbf{z}}_t^{(x)}/{\alpha }_t-\mathbf{R}{\widehat{\mathit{ϵ}}}_t^{(x)}{\sigma }_t/{\alpha }_t$, the distribution is equivariant.

Sampling from the model involves sampling ${\mathbf{z}}_T~𝒩(\mathbf{0},\mathbf{I})$ and then iteratively sampling z_t ~ p(z_t−1|z_t) for t = T, T − 1, …, 1. Lastly, we can sample x, h ~ p(x, h|z₀). For a high-level overview of the sampling algorithm for diffusion models such as GCDM, we refer interested readers to Hoogeboom et al. (2022).

GCDM Optimization Objective.

When we use the noise parametrization referred to in Eq. 8, the likelihood term of our model, ${ℒ}_t=-\mathrm{KL}\left(q\left({\mathbf{z}}_s\mid \mathbf{x},{\mathbf{z}}_t\right)\mid p\left({\mathbf{z}}_s,{\mathbf{z}}_t\right)\right)$, notably simplifies. Similar to Hoogeboom et al. (2022), with this parametrization, we observe that ${ℒ}_t$ reduces to

$${ℒ}_t={\mathbb{E}}_{{\mathit{ϵ}}_t~{𝒩}_{xh}(\mathbf{0},\mathbf{I})}\left[\frac{1}{2}w(t){\Vert {\mathit{ϵ}}_t-{\widehat{\mathit{ϵ}}}_t\Vert }^2\right],$$ (18)

where ${\widehat{\mathit{ϵ}}}_t=\varphi \left({\mathbf{z}}_t,t\right)$ and w(t) = (1 −SNR(t − 1)/SNR(t)). Following Ho et al. (2020), Hoogeboom et al. (2022), and others, practically speaking, we set w(t) = 1 to stabilize training and improve sample quality. For a high-level overview of the optimization algorithm for diffusion models such as GCDM, we refer interested readers once again to Hoogeboom et al. (2022).

To summarize, we have defined a diffusion process, a denoising model, and an optimization objective function between them. We now need to define the neural network model ϕ that we will reside within the denoising model.

4.2. Equivariant Dynamics

We use our previous definition of GCPNet in Section 3.2 to learn an SE(3)-equivariant dynamics function $[{\widehat{\mathit{ϵ}}}^{(x)},{\widehat{\mathit{ϵ}}}^{(h)}]=\varphi ({\mathbf{z}}_t^{(x)},{\mathbf{z}}_t^{(h)},t)$ as:

$${\widehat{\mathit{ϵ}}}_t^{(x)},{\widehat{\mathit{ϵ}}}_t^{(h)}=\mathrm{GCPNet}({\mathbf{z}}_t^{(x)},[{\mathbf{z}}_t^{(h)},t/T])-[{\mathbf{z}}_t^{(x)},\mathbf{0}],$$ (19)

where we inform our denoising model of the current time step by concatenating t/T as an additional node feature and where we subtract the coordinate representation outputs of GCPNet from its coordinate representation inputs after subtracting from the coordinate representation outputs their collective center of gravity. With the parametrization in Eq. 17, we have subsequently achieved rotation equivariance on ${\widehat{\mathbf{x}}}_i$.

4.3. Zeroth Likelihood Terms

For the zeroth likelihood terms corresponding to each type of input feature, we adopt the respective terms previously derived by Hoogeboom et al. (2022). In particular, for integer node features, we adopt the zeroth likelihood term:

$$p(\mathbf{h}\mid {\mathbf{z}}_0^{(h)})={\int }_{\mathbf{h}-\frac{1}{2}}^{\mathbf{h}+\frac{1}{2}}𝒩(\mathbf{u}\mid {\mathbf{z}}_0^{(h)},{\sigma }_0)\text{d}\mathbf{u},$$ (20)

where we use the CDF of a standard normal distribution, Φ, to compute Eq. 20 as $\mathrm{\Phi }((\mathbf{h}+\frac{1}{2}-{\mathbf{z}}_0^{(h)})/{\sigma }_0)-\mathrm{\Phi }((\mathbf{h}-\frac{1}{2}-{\mathbf{z}}_0^{(h)})/{\sigma }_0)\approx 1$ for reasonable noise parameters α₀ and σ₀. For categorical node features, we instead use the zeroth likelihood term:

$$p(\mathbf{h}\mid {\mathbf{z}}_0^{(h)})=C(\mathbf{h}\mid \mathbf{p}),\mathbf{p}\propto {\int }_{1-\frac{1}{2}}^{1+\frac{1}{2}}𝒩(\mathbf{u}\mid {\mathbf{z}}_0^{(h)},{\sigma }_0)\text{d}\mathbf{u},$$ (21)

where we normalize p to sum to one and where C is a categorical distribution. Lastly, for continuous node positions, we adopt the zeroth likelihood term:

$$p(\mathbf{x}\mid {\mathbf{z}}_0^{(x)})=𝒩\left(\mathbf{x}\mid {\mathbf{z}}_0^{(x)}/{\alpha }_0-{\sigma }_0/{\alpha }_0{\widehat{\mathit{ϵ}}}_0,{\sigma }_0^2/{\alpha }_0^2\mathbf{I}\right)$$ (22)

which gives rise to the log-likelihood component ${ℒ}_0^{(x)}$ as:

$${ℒ}_0^{(x)}={\mathbb{E}}_{{\mathit{ϵ}}^{(x)}~{𝒩}_x(\mathbf{0},\mathbf{I})}\left[\log Z^{-1}-\frac{1}{2}{\Vert {\mathit{ϵ}}^x-{\varphi }^{(x)}\left({\mathbf{z}}_0,0\right)\Vert }^2\right],$$ (23)

where d = 3 and the normalization constant $Z={\left(\sqrt{2\pi }\cdot {\sigma }_0/{\alpha }_0\right)}^{(N-1)\cdot d}$ - in particular, its (N−1)·d term - arises from the zero center of gravity trick mentioned in Section 4.1.

Scaling Node Features.

In line with Hoogeboom et al. (2022), to improve the log-likelihood of our model’s generated samples, we find it useful to train and perform sampling with our model using scaled node feature inputs as $\left[\mathbf{x},\frac{1}{4}{\mathbf{h}}^{\text{(}categorical)},\frac{1}{10}{\mathbf{h}}^{\text{(}integer)}\right]$.

Deriving The Number of Atoms.

Finally, to determine the number of atoms with which our model will generate a 3D molecule, we first sample N ~ p(N), where p(N) denotes the categorical distribution of molecule sizes over the model’s training dataset. Then, we conclude by sampling x, h ~ p(x, h|N).

5. Experiments

5.1. Unconditional 3D Molecule Generation - QM9

The QM9 dataset (Ramakrishnan et al. (2014)) contains molecular property descriptions and 3D atom coordinates for 130k small molecules. Each molecule in QM9 can contain up to 9 heavy atoms, that is, 29 atoms when including hydrogens. For the task of 3D molecule generation, we train GCDM to unconditionally generate molecules by producing atom types (H, C, N, O, and F), integer-valued atom charges, and 3D coordinates for each of the molecules’ atoms. Following Anderson et al. (2019), we split QM9 into training, validation, and test partitions consisting of 100k, 18k, and 13k molecule examples, respectively.

Metrics.

We adopt the scoring conventions of Satorras et al. (2021a) by using the distance between atom pairs and their respective atom types to predict bond types (single, double, triple, or none). Subsequently, we measure the proportion of generated atoms that have the right valency (atom stability) and the proportion of generated molecules for which all atoms are stable (molecule stability).

Baselines.

We compare GCDM to three existing E(3)-equivariant models: G-Schnet (Gebauer et al. (2019)), Equivariant Normalizing Flows (E-NF) (Satorras et al. (2021a)), and Equivariant Diffusion Models (EDM) (Hoogeboom et al. (2022)). For each of these three models, we report their results as reported in Hoogeboom et al. (2022). For comparison with models for this task that are not equivariant, we also report results from Hoogeboom et al. (2022) for Graph Diffusion Models (GDM) trained with random data rotations (GDM-aug) and without them (GDM). To the best of our knowledge, the force-guided molecule generation methods of Wu et al. (2022) are the most recent and performant state-of-the-art methods for 3D molecule generation, so we include their results for this experiment as well.

We further include two GCDM ablation models to more closely analyze the impact of certain key model components within GCDM. These two ablation models include GCDM without local geometric frames ${ℱ}_{ij}$ (i.e., GCDM w/o Frames) and GCDM without scalar message attention (SMA) applied to each edge message (i.e., GCDM w/o SMA) as m_ij = e_ijm_ij, where m_ij represents the scalar messages learned by GCPNet during geometric graph message passing and e_ij represents a 1 if an edge exists between nodes i and j (and 0 otherwise) via e_ij ≈ ϕ_inf(m_ij). Here, ${\varphi }_{inf}:{ℝ}^e\to {[0,1]}^1$ resembles a linear layer followed by a sigmoid function Satorras et al. (2021b). All GCDM models use 9 GCPConv layers; SiLU activations (Elfwing et al. (2018)); 256 and 64 scalar node and edge hidden features, respectively; and 32 and 16 vector-valued node and edge features, respectively. All GCDM models are also trained using the AdamW optimizer (Loshchilov & Hutter (2017)) with a batch size of 64, a learning rate of 10⁻⁴, and a weight decay rate of 10⁻¹².

Results.

In Table 1, we see that GCDM outperforms previous methods (E-NF, G-Schnet, EDM, Bridge, and Bridge + Force) as well as their non-equivariant counterparts (GDM and GDM-aug) for all metrics, except concerning Atom stable (%) compared to Bridge + Force. We note that the variance of this metric for Bridge + Force is notably larger than that of GCDM. Importantly, when evaluating this metric in conjunction with Mol stable (%), we see that GCDM generates a significantly larger proportion of realistic and stable molecules (negative log-likelihood & Mol stable (%)) compared to all other methods. It is specifically interesting to note how much lower the negative log-likelihood (NLL) of GCDM is compared to that of EDM, the previous NLL-based SOTA method for 3D molecule generation, indicating the generative distribution that GCDM learns from QM9 likely contains much sharper peaks compared to EDM.

Table 1:

Comparison of GCPNet with baseline methods for 3D molecule generation. The results are reported in terms of each method’s negative log-likelihood (NLL) - log p(x, h, N), atom stability, and molecule stability with standard deviations across three runs on QM9, each drawing 10,000 samples from the model. The top-1 (best) results for this task are in bold, and the second-best results are underlined.

Type	Method	NLL ↓	Atoms Stable (%) ↑	Mol Stable (%) ↑
Normalizing Flow	E-NF	−59.7	85.0	4.9
Graph Autoregression	G-Schnet	n/a	95.7	68.1
DDPM	GDM	−94.7	97.0	63.2
	GDM-aug	−92.5	97.6	71.6
	EDM	−110.7 ± 1.5	98.7 ± 0.1	82.0 ± 0.4
	Bridge		98.7 ± 0.1	81.8 ± 0.2
	Bridge + Force		98.8 ± 0.1	84.6 ± 0.3
DDPM - Ours	GCDM w/o Frames	−162.3 ± 0.3	98.4 ± 0.0	81.7 ± 0.5
	GCDM w/o SMA	−131.3 ± 0.8	95.7 ± 0.1	51.7 ± 1.4
	GCDM	−171.0 ± 0.2	98.7 ± 0.0	85.7 ± 0.4
Data			99.0	95.2

To offer additional insights into the behavior of each method for 3D molecule generation, we report as additional metrics the validity of a generated molecule as determined by RDKit (Landrum et al. (2013)) and the uniqueness of the generated molecules overall. Note that for G-Schnet, EDM, Bridge, and Bridge + Force, we directly derive the bonds from the distance between atom pairs. We see in Table 2 that GCDM generates the highest percentage of valid and unique molecules compared to all other methods, improving upon previous SOTA results in such measures by a notable margin. A possible explanation for why GCDM can achieve such results over equivariant methods such as EDM is that GCDM performs geometric (and geometry-complete) message-passing over each 3D input graph to remove the noise present therein, whereas other methods operate solely on scalar features. In future work, it would be interesting to investigate the impact of performing different kinds of geometric message-passing (e.g., type-2 tensor message-passing) on the performance of diffusion models for tasks such as 3D molecule generation.

Table 2:

Comparison of GCPNet with baseline methods for 3D molecule generation. The results are reported in terms of the validity and uniqueness of 10,000 samples generated by each method, with standard deviations across three runs on QM9. The best results for this task are in bold, and the second-best results are underlined.

Type	Method	Valid (%) ↑	Valid and Unique (%) ↑
Normalizing Flow	E-NF	40.2	39.4
Graph Autoregression	G-Schnet	85.5	80.3
DDPM	GDM-aug	90.4	89.5
	EDM	91.9 ± 0.5	90.7 ± 0.6
	Bridge		90.2
	Bridge + Force		90.7
DDPM - Ours	GCDM w/o Frames	93.9 ± 0.1	92.7 ± 0.1
	GCDM w/o SMA	83.1 ± 1.7	82.8 ± 1.7
	GCDM	94.8 ± 0.2	93.3 ± 0.0
Data		97.7	97.7

5.2. Conditional 3D Molecule Generation - QM9

Baselines.

Towards conditional generation of 3D molecules, we compare GCDM to an existing E(3)-equivariant model, EDM (Hoogeboom et al. (2022)), as well as to two naive baselines: ”Naive (Upper-bound)” where a property classifier ϕ_c predicts molecular properties given a method’s generated 3D molecules and shuffled (i.e., random) property labels; and ”# Atoms” where one uses the numbers of atoms in a method’s generated 3D molecules to predict their molecular properties. For each baseline method, we report its mean absolute error in terms of molecular property prediction by an EGNN classifier ϕ_c Satorras et al. (2021b) as reported in Hoogeboom et al. (2022). For GCDM, we train each conditional model by conditioning it on one of five distinct molecular properties - α, gap, homo, lumo, and μ - using the QM9 validation split of Hoogeboom et al. (2022) as the model’s training dataset and the QM9 training split of Hoogeboom et al. (2022) as the corresponding EGNN classifier’s training dataset. Consequently, one can expect the gap between a method’s performance and that of ”QM9 (Lower-bound)” to decrease as the method generates molecules that more accurately model a given molecular property.

Results.

We see in Table 3 that GCDM outperforms all other methods in conditioning on a given molecular property. In particular, GCDM improves upon the mean absolute error of the SOTA EDM method for the properties α, gap, homo, lumo, and μ by 28%, 8%, 3%, 18%, and 24%, respectively, demonstrating that, using geometry-complete message passing, GCDM can more accurately model important molecular properties for 3D molecule generation.

Table 3:

Comparison of GCPNet with baseline methods for property-conditional 3D molecule generation. The results are reported in terms of the mean absolute error for molecular property prediction by an EGNN classifier ϕ_c on a QM9 subset, GCDM-generated samples, and two different baselines ”Naive (Upper-bound)” and ”# Atoms”. The top-1 (best) results for this task are in bold, and the second-best results are underlined.

Task	α	Δϵ	ϵHOMO	ϵLUMO	μ
Units	Bohr 3	meV	meV	meV	D
Naive (Upper-bound)	9.01	1470	645	1457	1.616
# Atoms	3.86	866	426	813	1.053
EDM	2.76	655	356	584	1.111
GCDM	1.97	602	344	479	0.844
QM9 (Lower-bound)	0.10	64	39	36	0.043

6. Conclusion

In this work, we introduced GCDM, an SE(3)-equivariant geometry-complete denoising diffusion probabilistic model for 3D molecule generation. While previous equivariant methods for this task have had difficulty establishing sizeable performance gains over non-equivariant methods for this task, GCDM establishes a clear performance advantage over all other methods, generating more realistic, stable, valid, unique, and property-specific 3D molecules compared to existing approaches.