1 Introduction↩︎

Generative modeling for feature-rich geometries is an important task for many science fields. Typically, geometries can be represented as point clouds where each point is embedded in the Cartesian coordinates and labeled with rich features. Such structures are ubiquitous in scientific domains, e.g., we can represent molecules as atomic graphs in 3D [1] and proteins as proximity spatial graphs over amino acids [2]. Therefore, developing effective geometric generative models holds great promise for scientific discovery problems such as material and drug design [3]–[5]. Recently, considerable progress has been achieved with machine learning approaches, especially deep generative models. For example, [6], [7] and [8] proposed data-driven methods to generate 3D molecules (in silico) with autoregressive and flow-based models respectively. However, despite great potential, the results are still unsatisfactory with low chemical validity and small molecule size, due to the insufficient capacity of the underlying generative models [9].

Most recently, diffusion models (DMs) [10], [11] have emerged with surprising results on image synthesis [12] and beyond [13], [14]. DMs define a diffusion process that gradually perturbs the data, and learn neural networks to reverse this corruption by progressive denoising. Then the denoising network can conduct generation by iteratively cleaning data initialized from random noise. Several studies have also applied such frameworks to the geometric domain, especially molecular structures [15]–[17]. However, the existing models typically run DMs directly in the atomic feature space, which typically is composed of diverse physical quantities, e.g., charge, atom types, and coordinates. These features are multi-modal with discrete, integer, and continuous variables, making unified Gaussian diffusion frameworks sub-optimal [15], [16] or requiring sophisticated, decomposed modeling of different variables [17]. Besides, the high dimensionality of input features also increases DM modeling difficulty, since the model’s training and sampling require function forward and backward computation in the full input dimension. Therefore, the validity rate of generated molecules is still not satisfying enough, and an ideal approach would be a more flexible and expressive framework for modeling complex structures.

In this paper, we propose a novel and principled method to overcome the above limitations by utilizing a smoother latent space, named Geometric Latent Diffusion Models (GeoLDM). GeoLDM is set up as (variational) autoencoders (AEs) with DMs operating on the latent space. The encoder maps the raw geometries into a lower-dimensional representational space, and DMs learn to model the smaller and smoother distribution of latent variables. For modeling the 3D molecular geometry, our key innovation is constructing sufficient conditions for latent space to satisfy the critical 3D roto-translation equivariance constraints, where simply equipping latent variables with scalar-valued¹ (i.e., invariant) variables lead to extremely poor generation quality. Technically, we realize this constraint by building the latent space as point-structured latents with both invariant and equivariant variables, which in practice is implemented by parameterizing encoding and decoding functions with advanced equivariant networks. To the best of our knowledge, we are the first work to incorporate equivariant features, i.e., tensors, into the latent space modeling.

A unique advantage of GeoLDM is that unlike previous DM methods operating in the feature domain, we explicitly incorporate a latent space to capture the complex structures. This unified formulation enjoys several strengths. First, by mapping raw features into regularized latent space, the latent DMs learn to model a much smoother distribution. This alleviates the difficulty of directly modeling complex structures’ likelihood, and is therefore more expressive. Besides, the latent space enables GeoLDM to conduct training and sampling with a lower dimensionality, which can also benefit the generative modeling complexity. Furthermore, the use of latent variables also allows for better control over the generation process, which has shown promising results in text-guided image generation [18]. This enables the user to generate specific types of molecules with desired properties. Finally, our framework is very general and can be extended to various downstream molecular problems where DMs have shown promising results, i.e., target drug design [19] and antigen-specific antibody generation [20].

We conduct detailed evaluations of GeoLDM on multiple benchmarks, including both unconditional and property-conditioned molecule generation. Results demonstrate that GeoLDM can consistently achieve superior generation performance on all the metrics, with up to 7% higher valid rate for large biomolecules. Empirical studies also show significant improvement for controllable generation thanks to latent modeling. All the empirical results demonstrate that GeoLDM enjoys a significantly higher capacity to explore the chemical space and generate structurally novel and chemically feasible molecules.

2 Related Work↩︎

Latent Generative Models. To improve the generative modeling capacity, a lot of research [21], [22] has been conducted to learn more expressive generative models over the latent space. VQ-VAEs [9] proposed to discretize latent variables and use autoregressive models to learn an expressive prior there. [23] instead employed flow-based models as the latent prior, with applications on non-autoregressive text generation. Another line of research is inspired by variational autoencoder’s (VAE’s) problem that the simple Gaussian priors cannot accurately match the encoding posteriors and therefore generate poor samples, and [21], [24] therefore proposed to use VAEs and energy-based models respectively to learn the latent distribution. Most recently, several works successfully developed latent DMs with promising results on various applications, ranging from image [25], point clouds [26], to text [14] generation. Among them, the most impressive success is Stable Diffusion models [18], which show surprisingly realistic text-guided image generation results. Despite the considerable progress we have achieved, existing latent generative methods mainly work on latent space only filled with typical scalars, without any consideration for equivariance. By contrast, we study the novel and challenging task that latent space also contains equivariant tensors.

Molecule Generation in 3D. Although extensive prior work has focused on generating molecules as 2D graphs [27]–[29], interest has recently increased in 3D generation. G-Schnet and G-SphereNet [6], [7] employed autoregressive approaches to build molecules by sequential attachment of atoms or molecular fragments. Similar frameworks have also been applied to structure-based drug design [30]–[32]. However, this autoregressive approach requires careful formulation of a complex action space and action ordering. Other studies utilized atomic density grids, by which the entire molecule can be generated in “one step" by outputting a density over the voxelized 3D space [33]. However, these density grids lack the desirable equivariance property and require a separate fitting algorithm. In the past year, DMs have attracted attention for molecule generation in 3D [15], [16], with successful application in downstream tasks like target drug generation [19], antibody design [20], and protein design [17], [34]. However, existing models mainly still work on the original atomic space, while our method works on the fundamentally different and more expressive latent space.

3 Background↩︎

3.1 Problem Definition↩︎

In this paper, we consider generative modeling of molecular geometries from scratch. Let \(d\) be the dimension of node features, then each molecule is represented as point clouds \({\mathcal{G}}= \langle {\mathbf{x}}, {\mathbf{h}}\rangle\), where \({\mathbf{x}}= ({\mathbf{x}}_1, \dots, {\mathbf{x}}_N )\in \mathbb{R}^{N \times 3}\) is the atom coordinates matrix and \({\mathbf{h}}= ({\mathbf{h}}_1, \dots, {\mathbf{h}}_N )\in \mathbb{R}^{N \times d}\) is the node feature matrix, such as atomic type and charges. We consider the following two generation tasks:
(I) Unconditional generation. With a collection of molecules \({\mathcal{G}}\), learn parameterized generative models \(p_\theta({\mathcal{G}})\) which can generate diverse and realistic molecules \(\hat{{\mathcal{G}}}\) in 3D.
(II) Controllable generation. With molecules \({\mathcal{G}}\) labeled with certain properties \(s\), learn conditional generation models \(p_\theta({\mathcal{G}}| s)\) which can conduct controllable molecule generation given desired property value \(s\).

3.2 Equivariance↩︎

Equivariance is ubiquitous for geometric systems such as molecules, where vector features like atomic forces or dipoles should transform accordingly w.r.t. the coordinates [35]–[38]. Formally, a function \({\mathcal{F}}\) is defined as equivariant w.r.t the action of a group \(G\) if \(\mathcal{F} \circ S_g ({\mathbf{x}}) = T_g \circ \mathcal{F} ({\mathbf{x}}), \forall g \in G\) where \(S_g, T_g\) are transformations for a group element \(g\) [39]. In this work, we consider the Special Euclidean group SE(3), i.e., the group of rotation and translation in 3D space, where transformations \(T_g\) and \(S_g\) can be represented by a translation \({\boldsymbol{t}}\) and an orthogonal matrix rotation \({\mathbf{R}}\).

In molecules the features \({\mathbf{h}}\) are SE(3)-invariant while the coordinates will be affected² as \({\mathbf{R}}{\mathbf{x}}+ {\boldsymbol{t}}= ({\mathbf{R}}{\mathbf{x}}_1 + {\boldsymbol{t}}, \dots, {\mathbf{R}}{\mathbf{x}}_N + {\boldsymbol{t}})\). This requires our learned likelihood to be invariant to roto-translations. Such property has been shown important for improving the generalization capacity of 3D geometric modeling [8], [40].

3.3 Diffusion Models for Non-geometric Domains↩︎

Diffusion models [10], [41] are latent variable models that model the data \(\mathbf{x}_{0}\) as Markov chains \(\mathbf{x}_{T} \cdots \mathbf{x}_{0}\), with intermediate variables sharing the same dimension. DMs can be described with two Markovian processes: a forward diffusion process \(q(\mathbf{x}_{1:T} \mid \mathbf{x}_{0}) = \prod_{t=1}^T q(\mathbf{x}_{t} \mid \mathbf{x}_{t-1})\) and a reverse denoising process \(p_\theta(\mathbf{x}_{0:T}) = p(\mathbf{x}_{T}) \prod_{t=1}^T p_\theta(\mathbf{x}_{t-1} \mid \mathbf{x}_{t})\). The forward process gradually adds Gaussian noise to data \(\mathbf{x}_{t}\): \[\label{eq:ddpm95diffusion} q(\mathbf{x}_{t} \mid \mathbf{x}_{t-1}) = \mathcal{N} (\mathbf{x}_{t} ; \sqrt{1-\beta_t} \mathbf{x}_{t-1}, \beta_t {\boldsymbol{I}}),\tag{1}\] where the hyperparameter \(\beta_{1:T}\) controls the amount of noise added at each timestep \(t\). The \(\beta_{1:T}\) are chosen such that samples \(\mathbf{x}_{T}\) can approximately converge to standard Gaussians, i.e., \(q(\mathbf{x}_{T}) \approx \mathcal{N} (0, {\boldsymbol{I}})\). Typically, this forward process \(q\) is predefined without trainable parameters.

The generation process of DMs is defined as learning a parameterized reverse denoising process, which aims to incrementally denoise the noisy variables \(\mathbf{x}_{T:1}\) to approximate clean data \(\mathbf{x}_{0}\) in the target data distribution: \[\label{eq:ddpm95denoising} p_\theta(\mathbf{x}_{t-1}\mid \mathbf{x}_{t}) = \mathcal{N}(\mathbf{x}_{t-1}; {\boldsymbol{\mu}}_\theta(\mathbf{x}_{t}, t), \rho_t^2 {\boldsymbol{I}}),\tag{2}\] where the initial distribution \(p(\mathbf{x}_{T})\) is defined as \(\mathcal{N} (0, {\boldsymbol{I}})\). The means \({\boldsymbol{\mu}}_\theta\) typically are neural networks such as U-Nets for images or Transformers for text, and the variances \(\rho_t\) typically are also predefined.

As latent variable models, the forward process \(q(\mathbf{x}_{1:T} | \mathbf{x}_{0})\) can be viewed as a fixed posterior, to which the reverse process \(p_\theta(\mathbf{x}_{0:T})\) is trained to maximize the variational lower bound of the likelihood of the data \({\mathcal{L}}_\textit{vlb} = \mathbb{E}_{q(\mathbf{x}_{1:T} | \mathbf{x}_{0})} \Big[\log \frac{q(\mathbf{x}_{T} | \mathbf{x}_{0})}{p_\theta(\mathbf{x}_{T})} + \sum_{t=2}^T \log \frac{q(\mathbf{x}_{t-1} | \mathbf{x}_{0},\mathbf{x}_{t})}{p_\theta(\mathbf{x}_{t-1} | \mathbf{x}_{t}) } - \log p_\theta( \mathbf{x}_{0} | \mathbf{x}_{1})\Big].\) However, directly optimizing this objective is known to suffer serious training instability [42]. Instead, [10], [43] suggest a simple surrogate objective up to irrelevant constant terms: \[\label{eq:ddpm95loss} {\mathcal{L}}_\textit{DM} = \mathbb{E}_{\mathbf{x}_{0},{\boldsymbol{\epsilon}}\sim \mathcal{N} (0, {\boldsymbol{I}}), t} \big[ w(t) || {\boldsymbol{\epsilon}} - {\boldsymbol{\epsilon}}_\theta(\mathbf{x}_{t}, t)||^2 \big],\tag{3}\] where \(\mathbf{x}_{t} = \alpha_t \mathbf{x}_{0} + \sigma_t {\boldsymbol{\epsilon}}\), with \(\alpha_t=\sqrt{\prod_{s=1}^t(1-\beta_s)}\) and \(\sigma_t=\sqrt{1-\alpha_t^2}\) are parameters from the tractable diffusion distributions \(q({\mathbf{x}}_t|{\mathbf{x}}_0)={\mathcal{N}}({\mathbf{x}}_t;\alpha_t{\mathbf{x}}_0,\sigma_t^2{\boldsymbol{I}})\). \({\boldsymbol{\epsilon}}_\theta\) comes from the widely adopted parametrization of the means \(\mathbf{\mu}_\theta({\mathbf{x}}_t,t):=\tfrac{1}{\sqrt{1-\beta_t}}\big({\mathbf{x}}_t-\tfrac{\beta_t}{\sqrt{1-\alpha_t^2}}{\boldsymbol{\epsilon}}_\theta({\mathbf{x}}_t,t)\big)\). The reweighting terms are \(w(t)=\frac{\beta_t^2}{2\rho_t^2(1-\beta_t)(1-\alpha_t^2)}\), while in practice simply setting it as \(1\) often promotes the sampling quality. Intuitively, the model \({\boldsymbol{\epsilon}}_\theta\) is trained to predict the noise vector \({\boldsymbol{\epsilon}}\) to denoise diffused samples \({\mathbf{x}}_t\) at every step \(t\) towards a cleaner one \({\mathbf{x}}_{t-1}\). After training, we can draw samples with \({\boldsymbol{\epsilon}}_\theta\) by the iterative ancestral sampling: \[\label{eq:ddpm95sampling} {\mathbf{x}}_{t-1}=\tfrac{1}{\sqrt{1-\beta_t}}({\mathbf{x}}_t -\tfrac{\beta_t}{\sqrt{1-\alpha_t^2}}{\boldsymbol{\epsilon}}_\theta({\mathbf{x}}_t,t)) +\rho_t {\boldsymbol{\epsilon}},\tag{4}\] with \({\boldsymbol{\epsilon}}\sim {\mathcal{N}}(\boldsymbol{0}, {\boldsymbol{I}})\). The sampling chain is initialized from Gaussian prior \({\mathbf{x}}_T \sim p(x_T)={\mathcal{N}}({\mathbf{x}}_T;\boldsymbol{0}, {\boldsymbol{I}})\).

4 Method↩︎

In this section, we formally describe Geometric Latent Diffusion Models (GeoLDM). Our work is inspired by the recent success of stable (latent) diffusion models [18], but learning latent representations for the geometric domain is however challenging [44]. We address these challenges by learning a faithful point-structured latent space with both invariant and equivariant variables, and elaborate on the design details of geometric autoencoding and latent diffusion in 4.1 and 4.2 respectively. Finally, we briefly summarize the simple training and sampling scheme in 4.3, and further discuss extensions for conditioning mechanisms in 4.4. A high-level schematic is provided in 1.

4.1 Geometric Autoencoding↩︎

We are interested in first compressing the geometries \({\mathcal{G}}= \langle {\mathbf{x}}, {\mathbf{h}}\rangle \in \mathbb{R}^{N \times (3+d)}\) (see 3.1 for details) into lower-dimensional latent space. We consider the classic autoencoder (AE) framework, where the encoder \({\mathcal{E}}_\phi\) encodes \({\mathcal{G}}\) into latent domain \({\mathbf{z}}= {\mathcal{E}}_\phi ({\mathbf{x}}, {\mathbf{h}})\) and the decoder \({\mathcal{D}}_\xi\) learns to decode \({\mathbf{z}}\) back to data domain \(\tilde{{\mathbf{x}}}, \tilde{{\mathbf{h}}} = {\mathcal{D}}_\xi({\mathbf{z}})\). The whole framework can be trained by minimizing the reconstruction objective \({\boldsymbol{d}}({\mathcal{D}}({\mathcal{E}}({\mathcal{G}})), {\mathcal{G}})\), e.g., \(L_p\) norms.

However, this classic autoencoding scheme is non-trivial in the geometric domain. Considering we follow SE(3) group in this paper (see 3.2), the typical parameterization of latent space as invariant scalar-valued features [45] is very challenging:

The idea of this proposition is that Geometric AE requires an additional function \(\psi\) to represent appropriate group actions for encoding, and align output and input positions for decoding, to solve the reconstruction task. We leave a more detailed explanation with examples in 7. For euclidean groups SE(n), [46] suggests implementing \(\psi\) as equivariant ortho-normal vectors in the unit n-dimensional sphere \(S^n\).

In our method, instead of separately representing and applying the equivariance with \(\psi\), we propose to also incorporate equivariance into \({\mathcal{E}}\) and \({\mathcal{D}}\) by constructing latent features as point-structured variables \({\mathbf{z}}=\langle {\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}\rangle \in \mathbb{R}^{N \times (3+k)}\), which holds \(3\)-d equivariant and \(k\)-d invariant latent features \({\mathbf{z}}_{\textrm{x}}\) and \({\mathbf{z}}_{\textrm{h}}\) for each node. This in practice can be implemented by parameterizing \({\mathcal{E}}\) and \({\mathcal{D}}\) with equivariant graph neural networks (EGNNs) [47], which extract both invariant and equivariant embeddings with the property: \[\label{eq:equivariance:ae} \mathbf{R} {\mathbf{z}}_{\textrm{x}}+ {\boldsymbol{t}}, {\mathbf{z}}_{\textrm{h}}= {\mathcal{E}}_\phi(\mathbf{R}{\mathbf{x}}+ {\boldsymbol{t}}, {\mathbf{h}}); \mathbf{R} {\mathbf{x}}+ {\boldsymbol{t}}, {\mathbf{h}}= {\mathcal{D}}_\xi(\mathbf{R} {\mathbf{z}}_{\textrm{x}}+ {\boldsymbol{t}}, {\mathbf{z}}_{\textrm{h}}),\tag{5}\] for all rotations \(\mathbf{R}\) and translations \({\boldsymbol{t}}\). We provide parameterization details of EGNNs in 9. The latent points \({\mathbf{z}}_{\textrm{x}}\) can perform the role of \(\psi\) required in 1, to align the orientation of outputs towards inputs. Furthermore, this point-wise latent space follows the inherent structure of geometries \({\mathcal{G}}\), thereby achieving good reconstructions.

Then the encoding and decoding processes can be formulated by \(q_\phi ({\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}| {\mathbf{x}}, {\mathbf{h}}) = {\mathcal{N}}({\mathcal{E}}_\phi({\mathbf{x}}, {\mathbf{h}}), \sigma_0 {\boldsymbol{I}})\) and \(p_\xi ({\mathbf{x}}, {\mathbf{h}}| {\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}) = \prod_{i=1}^N p_\xi (x_i, h_i | {\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}})\) respectively. Following [15], [40] that linear subspaces with the center of gravity always being zero can induce translation-invariant distributions, we also define distributions of latent \({\mathbf{z}}_{\textrm{x}}\) and reconstructed \({\mathbf{x}}\) on the subspace that \(\sum_i {\mathbf{z}}_{{\textrm{x}},i} \text{ (or {\mathbf{x}}_i)} = 0\). The whole framework can be effectively optimized by: \[\begin{align} \label{eq:loss95ae} & {\mathcal{L}}_\textit{AE} = {\mathcal{L}}_\textit{recon} + {\mathcal{L}}_\textit{reg}, \\ & {\mathcal{L}}_\textit{recon} = - \mathbb{E}_{q_\phi ({\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}| {\mathbf{x}}, {\mathbf{h}})} p_\xi ({\mathbf{x}}, {\mathbf{h}}| {\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}), \end{align}\tag{6}\] which is a reconstruction loss combined with a regularization term. The reconstruction loss in practice is calculated as \(L_2\) norm or cross-entropy for continuous or discrete features. For the \({\mathcal{L}}_\textit{reg}\) terms we experimented with two variants: KL-reg [18], a slight Kullback-Leibler penalty of \(q_\phi\) towards standard Gaussians similar to variational AE; and ES-reg, an early-stop \(q_\phi\) training strategy to avoid a scattered latent space. The regularization prevents latent embeddings from arbitrarily high variance and is thus more suitable for learning the latent DMs (LDMs).

4.2 Geometric Latent Diffusion Models↩︎

With the equivariant autoencoding functions \({\mathcal{E}}_\phi\) and \({\mathcal{D}}_\xi\), now we can represent structures \({\mathcal{G}}\) using lower-dimensional latent variables \({\mathbf{z}}\) while still keeping geometric properties. Compared with the original atomic features which are high-dimensional with complicated data types and scales, the encoded latent space significantly benefits likelihood-based generative models since: (i) as described in 4.1, our proposed AEs can be viewed as regularized autoencoders [48], where the latent space is more compact and smoothed, thereby improving DM’s training; (ii) latent codes also enjoy lower dimensionality and benefit the generative modeling complexity, since DMs typically operate in the full dimension of inputs.

Existing latent generative models for images [25], [49] and texts [14] usually rely on typical autoregressive or diffusion models to model the scalar-valued latent space. By contrast, a fundamental challenge for our method is that the latent space \({\mathbf{z}}\) contains not only scalars (i.e,, invariant features) \({\mathbf{z}}_{\textrm{h}}\) but also tensors (i.e,, equivariant features) \({\mathbf{z}}_{\textrm{x}}\). This requires the distribution of latent DMs to satisfy the critical invariance: \[p_\theta({\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}) = p_\theta({\mathbf{R}}{\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}), \text{ \forall {\mathbf{R}}}.\] [40] proved that this can be achieved if the initial distribution \(p({\mathbf{z}}_{{\textrm{x}},T}, {\mathbf{z}}_{{\textrm{h}},T})\) is invariant while the transitions \(p_\theta({\mathbf{z}}_{{\textrm{x}}, t-1}, {\mathbf{z}}_{{\textrm{h}}, t-1} | {\mathbf{z}}_{{\textrm{x}},t}, {\mathbf{z}}_{{\textrm{h}},t})\) are equivariant: \[\begin{align} p_\theta({\mathbf{z}}_{{\textrm{x}}, t-1}, {\mathbf{z}}_{{\textrm{h}}, t-1} & | {\mathbf{z}}_{{\textrm{x}},t}, {\mathbf{z}}_{{\textrm{h}},t}) = \\ & p_\theta({\mathbf{R}}{\mathbf{z}}_{{\textrm{x}}, t-1}, {\mathbf{z}}_{{\textrm{h}}, t-1} | {\mathbf{R}}{\mathbf{z}}_{{\textrm{x}},t}, {\mathbf{z}}_{{\textrm{h}},t}), \text{ \forall {\mathbf{R}}}. \end{align}\] [15], [40] further show that this can be realized by implementing the denoising dynamics \({\boldsymbol{\epsilon}}_\theta\) with equivariant networks such that: \[\label{eq:equivariance:ldm} {\mathbf{R}}{\mathbf{z}}_{{\textrm{x}}, t-1} + {\boldsymbol{t}}, {\mathbf{z}}_{{\textrm{h}}, t-1} = {\boldsymbol{\epsilon}}_\theta({\mathbf{R}}{\mathbf{z}}_{{\textrm{x}},t} + {\boldsymbol{t}}, {\mathbf{z}}_{{\textrm{h}},t}, t), \text{ \forall {\mathbf{R}} and {\boldsymbol{t}}}.\tag{7}\] which in practice we parameterize as time-conditional EGNNs. More model details are also provided in 9. Similar to the encoding posterior, in order to keep translation invariance, all the intermediate states \({\mathbf{z}}_{{\textrm{x}},t}, {\mathbf{z}}_{{\textrm{h}},t}\) are also required to lie on the subspace by \(\sum_i {\mathbf{z}}_{{\textrm{x}},t,i} = 0\) by moving the center of gravity. Analogous to 3 , now we can train the model by: \[\label{eq:loss95ldm} {\mathcal{L}}_\textit{LDM} = \mathbb{E}_{{\mathcal{E}}({\mathcal{G}}),{\boldsymbol{\epsilon}}\sim \mathcal{N} (0, {\boldsymbol{I}}), t} \big[ w(t) || {\boldsymbol{\epsilon}} - {\boldsymbol{\epsilon}}_\theta({\mathbf{z}}_{{\textrm{x}},t}, {\mathbf{z}}_{{\textrm{h}},t}, t)||^2 \big],\tag{8}\] with \(w(t)\) simply set as \(1\) for all steps \(t\).

Theoretical analysis. The combined objective for the whole framework, i.e., \({\mathcal{L}}_\textit{AE} + {\mathcal{L}}_\textit{LDM}\), appears similar to the standard VAE objective with an additional regularization. We make the formal justification that considering neglecting the minor \({\mathcal{L}}_{reg}\) term, \({\mathcal{L}}= {\mathcal{L}}_\textit{recon} + {\mathcal{L}}_\textit{LDM}\) is theoretically an SE(3)-invariant variational lower bound of log-likelihood:

Furthermore, for the induced marginal distribution \(p_{\theta,\xi} ({\mathbf{x}}, {\mathbf{h}})\), we also hold the equivariance property that:

These theoretical analysis suggest that GeoLDM is parameterized and optimized in an SE(3)-invariant fashion, which is a critical inductive bias for geometric generative models [8], [40] and provides explanations as to why our framework can achieve better 3D geometries generation quality. We provide the full statements and proofs in 8

4.3 Training and Sampling↩︎

With the proposed formulation and practical parameterization, we now present the training and sampling schemes for GeoLDM. While objectives for training Geometric AEs and LDMs are already defined in 6 8 , it is still unclear whether the two components should be trained one by one, or optimized simultaneously by backpropagation through reparameterizing [45]. Previous work about latent DMs for image generation [18], [50] shows that the two-stage training strategy usually leads to better performance, and we notice similar phenomena in our experiments. This means we first train AE with regularization, and then train the latent DMs on the latent embeddings encoded by the pre-trained encoder. A formal description of the training process is provided in 2.

With GeoLDM we can formally define a residual generative distribution \(p_{\theta,\xi}({\mathbf{x}}, {\mathbf{h}}, {\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}) = p_\theta({\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}) p_\xi ({\mathbf{x}}, {\mathbf{h}}| {\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}})\), where \(p_\theta\) refers to the latent DM modeling the point-structured latent codes, and \(p_\xi\) denotes the decoder. We can generate molecular structures by first sampling equivariant latent embeddings from \(p_\theta\) and then translating them back to the original geometric space with \(p_\xi\). The pseudo-code of the sampling procedure is provided in 4.

For the number of nodes \(N\), in the above sections, we assume it to be predefined for each data point. In practice, we need to sample different numbers \(N\) for generating molecules of different sizes. We follow the common practice [8] to first count the distribution \(p(N)\) of molecular sizes on the training set. Then for generation, we can first sample \(N \sim p(N)\) and then generate latent variables and node features in size \(N\).

4.4 Controllable Generation↩︎

Similar to other generative models [45], [51], DMs are also capable of controllable generation with given conditions \(s\), by modeling conditional distributions \(p({\mathbf{z}}|s)\). This in DMs can be implemented with conditional denoising networks \({\boldsymbol{\epsilon}}_\theta({\mathbf{z}}, t, s)\), with the critical difference that it takes additional inputs \(s\). In the molecular domain, desired conditions \(s\) typically are chemical properties, which are much lower-dimensional than the text prompts for image generations [18], [52]. Therefore, instead of sophisticated cross-attention mechanisms used in text-guided image generation, we follow [15] and simply parameterize the conditioning by concatenating \(s\) to node features. Besides, as a whole framework, we also adopt similar concatenation methods for the encoder and decoder, i.e., \({\mathcal{E}}_\phi({\mathbf{x}}, {\mathbf{h}}, s)\) and \({\mathcal{D}}_\xi({\mathbf{z}}_{\textrm{x}}, {\mathbf{z}}_{\textrm{h}}, s)\), to further shift the latent codes towards data distribution with desired properties \(s\).

5 Experiments↩︎

In this section, we justify the advantages of GeoLDM with comprehensive experiments. We first introduce our experimental setup in 5.1. Then we report and analyze the evaluation results in 5.2 and 5.3, for unconditional and conditional generation respectively. We also provide further ablation studies in 11 to investigate the effect of several model designs. We leave more implementation details in 10.

5.1 Experiment Setup↩︎

Evaluation Task. Following previous works on molecule generation in 3D [6]–[8], [15], [16], we evaluate GeoLDM by comparing with the state-of-the-art approaches on three comprehensive tasks. Molecular Modeling and Generation measures the model’s capacity to learn the molecular data distribution and generate chemically valid and structurally diverse molecules. Controllable Molecule Generation concentrates on generating target molecules with desired chemical properties. For this task, we retrain the conditional version GeoLDM on molecular data with corresponding property labels.

Datasets. We first adopt QM9 dataset [53] for both unconditional and conditional molecule generation. QM9 is one of the most widely-used datasets for molecular machine learning research, which has also been adopted in previous 3D molecule generation studies [6], [54]. QM9 contains 3D structures together with several quantum properties for 130k small molecules, limited to 9 heavy atoms (29 atoms including hydrogens). Following [55], we split the train, validation, and test partitions, with 100K, 18K, and 13K samples. For the molecule generation task, we also test GeoLDM on the GEOM-DRUG (Geometric Ensemble Of Molecules) dataset. The DRUG dataset consists of much larger organic compounds, with up to 181 atoms and 44.2 atoms on average, in 5 different atom types. It covers 37 million molecular conformations for around 450,000 molecules, labeled with energy and statistical weight. We follow the common practice [15] to select the 30 lowest energy conformations of each molecule for training.

5.2 Molecular Modeling and Generation↩︎

Evaluation Metrics. We measure model performances by evaluating the chemical feasibility of generated molecules, indicating whether the model can learn chemical rules from data. Given molecular geometries, we first predict bond types (single, double, triple, or none) by pair-wise atomic distances and atom types. Then we calculate the atom stability and molecule stability of the predicted molecular graph. The first metric captures the proportion of atoms that have the right valency, while the latter is the proportion of generated molecules for which all atoms are stable. In addition, We report validity and uniqueness metrics, which are the percentages of valid (measured by RDKit) and unique molecules among all the generated compounds.

Baselines. We compare GeoLDM to several competitive baseline models. G-Schnet [6] and Equivariant Normalizing Flows (ENF) [8] are previous equivariant generative models for molecules, based on autoregressive and flow-based models respectively. Equivariant Graph Diffusion Models (EDM) with its non-equivariant variant (GDM) [15] are recent progress on diffusion models for molecule generation. Most recently, [16] proposed an improved version of EDM (EDM-Bridge), which further boosts the performance with well-designed informative prior bridges. To yield a fair comparison, all the baseline models use the same parameterization and training configurations as described in 5.1.

Results and Analysis. We generate \(10,000\) samples from each method to calculate the above metrics, and the results are reported in ¿tbl:tab:qm995results?. As shown in the table, GeoLDM outperforms competitive baseline methods on all metrics with an obvious margin. It is worth noticing that, for the DRUG dataset, even ground-truth molecules have \(86.5\%\) atom-level and nearly \(0\%\) molecule-level stability. This is because the DRUG molecules contain larger and more complex structures, creating errors during bond type prediction based on pair-wise atom types and distances. Furthermore, as DRUG contains many more molecules with diverse compositions, we also observe that unique metric is almost \(100\%\) for all methods. Therefore, we omit the molecule stability and unique metrics for the DRUG dataset. Overall, the superior performance demonstrates GeoLDM’s higher capacity to model the molecular distribution and generate chemically realistic molecular geometries. We provide visualization of randomly generated molecules in 3, and leave more visualizations in 12.

Ablation Study. Furthermore, to verify the benefits of incorporating equivariant latent features, we conduct ablation studies with only invariant variables in the latent space, called Graph Latent Diffusion Models (GraphLDM). We run GraphLDM with the same configuration as our method, except that all modules (i.e., encoder, decoder, and latent diffusion models) are instead equipped with typical non-equivariant graph networks. We also follow [15] to test GDM-aug and GraphLDM-aug, where models are trained with data augmented by random rotations. ¿tbl:tab:qm995results? shows the empirical improvement of GeoLDM over these ablation settings, which verifies the effectiveness of our latent equivariance design.

5.3 Controllable Molecule Generation↩︎

2pt

Evaluation Metrics. In this task, we aim to conduct controllable molecule generation with the given desired properties. This can be useful in realistic settings of material and drug design where we are interested in discovering molecules with specific property preferences. We test our conditional version of GeoLDM on QM9 with 6 properties: polarizability \(\alpha\), orbital energies \(\varepsilon_{\mathrm{HOMO}}\), \(\varepsilon_{\mathrm{LUMO}}\) and their gap \(\Delta \varepsilon\), Dipole moment \(\mu\), and heat capacity \(C_v\). For evaluating the model’s capacity to conduct property-conditioned generation, we follow [8] to first split the QM9 training set into two halves with \(50K\) samples in each. Then we train a property prediction network \(\omega\) on the first half, and train conditional models on the second half. Afterward, given a range of property values \(s\), we conditionally draw samples from the generative models and then use \(\omega\) to calculate their property values as \(\hat{s}\). We report the Mean Absolute Error (MAE) between \(s\) and \(\hat{s}\) to measure whether generated molecules are close to their conditioned property. We also test the MAE of directly running \(\omega\) on the second half QM9, named QM9 in ¿tbl:tab:conditional95results?, which measures the bias of \(\omega\). A smaller gap with QM9 numbers indicates a better property-conditioning performance.

Baselines. We incorporate existing EDM as our baseline model. In addition, we follow [15] to also list two baselines agnostic to ground-truth property \(s\), named Random and \(N_\textit{atoms}\). Random means we simply do random shuffling of the property labels in the dataset and then evaluate \(\omega\) on it. This operation removes any relation between molecule and property, which can be viewed as an upper bound of MAE metric. \(N_\text{atoms}\) predicts the molecular properties by only using the number of atoms in the molecule. The improvement over Random can verify the method is able to incorporate conditional property information into the generated molecules. And overcoming \(N_\text{atoms}\) further indicates the model can incorporate conditioning into molecular structures beyond the number of atoms.

Results and Analysis. We first provide a visualization of controlled molecule generation by GeoLDM in 5, as qualitative assessments. We interpolate the conditioning property with different Polarizability values \(\alpha\) while keeping the reparameterization noise \({\boldsymbol{\epsilon}}\) fixed. Polarizability refers to the tendency of matter, when subjected to an electric field, to acquire an electric dipole moment in proportion to that applied field. Typically, less isometrically molecular geometries lead to larger \(\alpha\) values. This is consistent with our observed phenomenon in 5.

We report the numerical results in ¿tbl:tab:conditional95results?. As shown in the table, GeoLDM significantly outperforms baseline models, including the previous diffusion model running on atomic features (EDM), on all the property metrics. The results demonstrate that by modeling in the latent space, GeoLDM acquired a higher capacity to incorporate given property information into the generation process.

6 Conclusion and Future Work↩︎

We presented GeoLDM, a novel latent diffusion model for molecular geometry generation. While current models operate directly on high-dimensional, multi-modal atom features, GeoLDM overcomes their limitations by learning diffusion models over a continuous, lower-dimensional latent space. By building point-structured latent codes with both invariant scalars and equivariant tensors, GeoLDM is able to effectively learn latent representations while maintaining roto-translational equivariance. Experimental results demonstrate its significantly better capacity for modeling chemically realistic molecules. For future work, as a general and principled framework, GeoLDM can be extended to various 3D geometric generation applications, e.g., apply GeoLDM in more realistic drug discovery scenarios with given protein targets, or scale up GeoLDM for more challenging 3D geometries such as peptides and proteins.

Acknowledgements↩︎

We thank Tailin Wu, Aaron Lou, Xiang Lisa Li, and Kexin Huang for discussions and for providing feedback on our manuscript. We also gratefully acknowledge the support of DARPA under Nos. HR00112190039 (TAMI), N660011924033 (MCS); ARO under Nos. W911NF-16-1-0342 (MURI), W911NF-16-1-0171 (DURIP); NSF under Nos. OAC-1835598 (CINES), OAC-1934578 (HDR), CCF-1918940 (Expeditions), NIH under No. 3U54HG010426-04S1 (HuBMAP), Stanford Data Science Initiative, Wu Tsai Neurosciences Institute, Amazon, Docomo, GSK, Hitachi, Intel, JPMorgan Chase, Juniper Networks, KDDI, NEC, and Toshiba. We also gratefully acknowledge the support of NSF (#1651565), ARO (W911NF-21-1-0125), ONR (N00014-23-1-2159), CZ Biohub, Stanford HAI. We also gratefully acknowledge the support of Novo Nordisk A/S. Minkai Xu thanks the generous support of Sequoia Capital Stanford Graduate Fellowship.

7 Explanation of 1↩︎

We first explain the intuition behind the theoretical justification of 1 by an example here. Considering given a input geometry \({\mathcal{G}}= \langle {\mathbf{h}}, {\mathbf{x}}\rangle\), the encoder \({\mathcal{E}}\) and decoder \({\mathcal{D}}\), such that \({\mathcal{G}}= {\mathcal{D}}({\mathcal{E}}({\mathcal{G}}))\). Then we can transform \({\mathcal{G}}\) by an action \(g\) from SE(3)-group to \(\hat{{\mathcal{G}}} = T_g {\mathcal{G}}= \langle {\mathbf{h}}, {\mathbf{R}}{\mathbf{x}}+ {\boldsymbol{t}}\rangle\) and input it into the autoencoders. Since the encoding function is invariant, we have \({\mathcal{E}}({\mathcal{G}}) = {\mathcal{E}}(\hat{{\mathcal{G}}})\), and thus the reconstructed geometry however will still be \({\mathcal{G}}= {\mathcal{D}}({\mathcal{E}}(\hat{{\mathcal{G}}}))\) instead of \(\hat{{\mathcal{G}}}\). This is problematic because we couldn’t calculate the reconstruction error based on \({\mathcal{G}}\) and \(\hat{{\mathcal{G}}}\), and a natural solution is that we need an additional function \(\psi\) to extract the group action \(g\). Then after decoding, we can apply the group action on generated \({\mathcal{G}}\) to recover \(\hat{{\mathcal{G}}}\), thereby solving the problem

Formally, the explanation is that all elements can be expressed in terms of coordinates with respect to a given basis. So we should consider a canonical basis for all orbits, and learn the equivariant function \(\psi\) to indicate to which orbit elements are decoded as “canonical”. For detailed theoretical analysis, we refer readers to [46].

8 Formal Statements and Proofs↩︎

8.1 Relationship to SE(3)-invariant Likelihood: 2↩︎

First, recall the informal theorem we provide in 4.2, which builds the connection between GeoLDM’s objective and SE(3)-invariant maximum likelihood:

Before providing the proof, we first present a formal version of the theorem:

As shown in the theorem, the conclusion is composed of two statements, i.e., 9 and 10 . The first equation states that \({\mathcal{L}}\) is a variational lower bound of the log-likelihood, and the second equation shows that the objective \({\mathcal{L}}\) is further SE(3)-invariant, i.e., invariant to any rotational and translational transformations. Here, we provide the full proofs of the two statements separately. We first present the proof for 9 :

8.2 Invariant Marginal Distribution: 3↩︎

We also include proofs for key properties of the equivariant probabilistic diffusion model here to be self-contained [15], [40]. Note that, since here we are interested in the equivariant properties, we omit the trivial scalar inputs \({\mathbf{h}}\) and focus on analyzing tensor features \({\mathbf{z}}\). The proof shows that when the the initial distribution \(p({\mathbf{z}}^{(T)})\) is invariant and transition distributions \(p({\mathbf{z}}_{\textrm{x}}^{(t-1)} | {\mathbf{z}}_{\textrm{x}}^{(t)})\) are equivariant, then the marginal distributions \(p({\mathbf{z}}_{\textrm{x}}^{(t)})\) will be invariant, importantly including \(p({\mathbf{z}}_{\textrm{x}}^{(0)})\). Similarly, with decoder \(p({\mathbf{x}}|{\mathbf{z}}_{\textrm{x}}^{(0)})\) also being equivariant, we can further have that our induced distribution \(p({\mathbf{x}})\) is invariant.

9 Model Architecture Details↩︎

In our implementation, all models are parameterized with EGNNs [47] as backbone. EGNNs are a class of Graph Neural Network that satisfies the equivariance property in 5 7 . In this work, we consider molecular geometries as point clouds, without specifying the connecting bonds. Therefore, in practice, we take the point clouds as fully connected graph \(G\) and model the interactions between all atoms \({\textrm{v}}_i \in \mathcal{V}\). Each node \({\textrm{v}}_i\) is embedded with coordinates \({\textrm{x}}_i \in \mathbb{R}^3\) and atomic features \({\textrm{h}}_i \in \mathbb{R}^d\). Then, EGNNs are composed of multiple Equivariant Convolutional Layers \({\mathbf{x}}^{l+1}, {\mathbf{h}}^{l+1} = \operatorname{EGCL}[{\mathbf{x}}^l, {\mathbf{h}}^l]\), with each single layer defined as: \[\begin{align} & {\mathbf{m}}_{ij} = \phi_{e}\left({\textrm{h}}_{i}^{l}, {\textrm{h}}_{j}^{l}, d_{ij}^2, a_{ij}\right), \\ & {\textrm{h}}_{i}^{l+1} = \phi_{h}({\textrm{h}}_{i}^l, { \sum_{j \neq i}} \tilde{e}_{ij} {\mathbf{m}}_{ij}), \\ & {\textrm{x}}_{i}^{l+1} = {\textrm{x}}_{i}^{l}+\sum_{j \neq i} \frac{{\textrm{x}}_{i}^{l}-{\textrm{x}}_{j}^{l}}{d_{ij}+ 1} \phi_{x}\left({\textrm{h}}_{i}^{l}, {\textrm{h}}_{j}^{l}, d_{ij}^2, a_{ij}\right), \label{eq:egnn-layer} \end{align}\tag{11}\] where \(l\) denotes the layer index. \(\tilde{e}_{ij}=\phi_{inf}({\mathbf{m}}_{ij})\) acts as the attention weights to reweight messages passed from different edges. \(d_{ij} = \| {\textrm{x}}_{i}^{l}-{\textrm{x}}_{j}^{l} \|_2\) represents the pairwise distance between atoms \({\textrm{v}}_i\) and \({\textrm{v}}_j\), and \(a_{ij}\) are optional edge features. We follow previous work [8], [15] to normalize the relative directions \({\textrm{x}}_i^l - {\textrm{x}}_j^l\) in 11 by \(d_{ij} + 1\), which empirically improved model stability. All learnable functions, i.e., \(\phi_e\), \(\phi_h\), \(\phi_x\) and \(\phi_{inf}\), are parameterized by Multi Layer Perceptrons (MLPs). Then a complete EGNN model can be realized by stacking \(L\) \(\operatorname{EGCL}\) layers such that \({\mathbf{x}}^L, {\mathbf{h}}^L = \operatorname{EGNN}[{\mathbf{x}}^0, {\mathbf{h}}^0]\), which can satisfy the required equivariant constraint in 5 7 .

10 Featurization and Implementation Details↩︎

We use the open-source software RDkit [56] to preprocess molecules. For QM9 we take atom types (H, C, N, O, F) and integer-valued atom charges as atomic features, while for Drugs we only use atom types. The results reported in [ssec:exp-molgen] [ssec:exp-condition] are based on the ES-reg regularization strategy (4.3), where the encoder is only optimized with \(1000\) iterations of warm-up training and then fixed. For the diffusion process (1 ), we use the polynomial noise schedule [15], [16], where \(\alpha\) linearly decays from \(10^3/T\) to \(0\) w.r.t. time step \(t\). And for the denoising process (2 ), the variances are defined as \(\rho_t = \sqrt{\frac{\sigma_{t-1}}{\sigma_t}}\beta_t\).

All neural networks used for the encoder, latent diffusion, and decoder are implemented with EGNNs [47] by PyTorch [57] package, as introduced in 9. We set the dimension of latent invariant features \(k\) to \(1\) for QM9 and \(2\) for DRUG, which extremely reduces the atomic feature dimension. For the training of latent denoising network \({\boldsymbol{\epsilon}}_\theta\): on QM9, we train EGNNs with \(9\) layers and \(256\) hidden features with a batch size \(64\); and on GEOM-DRUG, we train EGNNs with \(4\) layers and \(256\) hidden features, with batch size \(64\). For the autoencoders, we parameterize the decoder \({\mathcal{D}}_\xi\) in the same way as \({\boldsymbol{\epsilon}}_\theta\), but implement the encoder \({\mathcal{E}}_\phi\) with a \(1\) layer EGNN. The shallow encoder in practice constrains the encoding capacity and helps regularize the latent space. All models use SiLU activations. We train all the modules until convergence. For all the experiments, we choose the Adam optimizer  [58] with a constant learning rate of \(10^{-4}\) as our default training configuration. The training on QM9 takes approximately \(2000\) epochs, and on DRUG takes \(20\) epochs.

11 Ablation Studies↩︎

In this section, We provide additional experimental results on QM9 to justify the effect of several model designs. Specifically, we perform ablation studies on two key model designs: autoencoder regularization method and latent space dimension \(k\). The results are reported in ¿tbl:app:tab:ablation?.

We first discuss the effect of different autoencoder regularization methods, i.e., KL-reg and ES-reg (see details in 4.1), with the latent invariant feature dimension fixed as \(1\). Following previous practice in latent diffusion models of image and point clouds domains [18], [26], for KL-reg, we also weight the KL term with a small factor \(0.01\). However, during our initial experiments where we naturally first try the KL-reg method, we observed unexpected failure with extremely poor performance, as shown in the first row in ¿tbl:app:tab:ablation?. Note that, this reported result is already the best result we achieved for KL-reg, with searching over a large range of KL term weights and latent space dimension. In practice, we even notice the KL-reg is unstable for training, which often suffers from numerical errors during training. Our closer observation of the experimental results suggests that the equivariant latent feature part always tends to converge to highly scattered means and extremely small variances, which leads to the numerical issue for calculating KL term and also is not suitable for LDM training. Therefore, we turned to constraining the encoder, more precisely, constraining the value scale of encoded latent features, by early stopping the training encoder. This easy strategy turned out to work pretty well in practice as shown in ¿tbl:app:tab:ablation?, and we leave the further study of KL-reg as future work in this area.

We further study the effect of latent invariance feature dimension \(k\), and the results are also reported in ¿tbl:app:tab:ablation?. As shown in the table, we observe that generally GeoLDM shows better performance with lower \(k\). This phenomenon verifies our motivation that a lower dimensionality can alleviate the generative modeling complexity and benefit the training of LDM. Specifically, the performances of GeoLDM on QM9 with \(k\) set as \(1\) or \(2\) are very similar, so we only report \(k=1\) as representative in ¿tbl:app:tab:ablation?. In practice, we set \(k\) as \(1\) for QM9 dataset and \(2\) for DRUG which contains more atom types.

12 More Visualization Results↩︎

In this section, we provide more visualizations of molecules generated from GeoLDM. Samples drawn from models trained on QM9 and DRUG are provided in 6 and 7 respectively. These examples are randomly generated without any cherry pick. Therefore, the generated geometries might be difficult to see in some figures due to imperfect viewing direction.

As shown in the two figures, the model is always able to generate realistic molecular geometries for both small and large size molecules. An outlier case is that the model occasionally generates disconnected components, as shown in the rightest column of 7, which happens more often when trained on the large molecule DRUG dataset. However, this phenomenon actually is not a problem and is common in all non-autoregressive molecule generative models [59], [60], and can be easily fixed by just filtering the smaller components.

References↩︎