Source-linked AI summary
Junction Tree Variational Autoencoder for Molecular Graph Generation
Wengong Jin, Regina Barzilay, Tommi Jaakkola
TL;DR
The paper addresses the under-explored problem of generating molecular graphs from continuous representations for molecule design. It introduces a junction tree variational autoencoder that first generates a scaffold of valid chemical substructures and then assembles them into a molecular graph. Across generation and optimization tasks, the method outperforms prior work, including 100% validity from prior sampling and a 30% relative gain over baselines for desired-property discovery.
Problem
Molecular graph generation from latent representations is under-explored, while SMILES can represent similar molecules with markedly different strings.
Method
A junction tree variational autoencoder generates a tree-structured scaffold of valid substructures before assembling those components into a molecular graph.
Results
Across molecule generation and optimization tasks, the model outperforms previous work, producing 100% valid prior samples and achieving a 30% relative gain over baselines for desired-property discovery.
Takeaways & Limitations
Using valid substructures as building blocks supports molecular graph generation while maintaining validity during incremental construction.
Takeaways & Limitations
The predefined cluster vocabulary cannot guarantee that a junction tree exists for every arbitrary molecule, although it covers most test molecules empirically.
Abstract
from arXiv · showhide
We seek to automate the design of molecules based on specific chemical properties. In computational terms, this task involves continuous embedding and generation of molecular graphs. Our primary contribution is the direct realization of molecular graphs, a task previously approached by generating linear SMILES strings instead of graphs. Our junction tree variational autoencoder generates molecular graphs in two phases, by first generating a tree-structured scaffold over chemical substructures, and then combining them into a molecule with a graph message passing network. This approach allows us to incrementally expand molecules while maintaining chemical validity at every step. We evaluate our model on multiple tasks ranging from molecular generation to optimization. Across these tasks, our model outperforms previous state-of-the-art baselines by a significant margin.
1. Introduction
Molecular graph generation remains a harder, under-explored challenge than graph encoding, while SMILES-based approaches can obscure molecular similarity and produce invalid intermediate states. The paper proposes a junction tree variational autoencoder that generates valid substructure scaffolds before assembling molecular graphs, and reports strong results across generation and optimization tasks.
- Molecular graph generation from latent representations remains under-explored despite extensive work on molecular graph encoding.
- Similar molecules can have markedly different SMILES strings, limiting the usefulness of string representations for learning molecular similarity.Figure 1 reports an edit distance of 22, or 50.5% of the sequence, between two almost identical molecules.
- Node-by-node graph generation can create chemically invalid intermediate molecules before the completed graph can be validated.The paper contrasts this with structure-by-structure generation, which avoids the invalid intermediate states shown in Figure 2.
- The proposed junction tree variational autoencoder first generates a tree-structured scaffold of valid chemical substructures and then assembles those components into a molecular graph.The substructures are automatically extracted from the training set and used as building blocks.
- Across molecular generation and optimization tasks, the model produces 100% valid molecules from a prior distribution and achieves a 30% relative gain over baselines for discovering molecules with desired properties.
2. Junction Tree Variational Autoencoder
JT-VAE represents molecules with complementary junction-tree and graph structures, then decodes them in two stages: tree reconstruction followed by graph assembly. Valid chemical components and tree-dependent context support chemically feasible, structured decoding.
- Model representation: JT-VAE extends variational autoencoders to molecular graphs using valid subgraphs as reusable building blocks.The same component vocabulary supports encoding molecules and decoding latent vectors into molecular graphs.
- Model representation: A molecule is encoded into z = [zT, zG], where zT captures tree structure and cluster identities while zG captures fine-grained connectivity.The two latent parts are produced by separate tree and graph encoders.
- Molecular tree decomposition: Tree decomposition contracts molecular structures into a cycle-free junction tree whose nodes are induced subgraphs satisfying union and running-intersection constraints.The molecular cluster vocabulary contains cycles and single edges; for a standard 250K-molecule dataset, |X| = 780.
- Encoders: The tree encoder uses bottom-up and top-down message passing, retaining the rooted-tree representation rather than averaging node states.The schedule follows tree belief propagation, and root pooling preserves the order needed by the tree decoder.
- Tree decoder: The decoder reconstructs a junction tree top-down in depth-first order, predicting child topology and labels before backtracking when branches end.During sampling, chemically incompatible labels are masked from each node’s candidate set.
- Graph decoder: Graph decoding selects candidate graphs compatible with the predicted tree using a cluster-decomposed scoring function and tree-dependent positional messages.These messages distinguish candidate assemblies that share local cluster structure but differ in their placement within the tree.
- Complexity: The junction-tree decomposition makes assembly linear in the number of clusters, with approximately four candidates per node on average after pruning and graph deduplication.Any two clusters share at most two atoms or one bond, limiting the merge operations.
3. Experiments
The experiments evaluate JT-VAE on reconstruction and validity, Bayesian optimization, and constrained molecule optimization. Across these tasks, the model produces valid molecules, discovers higher-scoring molecules, and improves properties while meeting similarity constraints.
- JT-VAE is evaluated on molecule reconstruction and validity, Bayesian optimization, and constrained molecule optimization.The constrained task modifies existing molecules while limiting deviation from the original.
- Molecule Reconstruction and Validity: JT-VAE always produces valid molecules when sampled from the prior distribution, whereas atom-by-atom generation achieves 89.2% validity.The model uses valid substructures as building blocks, avoiding invalid intermediate states.
- Molecule Reconstruction and Validity: JT-VAE’s learned molecular neighborhoods avoid huge rings and include closely resembling molecules with identical tree structures but different cluster attachments.These observations demonstrate smoothness of the learned molecular embeddings.
- Bayesian Optimization: JT-VAE finds significantly better molecular property scores than previous methods and discovers over 50 molecules with scores above 3.50.Its sparse Gaussian process also has better predictive performance when trained on JT-VAE embeddings.
- Constrained Optimization: With similarity threshold δ = 0.4, JT-VAE finds similar molecules about 80% of the time with an average improvement of 0.84.The unconstrained setting has the best average improvement but often proposes dissimilar molecules.
4. Related Work
Related work includes SMILES-based molecule generation, direct graph generation, graph encoders, and tree-structured models. JT-VAE differs by generating chemically valid substructures in a coarse-to-fine junction-tree framework.
- Molecule Generation: Earlier molecule-generation methods primarily generate SMILES strings, which can produce invalid molecules and do not capture molecular similarity well.Grammar-based approaches add syntactic and semantic constraints to SMILES generation.
- Molecule Generation: JT-VAE generates molecular graphs directly and enforces chemical validity through coarse-to-fine generation from valid substructures.This contrasts with atom-by-atom graph generation, which can pass through chemically invalid intermediate states.
- Graph- and Tree-Structured Encoders: The graph and tree encoders relate to message-passing and graphical-model perspectives, while the tree encoder propagates information both bottom-up and top-down.This differs from tree models that use only bottom-up transformations.
- Tree-Structured Models: JT-VAE’s tree decoder separates topology and label prediction, generates nodes depth-first, and adds bottom-up information propagation during decoding.Its message-passing formulation differs from node-based forward-backward tree models.
5. Conclusion
The paper presents a junction tree variational autoencoder for molecular graph generation, reporting strong performance in molecule generation and optimization while identifying general low-treewidth graphs as future work.
- The paper presents a junction tree variational autoencoder for generating molecular graphs.
- The method significantly outperforms previous work in molecule generation and optimization.
- Future work aims to generalize the method to general low-treewidth graphs.
A. Tree Decomposition
The tree decomposition represents molecules through non-ring bonds, simple rings, and multi-cluster intersection atoms, then constructs a maximum spanning tree over their cluster graph.
- Non-ring bonds and simple rings form the initial cluster sets for molecular tree decomposition.Simple rings are extracted with RDKit’s GetSymmSSSR function.
- The junction tree is not unique when the molecular cluster graph contains cycles, introducing uncertainty for probabilistic modeling.The method reduces this variation through additional decomposition choices described in the passage.
- Rings sharing more than two atoms are merged because they form bridged compounds.
- Atoms belonging to the intersection of three or more clusters form the V0 set.
- The cluster graph connects overlapping clusters and returns their maximum spanning tree.Edges receive weight ∞ when either endpoint is in V0 and weight 1 otherwise.
B. Stereochemistry
Because molecules include three-dimensional stereochemical configurations beyond their 2D graphs, JT-VAE generates the 2D structure first and ranks enumerated stereoisomers afterward.
- Stereoisomers share a 2D structure but differ in the three-dimensional orientations of their atoms.Stereochemical feasibility cannot be simply encoded as context-free or attribute grammars.
- JT-VAE first generates a molecule’s 2D structure, then enumerates its stereoisomers with RDKit.
- The model reconstructs the final 3D structure by selecting the stereoisomer with the highest scoring function value.
- Post-ranking is efficient because only a few atoms typically have stereochemical variations.
C. Training Details
Training uses a 780-component vocabulary extracted from 240K ZINC molecules, 450-dimensional hidden states, and a 56-dimensional latent bottleneck.
- The component vocabulary contains 780 entries collected by decomposing 240K molecules from the ZINC dataset.
- JT-VAE uses hidden states of dimension 450 and a latent bottleneck of dimension 56.
- Atom features include atom type, degree, formal charge, and chiral configuration, while bond features include bond type, ring membership, and cis-trans configuration.
- The graph encoder and decoder each run three iterations of neural message passing.
- PyTorch implements the neural components and RDKit processes the molecules.
D. More Experimental Results
The experiments examine prior sampling, latent-space neighborhoods, Bayesian optimization, and constrained molecular modification. Results show diverse generated structures and that stricter similarity constraints preserve more of the original molecule.
- Sampled Molecules: 250 molecules sampled from prior distribution N(0, I) showed rich variety and structural complexity rather than trivial chain-like solutions.The authors use this sample to support the soundness of the model’s prior-validity improvement.
- Neighborhood Visualization: Neighborhood visualization decodes latent-space grids generated by moving along combinations of two random orthogonal unit vectors.The procedure encodes a molecule, defines two grid axes in latent space, and decodes the resulting latent vectors.
- Bayesian Optimization: Bayesian optimization uses a sparse Gaussian process with 500 inducing points and five batch iterations to propose new latent vectors.Each iteration proposes 50 latent vectors, decodes them into molecules, and adds them to the training set.
- Constrained Optimization: Constrained optimization trains a property predictor jointly with the VAE to estimate y(m) = logP(m) − SA(m) from molecular latent embeddings.The predictor is a one-hidden-layer feed-forward network with hidden dimension 450 and tanh activation.
- Constrained Optimization: Tighter similarity constraints force modified molecules to preserve more of the original structure, while inaccurate property prediction can produce decreased-property examples.The comparisons use thresholds δ = 0.2, 0.4, and 0.6 and report improved and decreased-property pairs for each group.