Geometric Latent Diffusion Models for 3D Molecule Generation

Minkai Xu; Alexander Powers; Ron Dror; Stefano Ermon; Jure Leskovec

Source-linked AI summary

Geometric Latent Diffusion Models for 3D Molecule Generation

Minkai Xu, Alexander Powers, Ron Dror, Stefano Ermon, Jure Leskovec

arXiv:2305.01140v1May 2, 2023cs.LGq-bio.QM

TL;DR

Existing 3D molecular generative models have limited chemical validity and molecule size, while diffusion models usually operate directly on heterogeneous atomic features. GeoLDM learns a continuous, lower-dimensional latent space with invariant scalars and equivariant tensors, then performs diffusion there. It consistently improves generation across multiple benchmarks, including a reported 7% gain in valid rate for large biomolecules, and improves controllable generation.

Problem
Existing 3D molecule generators achieve low chemical validity and small molecule size, while current geometric diffusion models typically operate directly on diverse atomic features.
Method
GeoLDM uses autoencoders to encode molecular geometries into continuous point-structured latents containing invariant scalars and equivariant tensors, with diffusion operating in latent space.
Results
GeoLDM consistently achieves superior performance across unconditional and property-conditioned benchmarks, including up to 7% higher valid rate for large biomolecules.
Takeaways & Limitations
Latent modeling gives GeoLDM higher reported capacity for controllable generation and for producing structurally novel, chemically feasible molecules.
Takeaways & Limitations
The learned likelihood must remain invariant to molecular roto-translations under SE(3).

Abstract

from arXiv · show

Generative models, especially diffusion models (DMs), have achieved promising results for generating feature-rich geometries and advancing foundational science problems such as molecule design. Inspired by the recent huge success of Stable (latent) Diffusion models, we propose a novel and principled method for 3D molecule generation named Geometric Latent Diffusion Models (GeoLDM). GeoLDM is the first latent DM model for the molecular geometry domain, composed of autoencoders encoding structures into continuous latent codes and DMs operating in the latent space. Our key innovation is that for modeling the 3D molecular geometries, we capture its critical roto-translational equivariance constraints by building a point-structured latent space with both invariant scalars and equivariant tensors. Extensive experiments demonstrate that GeoLDM can consistently achieve better performance on multiple molecule generation benchmarks, with up to 7\% improvement for the valid percentage of large biomolecules. Results also demonstrate GeoLDM's higher capacity for controllable generation thanks to the latent modeling. Code is provided at \url{https://github.com/MinkaiXu/GeoLDM}.

1. Introduction

Geometric generative models support scientific discovery, but existing 3D molecule generators struggle with chemical validity and molecule size. GEOLDM addresses these limitations with latent diffusion over equivariant point-structured representations and reports stronger benchmark and controllable-generation performance.

Motivation: Feature-rich geometric generative models represent scientific objects as point clouds with coordinates and rich point-level features.Molecules can be represented as 3D atomic graphs, while proteins can be represented as proximity spatial graphs over amino acids.
Limitations of prior models: Existing autoregressive and flow-based 3D molecule generators show low chemical validity and small molecule size, reflecting limited generative-model capacity.
Limitations of prior models: Diffusion models improve generative modeling, but existing geometric approaches typically operate directly on heterogeneous atomic features such as charge and atom type.
GEOLDM: GEOLDM combines autoencoders with latent-space diffusion, mapping raw geometries into a lower-dimensional, smoother representational space.Its encoder produces latent variables, while diffusion models learn the latent distribution.
GEOLDM: GEOLDM constructs point-structured latents containing invariant scalars and equivariant tensors to preserve 3D roto-translational constraints.The encoding and decoding functions are parameterized with equivariant networks; scalar-only latents produced extremely poor generation quality.
Results: GEOLDM consistently outperforms competing methods across unconditional and property-conditioned benchmarks, including up to 7% higher valid rate for large biomolecules.Empirical studies also report improved controllable generation attributed to latent modeling.

2. Related Work

Prior work has explored latent generative modeling and 3D molecule generation through autoregressive, flow-based, and related approaches. These methods motivate latent-space modeling while leaving challenges in complex geometric molecule generation.

Latent Generative Models: Latent generative models improve modeling capacity by learning expressive distributions over latent spaces.Examples include discrete VQ-VAE latents with autoregressive priors and flow-based latent priors for non-autoregressive text generation.
Latent Generative Models: GEOLDM’s latent space also contains equivariant tensors, extending latent modeling to geometric features.
Molecule Generation in 3D: Recent 3D molecule-generation methods use autoregressive construction by sequentially attaching atoms or molecular fragments.Related frameworks have also been applied to structure-based drug design.

3. Background

The paper formulates 3D molecules as point clouds and defines unconditional and property-conditioned generation tasks. It reviews SE(3) equivariance and diffusion models, which add noise in a forward process and learn reverse denoising from Gaussian initialization.

Problem Definition: A molecule is represented as point clouds G = ⟨x, h⟩, with atom coordinates x ∈ R^N×3 and node features h ∈ R^N×d such as atomic types and charges.
Problem Definition: The two tasks are unconditional generation of diverse realistic 3D molecules and controllable generation conditioned on a desired property value.
Equivariance: SE(3) equivariance requires geometric outputs to transform consistently under 3D rotations and translations, while molecular features h remain invariant.Coordinates transform as Rx + t, and the learned likelihood must be invariant to roto-translations.
Diffusion Models: Diffusion models define forward and reverse Markov processes with intermediate variables sharing the data dimension.The forward process progressively perturbs data, while the reverse process learns to reconstruct clean samples.
Diffusion Models: The forward process uses a predefined noise schedule so samples approximately converge to a standard Gaussian distribution.The reverse generation process starts from Gaussian noise and incrementally denoises toward the target data distribution.
Diffusion Models: Diffusion training can maximize a variational lower bound, but directly optimizing that objective is known to suffer serious training instability.Simplified objectives are therefore commonly used in practice.

4. Method

GEOLDM combines geometric autoencoding with latent diffusion, using point-structured latents that preserve invariant and equivariant information under 3D transformations. Its training and sampling procedures operate in this latent space while maintaining SE(3)-consistent molecular generation.

Geometric Latent Diffusion Models: GEOLDM encodes molecular geometries into lower-dimensional latent variables and trains diffusion models to generate in the latent space.The decoder maps sampled latent codes back to molecular coordinates and features.
Geometric Autoencoding: Scalar-only latent variables perform poorly for 3D molecular generation because the latent representation must preserve roto-translational equivariance.The point-structured latent coordinates can also represent the group actions needed to align decoded outputs with inputs.
Geometric Autoencoding: Point-structured latents combine 3-dimensional equivariant coordinates with k-dimensional invariant features for each node.Equivariant graph neural networks parameterize the encoder and decoder to extract both types of embeddings.
Geometric Autoencoding: Regularization combines reconstruction loss with KL-based or early-stopping variants to prevent arbitrarily high latent variance and support latent diffusion.The autoencoder objective uses reconstruction and regularization terms, with continuous and discrete features handled by L2 loss or cross-entropy.
Geometric Latent Diffusion Models: The latent diffusion process uses equivariant denoising dynamics and center-of-gravity constraints so intermediate latent coordinates remain translation-invariant.The initial latent distribution is invariant, transitions are equivariant, and intermediate coordinate states satisfy a zero-sum constraint.

5. Experiments

The experiments evaluate GEOLDM on unconditional and property-conditioned 3D molecule generation using standard datasets, metrics, and competitive baselines. GEOLDM outperforms the baselines across reported metrics, including improved controllable generation and up to 7% higher validity for large molecules.

Experimental Setup: GEOLDM is evaluated on molecular modeling, generation, and controllable molecule generation tasks using QM9 and larger DRUG molecules.QM9 contains 130k small molecules with 3D structures and quantum properties; the DRUG dataset contains larger, more complex structures.
Molecular Modeling and Generation: Up to 7% higher validity is achieved for large molecule generation with latent variables that combine invariant and equivariant features.Table 1 reports atom stability, molecule stability, validity, and validity×uniqueness, with higher values indicating better quality.
Experimental Setup: Evaluation measures atom stability, molecule stability, validity, uniqueness, and validity×uniqueness for generated compounds.Stability assesses valency, while validity and uniqueness measure chemically valid and non-duplicate molecules among generated samples.
Molecular Modeling and Generation: GEOLDM outperforms competitive baselines on all reported molecular-generation metrics, demonstrating stronger modeling of chemically realistic geometries.The comparison uses 10,000 generated samples per method; DRUG molecule-stability and uniqueness metrics are omitted because they are nearly 0% and 100%, respectively, across methods.
Controllable Molecule Generation: Interpolating polarizability while fixing reparameterization noise produces geometry changes consistent with the relationship between molecular isometry and polarizability.The qualitative assessment varies polarizability α; less isometrically molecular geometries typically have larger α values.
Controllable Molecule Generation: GEOLDM improves property-conditioned generation across all six QM9 property metrics relative to EDM and other baselines.The evaluation uses MAE between desired and predicted properties, with smaller gaps to the predictor’s QM9 reference indicating better conditioning.

6. Conclusion and Future Work

GEOLDM is presented as a lower-dimensional latent diffusion framework for molecular geometry generation that preserves roto-translational equivariance through structured latent codes. The authors report stronger modeling of chemically realistic molecules and identify extensions to drug discovery and larger 3D geometries.

6. Conclusion and Future Work: GEOLDM applies diffusion models in a continuous, lower-dimensional latent space rather than directly modeling high-dimensional multimodal atom features.Autoencoder-based latent representations are used to support diffusion modeling of molecular geometries.
6. Conclusion and Future Work: Point-structured latent codes combine invariant scalars and equivariant tensors to maintain roto-translational equivariance.The conclusion identifies this construction as enabling effective latent representations while preserving the relevant geometric symmetry.
6. Conclusion and Future Work: The authors propose extending GEOLDM to protein-targeted drug discovery and more challenging geometries such as peptides and proteins.These are stated as future applications of the general framework rather than demonstrated results in this section.
6. Conclusion and Future Work: An invariant encoder alone can reconstruct a transformed geometry as the original geometry, motivating an equivariant function that selects canonical decoded representatives.For a transformation in the SE(3) group, encoding is unchanged, so decoding does not recover the transformed input without an appropriate equivariant construction.

B. Formal Statements and Proofs

The formal appendix analyzes how equivariant latent diffusion and decoding yield invariant marginal distributions over generated molecular geometries. The argument assumes an invariant initial latent distribution and equivariant transition and decoder distributions.

B.2. Invariant Marginal Distribution: Proposition 4.3: The proof focuses on tensor features z and omits scalar inputs h because the analysis concerns equivariant properties.
B.2. Invariant Marginal Distribution: Proposition 4.3: The formal justification is presented as a self-contained proof of invariant marginal distributions for the equivariant probabilistic diffusion model.
B.2. Invariant Marginal Distribution: Proposition 4.3: Assuming equivariant transition distributions, the marginal latent distributions p(z_t) remain invariant, including the final latent distribution p(z_0).
B.2. Invariant Marginal Distribution: Proposition 4.3: An equivariant decoder distribution p(x|z_0) transfers latent invariance to the induced geometry distribution p(x).
B.2. Invariant Marginal Distribution: Proposition 4.3: The construction begins with an invariant initial distribution p(z_T), specified as N(0, I) under rotations.
B.2. Invariant Marginal Distribution: Proposition 4.3: The argument considers orthogonal transformations R and requires the relevant distributions to be invariant for all such R.
B.2. Invariant Marginal Distribution: Proposition 4.3: The proof’s change-of-variables step tracks how transformed tensor coordinates and the determinant of R affect the distribution.

B.1. Relationship to SE(3)-invariant Likelihood: Theorem 4.2

Theorem 4.2 establishes that GEOLDM’s objective is a variational lower bound to log-likelihood that remains invariant under 3D rotations and translations.

Theorem 4.2 states that L = Lrecon + LLDM is an SE(3)-invariant variational lower bound to log-likelihood.The result holds with certain timestep weights w(t).
The objective satisfies L(x, h) = L(Rx + t, h) for every rotation R and translation t.
The reconstruction and latent diffusion terms are separately shown to be invariant under the relevant equivariant transformations.The reconstruction term follows from equivariant encoder and decoder distributions, while the latent diffusion term follows from equivariant transition distributions.
Translation invariance follows naturally because the model operates on the zero-mean subspace.

C. Model Architecture Details

GEOLDM uses fully connected EGNNs to process molecular point clouds while preserving the required equivariant structure through coordinate and feature updates.

Molecular geometries are represented as fully connected graphs whose nodes contain atom coordinates and atomic features.The model does not specify connecting bonds and instead models interactions between all atoms.
EGNN layers update coordinates and node features using learned edge, node, coordinate, and interaction functions.Pairwise distances and optional edge features contribute to message reweighting and propagation.
Stacking EGCL layers produces an EGNN satisfying the required equivariant constraints.All learnable functions are parameterized by multilayer perceptrons.

D. Featurization and Implementation Details

The implementation uses dataset-specific atomic features, EGNN backbones, an early-stopped encoder, and low-dimensional latent invariant features.

QM9 uses atom types H, C, N, O, F and integer-valued atom charges, whereas Drugs uses only atom types.Molecules are preprocessed with the open-source RDKit software.
Reported results use ES-reg, with the encoder trained for 1000 warm-up iterations and then fixed.
All encoder, latent diffusion, and decoder networks use EGNNs implemented in PyTorch.
The latent invariant feature dimension k is set to 1 for QM9 and 2 for DRUG.These settings reduce the atomic feature dimension.

E. Ablation Studies

Ablations examine autoencoder regularization and latent invariant dimension, showing that ES-reg is preferable to KL-reg and that lower latent dimensionality generally performs better.

The ablation study evaluates autoencoder regularization methods and latent invariant feature dimension k using 10000 generated samples per setting.
Autoencoder regularization: KL-reg produces extremely poor performance and unstable training, whereas constraining encoded latent scale through encoder early stopping works well in practice.The authors attribute KL-reg difficulties to scattered equivariant latent means and extremely small variances.
Latent dimension: GEOLDM generally performs better with lower k, supporting the motivation that lower dimensionality reduces latent diffusion modeling complexity.Performance for k = 1 and k = 2 is very similar on QM9, so k = 1 is reported there.

F. More Visualization Results

The section visualizes randomly generated GeoLDM molecules trained on QM9 and DRUG, without cherry-picking. The examples show realistic geometries across small and large molecules, while DRUG samples occasionally contain disconnected components that can be filtered.

Randomly generated molecules from GeoLDM trained on QM9 and DRUG are shown without cherry-picking.Viewing direction may make some geometries difficult to see.
The generated samples show realistic molecular geometries for both small and large molecules.
Disconnected components occasionally appear, more often for models trained on the large-molecule DRUG dataset.The authors state that filtering smaller components can easily fix this issue.