[Paper Review] Geometric Latent Diffusion Models for 3D Molecule Generation

Minkai Xu, Alexander K. Powers+3 more|arXiv (Cornell University)|May 2, 2023

Model Reduction and Neural Networks32 citations

TL;DR

GeoLDM introduces a latent diffusion framework with SE(3)-equivariant latent spaces (invariant scalars and equivariant tensors) for 3D molecule generation, achieving higher validity and controllability than prior methods.

ABSTRACT

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}.

Motivation & Objective

• Advance 3D molecule generation by modeling in a latent, geometry-aware space that respects roto-translation equivariance.
• Introduce a geometric autoencoder that produces latent codes with both invariant and equivariant components.
• Train diffusion models in the latent space to improve validity, diversity, and controllable generation of 3D molecules.
• Show that latent GeoLDM improves generation metrics on QM9 and DRUG benchmarks and enables controllable generation.
• Provide public code for replication and further research.

Proposed method

• Construct a geometric autoencoder that maps 3D molecular geometries to a latent space with per-node invariant (z_h) and equivariant (z_x) components using equivariant graph neural networks (EGNNs).
• Ensure SE(3) invariance by designing latent transitions to be equivariant and latent priors to be invariant, so the resulting marginal is SE(3)-invariant.
• Train a latent diffusion model (LDM) on the structured latent space with a denoising network epsilon_theta implemented as time-conditional EGNNs.
• Adopt a two-stage training: first train the geometric AE with reconstruction and regularization, then fix the encoder and train the latent DM; use p_theta(z_x,z_h) for generation followed by decoding p_xi(x,h|z_x,z_h).
• Support controllable generation by conditioning the denoiser and encoder/decoder on property vectors s (concatenation rather than cross-attention).
• Optionally sample molecule sizes N from the training distribution p(N) and generate accordingly.

Figure 1: Illustration of GeoLDM . The encoder ${\mathcal{E}}_{\phi}$ encodes molecular features ${\mathbf{x}},{\mathbf{h}}$ into equivariant latent variables ${\mathbf{z}}_{\textnormal{x}},{\mathbf{z}}_{\textnormal{h}}$ , and the latent diffusion transitions $q({\mathbf{z}}_{{\textnormal{x}},t},{\m

Experimental results

Research questions

• RQ1 Can a latent diffusion model operating on a geometry-aware latent space (with both invariant and equivariant components) improve 3D molecule generation quality?
• RQ2 Does enforcing SE(3) equivariance in the latent space lead to higher chemical validity and structural diversity compared to latent models that operate in the raw atom feature space?
• RQ3 To what extent can GeoLDM enable controllable generation for target chemical properties, and how does it perform on standard 3D molecule benchmarks?
• RQ4 How does GeoLDM perform on unconditional vs. conditional generation across QM9 and DRUG datasets?

Key findings

• GeoLDM achieves higher generation quality across benchmarks, with up to 7% higher valid rate for large biomolecules on QM9/DRUG compared to baselines.
• Latent modeling with both invariant and equivariant latent variables yields superior validity and diversity relative to models using only scalars.
• Conditional generation experiments show improved controllability due to latent space design and simple property-conditioning via concatenation.
• GeoLDM demonstrated strong SE(3)-invariant properties in the learned marginal distribution, aiding generalization to roto-translations.
• Empirical results show GeoLDM surpasses several prior 3D molecule generation methods on multiple metrics (validity, stability, and validity×uniqueness).
• The authors provide code at the project site for replication and further exploration.

Figure 2: Molecules generated by GeoLDM trained on QM9 (left three) and DRUG (right four).