NeVAE: A Deep Generative Model for Molecular Graphs

Introduction

Drug design aims to identify (new) molecules with a set of specified properties, which in turn results in a therapeutic benefit to a group of patients. However, drug design is still a lengthy, expensive, difficult, and inefficient process with a low rate of new therapeutic discovery , in which candidate molecules are produced through chemical synthesis or biological processes. In the context of computer-aided drug design , there is a great interest in developing automated, machine learning techniques to discover sizeable numbers of plausible, diverse, and novel candidate molecules with various desirable properties in the vast ($10^{23}-10^{60}$) and unstructured molecular space .

In recent years, there has been a flurry of work devoted to developing deep generative models for automatic molecule design , which has predominantly followed two strategies. The first strategy consists of representing molecules using a domain specific textual representation—SMILES strings—and then leveraging deep generative models for text generation for molecule design. Unfortunately, SMILE strings do not capture the structural similarity between molecules and, moreover, a molecule can have multiple SMILES representations. As a consequence, the generated molecules lack in terms of diversity and validity, as shown in Tables 1–2 and Figure 8. The second strategy consists of representing molecules using molecular graphs, rather than SMILES representations, and then developing deep generative models for molecular graphs, in which atoms correspond to nodes and bonds correspond to edges. However, current generative models for molecular graphs share one or more of the following limitations, which preclude them from realizing all their potential:

They can only generate (and be trained on) molecules with the same number of atoms while, in practice, molecules having similar properties often come with a different number of atoms and bonds.

They are not invariant to permutations of their node labels, however, molecular graphs remain isomorphic under permutation of their node labels.

Their training procedure suffers from a quadratic complexity with respect to the number of nodes in the graph, which makes it difficult to leverage a sizeable number of large molecules during training.

They generate molecular graphs by combining a small set of molecular graphlets (or subgraphs), which constrain the diversity of the generated molecules, as shown in Table 1 and Figure 8.

They do not provide the spatial coordinates of the atoms they generate, whereas in practice, a molecule is a three-dimensional object in which the coordinates of its atoms significantly influence its chemical properties, as shown in Figure 5.

To identify molecules that maximize the value of certain property (e.g., solubility in water), they resort to either traditional Bayesian optimization or reinforcement learning over the continuous latent representation of molecules they find. However, such procedures are unable to discover a sizeable set of candidate molecules with high property values, as shown in Table 3 and Figure 8.

To address the first five shortcomings (I-V), we develop NeVAE, a deep generative model for molecular graphs based on variational autoencoders. Our model relies on several technical innovations, which distinguish us from previous work:

Our probabilistic encoder learns to aggregate information (e.g., bond features, atoms and their coordinates) from a different number of hops away from a given atom and then map this aggregate information into a continuous latent space, as in inductive graph representation learning . However, in contrast with inductive graph representation learning, the aggregator functions are learned via variational inference so that the resulting aggregator functions are especially well suited to enable the probabilistic decoder to generate new molecules rather than other machine learning tasks such as, e.g., link prediction. Moreover, by using (symmetric) aggregator functions, it is invariant to permutations of the node labels and can encode graphs with a variable number of atoms, as opposed to existing graph generative models, with a few the notable exception of those based on GCNs .

Our probabilistic decoder jointly represents all edges as an unnormalized log probability vector (or ‘logit’), which then feeds a single multinomial edge distribution. Such scheme allows for an efficient inference algorithm with $O(l)$ complexity, where $l$ is the number of true edges in the molecules, which is also invariant to permutations of the node labels. In contrast, previous work typically models the presence and absence of each potential edge using a Bernoulli distribution and this leads to inference algorithms with $O(n^{2})$ complexity, where $n$ is the number of nodes, which are not permutation invariant.

Our probabilistic decoder is able to guarantee a set of local structural and functional properties in the generated molecules by using a mask in the edge distribution definition, which can prevent the generation of certain undesirable edges during the decoding process. While masking have been increasingly used to account for prior (expert) knowledge in generative models based on SMILES, their use in generative models for molecular graphs has been lacking.

Our probabilistic decoder is able to provide the spatial coordinates of the atoms of the molecules it generates. To do so, it models the position of each atom using a Gaussian distribution whose mean and variance depend on its latent representation as well as that of each of its neighbors.

To address the last shortcoming (VI), we develop a gradient-based algorithm to optimize the decoder of our model for property oriented molecule generation, i.e., to optimize the decoder so that it learns to generate molecules that maximize the value of certain property (e.g., solubility in water). Note that, in contrast with recent reinforcement learning methods for property oriented molecule generation , our gradient-based algorithm benefits from the inductive bias provided by the original decoder, which in turns enable us to identify better molecules, as shown in Table 3. Moreover, given a molecule of interest, our gradient-based algorithm can also be used to optimize the spatial configuration of its atoms for greater stability. We believe our algorithm is of independent interest since it may be adapted to other deep generative models designed for other data types such as graphs, images, text or audio, that maximizes certain property value.

We experiment with molecules from two publicly available datasets, ZINC and QM9 . First, we show that NeVAE beats the state of the art in terms of several relevant quality metrics, i.e., validity, novelty and uniqueness, and the resulting latent space representation of molecules exhibits powerful semantics—we can smoothly interpolate between molecules—and generalization ability—we can generate (valid) molecules that are larger than any of the molecules in the datasets. Then, we demonstrate that, for several properties of interest (e.g., solubility in water), our gradient-based algorithm is able to successfully optimize NeVAE’s decoder for property oriented molecule generation. In particular, the optimized decoder is able to identify molecules with property values $121$% higher than those identified by several state of the art methods based on Bayesian optimization and reinforcement learning and, given a molecule of interest, it is able to optimize the spatial configuration of its atoms for greater stability, i.e., lower potential energy. To facilitate research in this area, we are releasing an open source implementation of our model in Tensorflow as well as synthetic and real-world data used in our experiments https://github.com/Networks-Learning/nevae.

Background on Variational Autoencoders

Finally, note that the quality of this variational lower bound depends on the expressive ability of the approximate inference model $q_{\phi}(\bm{z}|\bm{x})$, which is typically assumed to be a normal distribution whose mean and variance are parametrized by a neural network $\phi$ with the observed data $\bm{x}$ as input.

NeVAE: A Variational Autoencoder for Molecular Graphs

In this section, we first give a high-level overview of the design of NeVAE, our variational autoencoder for molecular graphs, starting from the data it is designed for. Then, we describe more in-depth the key technical aspects of its individual components. Finally, we elaborate on the training procedure, scalability and implementation details.

High-level overview. We observe a collection of $N$ molecular graphs $\{\mathcal{G}_{i}=(\mathcal{V}_{i},\mathcal{E}_{i})\}_{i\in[N]}$, where $\mathcal{V}_{i}$ and $\mathcal{E}_{i}$ denote the corresponding set of nodes (atoms) and edges (bonds), respectively, and this collection may contain graphs with a different number of nodes and edges. Moreover, for each molecular graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$, we also observe a set of node features $\mathcal{F}=\{\bm{t}_{u},\bm{x}_{u}\}_{u\in\mathcal{V}}$ and edge weights $\mathcal{Y}=\{y_{uv}\}_{(u,v)\in\mathcal{E}}$. More specifically, $\bm{t}_{u}$ are one-hot representations of the type of the atoms (i.e., $C$, $H$, $N$ or $O$), $\bm{x}_{u}$ are the coordinates of the atoms in three dimensional space, and $y_{uv}$ are the bond types (i.e., single, double, triple). Our goal is then to design a variational autoencoder for molecular graphs that, once trained on this collection of graphs, has the ability of creating new plausible molecular graphs, including node features and edge weights. In doing so, it will also provide a latent representation of any graph in the collection (or elsewhere) with meaningful semantics.

Following the above background on variational autoencoders, we characterize NeVAE by means of:

In the above characterization, note that we define one latent variable per node, i.e., we have a node-based latent representation, and the number of nodes is a random variable. As a consequence, both the latent representation as well as the graph can vary in size. Next, we formally define the functional form of the inference model, the generative model, and the prior.

Inference model (probabilistic encoder). Given a graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ with node features $\mathcal{F}$ and edge weights $\mathcal{Y}$, our inference model $q_{\phi}$ defines a probabilistic encoding for each node in the graph by aggregating information from different distances. More formally, for each node $u$, the inference model is defined as follows:

In the above, $\bm{W}^{{\mathcal{T}}}_{k}$ and $\bm{W}^{\mathcal{X}}_{k}$ are trainable weight matrices, which propagate information between different search depths, $\bm{\Lambda}(.)$ is a (possibly nonlinear) symmetric aggregator function in its arguments, $\bm{g}(\cdot)$ and $\bm{r}(\cdot)$ are (possibly nonlinear) differentiable functions, $\phi^{enc}$ is a neural network, and $\odot$ denotes pairwise product. Figure 1 describes our encoder architecture.

The above node embeddings, defined by Eq. 2, are very similar to the ones used in several graph representation learning algorithms such as GraphSAGE , column networks , and GCNs , the main difference with our work is the way we will train the weight matrices $\bm{W}^{\bullet}_{k}$. Here, we will use variational inference so that the resulting embeddings are especially well suited to enable our probabilistic decoder to generate new, plausible molecular graphs. In contrast, the above algorithms use non variational approaches to compute general purpose embeddings to feed downstream machine learning tasks.

The following proposition highlights several desirable theoretical properties of our probabilistic encoder, which distinguishes our design from most existing generative models of graphs :

The probabilistic encoder defined by Eqs. 1 and 2 has the following properties:

For each node $u$, its corresponding embedding $\mathbf{c}_{u}(k)$ is invariant to permutations of the node labels of its neighbors.

The weight matrices $\bm{W}^{{\mathcal{T}}}_{1},\ldots,\bm{W}^{{\mathcal{T}}}_{K}$ and $\bm{W}^{\mathcal{X}}_{1},\ldots,\bm{W}^{\mathcal{X}}_{K}$ do not depend on the number of nodes and edges in the graph and thus a single encoder allows for graphs with a variable number of nodes and edges.

Consider $\pi$, a permutation of the node labels, i.e. for each $u$, we have $\pi(u)\in\mathcal{V}$; and the set of all shuffled labels $\{\pi(w)|w\in\mathcal{V}\}=\mathcal{V}$. Let us denote $\widetilde{u}:=\pi(u)$. Now we need to prove

We proof this by induction. Since the features $\bm{t}_{u}$ and $\bm{x}_{u}$ are independent of the node label of $u$, we have that $\bm{t}_{u}=\bm{t}_{\widetilde{u}}$ and $\bm{x}_{u}=\bm{x}_{\widetilde{u}}$, which proves Eq. 3 for $k=1$, $\forall u\in\mathcal{V}$. Now assume that Eq. 3 is true for $k\leq k^{\prime}-1$, with $k^{\prime}>1$. That is, we have,

Also, since the edge-weight $y_{uv}$ between nodes does not depend on their labels, we have

This, along with Eq. 4 gives $\{\cup_{v\in\mathcal{N}(u)}y_{uv}\,\bm{g}(\mathbf{c}_{v}(k^{\prime}-1)\}=\{\cup_{\widetilde{v}\in\mathcal{N}(\widetilde{u})}y_{\widetilde{u}\widetilde{v}}\,\bm{g}(\mathbf{c}_{\widetilde{v}}(k^{\prime}-1)\}$ which, due to the symmetric property of $\bm{\Lambda}(.)$, implies

The above equation, together with the fact that $\bm{t}_{u}=\bm{t}_{\widetilde{u}}$ and $\bm{x}_{u}=\bm{x}_{\widetilde{u}}$ proves Eq. 3 for $k=k^{\prime}$.

Generative model (probabilistic decoder). Given a set of of $n$ nodes with latent variables $\mathcal{Z}=\{\bm{z}_{u}\}_{u\in[n]}$, our generative model $p_{\theta}$ is defined as follows:

where the ordering for the edge and edge weights is independent of node labels and hence permutation invariant, $e_{k}$ and $y_{u_{k}v_{k}}$ denote the $k$-th edge and edge weight under the chosen order, and $\mathcal{E}_{k-1}=\{e_{1},\ldots,e_{k-1}\}$ and $\mathcal{Y}_{k-1}=\{y_{u_{1}v_{1}},\ldots,y_{u_{k-1}v_{k-1}}\}$ denote the $k-1$ previously generated edges and edge weights respectively.

Moreover, the model characterizes the conditional probabilities in the above formulation as follows. For each node, it represents all potential values for the atom types $\bm{t}_{u}=\bm{q}$ as an unnormalized log probability vector (or ‘logits’), feeds this logit into a softmax distribution and samples the node features. Then, it represents the average number of edges as a logit, feeds this logit into a Poisson distribution and samples the number of edges. Next, it represents all potential edges as logits and, for each edge, all potential edge weights as another logit, and it feeds the former vector into a single softmax distribution and the latter vectors each into a different softmax distribution. Moreover, the edge distribution and the corresponding edge weight distributions depend on a set of binary masks, which may depend on the sampled node features and also get updated every time a new edge and edge weight are sampled. By doing so, it prevents the generation of certain undesirable edges and edges weights, allowing for the generated graph to fulfill a set of predefined local structural and functional properties. Finally, for each atom, it samples its coordinates $\bm{x}_{u}$ from a multidimensional Gaussian distribution whose mean and variance depends on the latent vectors of the corresponding atom as well as its neighbors and the underlying chemical bonds.

More formally, the distributions of each node feature, the number of edges, each edge and edge weight are given by:

Prior. Given a set of $n$ nodes with latent variables $\mathcal{Z}=\{\bm{z}_{u}\}_{u\in[n]}$, $p_{z}(\mathcal{Z})\sim\mathcal{N}(\mathbf{0},\bm{I})$.

Training. Given a collection of $N$ molecular graphs $\{\mathcal{G}_{i}=(\mathcal{V}_{i},\mathcal{E}_{i})\}_{i\in[N]}$, each with $n_{i}$ nodes, a set of node features $\mathcal{F}_{i}$, set of node coordinates $\mathcal{X}_{i}$, set of edge weights $\mathcal{Y}_{i}$, we train our variational autoencoder for graphs by maximizing the evidence lower bound (ELBO), as described in the previous section, plus the log-likelihood of the Poisson distribution $p_{\lambda_{n}}$ modeling the number of nodes in each graph. Hence we aim to solve:

Therefore, to train our model, we maximize

The following theorem points out the key property of our objective function.

If the source distribution $\zeta$ does not depend on the node labels, then the parameters learned by maximizing the objective in Eq. 10 are invariant to the permutations of the node labels.

Proof For each training graph $\mathcal{G}_{i}=(\mathcal{V}_{i},\mathcal{E}_{i})$, we denote each corresponding component in the objective function as

where $\Theta$ is the set of trainable parameters.

$\zeta(\mathcal{V}_{i})$ does not depend on node labels (e.g. uniform sampling, degree based sampling, etc.);

the edge sequence of $\mathcal{E}_{i}$ is determined by BFS with randomized tie breaking;

$\epsilon_{u}$ does not depend on $u$ since it is sampled from $\mathcal{N}(\bm{0},\bm{I})$.

Scalability and implementation details. In terms of scalability, the major bottleneck is computing the gradient of the first term in Eq. 10 during training, rather than encoding and decoding graphs once the model is trained. More specifically, given a source node for a network without masks, an exact computation of the per edge partition function of the log-likelihood of the edges, i.e., $\sum_{e^{\prime}=(u^{\prime},v^{\prime})\notin\mathcal{E}_{k-1}}\exp(\theta^{dec}_{\bm{\alpha}}(\bm{z}_{u^{\prime}},\bm{z}_{v^{\prime}}))$, requires $O(|\mathcal{V}|^{2})$ computations, similarly as in most inference algorithms for existing generative models of graphs, and hence is costly to compute even for medium networks. Fortunately, in practice, we can approximate such partition function using negative sampling which reduces the likelihood computation to $O(l)$, where $l=|\mathcal{E}|$ is the number of (true) edges in the graph. Therefore, for $S$ samples of source nodes, the complexity becomes $O(Sl)$. Here, note that most real-world graphs are sparse and thus $l\ll|\mathcal{V}|^{2}$.

Property Oriented Molecule Generation

In this section, we aim to optimize the probabilistic decoder of our variational autoencoder, described in Section 3, so that it learns to generate molecules that maximize certain molecular property (e.g., solubility in water). To this aim, we approach the problem from the perspective of variational inference and show that the optimal property-oriented decoder can be expressed in terms of the original decoder and the value of the molecular property. This result means that we can obtain molecules from the optimal property-oriented decoder just by applying rejection sampling on the molecules samples from the original decoder. However, in practice, such a naive sampling strategy will be inefficient and impractical, specially given the high-dimensional nature of the data. Therefore, we design a practical method for approximating the optimal property-oriented decoder, which iteratively adapts the parameters of a (parameterized) property-oriented decoder using a stochastic gradient-based algorithm.

where the inner expectation is taken over all molecules generated using the property-oriented decoder $p(\cdot|\mathcal{Z})$ given the latent vectors $\mathcal{Z}$, the outer expectation is taken over all possible latent vectors under the prior distribution If a molecule is given, instead of the prior distribution, one may also consider using the posterior $q_{\phi}(\cdot|\mathcal{V},\mathcal{E},\mathcal{F},\mathcal{Y})$. $p_{z}(\bm{z}_{1},\ldots,\bm{z}_{n})$ with $n\sim\text{Poisson}(\lambda_{n})$, and we do not assume any specific parametric form for the property-oriented decoder $p$. In Eq. 14, the first term penalizes molecules with a low value of the property of interest, the second term penalizes property-oriented decoders whose generated molecules differ more from those that the original decoder would generate and the parameter $\rho$ controls the trade off between both terms. Here, note that the second term provides an inductive bias that ensures that the molecules generated by the property-oriented decoder are plausible. Moreover, we can rewrite the inner expectation of the second term in terms of Kullback-Leibler (KL) divergence , i.e.,

which is commonly used as a distance measure between distributions.

Then, it is straightforward to show that the above optimization problem is equivalent to the following problem:

The above objective function achieves its global minimum of zero if the numerator and the denominator are equal. Thus, the optimal property-oriented decoder is just given by:

The above result has an important implication. It means that we can use sampling methods to obtain (unbiased) samples from the optimal property-oriented decoder. For example, we can apply rejection sampling on the molecules generated by the original decoder $p_{\theta}$, where we accept or reject them according to the (exponentiated) property value of interest. However, in practice, these sampling methods may be inefficient if the generated molecules under the original decoder $p_{\theta}$ have low probability under the optimal property-oriented decoder model. Given that molecules are high dimensional objects, this is specially problematic due to the curse of dimensionality. Next, we will design a practical method for approximating $p^{*}(\mathcal{E},\mathcal{Y},\mathcal{F}|\mathcal{Z})$, which iteratively adapts the parameters of a (parameterized) property-oriented model using a stochastic gradient-based algorithm.

A stochastic gradient-based algorithm. In this section, we aim to find a property-oriented decoder $p_{\theta^{\prime}}$ within the class of parameterized probabilistic decoders defined by Eq. 7 that approximates well the optimal property-oriented decoder $p^{*}$ that minimizes the objective function in Eq. 13, i.e.,

To this aim, we introduce a general gradient-based algorithm, which iteratively update the parameters $\theta^{\prime}$ of the parameterized property-oriented decoder $p_{\theta^{\prime}}$ using stochastic gradient descent (SGD) , i.e.,

$\alpha_{j}>0$ is the learning rate at step $j$, and $\theta^{\prime}_{0}=\theta$.

The above expression readily yields the following unbiased finite sample Monte Carlo estimator:

where $M$ is total number of sampled molecules generated from the property-oriented decoder $p_{\theta^{\prime}}$. Algorithm 1 summarizes the overall procedure.

Experiments

In this section, we first show that NeVAE beats several state of the art machine learning models for molecule design in terms of several relevant quality metrics, i.e., validity, novelty and uniqueness. Then, we also show that the continuous latent representations of molecules that our model finds are smooth. Finally, we demonstrate that the property-oriented decoder provided by Algorithm 1 is able to generate molecules that maximize certain desirable properties more effectively than several baselines based on Bayesian optimization and reinforcement learning. Appendix C contains additional experiments on synthetic data.

2 Quality of the generated molecules

We first make a quantitative analysis of our model by comparing the quality of the molecules generated by our trained models against the molecules generated by several state of the art competing methods and then provide a qualitative analysis by demonstrating that the latent space of the molecules inferred by our model is smooth. For the quantitative analysis, we use eight baselines for comparison: (i) GraphVAE , (ii) GrammarVAE , (iii) CVAE , (iv) SDVAE , (v) JTVAE , (vi) CGVAE , (vii) MOLGAN , (viii) ORGAN , and (ix) GCPN . Among them, GraphVAE, JTVAE, CGVAE, MOLGAN and GCPN use molecular graphs and GrammarVAE, CVAE, SDVAE, JTVAE and ORGAN use SMILES strings, a domain specific textual representation of molecules.

Moreover, we use the following evaluation metrics for performance comparison:

Novelty: we use this metric to evaluate to which degree a method generates novel molecules, i.e., molecules which were not present in the (training) dataset, i.e. $\text{Novelty}=1-{|\mathcal{C}_{s}\cap\mathcal{D}|}/{|\mathcal{C}_{s}|}$, where $\mathcal{C}_{s}$ is the set of generated molecules which are chemically valid, $\mathcal{D}$ is the training dataset, and $\text{Novelty}\in$.

Uniqueness: we use this metric to evaluate to what extent a method generates unique chemically valid molecules. We define, $\text{Uniqueness}={|\text{set}(\mathcal{C}_{s})|}/{n_{s}}$ where $n_{s}$ is the number of generated molecules and $\text{Unique}\in$.

Validity: we use this metric to evaluate to which degree a method generates chemically valid molecules We used the opensource cheminformatics suite RDkit (http://www.rdkit.org) to check the validity of a generated molecule.. That is, $\text{Validity}={|{\mathcal{C}_{s}}|}/{n_{s}}$ where $n_{s}$ is the number of generated molecules, $\mathcal{C}_{s}$ is the set of generated molecules which are chemically valid, and note that $\text{Validity}\in$.

Tables 1–2 compare our trained models to the state of the art methods above in terms of novelty, uniqueness, and validity. For GraphVAE and CGVAE we report the results reported in the paper and, for SDVAE, since there is no public domain implementation of these methods at the time of writing, we have used the sampled molecules from the prior provided by the authors for the ZINC dataset. For CVAE, GrammarVAE, JTVAE, ORGAN, MOLGAN and GCPN, we run their public domain implementations in the same set of molecules that we used. For MOLGAN, ORGAN and GCPN, we only report the validity on the discovered molecules and refrain comparing their performance in terms of novelty and uniqueness given that their focus is on generating molecules maximizing certain property value.

We find that, in terms of novelty, both our trained models and all competing methods except for the GraphVAE, which assumes a fixed number of nodes, are able to (almost) always generate novel molecules. However, we would also like to note that novelty is only defined over chemically valid molecules. Therefore, despite having (almost) perfect novelty scores, GraphVAE, GrammarVAE, CVAE and SDVAE generate significantly fewer novel molecules than our method. In terms of uniqueness, which is defined over the set of sampled molecules, we observe that all baseline methods, except CGVAE (for ZINC and QM9) and JTVAE (for ZINC), perform very poorly in both datasets in comparison with our method. In terms of validity, our trained models significantly outperform four competing methods— GraphVAE, GrammarVAE, CVAE, SDVAE and ORGAN— even without the use of masking, and achieve comparable performance to JTVAE, CGVAE and GCPN.

We would like to highlight that, in contrast to our model, GrammarVAE, CVAE and SDVAE use SMILES, a domain specific string based representation, and thus they may be constrained by its limited expressiveness. Among them, GrammarVAE and SDVAE achieve better performance by using grammar to favor valid molecules. GraphVAE generates molecular graphs, as our model, however, its performance is inferior to our method because it assumes a fixed number of nodes, it samples edges independently from a Bernoulli distribution, and is not permutation invariant.

Next, we qualitatively demonstrate that the latent space of molecules inferred by our model is smooth. To that aim, given a molecule, along with its associated graph $\mathcal{G}$, node features $\mathcal{F}$ and edge weights $\mathcal{Y}$, we first sample its latent representation $\mathcal{Z}$ using our probabilistic encoder, i.e., $\mathcal{Z}\sim q_{\phi}(\mathcal{Z}|G,\mathcal{F},\mathcal{Y})$. Then, given this latent representation, we generate various molecular graphs by sampling from our probabilistic decoder, i.e., $\mathcal{G}_{i}\sim p_{\theta}(\mathcal{G}|\mathcal{Z})$. Figure 3 summarizes the results for one molecule from ZINC dataset, which shows that the sampled molecules are topologically similar to the given molecule. Finally, we also show that our encoder, once trained, creates a latent space representation of molecules with powerful semantics. In particular, since each node in a molecule has a latent representation, we can make fine-grained changes to the structure of a molecule by perturbing the latent representation of single nodes. To this aim, we proceed as follows.

First, we select one molecule with $n$ nodes from the ZINC dataset. Given its corresponding graph, node features and edge weights, $\mathcal{G}$, $\mathcal{F}$ and $\mathcal{Y}$, we sample its latent representation $\mathcal{Z}_{0}$. Then, we sample new molecular graphs $\mathcal{G}$ from the probabilistic decoder $\mathcal{G}\sim p_{\theta}(\mathcal{G}|\mathcal{Z})$, where $\mathcal{Z}=\{\bm{z}_{i}+a_{i}\bm{z}_{i}|\,\bm{z}_{i}\in\mathcal{Z}_{0},a_{i}\geq 0\}$ and $a_{i}$ are given parameters. Figure 4 provides several examples across both datasets, which show that the latent space representation is smooth and, as the distance from the initial molecule increases in the latent space, the resulting molecule differs more from the original.

3 Property oriented molecule generation

In this section, we first use our gradient-based algorithm (refer to Algorithm 1) to design property-oriented decoders that maximize the following two properties:

the octanol-water partition coefficient, penalized by synthetic accessibility (SA) score and number of long cycles (penalized logP, $y_{1}(m)$); and,

the quantitative estimation of drug-likeness (QED, $y_{2}(m)$).

Then, we use our gradient-based algorithm to design property-oriented decoders that, given a molecule of interest, are able to optimize the spatial configuration of its atoms for greater stability, i.e., lower potential energy.

Table 3 shows the values of penalized logP and QED for the best three molecules generated by each method. The results show that our property oriented decoder, NeVAE (Algorithm 1), is able to identify molecules with property values $121$% higher than those identified by the best performing competitor, i.e., Bayesian optimization over the latent spaced of molecules generated by JTVAE. Finally, note that, in contrast with all the competitors, our property oriented decoder is also able to provide a (plausible) spatial configuration for the atoms of each of the identified molecules, as shown in Figure 5.

Discovering molecules with low potential energy values. The stability of a molecule depends on its potential energy—a lower value of potential energy indicates higher stability. In this section, our goal is to generate the most stable three dimensional structure for a given two dimensional molecular graph $\mathcal{G}=(\mathcal{V},\mathcal{E})$ with atoms $\{\bm{t}_{u}\}_{u\in\mathcal{V}}$, initial spatial coordinates $\bm{x}^{\prime}_{u}$ and bond-types $\mathcal{Y}$.

To this aim, we only need to optimize the part of the decoder that generates the spatial coordinates rather than the entire decoder. Therefore, given a molecular graph, we solve the following optimization problem using the same gradient-descent algorithm described in Section 4:

Figures 6 and 7 summarize the results for $100$ molecular graphs The molecular graphs were not present in the training set used to train NeVAE.. The results show that, for a majority of the molecular graphs, we are able to find spatial configurations that increase their stability (i.e., decrease their potential energy).

Conclusion

In this work, we have introduced a variational autoencoder for molecular graphs, that is permutation invariant of the nodes labels of the graphs they are trained with, and allow for graphs with different number of nodes and edges as well as three dimensional spatial coordinates for atoms. Moreover, the decoder is able to guarantee a set of local structural and functional properties in the generated graphs through masking. Then, we have developed a gradient based algorithm to optimize the decoder of our model so that it learns to generate molecules that maximize the value of certain property of interest. Finally, we have shown that our variational autoencoder is able to discover plausible, diverse and novel molecules more effectively than several state of the art methods and, for several properties of interest, our optimized decoder is able to identify molecules with property values $121$% higher than those identified by several state of the art methods.

Our work also opens many interesting venues for future work. For example, in the design of our variational autoencoder, we have assumed graphs to be static, however, it would be interesting to augment our design to dynamic graphs by, e.g., incorporating a recurrent neural network or long short-term memory (LSTM) units. Moreover, we have focused on molecular graphs, however, we believe our methodology could be adapted to other real-world graphs. Finally, there are other problems related to molecular design, such as retro synthesis , where machine learning may advance the state of the art.

References

Appendix A Implementation Details

Architecture details. Table 4 provides additional details on the architecture of our variational autoencoder for graphs, where it is important to notice that the parameters to be learned do not depend on the size of the graphs (i.e., the number of nodes and edges). Note that, $\bm{r}$ and $\bm{g}$ are linear forms and the aggregator function $\bm{\Lambda}$ is a sum, which is a symmetric function, for simplicity We did experiment with other symmetric aggregator functions such as pooling, as in the inductive graph representation learning , and did not notice significant gains in practice..

Hyperparameter tuning. At the very outset, to train NeVAE, we implemented stochastic gradient descent (SGD) using the Adam optimizer. Therein, we had to specify four hyperparameters: (i) $D$ – the dimension of $z_{u}$, (ii) $K$ – the maximum number of hops used in encoder to aggregate information, (iii) $L$ – the number of negative samples, (iv) $l_{r}$– the learning rate. Note that, all the parameters $\bm{W}_{\bullet}$’s and $\bm{b}_{\bullet}$’s in the input, hidden and output layers depend on $D$ and $K$. We selected these hyperparameters using cross validation. More specifically, we varied $l_{r}$ in a logarithmic scale, i.e., $\{0.0005,0.005,0.05,0.5\}$, and the rest of the hyperparameters in an arithmetic scale, and chose the hyperparameters maximizing the value of the objective function in the validation set. For synthetic (real) data, the resulting hyperparameter values were $D=7(5)$, $K=3(5)$, $L=10(10)$ and $l_{r}=0.005(0.005)$. To run the baseline algorithms, we followed the instructions in the corresponding repository (or paper).

Training with minibatch. We implemented stochastic gradient descent (SGD) using minibatches, where each batch contained graphs with the same number of nodes. More specifically, we first group the training graphs $\mathcal{G}_{i}$’s into batches $\mathcal{B}=\{\mathcal{B}_{k}\}$ such that $|\mathcal{V}_{i}|=|\mathcal{V}_{j}|$ for all $\mathcal{G}_{i},\mathcal{G}_{j}\in\mathcal{B}_{k}$. Then, at each iteration, we select a batch at random, build a computation graph for the number of nodes corresponding to the batch using the parameters estimated in the previous iteration, and update the parameters using the computation graph and the batch of graphs. Such a procedure helps to reduce the overhead time for building the computational graph, from per sample to per batch. This batching and training process is summarizedd in Algorithm 2, where “CreateBatches(…)” group the training graphs into batches, “BuildComputationalGraph(…)” builds the computation graph “NeVAE” using the parameters from the previous iteration and a given number of nodes, “Nodes(…)” returns the number of nodes of the graphs in a batch, and “Train(…)” updates the parameters given the computation graph and the parameters from the previous iteration.

Hardware and software specifications. We carried out all our experiments for NeVAE using Tensorflow 1.4.1, on a 64 bit Debian distribution with 16 core Intel Xenon CPU (E5-2667 v4 @3.20 GHz) and 512GB RAM.

Appendix B Additional details on Bayesian optimization

To implement Bayesian optimization (BO) for property-oriented molecule generation, we proceed similarly as in previous work . More specifically, we first sample $3{,}000$ molecules from our ZINC dataset, which we split into training (90%) and test (10%) sets. Then, for our model and each competing model with public domain implementations, we train a sparse Gaussian process (SGP) with the latent representations and $y(m)$ values of $100$ inducing points sampled from the training set. The SGPs allow us to make predictions for the property values of new molecules in the latent spaces. Then, we run $5$ iterations of batch Bayesian optimization (BO) using the expected improvement (EI) heuristic , with $50$ (new) latent vectors (molecules) per iteration.

In this section, we complement the performance comparison between NeVAE, GrammarVAE, CVAE, and JTVAE from Table 3 using two additional quality measures:

the predictive performance of the trained SGPs in terms of log-likelihood (LL) and root mean square error (RMSE) on the test set; and,

Appendix C Additional Experiments on Synthetic Graphs

In this section, we first demonstrate that our original model NeVAE is able to generate graphs with a predefined local topological property, i.e., graphs without triangles. Then, we show that our model is able to learn smooth latent representations of a popular type of random graphs, Kronecker graphs . Then, we present additional quantitative results on the ability of our model to learn and mimic the generative processes that determine the absence or presence of nodes and edges in of Kronecker graphs and Barabási-Albert graphs , a scalability analysis and finally illustrate the effect of node label permutations on the decoder parameter estimation. Finally, we show that the optimal property-oriented decoder designed using variational inference is able to generate synthetic graphs with certain structural properties.

We first generate two sets of synthetic networks, each containing $100$ graphs, with up to $n=1000$ number of nodes. The first set contains triangle free graphs and the second set contains a 50%-50% mixture of Kronecker graphs with initiator matrices: $\bm{\Theta}_{1}=[0.9,0.6;0.3,0.2]$, and $\bm{\Theta}_{2}=[0.6,0.6;0.6,0.6]$. For each dataset, we train our variational autoencoder for graphs by maximizing the corresponding evidence lower bound (ELBO). Then, we use the trained models to generate three sets of $1000$ graphs by sampling from the decoders, i.e., $\mathcal{G}\sim p_{\theta}(\mathcal{G}|\mathcal{Z})$, where $\mathcal{Z}\sim p(\mathcal{Z})$.

C.2 Quality of the generated graphs

Next, we evaluate the ability of our model to learn smooth latent representations of Kronecker graphs as follows. First, we select two graphs ($\mathcal{G}_{0}$ and $\mathcal{G}_{1}$) from the training set, one generated using an initiator matrix $\bm{\Theta}_{0}=[0.9,0.6;0.5,0.1]$ and the other using $\bm{\Theta}_{1}=[0.6,0.6;0.6,0.6]$. Then, we sample the latent representations $\mathcal{Z}_{0}$ and $\mathcal{Z}_{1}$ for $\mathcal{G}_{0}$ and $\mathcal{G}_{1}$, respectively, and sample new graphs from latent values $\mathcal{Z}$ in between these latent representations (using a linear interpolation), i.e., $\mathcal{G}\sim p_{\theta}(\mathcal{G}|\mathcal{Z})$, where $\mathcal{Z}=a\mathcal{Z}_{0}+(1-a)\mathcal{Z}_{1}$ and $a\in$, and the node labels, which define the matching between pairs of nodes in both graphs, are arbitrary. Figure 11 provides an example, which shows that, remarkably, as $\mathcal{Z}$ moves towards $\mathcal{Z}_{0}$ ($\mathcal{Z}_{1}$), the sampled graph becomes similar to that of $\mathcal{G}_{0}$ ($\mathcal{G}_{1}$) and the inferred initiator matrices along the way smoothly interpolate between both initiator matrix. Here, we infer the initiator matrices of the graphs generated by our trained decoder using the method by Leskovec et al. . Table 6 provides a quantitative evaluation of the quality of the generated graphs, i.e., it shows that the graphs our model generates are indistinguishable from true Kronecker graphs.

Finally, we create a set of $100$ graphs with up to $n=1000$ number of nodes sampled from the Barabási-Albert graph model with generation parameter $m=1$. For both Barabási-Albert and Kronecker graphs, we evaluate the quality of the generated graphs using two quantitative evaluation metrics:

where $\rho_{p_{\theta}},\rho_{\text{ELBO}}\in$.

where $\gamma^{\uparrow},\gamma^{\downarrow}\in$ and ${\mathcal{T}}_{x}^{\uparrow}$ (${\mathcal{T}}_{x}^{\downarrow}$) is the top (bottom) half of either $x={\mathcal{T}}_{p_{\theta}}$ or $x={\mathcal{T}}_{\text{ELBO}}$.

Table 6 summarizes the results, which show that our model is able to learn the generative process of Barabási-Albert more accurately than Kronecker graphs. This may be due to the higher complexity of the generative process Kronecker graph use. That being said, it is remarkable that our model is able to achieve correlation and precision values over $0.4$ in both cases.

C.3 Effect of permuting node labels on decoder parameter estimation

Figure 12 summarizes the results which show that degree based methods perform best in case of Barabási-Albert graph and uniform distribution performs best in Kronecker graph. This is because, the degree distribution of Barabási-Albert graph is skewed and as a result, a very few source nodes are sampled again and again, thereby giving similar parameter values. On the other hand, for homogeneous Kronecker graph, the degree distribution is more or less uniform. Consequently, the degree based methods perform worse in that case.

C.4 Effect of KK (search-depth in encoder) on model performance

In this section, we investigate the behavior of our model with respect to the search depths $K$ used in the decoder. Figure 13 summarizes the results, which show that, for Barabási-Albert graphs, our model performs consistently well for low values of $K$, however, for Kronecker graphs, the performance is better for high values of $K$. A plausible explanation for this is that Barabási-Albert networks are generated sequentially using only local topological features (only node-degrees), whereas the generation process of Kronecker graphs incorporates global topological features.

C.5 Scalability

We first compute the running time of our variational inference procedure against the size of the graphs in the training set and then compute the running time of our probabilistic decoder against the size of the sampled (generated) graphs. Figure 14 summarizes the results, which show that both in terms of inference and sampling, our model easily scales to $\sim$$1{,}000$ nodes. For example, for graphs with $1000$ nodes (average degree $3$), our inference procedure takes $67+20$ seconds to run one iteration of SGD with a batch size of $10$ graphs and, for graphs with $50$ nodes, our inference procedure takes less than $10$ seconds per iteration. Moreover, our probabilistic decoder can sample a graph with $1000$ ($50$) nodes (average degree $3$) in only $5$ ($0.5$) seconds.

C.6 Property oriented graph generation

Table 7 summarizes the results for different $\rho$, which shows that, the lower the value of $\rho$, the higher the success rate.