Equivariant Diffusion for Molecule Generation in 3D

NB: This post is just a recap of my work, but you can get my full report here.

Project Overview

The E(3) Equivariant Diffusion Model (EDM) presented by Hoogeboom et al. ¹ introduces a novel paradigm for 3D molecular generation. It is a diffusion model which goes from an initial noise $\mathbf{x}_T\sim\mathcal{N}(0,I)$ in $(\mathbb{R}^{\text{nf}+3})^M$ (where $M$ is the number of atoms and $\text{nf}$ the atom features) to a final denoised object $\mathbf{x}_0$. A final reconstruction step then transforms $\mathbf{x}_0$ to a 3D molecule with $M$ atoms. The novelty in this approach comes from the utilization of an E(3) invariant Graph Neural Network as the denoising network, which effectively creates an inductive bias that helps the EDM generalize better.

This project discusses the geometric underpinnings of the E(3) EDM. We demonstrate the benefits of incorporating E(3)-equivariance in Graph Neural Networks (GNNs) through a toy model experiment.

E(3) Equivariant Diffusion Model

The E(3) EDM is a complicated model so we break it up into chunks to understand it better.

Diffusion

Diffusion models are a class of generative models that operate by iteratively denoising a noised signal. In essence, diffusion models define a (forward) noising process which diffuses the complex target distribution $p_\text{data}(\mathbf{x})$ into a simple distribution we know how to sample, usually a Gaussian distribution. The training consists in learning the reverse process, usually by leveraging neural networks which given the noisy signal at step $t$ predict the noise increment so that we can recover the original signal at step $t-1$. Once the reverse process is learned, one can sample from the target distribution by starting from the simple distribution $\mathbf{x} _ T \sim\mathcal{N}(0,I)$ and iteratively denoising it $T$ times to get a sample from the target distribution $\mathbf{x} _ 0 \sim p _ \text{data}(\mathbf{x})$.

Diffusion models were first introduced in the context of image generation ² but they can also be applied to generate other modalities such as molecules. One simple approach to do so is to represent the molecule as a vector $\left[\mathbf{x}, h \right]$ where $\mathbf{x} \in (\mathbb{R}^3)^M$ is the position of the atoms in 3D space and $h \in (\mathbb{R}^\text{nf})^M$ is the atom features. Here $M$ is the number of atoms in the molecule and $\text{nf}$ is the number of features used to embed the atoms. In the paper, the atom features are the atom type (H, C, O, N, F one-hot encoded in a 5-dimensional vector) and charge (integer) such that $\text{nf} = 6$.

It makes sense to treat differently the atom positions $\mathbf{x}$ and the atom features $h$ since the former live in $\mathbb{R}^3$ and are subject to the symmetries of the Euclidean group $E(3)$ while the latter live in $\mathbb{R}^\text{nf}$ and are not subject to these symmetries. For this reason, the latent space distingues between the atom positions and the atom features by representing the molecule as a vector concatenation $\left[\mathbf{z}^{(\mathbf{x})} _ t, \mathbf{z}^{(\mathbf{h})} _ t \right]$. However, the two vectors do interact in the reverse diffusion process i.e. there is only one diffusion process for the whole molecule.

Graph Neural Networks

Graph Neural Networks (GNNs) are a class of neural networks that operate on graph-structured data. They are designed to capture the structure of the graph and the interactions between its nodes. More precisely, a GNN is an optimizable transformation on all attributes of the graph (nodes, edges, and global attributes) that preserves graph symmetries. (permutation invariance) ³

Equivariance, Invariance and the E(3) Group

Equivariance is a form of symmetry for functions. Roughly speaking, a function $f$ is equivariant to a group of transformations $G$ if and only if $\forall T\in G, f(T(\mathbf{x})) = T(f(\mathbf{x}))$. Likewise, invariance is when the function is constant under the action of the group i.e. $\forall T\in G, f(T(\mathbf{x})) = f(\mathbf{x})$. In the context of molecule generation, we are interested in the spatial symmetries of the molecules. Indeed, only the relative positions of the atoms matter, not their absolute positions. Therefore, we would want our generative model’s final output to be invariant to the Euclidean group $E(3)$, which is the group of translations, rotations and reflections in 3D space. Likewise, since we use diffusion to gradually denoise the molecule, we would want the denoising process to be equivariant to $E(3)$. This inductive bias is expected to help the model learn the spatial structure of the molecules more efficiently i.e. the model will be able to generalize better and with less training data.

E(3) Equivariant GNN

More precisely, the EGCL computes $\mathbf{x}^{t+1}, \mathbf{h}^{t+1} = \text{EGCL}(\mathbf{x}^t, \mathbf{h}^t)$ as follows: \(\begin{gather} m_{ij} = \phi_e(\mathbf{h}_i^t, \mathbf{h}_j^t, d_{ij}^2), \quad \tilde{e}_{ij} = \phi_{inf}(m_{ij}) \\ \mathbf{h}_i^{t+1} = \phi_h(\mathbf{h}_i^t, \sum_{j\neq i} \tilde{e}_{ij} m_{ij}) \\ \mathbf{x}_i^{t+1} = \mathbf{x}_i^t + \sum_{j\neq i} \frac{\mathbf{x}_i^t - \mathbf{x}_j^t}{d_{ij}+1} \phi_x(\mathbf{h}_i^t, \mathbf{h}_j^t, d_{ij}^2, a_{ij}) \end{gather}\) where $\phi_e, \phi_{inf}, \phi_h, \phi_x$ are neural networks and $d_{ij}$ is the Euclidean distance between atoms $i$ and $j$.

Importantly, the EGCL is equivariant to actions the Euclidean group $E(3)$ on the atom positions $\mathbf{x}$ since replacing $\mathbf{x}_i^t$ by $R \mathbf{x}_i^t + t$ for any $R\in SO(3)$ and $t\in \mathbb{R}^3$ results in the same transformation on $\mathbf{x}_i^{t+1}$. Besides, one can easily show by induction that the EGNN is also equivariant to $E(3)$ since it is a stack of EGCLs.

Wrapping it together: EDM

Finally, the Equivariant Diffusion Model (EDM) is obtained by combining the diffusion model and the EGNN. The training objective is $\mathbb{E}_{t, \epsilon_t}[|| \epsilon_t - \hat{\epsilon}_t ||^2]$ i.e. we want to predict the noise. We use the EGNN to predict the noise increment $\hat{\epsilon}_t$ given the noisy molecule $\left[\mathbf{x}_t, \mathbf{h}_t \right]$ and the time step $t$. To be precise, we set $\hat{\epsilon}_t = \left[ \hat{\epsilon}_t^{(\mathbf{x})}, \hat{\epsilon}_t^{(\mathbf{h})} \right] = \text{EGNN}(\mathbf{z}_t^{(\mathbf{x})}, \mathbf{z}_t^{(\mathbf{h})}) - \left[ \mathbf{z}_t^{(\mathbf{x})}, 0 \right]$. This recentering trick is necessary for translational equivariance of the EGNN. More details can be found in the original paper¹.

Once the EGNN is trained, we can sample from the target distribution by starting from the simple distribution $\mathbf{z}_T \sim \mathcal{N}(0,I)$ and iteratively denoising it $T$ times to get $\mathbf{z}_0$. We must then decode $\mathbf{z}_0 = \left[\mathbf{z}_0^{(\mathbf{x})}, \mathbf{z}_0^{(\mathbf{h})} \right]$ to get the molecule $\left[\mathbf{x}_0, \mathbf{h}_0 \right]$. Doing so is non-trivial since we have continuous (positions), categorical (atom type) and ordinal (charge) variables. In the paper, the authors use Bayes rule and some approximations to decode $\mathbf{z}_0$.

• The atom positions $\mathbf{x}_0$ are obtained by sampling from a gaussian distribution centered around $\mathbf{z}_0^{(\mathbf{x})}$ (plus a correction term).
• The atom type is obtained by sampling from a categorical distribution which essentially amounts to taking the nearest one-hot encoded vector to $(\mathbf{z} _ 0^{(\mathbf{h})}) _ {1:5}$.
• The charge is obtained by sampling from a gaussian distribution centered around $(\mathbf{z} _ 0^{(\mathbf{h})}) _ 6$ which essentially amounts to taking the nearest integer to $(\mathbf{z} _ 0^{(\mathbf{h})}) _ 6$. In addition, the authors use a heuristic distance-based method to create the edges of the molecule. That is, for each pair of atoms $i$ and $j$, they compute the distance $d _ {ij}$ and based on that distance and a table of bond lengths for these atoms they decide whether there is a bond between $i$ and $j$ and if so what type of bond it is.

Finally, note that since the initial distribution $\mathbf{z} _ T \sim \mathcal{N}(0,I)$ is E(3)-invariant and the EGNN is E(3)-equivariant, the final distribution of the model is also E(3)-invariant. (this can be shown easily by induction)

Experiments on toy model

The key contribution of Hoogeboom’s paper is the introduction of E(3)-equivariance. Their choice is motivated by the spatial symmetries of molecules. We aim to demonstrate the benefits of incorporating E(3)-equivariance in Graph Neural Networks (GNN) through a toy model experiment. We will focus on a regression task, which is less computationally intensive than generation and yet sufficient to highlight the advantages of equivariant models.

Dataset

Our dataset is QM9: it only contains small molecules with atoms in H, C, O, N, F. This allows us to focus on the model’s ability to capture the spatial structure of the molecules without being overwhelmed by the complexity of larger molecules. Additionnally, since QM9 has over 100,000 molecules, we filter to keep only molecules with 12 atoms or less. This brings us down to 4005 molecules.

Each sample in the dataset is a graph with several properties:

• Number of atoms $M$
• Atom positions $x \in (\mathbb{R}^3)^M$
• Atom features $h \in (\mathbb{R}^\text{nf})^M$ where $\text{nf} = 11$
  • 1-5: atom type (one-hot: H, C, O, N, F)
  • 6: atomic number (number of protons)
  • 7: aromaticity (binary)
  • 8-10: electron orbital hybridization (one-hot: sp, sp2, sp3)
  • 11: number of hydrogens
• Graph edges (bonds) $\mathcal{E} \in \mathbb{R}^{2\times\text{#edges}}$
• Edge features $\mathcal{E}_\text{type} \in \mathbb{R}^{\text{#edges}\times 4}$ (one-hot: single, double, triple, aromatic)

The QM9 dataset has 19 regression targets which are various physical properties of the molecules. We will focus on the first target which is the dipole moment of the molecule. We chose this target because it is typically geometric so we expect to clearly see the benefits of incorporating E(3)-equivariance in the model.

Models

Taking inspiration from⁴, we consider the following models:

• MPNN: a Message Passing Neural Network (MPNN) built with stacked Graph Convolutional Layers (GCL)
• EGNN: an E(3)-equivariant GNN built with stacked E(3)-Equivariant Graph Convolutional Layers (EGCL)
• EGNN_edge: same architecture as EGNN but using the edge features $\mathcal{E}_\textnormal{type}$ as well when computing the messages

All models have $L=4$ layers and $d_\text{embed}=11$ (i.e. the same as the input dimension).

We also consider LinReg, a linear regression model (projects each atom to a scalar and sums them up to get the final prediction), for baseline comparison with the geometric models. We expect poor performance as it doesn’t leverage the graph structure of the molecules.

Table: Symmetry Properties of the Different Models

Results

The results are reported on two figures: the training loss and the validation loss.

To begin with, the LinReg model performs very poorly as expected. Unlike its geometric counterparts, it cannot leverage the graph structure of the molecules and thus fails to learn anything meaningful.

Then, MPNN and EGNN_edge have similar training losses but EGNN_edge drastically outperforms MPNN on the validation set. This gap shows that EGNN_edge’s inductive bias helps it generalize better than MPNN.

Finally, EGNN is also significantly better than MPNN but learns much slower than EGNN_edge. This suggests that incorporating edge features is beneficial to the model’s performance. Nonetheless, edge features can intuitively be inferred from the atom positions and types (using a distance-based heuristic like in the original paper), so the improvement is not as significant as the one brought by E(3)-equivariance. That would explain why the EGNN is still learning and could potentially be as good as EGNN_edge with more training. To verify this hypothesis, we trained the model for 1000 epochs, but the validation loss and training loss actually plateau around 500 epochs. There could be two reasons for that: either EGNN doesn’t have enough capacity to infer the edge types itself (need for more layers or higher embedding dimension) or EGNN got stuck in a local minimum during training.

Finally, given the (intuitively) low complexity of our regression task, we expect that scaling up and fine-tuning the EGNN would lead to much better performance i.e. a much lower validation loss. However, the goal of our toy experiment was solely to demonstrate the superiority of E(3) equivariant models over architectures which do not leverage the geometry of the problem. We leave the task of maximizing performance to researchers with professional computational resources.

References: