[Paper Review] Learning Latent Space Energy-Based Prior Model for Molecule Generation

Bo Pang, Tian Han+1 more|arXiv (Cornell University)|Oct 19, 2020

Machine Learning in Materials Science9 references7 citations

TL;DR

This paper proposes a latent space energy-based prior model that improves molecule generation using character-level SMILES strings, enabling valid and unique molecular generation without explicit chemical rule enforcement. By learning a data-driven, expressive prior in the latent space, the model achieves state-of-the-art performance in validity, uniqueness, and molecular property distribution matching, demonstrating that complex chemical rules can be implicitly learned from data.

ABSTRACT

Deep generative models have recently been applied to molecule design. If the molecules are encoded in linear SMILES strings, modeling becomes convenient. However, models relying on string representations tend to generate invalid samples and duplicates. Prior work addressed these issues by building models on chemically-valid fragments or explicitly enforcing chemical rules in the generation process. We argue that an expressive model is sufficient to implicitly and automatically learn the complicated chemical rules from the data, even if molecules are encoded in simple character-level SMILES strings. We propose to learn latent space energy-based prior model with SMILES representation for molecule modeling. Our experiments show that our method is able to generate molecules with validity and uniqueness competitive with state-of-the-art models. Interestingly, generated molecules have structural and chemical features whose distributions almost perfectly match those of the real molecules.

Motivation & Objective

• Address the challenge of generating chemically valid and unique molecules using simple character-level SMILES representations, which often lead to invalid or duplicate outputs in standard language model-based approaches.
• Overcome the limitations of prior VAE-based models that rely on a simple isotropic Gaussian prior, which may hinder effective learning of complex chemical rules.
• Demonstrate that an expressive latent space prior can implicitly learn chemical rules from data, eliminating the need for explicit valency checks or fragment-based representations.
• Achieve performance competitive with state-of-the-art graph-based and fragment-based models in validity, uniqueness, and molecular property distribution matching.
• Show that the model can generate molecules with structural and chemical features (e.g., logP, QED, SAS) that closely mirror real data distributions without explicit supervision on these properties.

Proposed method

• Propose a hybrid generative model combining a conditional autoregressive generator (LSTM-based) with a learned latent space energy-based prior, replacing the standard isotropic Gaussian prior in VAEs.
• Formulate the prior as a Gibbs distribution: $ p_{ar{\alpha}}(z) = \frac{1}{Z(\alpha)} \exp(f_{\alpha}(z)) p_0(z) $, where $ f_{\alpha}(z) $ is a neural network that learns data-driven corrections to the base Gaussian distribution.
• Train the model using a variational inference framework with a reparameterization trick and short-run Markov Chain Monte Carlo (MCMC) to approximate posterior expectations.
• Use a two-stage learning process: first approximate the posterior $ p_{\theta}(z|x) $ via MCMC, then update the model parameters $ \alpha $ and $ \beta $ using gradient estimates based on the approximate posterior.
• Leverage the low-dimensional latent space to efficiently perform MCMC sampling, improving mixing and convergence during training.
• The generator network is a single-layer LSTM that autoregressively predicts one-hot encoded SMILES tokens conditioned on the latent code, enabling efficient and effective molecule generation.

Experimental results

Research questions

• RQ1 Can a latent space energy-based prior model generate chemically valid and unique molecules using only character-level SMILES strings, without explicit chemical rule enforcement?
• RQ2 Does learning a complex, data-driven prior in the latent space enable implicit capture of complex chemical rules, such as valency and aromaticity, from raw SMILES data?
• RQ3 How well do the generated molecules match the distribution of real molecular properties (e.g., logP, QED, SAS) compared to real data, especially without explicit supervision on these properties?
• RQ4 Can this approach achieve performance on par with state-of-the-art graph-based and fragment-based models in validity, uniqueness, and property matching?
• RQ5 To what extent does the model’s expressivity surpass that of standard VAEs with isotropic Gaussian priors in generating realistic and diverse molecules?

Key findings

• The proposed model achieves 95.5% validity and 100% uniqueness in generating 10,000 molecules from the ZINC dataset, outperforming standard LM-based models and matching the best-performing fragment-based and graph-based models.
• The model generates molecules with molecular property distributions (logP, QED, SAS) that closely match those of real molecules in the training data, even without explicit supervision on these properties.
• Compared to FragmentVAE, the proposed model shows superior performance in matching molecular property distributions, indicating better generalization and expressivity.
• The model achieves 100% novelty, meaning all generated molecules are unique and not present in the training set, demonstrating strong generalization and diversity.
• The results suggest that complex chemical rules can be implicitly learned from data via a flexible latent prior, reducing reliance on explicit constraints like valency checks in graph-based models.
• The model’s performance demonstrates that expressive latent priors can compensate for the simplicity of character-level SMILES representation, enabling high-quality molecule generation without architectural complexity.