[Paper Review] GraphEBM: Molecular Graph Generation with Energy-Based Models

Meng Liu, Keqiang Yan+2 more|arXiv (Cornell University)|Jan 31, 2021

Machine Learning in Materials Science20 citations

TL;DR

GraphEBM proposes a permutation-invariant energy-based model for molecular graph generation using graph neural networks to parameterize the energy function, enabling training via Langevin dynamics and flexible, goal-directed generation by adjusting energy scores based on molecular properties. It achieves state-of-the-art performance in validity, uniqueness, and novelty across multiple benchmarks, including 100% validity and 97.01% novelty on QM9 and 99.96% validity and 100% novelty on ZINC250k.

ABSTRACT

We note that most existing approaches for molecular graph generation fail to guarantee the intrinsic property of permutation invariance, resulting in unexpected bias in generative models. In this work, we propose GraphEBM to generate molecular graphs using energy-based models. In particular, we parameterize the energy function in a permutation invariant manner, thus making GraphEBM permutation invariant. We apply Langevin dynamics to train the energy function by approximately maximizing likelihood and generate samples with low energies. Furthermore, to generate molecules with a desirable property, we propose a simple yet effective strategy, which pushes down energies with flexible degrees according to the properties of corresponding molecules. Finally, we explore the use of GraphEBM for generating molecules with multiple objectives in a compositional manner. Comprehensive experimental results on random, goal-directed, and compositional generation tasks demonstrate the effectiveness of our proposed method.

Motivation & Objective

• To address the lack of permutation invariance in existing molecular graph generation methods, which introduces bias in generative models.
• To develop a novel energy-based model (GraphEBM) that inherently preserves permutation invariance through a graph neural network parameterization of the energy function.
• To enable goal-directed molecular generation by modifying energy scores based on desired molecular properties with flexible control.
• To explore compositional generation of molecules with multiple objectives simultaneously, a capability not supported by prior methods.
• To demonstrate superior performance in random, goal-directed, and multi-objective molecular generation tasks.

Proposed method

• Parameterize the energy function Eθ(X,A) using a graph neural network that processes node features X and adjacency tensors A, ensuring permutation invariance through symmetric aggregation.
• Train the energy function using Langevin dynamics to approximate maximum likelihood, sampling from the unnormalized Boltzmann distribution e^(-Eθ(X,A)).
• Introduce a flexible strategy to lower energy scores based on target molecular properties (e.g., QED, pLogP), enabling goal-directed generation by modifying the energy function with property-dependent terms.
• Enable compositional generation by combining multiple property-specific energy functions into a single composite energy function, allowing simultaneous optimization for multiple objectives.
• Use virtual nodes and virtual edges to standardize graph dimensions across molecules with varying sizes, preserving structural consistency during training and generation.
• Leverage existing molecules from the training set as initialization points for Langevin dynamics to improve sample quality and convergence in goal-directed generation.

Experimental results

Research questions

• RQ1 Can an energy-based model be designed to preserve permutation invariance in molecular graph generation, eliminating bias from node ordering?
• RQ2 Can Langevin dynamics effectively train an energy-based model for molecular graph generation while maintaining high validity and novelty?
• RQ3 Can the energy function be modified to enable flexible, property-directed generation with adjustable degrees of influence on target properties?
• RQ4 Can GraphEBM perform compositional generation of molecules with multiple desired properties simultaneously, beyond what existing methods can achieve?
• RQ5 How does GraphEBM compare to state-of-the-art models in terms of validity, uniqueness, and novelty across diverse molecular generation benchmarks?

Key findings

• GraphEBM achieves 100.00% validity, 97.90% uniqueness, and 97.01% novelty on the QM9 dataset, outperforming existing methods in all metrics.
• On the ZINC250k dataset, GraphEBM attains 99.96% validity, 98.79% uniqueness, and 100.00% novelty, demonstrating high-quality generation with near-perfect validity and full novelty.
• In goal-directed generation for QED, GraphEBM discovers novel molecules with QED scores up to 0.92, significantly improving over baseline methods.
• For constrained property optimization with Tanimoto similarity ≥0.8, GraphEBM successfully generates molecules with improved pLogP scores while maintaining structural similarity to initial molecules.
• GraphEBM enables compositional generation by combining multiple property-specific energy functions, producing molecules that satisfy multiple objectives simultaneously, a capability absent in prior approaches.
• Visualization of the generation process confirms that GraphEBM learns meaningful molecular structures through iterative Langevin sampling, with stable convergence and plausible chemical connectivity.