Source-linked AI summary
Constrained Graph Variational Autoencoders for Molecule Design
Qi Liu, Miltiadis Allamanis, Marc Brockschmidt, Alexander L. Gaunt
TL;DR
Graph generation for molecule design requires modeling correlations among labeled nodes and edges while handling chemical validity. The paper proposes CGVAE, a VAE with GGNN encoder and decoder, sequential graph construction, and constraint masking. Experiments report state-of-the-art generation and optimization results, including valid molecule generation and latent-space property optimization.
Problem
Direct sampling from the joint distribution over labeled graph nodes and edges is intractable for reasonably sized graphs, motivating structured decompositions for molecule generation.
Method
CGVAE uses GGNNs in a VAE encoder and decoder, sequentially expands graphs, masks choices violating chemical rules, and shapes latent space for property optimization.
Results
CGVAE achieved state-of-the-art molecule generation and optimization results, matching chemically relevant statistics while generating valid, novel, and unique molecules.
Takeaways & Limitations
The model supports molecule design by combining chemically constrained generation with continuous latent-space optimization of numerical properties.
Takeaways & Limitations
The decoder objective estimates expectations from a very small set of enumerated generation traces, which can result in high-variance estimates; a set2set model could improve the bound.
Abstract
from arXiv · showhide
Graphs are ubiquitous data structures for representing interactions between entities. With an emphasis on the use of graphs to represent chemical molecules, we explore the task of learning to generate graphs that conform to a distribution observed in training data. We propose a variational autoencoder model in which both encoder and decoder are graph-structured. Our decoder assumes a sequential ordering of graph extension steps and we discuss and analyze design choices that mitigate the potential downsides of this linearization. Experiments compare our approach with a wide range of baselines on the molecule generation task and show that our method is more successful at matching the statistics of the original dataset on semantically important metrics. Furthermore, we show that by using appropriate shaping of the latent space, our model allows us to design molecules that are (locally) optimal in desired properties.
1 Introduction
The paper addresses graph generation for molecules by using a sequential probabilistic model that captures edge correlations while enforcing hard chemical constraints. It combines GGNNs with a VAE and shapes the latent space for molecule-property optimization.
- CGVAE integrates GGNNs into both the encoder and decoder of a VAE and incorporates hard domain-specific constraints for molecule generation.
- Experiments on pharmaceutical and photovoltaic molecule datasets show benefits in matching chemically relevant statistics and optimizing molecule properties.
- Graph generation is difficult because directly sampling the joint distribution over labeled nodes and edges is intractable for reasonably sized graphs.
- The model factors graph generation into sequential construction decisions because correlations between edges are important in real applications.
- The model conditions its learned component on the current generation state rather than the arbitrary path used to reach that state.
2 Related Work
Related work spans uncorrelated graph generation, sequential construction, and molecule-design methods. The paper situates its approach among learned graph generators and domain-specific molecular techniques.
- Uncorrelated generation: Uncorrelated methods sample edges independently, while learned GraphVAEs independently predict edge and node existence and labels.
- Sequential generation: Sequential graph generators use ordered construction decisions, often embedding the current graph with GGNNs and driving decisions with neural networks.
- Sequential generation: Dense GGNN implementations can scale as O(eN^2), whereas sparse implementations reduce the stated complexity to O(e^2).
- Molecule design: Molecule-design approaches include domain-knowledge-based simulation and heuristic search, as well as deep models that generate frequent molecular fragments.
- Molecule design: The paper’s masking approach helps ensure generated molecules can be converted to syntactically valid SMILES strings.
3 Generative Model
The generative model initializes node-level latent specifications, then grows a connected graph through focus and expansion decisions with GGNN-based state updates. Constraint masks restrict illegal edges and guide termination.
- Initialization and focus: Generation starts from N node vectors that form a latent specification, with N serving as an upper bound on the final graph’s node count.
- Initialization and focus: A focus function selects a node and an expand function selects edges to add, with focus implemented as a deterministic queue after a random initial choice.
- State representation: The expand function conditions only on the partial graph structure, avoiding dependence on the arbitrary generation trace and reducing overly deep computation graphs.
- Initialization and focus: Each node latent vector is paired with a sampled one-hot node-type representation, whose interpretable component supports hard-constraint enforcement.
- State representation: GGNN decoder updates recompute node representations after graph changes, using current neighborhoods and global representations of the initial and current graphs.
- Edge selection and termination: Edge candidates are represented using focus-node information, candidate-node information, graph distance, the initial specification, and the current graph state.
- Edge selection and termination: Binary masks forbid constraint-violating edges, while newly connected nodes enter the focus queue and stop-node selection closes the current node.
4 Training the Generative Model
The model trains a graph-structured VAE by encoding graphs into latent node distributions and decoding them through supervised breadth-first generation traces. Its objective handles trace multiplicity and shapes the latent space for local optimization of numerical graph properties.
- 4.1 Encoder: The encoder GGNN maps each graph node to a diagonal normal latent distribution, whose samples support decoder conditioning and KL regularization toward a standard Gaussian prior.
- 4.2 Decoder: Training supervises edge additions using breadth-first traces, although exact marginalization over all valid traces is computationally intractable.
- 4.2 Decoder: The objective reorganizes trace likelihoods into an expectation over valid partial-graph and focus-node states, weighted by each state’s trace multiplicity.
- 4.2 Decoder: A Monte Carlo estimate over enumerated traces avoids penalizing the model for choosing any valid expansion at a state, rather than requiring one exact sampled path.
- 4.3 Optimizing Graph Properties: The regression loss shapes the continuous latent space so gradient ascent can reach locally optimal points for a target property while remaining near the VAE’s standard-normal prior.
5 Application: Molecule Generation
For molecule generation, the model applies chemical preprocessing and domain-specific masks during sequential graph construction. Valency-aware masking prevents syntactically invalid molecular graphs while the evaluation compares generated statistics across datasets and models.
- 5.1 Datasets: The study uses QM9, ZINC, and CEPDB, with ZINC containing larger and more structurally complex drug-like compounds than QM9.
- 5.1 Datasets: Molecules are kekulized, hydrogens are removed, and remaining edge types are single, double, and triple covalent bonds.
- 5.3 Evaluation: The evaluation compares atom and bond statistics, ring counts, syntactic validity, novelty, and uniqueness across sampled molecules and generative models.
- 5.2 Valency masking: Valency masks use atom-specific chemical limits, counting double and triple bonds as two and three bonds, respectively.
- 5.2 Valency masking: The masks prevent excess valency, duplicate edges, self-loops, and connections to closed nodes during graph construction.
6 Experiments
Experiments compare CGVAE with baselines and ablations on molecule generation and latent-space optimization. CGVAE matches graph statistics while producing valid, novel, and unique molecules, and its design choices affect performance.
- Baseline comparison: CGVAE is evaluated against DeepGAR, an LSTM language model, ChemVAE, GrammarVAE, GraphVAE, and another graph model across three datasets.The models are sampled to obtain 20k molecules for comparison.
- Generation quality: CGVAE matches graph statistics while generating valid, novel, and unique molecules across all considered datasets.The evaluation includes syntactic validity, novelty, uniqueness, atom and bond counts, and ring counts.
- Generation quality: More than 10% of ZINC-like molecules generated by DeepGAR are invalid, whereas CGVAE’s masking mechanism guarantees molecule validity.The LSTM generates valid molecules effectively but does not support meaningful domain-specific constraints on SMILES.
- Ablations: Removing the distance feature harms performance on larger ZINC molecules, while independent edge generation and removing the decoder GGNN produce poor results.These ablations identify sequential decoding with GGNN node updates before each decision as key design choices.
- Directed molecule generation: CGVAE directs molecule generation toward high QED values by predicting QED from latent space and applying gradient ascent.Figure 5 depicts interpolation from low QED to a local maximum, showing generated molecules and predicted and RDKit-computed QED values.
7 Conclusion
The paper concludes that CGVAE combines sequential graph generation, GGNN-based encoding and decoding, chemical-rule masks, and latent-space optimization. It reports state-of-the-art molecule generation and optimization results while identifying richer chemistry metrics as future work.
- Conclusion: CGVAE is a sequential graph-generative VAE with GGNNs in its encoder and decoder, specialized to molecules using masks for chemical rules.The model also shapes its latent space for numerical-property optimization.
- Conclusion: The authors report state-of-the-art generation and optimization results for molecule design.They introduce basic statistics for validating generated molecules.
- Conclusion: Future work should define additional chemistry metrics to guide models and datasets for real-world molecule design tasks.
A Molecule Samples
The paper provides random samples from CGVAE for qualitative comparison with molecules from each training dataset. The sample display labels the two groups as data samples and CGVAE samples.
- Molecule Samples: The authors provide 25 random CGVAE samples for qualitative comparison with samples from each training dataset.
- Molecule Samples: The sample comparison distinguishes training-data examples from generated CGVAE examples.The labels are “Data Samples” and “CGVAE Samples.”
B Effect of multiple training paths
The experiments compare CGVAE trained with one versus 50 generation traces. Using more traces slightly improves ring-statistic matching but adds considerable preparation cost.
- Multiple training paths: The standard experiments use one breadth-first generation trace per molecule, while CGVAE (50) enumerates 50 traces and samples the same number of state transitions.The final dataset size remains the same in both settings.
- Multiple training paths: The benefit of 50 traces is not clearly worth the considerable computational overhead required to prepare the dataset.
C Additional Molecular Properties
This section presents molecular-property histograms for sampled molecules from the proposed method and two baselines, alongside additional QED optimization trajectories on ZINC-trained models.
- Histograms compare sampled-molecule properties for the proposed method, DeepGAR, and LSTM baselines.
- The reported properties include molecular weight, Bertz complexity, hydrogen-bond donor count, and hydrogen-bond acceptor count.Molecular weight is isotope-averaged; Bertz complexity is a topological molecular-complexity index.
- Additional QED optimization trajectories are provided for the model trained on the ZINC dataset.