Constrained Graph Variational Autoencoders for Molecule Design

Introduction

Structured objects such as program source code, physical systems, chemical molecules and even 3D scenes are often well represented using graphs . Recently, considerable progress has been made on building discriminative deep learning models that ingest graphs as inputs . Deep learning approaches have also been suggested for graph generation. More specifically, generating and optimizing chemical molecules has been identified as an important real-world application for this set of techniques .

In this paper, we propose a novel probabilistic model for graph generation that builds gated graph neural networks (GGNNs) into the encoder and decoder of a variational autoencoder (VAE) . Furthermore, we demonstrate how to incorporate hard domain-specific constraints into our model to adapt it for the molecule generation task. With these constraints in place, we refer to our model as a constrained graph variational autoencoder (CGVAE). Additionally, we shape the latent space of the VAE to allow optimization of numerical properties of the resulting molecules. Our experiments are performed with real-world datasets of molecules with pharmaceutical and photo-voltaic applications. By generating novel molecules from these datasets, we demonstrate the benefits of our architectural choices. In particular, we observe that (1) the GGNN architecture is beneficial for state-of-the-art generation of molecules matching chemically relevant statistics of the training distribution, and (2) the semantically meaningful latent space arising from the VAE allows continuous optimization of molecule properties .

The key challenge in generating graphs is that sampling directly from a joint distribution over all configurations of labeled nodes and edges is intractable for reasonably sized graphs. Therefore, a generative model must decompose this joint in some way. A straightforward approximation is to ignore correlations and model the existence and label of each edge with independent random variables . An alternative approach is to factor the distribution into a sequence of discrete decisions in a graph construction trace . Since correlations between edges are usually crucial in real applications, we pick the latter, sequential, approach in this work. Note that for molecule design, some correlations take the form of known hard rules governing molecule stability, and we explicitly enforce these rules wherever possible using a technique that masks out choices leading to illegal graphs . The remaining “soft” correlations (e.g. disfavoring of small cycles) are learned by our graph structured VAE.

By opting to generate graphs sequentially, we lose permutation symmetry and have to train using arbitrary graph linearizations. For computational reasons, we cannot consider all possible linearizations for each graph, so it is challenging to marginalize out the construction trace when computing the log-likelihood of a graph in the VAE objective. We design a generative model where the learned component is conditioned only on the current state of generation and not on the arbitrarily chosen path to that state. We argue that this property is intuitively desirable and show how to derive a bound for the desired log-likelihood under this model. Furthermore, this property makes the model relatively shallow and it is easy scale and train.

Related Work

Generating graphs has a long history in research, and we consider three classes of related work: Works that ignore correlations between edges, works that generate graphs sequentially and works that emphasize the application to molecule design.

The Erdős-Rényi $G(n,p)$ random graph model is the simplest example of this class of algorithms, where each edge exists with independent probability $p$. Stochastic block models add community structure to the Erdős-Rényi model, but retain uncorrelated edge sampling. Other traditional random graph models such as those of Albert and Barabási , Leskovec et al. do account for edge correlations, but they are hand-crafted into the models. A more modern learned approach in this class is GraphVAEs , where the decoder emits independent probabilities governing edge and node existence and labels.

Sequential generation

Johnson sidesteps the issue of permutation symmetry by considering the task of generating a graph from an auxiliary stream of information that imposes an order on construction steps. This work outlined many ingredients for the general sequential graph generation task: using GGNNs to embed the current state of generation and multi-layer perceptrons (MLPs) to drive decisions based on this embedding. Li et al. uses these ingredients to build an autoregressive model for graphs without the auxiliary stream. Their model gives good results, but each decision is conditioned on a full history of the generation sequence, and the authors remark on stability and scalability problems arising from the optimization of very deep neural networks. In addition, they describe some evidence for overfitting to the chosen linearization scheme due to the strong history dependence. Our approach also uses the ingredients from Johnson , but avoids the training and overfitting problems using a model that is conditioned only on the current partial graph rather than on full generation traces. In addition, we combine Johnson’s ingredients with a VAE that produces a meaningful latent space to enable continuous graph optimization .

An alternative sequential generation algorithm based on RNNs is presented in You et al. . The authors point out that a dense implementation of a GGNN requires a large number $\mathcal{O}(eN^{2})$ of operations to construct a graph with $e$ edges and $N$ nodes. We note that this scaling problem can be mitigated using a sparse GGNN implementation , which reduces complexity to $\mathcal{O}(e^{2})$.

Molecule design

Traditional in silico molecule design approaches rely on considerable domain knowledge, physical simulation and heuristic search algorithms (for a recent example, see Gómez-Bombarelli et al. ). Several deep learning approaches have also been tailored to molecule design, for example is a very promising method that uses a library of frequent (ring-containing) fragments to reduce the graph generation process to a tree generation process where nodes represent entire fragments. Alternatively, many methods rely on the SMILES linearization of molecules and use RNNs to generate new SMILES strings . A particular challenge of this approach is to ensure that the generated strings are syntactically valid under the SMILES grammar. The Grammar VAE uses a mask to impose these constraints during generation and a similar technique is applied for general graph construction in Samanta et al. . Our model also employs masking that, among other things, ensures that the molecules we generate can be converted to syntactically valid SMILES strings.

Generative Model

Our generative procedure is illustrated in Fig. 1. The process is seeded with $N$ vectors $\mathbf{z}_{v}$ that together form a latent “specification” for the graph to be generated ($N$ is an upper bound on the number of nodes in the final graph). Generation of edges between these nodes then proceeds using two decision functions: focus and expand. In each step the focus function chooses a focus node to visit, and then the expand function chooses edges to add from the focus node. As in breadth-first traversal, we implement focus as a deterministic queue (with a random choice for the initial node).

Our task is thus reduced to learning the expand function that enumerates new edges connected to the currently focused node. One design choice is to make expand condition upon the full history of the generation. However, this has both theoretical and practical downsides. Theoretically, this means that the learned model is likely to learn to reproduce generation traces. This is undesirable, since the underlying data usually only contains fully formed graphs; thus the exact form of the trace is an artifact of the implemented data preprocessing. Practically, this would lead to extremely deep computation graphs, as even small graphs easily have many dozens of edges; this makes training of the resulting models very hard as mentioned in mentioned in Li et al. . Hence, we condition expand only upon the partial graph structure $\smash{\mathcal{G}^{(t)}}$ generated so far; intuitively, this corresponds to learning how to complete a partial graph without using any information about how the partial graph was generated. We now present the details of each stage of this generative procedure.

We associate a state $\smash{\mathbf{h}^{(t=0)}_{v}}$ with each node $v$ in a set of initially unconnected nodes. Specifically, $\mathbf{z}_{v}$ is drawn from the $d$-dimensional standard normal $\mathcal{N}\left(\mathbf{0},\mathbf{I}\right)$, and $\smash{\mathbf{h}^{(t=0)}_{v}}$ is the concatenation $[\mathbf{z}_{v},\bm{\tau}_{v}]$, where $\bm{\tau}_{v}$ is an interpretable one-hot vector indicating the node type. $\bm{\tau}_{v}$ is derived from $\mathbf{z}_{v}$ by sampling from the softmax output of a learned mapping $\bm{\tau}_{v}\sim f(\mathbf{z}_{v})$ where $f$ is a neural networkWe implement $f$ as a linear classifier from the $100$ dimensional latent space to one of the node type classes.. The interpretable component of $\smash{\mathbf{h}^{(t=0)}_{v}}$ gives us a means to enforce hard constraints during generation.

From these node-level variables, we can calculate global representations $\smash{\mathbf{H}^{(t)}}$ (the average representation of nodes in the connected component at generation step $t$), and $\mathbf{H}_{\rm init}$ (the average representation of all nodes at $t=0$). In addition to $N$ working nodes, we also initialize a special “stop node” to a learned representation $\mathbf{h}_{\oslash}$ for managing algorithm termination (see below).

Node Update

Whenever we obtain a new graph $\mathcal{G}^{(t+1)}$, we discard $\smash{\mathbf{h}^{(t)}_{v}}$ and compute new representations $\smash{\mathbf{h}^{(t+1)}_{v}}$ for all nodes taking their (possibly changed) neighborhood into account. This is implemented using a standard gated graph neural network (GGNN) $G_{\rm dec}$ for $S$ stepsOur experiments use $S=7$., which is defined as a recurrent operation over messages $\smash{\mathbf{m}_{v}^{(s)}}$.

Edge Selection and Labelling

Termination

We keep adding edges to a node $v$ using expand and $G_{\rm dec}$ until an edge to the stop node is selected. Node $v$ then loses focus and becomes “closed” (mask $M$ ensures that no further edges will ever be made to $v$). The next focus node is selected from the focus queue. In this way, a single connected component is grown in a breadth-first manner. Edge generation continues until the queue is empty (note that this may leave some unconnected nodes that will be discarded).

Training the Generative Model

The encoder of our VAE is a GGNN $G_{\rm enc}$ that embeds each node in an input graph $\mathcal{G}$ to a diagonal normal distribution in $d$-dimensional latent space parametrized by mean $\bm{\mu}_{v}$ and standard deviation $\bm{\sigma}_{v}$ vectors. The latent vectors $\mathbf{z}_{v}$ are sampled from these distributions, and we construct the usual VAE regularizer term measuring the KL divergence between the encoder distribution and the standard Gaussian prior: $\mathcal{L}_{\rm latent}=\sum_{v\in\mathcal{G}}{\rm KL}(\mathcal{N}\left(\bm{\mu}_{v},{\rm diag}(\bm{\sigma}_{v})^{2}\right)||\;\mathcal{N}\left(\mathbf{0},\mathbf{I}\right))$.

2 Decoder

The decoder is the generative procedure described in Sect. 3, and we condition generation on a latent sample from the encoder distribution during training. We supervise training of the overall model using generation traces extracted from graphs in $\mathcal{D}$.

To obtain initial node states $\smash{\mathbf{h}^{(t=0)}_{v}}$, we first sample a node specification $\mathbf{z}_{v}$ for each node $v$ and then independently for each node we generate the label $\bm{\tau}_{v}$ using the learned function $f$. The probability of re-generating the labels $\bm{\tau}^{*}_{v}$ observed in the encoded graph is given by a sum over node permutations $\mathcal{P}$:

This inequality provides a lower bound given by the single contribution from the ordering used in the encoder (recall that in the encoder we know the node type $\bm{\tau}^{*}_{v}$ from which $\mathbf{z}_{v}$ was generated). A set2set model could improve this bound.

Edge Selection and Labelling

During training, we provide supervision on the sequence of edge additions based on breadth-first traversals of each graph in the dataset $\mathcal{D}$. Formally, to learn a distribution over graphs (and not graph generation traces), we would need to train with an objective that computes the log-likelihood of each graph by marginalizing over all possible breadth-first traces. This is computationally intractable, so in practice we only compute a Monte-Carlo estimate of the marginal on a small set of sampled traces. However, recall from Sect. 3 that our expand model is not conditioned on full traces, and instead only considers the partial graph generated so far. Below we outline how this intuitive design formally affects the VAE training objective.

Given the initial collection of unconnected nodes, $\mathcal{G}^{(0)}$, from the initialization above, we first use Jensen’s inequality to show that the log-likelihood of a graph $\mathcal{G}$ is loosely lower bounded by the expected log-likelihood of all the traces $\Pi$ that generate it.

The first term corresponds to the choice of $v$ as focus node at step $t$ of trace $\pi$. As our focus function is fixed, this choice is uniform in the first focus node and then deterministically follows a breadth-first queuing system. A summation over this term thus evaluates to the constant $\log(1/N)$.

As discussed above, the second term is only conditioned on the current graph (and not the whole generation history $\mathcal{G}^{(0)}\ldots\mathcal{G}^{(t-1)}$). To evaluate it further, we consider the set of generation states $\mathcal{S}$ of all valid state pairs $s=(\mathcal{G}^{(t)},v)$ of a partial graph $\mathcal{G}^{(t)}$ and a focus node $v$.

We use $|s|$ to denote the multiplicity of state $s$ in $\Pi$, i.e., the number of traces that contain graph $\mathcal{G}^{(t)}$ and focus on node $v$. Let $\mathcal{E}_{s}$ denote all edges that could be generated at state $s$, i.e., the edges from the focus node $v$ that are present in the graph $\mathcal{G}$ from the dataset, but are not yet present in $\mathcal{G}^{(t)}$. Then, each of these appears uniformly as the next edge to generate in a trace for all $|s|$ occurrences of $s$ in a trace from $\Pi$,

and therefore, we can rearrange a sum over paths into a sum over steps:

Here we use that $|s|/|\Pi|$ is the probability of observing state $s$ in a random draw from all states in $\Pi$. We use this expression in Eq. 2 and train our VAE with a reconstruction loss $\mathcal{L}_{\rm recon.}=\sum_{\mathcal{G}\in\mathcal{D}}\log\left[p(\mathcal{G}\mid\mathcal{G}^{(0)})\cdot p(\mathcal{G}^{(0)}\mid\mathbf{z})\right]$ ignoring additive constants.

We evaluate the expectation over states $s$ using a Monte Carlo estimate from a set of enumerated generation traces. In practice, this set of paths is very small (e.g. a single trace) resulting in a high variance estimate. Intuitively, Fig. 2 shows that rather than requiring the model to exactly reproduce each step of the sampled paths (orange) our objective does not penalize the model for choosing any valid expansion at each step (black).

3 Optimizing Graph Properties

So far, we have described a generative model for graphs. In addition, we may wish to perform (local) optimization of these graphs with respect to some numerical property, $Q$. This is achieved by gradient ascent in the continuous latent space using a differentiable gated regression model

where $g_{1}$ and $g_{2}$ are neural networksIn our experiments, both $g_{1}$ and $g_{2}$ are implemented as linear transformations that project to scalars. and $\sigma$ is the sigmoid function. Note that the combination of $R$ with $G_{\rm enc}$ (i.e., $R(G_{\rm enc}(\mathcal{G}))$) is exactly the GGNN regression model from Gilmer et al. . During training, we use an $L_{2}$ distance loss $\mathcal{L}_{Q}$ between $R(\mathbf{z}_{v})$ and the labeled properties $Q$. This regression objective shapes the latent space, allowing us to optimize for the property $Q$ in it. Thus, at test time, we can sample an initial latent point $\mathbf{z}_{v}$ and then use gradient ascent to a locally optimal point $\mathbf{z}_{v}^{*}$ subject to an $L_{2}$ penalty that keeps the $\mathbf{z}_{v}^{*}$ within the standard normal prior of the VAE. Decoding from the point $\mathbf{z}_{v}^{*}$ then produces graphs with an optimized property $Q$. We show this in our experiments in Sect. 6.2.

4 Training objective

The overall objective is $\mathcal{L}=\mathcal{L}_{\rm recon.}+\lambda_{1}\mathcal{L}_{\rm latent}+\lambda_{2}\mathcal{L}_{Q}$, consisting of the usual VAE objective (reconstruction terms and regularization on the latent variables) and the regression loss. Note that we allow deviation from the pure VAE loss ($\lambda_{1}=1$) following Yeung et al. .

Application: Molecule Generation

In this section, we describe additional specialization of our model for the application of generating chemical molecules. Specifically, we outline details of the molecular datasets that we use and the domain specific masking factors that appear in Eq. 1.

We consider three datasets commonly used in the evaluation of computational chemistry approaches:

QM9 , an enumeration of $\sim{}134$k stable organic molecules with up to 9 heavy atoms (carbon, oxygen, nitrogen and fluorine). As no filtering is applied, the molecules in this dataset only reflect basic structural constraints.

ZINC dataset , a curated set of 250k commercially available drug-like chemical compounds. On average, these molecules are bigger ($\sim{}23$ heavy atoms) and structurally more complex than the molecules in QM9.

CEPDB , a dataset of organic molecules with an emphasis on photo-voltaic applications. The contained molecules have $\sim{}28$ heavy atoms on average and contain six to seven rings each. We use a subset of the full database containing 250k randomly sampled molecules.

For all datasets we kekulize the molecules so that the only edge types to consider are single, double and triple covalent bonds and we remove all hydrogen atoms. In the encoder, molecular graphs are presented with nodes annotated with onehot vectors $\bm{\tau}^{*}_{v}$ indicating their atom type and charge.

2 Valency masking

Valency rules impose a strong constraint on constructing syntactically valid moleculesNote that more complex domain knowledge e.g. Bredt’s rule could also be handled in our model but we do not implement this here.. The valency of an atom indicates the number of bonds that that atom can make in a stable molecule, where edge types “double” and “triple” count for 2 and 3 bonds respectively. In our data, each node type has a fixed valency given by known chemical properties, for example node type “O” (an oxygen atom) has a valency of 2 and node type “O-” (an oxygen ion) has valency of 1. Throughout the generation process, we use masks $M$ and $m$ to guarantee that the number of bonds $b_{v}$ at each node never exceeds the valency $b_{v}^{*}$ of the node. If $b_{v}<b_{v}^{*}$ at the end of generation we link $b_{v}^{*}-b_{v}$ hydrogen atoms to node $v$. In this way, our generation process always produces syntactically valid molecules (we define syntactic validity as the ability to parse the graph to a SMILES string using the RDKit parser ). More specifically, $\smash{M_{v\mathrel{\leftrightarrow}u}^{(t)}}$ also handles avoidance of edge duplication and self loops, and is defined as:

Experiments

We evaluate baseline models, our model (CGVAE) and a number of ablations on the two tasks of molecule generation and optimizationOur implementation of CGVAE can be found at https://github.com/Microsoft/constrained-graph-variational-autoencoder..

As baseline models, we consider the deep autoregressive graph model (that we refer to as DeepGAR) from , a SMILES generating LSTM language model with 256 hidden units (reduced to 64 units for the smaller QM9 dataset), ChemVAE , GrammarVAE , GraphVAE , and the graph model from . We train these and on our three datasets and then sample 20k molecules from the trained models (in the case of , we obtained sets of sampled molecules from the authors).

We analyze the methods using two sets of metrics. First in Fig. 3(a) we show metrics from existing work: syntactic validity, novelty (i.e. fraction of sampled molecules not appearing in the training data) and uniqueness (i.e. ratio of sample set size before and after deduplication of identical molecules). Second, in Fig. 3(b) we introduce new metrics to measure how well each model captures the distribution of molecules in the training set. Specifically, we measure the average number of each atom type and each bond type in the sampled molecules, and we count the average number of 3-, 4-, 5-, and 6-membered cycles in each molecule. This latter metric is chemically relevant because 3- and 4-membered rings are typically unstable due to their high ring strain. Fig. 3(c) shows 2 samples from our model for each dataset and we show more samples of generated molecules in the supplementary material.

The results in Fig. 3 show that CGVAE is excellent at matching graph statistics, while generating valid, novel and unique molecules for all datasets considered (additional details are found in supplementary material B and C). The only competitive baselines are DeepGAR from Li et al. and an LSTM language model. Our approach has three advantages over these baselines: First, whereas >10% of ZINC-like molecules generated by DeepGAR are invalid, our masking mechanism guarantees molecule validity. An LSTM is surprisingly effective at generating valid molecules, however, LSTMs do not permit the injection of domain knowledge (e.g. valence rules or requirement for the existance of a particular scaffold) because meaningful constraints cannot be imposed on the flat SMILES representation during generation. Second, we train a shallow model on breadth-first steps rather than full paths and therefore do not experience problems with training instability or overfitting that are described in Li et al. . Empirical indication for overfitting in DeepGAR is seen by the fact that Li et al. achieves the lowest novelty score on the ZINC dataset, suggesting that it more often replays memorized construction traces. It is also observed in the LSTM case, where on average 60% of each generated SMILES string is copied from the nearest neighbour in the training set. Converting our generated graphs to SMILES strings reveals only 40% similarity to the nearest neighbour in the same metric. Third we are able to use our continuous latent space for molecule optimization (see below).

We also perform an ablation study on our method. For brevity we only report results using our ring count metrics, and other statistics show similar behavior. From all our experiments we highlight three aspects that are important choices to obtain good results, and we report these in ablation experiments A, B and C in Fig. 4. In experiment A we remove the distance feature $d_{v,u}$ from $\bm{\phi}$ and see that this harms performance on the larger molecules in the ZINC dataset. More interestingly, we see poor results in experiment B where we make an independence assumption on edge generation (i.e. use features $\bm{\phi}$ to calculate independent probabilities for all possible edges and sample an entire molecule in one step). We also see poor results in experiment C where we remove the GGNN from the decoder (i.e. perform sequential construction with $\smash{\mathbf{h}^{(t)}_{v}=\mathbf{h}^{(0)}_{v}}$). This indicates that the choice to perform sequential decoding with GGNN node updates before each decision are the keys to the success of our model.

2 Directed molecule generation

Finally, we show that we can use the VAE structure of our method to direct the molecule generation towards especially interesting molecules. As discussed in Sect. 4.3 (and first shown by Gómez-Bombarelli et al. in this setting), we extend our architecture to predict the Quantitative Estimate of Drug-Likeness (QED) directly from latent space. This allows us to generate molecules with very high QED values by performing gradient ascent in the latent space using the trained QED-scoring network. Fig. 5 shows an interpolation sequence from a point in latent space with an low QED value (which ranges between 0 and 1) to the local maximum. For each point in the sequence, the figure shows a generated molecule, the QED value our architecture predicts for this molecule, as well as the QED value computed by RDKit.

Conclusion

We proposed CGVAE, a sequential generative model for graphs built from a VAE with GGNNs in the encoder and decoder. Using masks that enforce chemical rules, we specialized our model to the application of molecule generation and achieved state-of-the-art generation and optimization results. We introduced basic statistics to validate the quality of our generated molecules. Future work will need to link to the chemistry community to define additional metrics that further guide the construction of models and datasets for real-world molecule design tasks.

References

Appendix A Molecule Samples

We provide 25 random samples from our model for qualitative comparison with samples from each training dataset.

Appendix B Effect of multiple training paths

Section 4.2 describes how we should enumerate all breadth first graph generation traces, break these traces into state transitions and then randomly sample state transitions to give the Monte Carlo estimate of the reconstruction loss. However, for computational efficiency, in the presented experiments we provide only a single trace containing $E$ transitions (where $E$ is the number of edges in the final molecule including edges to the stop node). Figure 6 shows an additional experiment (CGVAE (50)) where we enumerate 50 traces for each molecule and sample $E$ transitions from this enumeration (so the final dataset size is the same). While increasing the number of traces considered produces a small improvement in the matching of ring statistics in the sampled molecules, it is not clear that this benefit is worth the considerable computational overhead required in preparing the dataset.

Appendix C Additional Molecular Properties

Here we provide histograms of the following molecular properites of the sampled molecules for our method and the DeepGAR and LSTM baselines:

Appendix D Optimization trajectories

We provide additional QED optimization trajectories for our model trained on the ZINC dataset.