Title: DEFactor: Differentiable Edge Factorization-based Probabilistic Graph Generation

Abstract: Generating novel molecules with optimal properties is a crucial step in many industries such as drug discovery. Recently, deep generative models have shown a promising way of performing de-novo molecular design. Although graph generative models are currently available they either have a graph size dependency in their number of parameters, limiting their use to only small graphs or are formulated as a sequence of discrete actions needed to construct a graph, making the output graph non-differentiable w.r.t the model parameters, therefore preventing them to be used in scenarios such as conditional graph generation. In this work we propose a model for conditional graph generation that directly optimises properties of the graph, and generates a probabilistic graph, making the decoding process differentiable. We demonstrate favourable performance of our model on prototype-based molecular graph conditional generation task.

1. Introduction:

In this paper we address the problem of learning probabilistic generative graph models for tasks such as the conditional generation of molecules with optimal properties. More precisely we focus on generating realistic molecular graphs, similar to a target molecule (the prototype).

The main challenge stems from the discrete nature of molecules. This particularly prevents us from using global discriminators that assess generated samples and back-propagate their gradients to guide the optimisation of the generator. This becomes very important in cases where we want to either optimise the property of a graph or explore the vicinity of an input graph (prototype) for conditional optimal generations, an approach that has proven successful in controlled image generation [Chen et al. (2016) ; Odena et al. (2016) ].

A number of recent approaches aim to address this limitation by performing indirect optimisation (You et al., 2018a; Li et al., 2018a) . You et al. You et al. (2018a) formulate the molecular graph optimisation task in a reinforcement learning setting, and optimise the loss with policy gradient Yu et al. (2016) . However policy gradient tends to suffer from high variance during training. Kang and Cho Kang & Cho (2018) suggest a reconstruction-based formulation which is directly applicable to discrete structures and does not require gradient estimation. However, it is limited by the number of samples available. Moreover, there is always a risk that the generator simply ignores the part of the latent code containing the property that we want to optimise. Finally, Jin et al, Jin et al. (2018) utilise Bayesian optimisation to optimise a proxy (the latent code) of the molecular graph, rather than the graph itself.

In contrast, Simonovsky and Komodakis Simonovsky & Komodakis (2018) and De Cao and Kipf Cao & Kipf (2018) have proposed decoding schemes that output graphs (adjacencies and node/edge features tensors) in a single step, and so are able to perform direct optimisation on the probabilistic continuous approximation of a graph. However, both decoding schemes make use of fixed size MLP layers which restricts their use to very small graphs of a predefined maximum size.

In this work, we address these issues directly by:

• Proposing a novel probabilistic graph decoding scheme that is computationally efficient and capable of generating arbitrary sized graphs (Section 3).

• Evaluating the generator's capacity to modify molecular graph attributes in the context of prototype-based molecular graph generation (Section 4).

Our approach (DEFactor) depicted in Figure 2 aims to directly address these issues with a probabilistic graph decoding scheme that is end-to-end differentiable, computationally efficient w.r.t the number of parameters in the model and capable of generating arbitrary sized graphs. We evaluate DEFactor on the task of constrained molecule property optimisation Jin et al. (2018) ; You et al. (2018a) and demonstrate that our results are competitive with recent results.

2. Related work:

Lead-based Molecule Optimisation. The aim here is to obtain molecules that satisfy a target set of objectives, for example activity against a biological target while not being toxic or maintaining certain properties, such as solubility. Currently a popular strategy is to fine-tune an pretrained generative model to produce/select molecules that satisfy a desired set of properties Segler et al. (2017) .

Bayesian optimisation is proposed to explore the learnt latent spaces for molecules in GÃşmez-Bombarelli et al. (2016) , and is shown to be effective at exploiting feature rich latent representations Kusner et al. (2017) ; Dai et al. (2018) ; Jin et al. (2018) . In Li et al. (2018b; a) sequential graph decoding schemes whereby conditioning properties can be added to the input are proposed. However these approaches are unable to perform direct optimisation for objectives. Finally You et al. (2018a) reformulates the problem in a reinforcement learning setting, and objective optimisation is performed while keeping an efficient sequential-like generative scheme You et al. (2018b) .

Graph Generation Models.:

Current work on graph generation can be divided into two type of graph decoding schemes. Sequential methods to graph generation [ You et al. (2018b) ; Li et al. (2018a) ; You et al. (2018a) ; Li et al. (2018b) ] aim to construct a graph by predicting a sequence of addition/edition actions of nodes/edges. Starting from a sub-graph (normally empty), at each time step a discrete transition is predicted and the sub-graph is updated. Although sequential approaches enable us to decouple the number of parameters in models from the the maximum size of the graph processed, due to the discretisation of the final outputs, the graph is still non-differentiable w.r.t. to the decoder's parameters. This again prevents us from directly optimising for the objectives we are interested in.

In contrast to the sequential process Cao & Kipf (2018); Simonovsky & Komodakis (2018) reconstruct probabilistic graphs. These methods however make use of fixed size MLP layers when decoding to predict the graph adjacency and node tensors. This however limits their use to very small graphs of a pre-chosen maximum size. They therefore restrict study and application to small molecular graphs; a maximum number of 9 heavy atoms, compared to approximately 40 in sequential models.

We propose to tackle these drawbacks by designing a graph decoding scheme that is:

• Efficient: so that the number of parameters of the decoder does not depend on a fixed maximum graph size.

• Differentiable: in particular we would like the final graph to be differentiable w.r.t the decoder's parameters, so that we are able to directly optimise the graph for target objectives.

Edge-factorization.:

The idea of using tensor factorization methods to decode edges is not new and has been extensively used in relational inference tasks (Nickel & Tresp (2013) ). Kipf & Welling (2016a) recently proposed an graph autoencoder for rich node embeddings learning and link prediction. However their formulation suppose a fixed size graph in which we want to predict existing and potential novel links between its nodes. Such assumption is not applicable to the case of molecular graph optimization where given a property we want to be able to generate a graph of a priori unknown size. To that extent we suggest a novel autoregressive model for nodes embeddings generation which, when combined with edge-factorization graph decoding constitute our full graph decoder. We describe the model more extensively in the next section. is the conditional setting a discriminator D that assesses the outputs and gives its feedback to the generator G. L rec (resp. L cond ) refers to the reconstruction (resp. condtional) loss described in section 3.3.

3. DEFactor:

Molecules can be represented as graphs G = (V, E) where atoms and bonds correspond to the nodes and edges respectively. Each node in V is labeled with its atom type which can be considered as part of its features. Equally, a molecular graph can be defined by its adjacency tensor E ∈ {0, 1} n×n×e where n is the number of nodes (atoms) in the graph and e is the number of edge (bond) types that we can have between two atoms and the node types are represented by a node feature tensor N ∈ {0, 1}

n×d which is composed of several one-hot-encoded properties. In the following sections we will describe our proposed DEFactor model applied to molecular graphs.

3.1. Graph Construction Process:

Given a molecular graph defined as G = (N, E) we propose to leverage the edge-specific information propagation framework described in Simonovsky & Komodakis (2017) to learn a set of informative embeddings from which we can directly infer a graph. More specifically, our graph construction process is composed of two parts:

• An Encoder that in -step 1 performs several spatial graph convolutions on the input graph, and in -step 2 aggregates those embeddings into a single graph latent representation.

• A Decoder that in -step 3 autoregressively generates a set of continuous node embeddings conditioned on the learnt latent representation, and in -step 4 reconstruct the whole graph in an edge-factorization fashion.

Figure 2 (a) and (b) provides a summary of those 4 steps.

Steps 1 and 2: Graph Representation Learning. We propose an encoder that makes an efficient use of the information contained in the bonds by having a separate information propagation channel for each bond type. The information is propagated in the graph using a Graph Convolutional Network (GCN) update rule (Kipf & Welling, 2016b ) so that each node embedding can be written as a weighted sum of the edge-conditioned information of its neighbors in the graph. Namely for each l-th layer of the encoder, the representation is given by:

EQUATION

where E e is the e-th frontal slice of the adjacency tensor, D e the corresponding degree tensor and W l e and W l s are learned parameters of the layer. Once we have those node embeddings we further aggregate them to obtain a fixed-length latent representation of the graph. We propose to parametrize this aggregation step by an LSTM and we compute the graph latent representation by a simple linear transformation of the last hidden state of this Aggregator:

EQUATION

Even though the use of an LSTM makes the aggregation non permutation-invariant, similar to GraphSAGE Hamilton et al. (2017) , we adapt the LSTM aggregator to work on a randomly permuted set of embeddings and noticed that it did not affect the performance of the model.

In the subsequent steps we are interested in designing a graph decoding scheme from the latent code that is both scalable and powerful enough to model the interdependencies between the nodes and edges in the graph.

Step 3: Autoregressive Embeddings Generation. As specified above, we are interested in building a graph decoding scheme that models the nodes and their connectivity (represented by continuous embeddings S) in an autoregressive fashion so that the latent code on which is conditioned the decoding has enough dimensions to encode more highlevel features. Like stated previously, methods that suggest to compute the graph all at once (Simonovsky & Komodakis, 2018; Cao & Kipf, 2018) model each node and edge as conditionally independent given the latent code z: this means that every detail of their dependencies within the graph has to be encoded in this latent variable: that does not leave much room for high-level feature learning if the size of the latent code is not chosen carefully. We propose to tackle this drawback by autoregressive generation of the continuous embeddings s = [s 0 , s 1 , ..., s n ] for n nodes, with the embedding containing enough information about the node itself and its neighbourhood. More precisely we model the generation of node embeddings such that:

p(s|z) = n i=1 p(s i |s <i , z).

(3)

In our model, the autoregressive generation of embeddings is parametrized by a simple Long Short-Term Memory (LSTM, ? and is completely deterministic such that at each time step t the LSTM decoder takes as input the previously generated embeddings and the latent code z which captured node-invariant features of the graph. Each embedding is computed as a function of the concatenation of the current hidden state and the latent code z such that:

EQUATION

)

where f d LST M corresponds to the LSTM recurrence operation and g in and f embed are parametrized as simple MLP layers to perform nonlinear feature extraction.

Step 4: Graph Decoding from Node Embeddings. As stated previously, we want to drive the generation of the continuous embeddings s towards latent factors that contains enough information about the node they represent (i.e. we can easily retrieve the one-hot atom type performing a linear transformation of the continuous embedding) and its neighbourhood (i.e. the adjacency tensor can be easily retrieved by comparing those embeddings in a pair-wise manner). For those reasons, we suggest to factorize each bond type that we can have between two atoms in a relational inference fashion (Zitnik et al., 2018; Kipf et al., 2018) .

Let S ∈ R n×p be the concatenated continuous node embeddings generated in the previous step. We reconstruct the adjacency tensor E by learning edge-specific similarity measure as follows for k-th edge type:

EQUATION

)

We model this by a set of edge-specific factors U = (u 1 , • • • , u e ) ∈ R e×p such that we can reconstruct the adjacency tensor as :

EQUATION

)

where σ is the logistic sigmoid function, D k the corresponding diagonal matrix of the vector u k and the factors (u i ) ∈ R e×p are learned parameters.

We reconstruct the node features (i.e. the atom type) with a simple affine transformation such that:

EQUATION

where W ∈ R p×d is a learned parameter.

These four steps define our proposed graph autoencoder.

Generating Graphs of Arbitrary Sizes. In order to generate graphs of different sizes we need to add what we call here an Existence module that retrieves a probability of belonging to the final graph from each continuous embedding generated by the embeddings generator (in step 3). This module is parametrized as a simple MLP layer followed by a sigmoid activation and stops the unrolling of the embedding LSTM generator whenever we encounter a non-informative embedding. This module can be interpreted as an < eos > translator.

3.2. Training:

Teacher forcing. In order to make the model converge in reasonable time we used a trick similar to the teacher-forcing based training of language models (Williams & Zipser, 1989 ).

The training is thus done in 3 phases:

• We first pre-train the GCN part along with the embedding decoder (factorization, nodes and existence modules) to reconstruct the graphs. This corresponds to the training of a simple Graph AE as in Kipf & Welling (2016a) except that we also want to reconstruct the nodes' one-hot features (and not just the relations).

• We then append those two units to the embedding aggregator and generator while keeping them fixed. In this second phase, the embedding generator is trained in a teacher forcing fashion where at each time step t the LSTM decoder does not take as input the previously generated embedding but the true one that is the direct output of the pretrained GCN embedding encoder.

• Finally in order to transition from a teacher-forcing to a fully autoregressive state we increasingly (Bengio et al., 2015) feed the LSTM generator more of its own predictions. When that fully autoregressive state is reached the pre-trained units are not frozen anymore and the whole model continues training end-to-end.

Log-Likelihood Estimates:

We train this first autoencoder framework in a MLE fashion with the following negative log-likelihood estimate L rec = L X + L X + L N corresponding to the existing edges (X), the non-existing edges ( X) and the node features (N ) reconstruction terms:

L X = - 1 |X| (i,j)∈X

E T i,j,: log( Ẽi,j,: ) + (1 -E i,j,: ) T log(1 -Ẽi,j,: ) ( 9)

EQUATION

)

where X is the set of existing edges, X the set of non existing edges, E the adjacency tensor, N the node features tensor, n the number of nodes in the graph. As molecular graphs are sparse we found that such separate normalisations were crucial for the training.

3.3. Conditional Generation and Optimisation:

Model overview. In this part we discuss our approach to build a conditional framework starting from the previous autoencoder architecture where the construction of a probabilistic graph G is conditioned on some unregularized latent code z derived from a given input graph. We then augment this unstructured z with a set of structured attributes y that represent some physico-chemical properties of interest so that the decoder is conditioned on the joint (z, y). At the end of a successful training we expect this decoder to generate samples that have the properties specified in y and to be similar (in terms of information contained in z) to the original query molecular graph (encoded as z). To do so we choose a mutual information maximization approach (detailed in the appendix) that involves the use of discriminators that assess the properties ỹ of the generated samples.

Discriminator Pre-Training:

In this phase we pre-train a discriminator to assess the property y of a generated sample so that we can backpropagate its feedback to the generator (the discriminator can be trained on another dataset and we can have several discriminators for several attributes of interest). In order to have informative gradients in the early stages of the training we have trained the discriminator on continuous approximations of the discrete training graphs (details in the appendix A.1) so that our objective becomes:

EQUATION

)

where the graphs sampled from pdata (x) are the probabilistic approximations of the discrete ones.

The next step is to incorporate the feedback signal of the trained discriminator in order to formulate the property attribute constraint. The training is decomposed in two phases in which we learn to reconstruct graphs of the dataset (MLE phase) and to modify chemical attributes (Variational MI maximization phase). 

Model:

EQUATION

)

Generator Learning. The generator is updated in both reconstruction and conditional phases. In the MLE phase the generator is trained with same loss L enc as the encoder so that it is pushed towards generating realistic molecular graphs. In the MI maximization phase we sample the attributes from a prior p(y) s.t. we minimize the following objective:

EQUATION

)

In this phase the only optimisation signal comes from the trained discriminator. Consequently there is a risk of falling off the manifold of the molecular graphs as no realism constraint is specified. A good way to make sure that this does not happen is to add a similar discriminator trained to distinguish between the real probabilistic graph and the generated ones so that when trying to satisfy the attribute constraint the generator is forced to produce valid molecular graphs. We leave that additional feature for future work.

4. Experiments:

To compare with recent results in constrained molecular graph optimization Jin et al. (2018) ; You et al. (2018a) , we present the following experiments :

• Molecular Graph Reconstruction: We test the autoencoder framework on the task of reconstructing input molecules from their latent representations.

• Conditional Generation: We test our conditional generation framework on the task of generating novel molecules that satisfy a given input property. Here, we are interested in the octanol-water partition coefficient (LogP) optimization used as benchmark in Kusner et al. (2017) ; Jin et al. (2018) ; You et al. (2018a) .

• Constrained Property Optimization: Finally, we test our conditional autoencoder on the task of modifying a given molecule to improve a specified property, while constraining the degree of deviation from the original molecule. Again we use the LogP benchmark for the experiment.

Molecular graph reconstruction::

In this task we evaluate the exact reconstruction error from encoding and decoding a given molecular graph from the test set. We report in Table 1 the ratio of exactly reconstructed graphs, where we see that the our autoencoder outperforms the Junction Tree-VAE (JT-VAE) Jin et al. (2018) which has the current stateof-the-art performance in this task. Appendix B.3 reports the reconstruction ratio as a function of the molecule size (number of heavy atoms).

Conditional Generation::

In this task we evaluate the conditional generation formulation described in Section 3.3. For a given molecule m i with an observed property value y i , the goal here is to modify the molecule to generate a new molecule with the given target property value; (m i , y * i ). New molecules are generated by conditioning the decoder on (z i ; y * i ), where z i is the latent code for m i . The decoded new molecule m i , is ideally best suited to satisfy the target property. This is evaluated by comparing the property value of the new molecule with the target property value. A generator that performs well at this task will produce predicted molecules with property values that are close to the target. In these experiments, LogP was chosen as the desired property, and we use RDKit, online to calculate the LogP values of generated molecules.

The scatter plots in Figure 2 give for a randomly selected set of test molecules, the correlation of target property values against the evaluated property value of the correctly decoded molecules.

Constrained Property Optimization::

In this section we follow the evaluation methodology outlined in Jin et al. (2018) ; You et al. (2018a) , and evaluate our model in a constrained molecule property optimization. In contrast to Jin et al. (2018) , because of the conditional formulation, our autoencoder is already suited for that task without the need for retraining.

Given the 800 molecules with the lowest penalized LogP1 property scores from the test set, we evaluate the decoder by providing pairs of (z i , y * i ) with increasing property scores, and among the valid decoded graphs we compute:

• Their similarity scores (Sim.) to the encoded target molecule (called the prototype);

• Their penalized LogP scores. Note that in this setting the conditioning property values (y * i ) are the unpenalized LogP scores. However, to evaluate the model we compute the penalized LogP scores to assess the model's ability to decode synthetically accessible molecules.

• While varying the similarity threshold values (δ), we compute the success rate (Suc.) for all 800 molecules. This measures how often we to get a novel molecule with an improved penalized LogP score.

• Finally, for different similarity thresholds, for successfully decoded molecules, we report the average improvements (Imp.) and the similarity (Sim.) for the molecule that is most improved. We compare our results with Jin et al. (2018) ; You et al. (2018a) .

The final results are reported in Table 2 . As can be seen, although slightly behind GCPN You et al. (2018a) w.r.t. success rates (Suc.), DEFactor significantly outperforms other models in terms of improvements (Imp.) achieved (by between 130% and 195% for thresholds 0.2 and 0.6 respectively, with respect to the next best model GCPN).

5. Future work:

In this paper, we designed a new way of generating molecular graphs in a conditional optimisation setting . We believe that our DEFactor model is a significant step forward (2018a) .

to build ML-driven applications for de-novo drug design or generation of molecules with optimal properties without resorting to methods that do not directly optimise the desired properties.

Note that a drawback of our model is that it uses an MLE training process which forces us to either fix the ordering of nodes or to perform a computationally expensive graph matching operation to compute the loss. Moreover in our fully deterministic conditional formulation we assume that chemical properties optimisation is a one-to-one mapping but in reality there may exist many suitable way of optimizing a molecule to satisfy one property condition while staying similar to the query molecule. To that extent it could be interesting to augment our model to include the possibility of a one-to-many mapping. Another way of improving the model could also be to include a validity constraint formulated as training a discriminator that discriminates between valid and generated graphs.

B.2 Mutual information maximization:

For the conditional setting we choose a simple mutual information maximization formulation. The objective is to maximize the MI I(X; Y ) between the target property Y and the decoder's output X = G θ (Y ) under the joint p θ (X, Y ) defined by the decoder G θ . In the conditional setting G θ is also conditioned on the encoded molecule z but for simplicity we treat it as a parameter of the decoder (and thus reason with one target molecule from which we want to modify attributes). We define the MI as: In our conditional setting we pre-trained the discriminators (parametrized by Q in the lower bound derivation) to approximate p data (y|x) which makes the bound tight only when p θ (y paired |x) is close to p data (y|x) and this corresponds to a stage where the decoder has maximized the log-likelihood of the data well enough (i.e. when it is able to reconstruct input graphs properly when z and y are paired). Thus, in the conditional setting we are maximizing the following objective: Notice that as we make use of a simple LSTM to encode a graph representation, there is a risk that for the largest molecules the long term dependencies of the embeddings are not captured well resulting in a bad reconstruction error. We capture this observation in figure 4 . One possible amelioration could be to add other at each step other non-sequential aggregation of the embeddings (average pooling of the emebeddings for example) or to make the encoder more powerful by adding some attention mechanisms. We leave those for future work.

L cond = E x,

B.4 Visual similarity samples:



:

The penalized logP is octanol-water partition coefficient (logP) penalized by the synthetic accessibility (SA) score and the number of long cycles, seeJin et al. (2018)

Figures and Tables:

Figure 1: (a) is the full autoencoder (step 1 to 4) , (b) is expanding the steps (3 for the LSTM and 4 for the factorization and node decoding) of the generator G of the autoencoder and (c) is the conditional setting a discriminator D that assesses the outputs and gives its feedback to the generator G. L rec (resp. L cond ) refers to the reconstruction (resp. condtional) loss described in section 3.3.

Figure 2: Conditional generation: The initial LogP value of the query molecule is specified as IL and the Pearson correlation coefficient is specified as c. We report on the y-axis the conditional value given as input and on the x-axis the true LogP of the generated graph when translated back into molecule (we have of course only reported the values when the decoded graphs correspond to valid molecules, which explains the difference in the total number of points for each query molecule).

(y; G θ (y)) = E x∼G θ (y) [E y ∼p θ (y|x) [log p θ (y |x)]] + H(y) = E x∼G θ (y) [D KL (p θ (.|x)||Q(.|x)) + E y ∼p θ (y|x) [log Q(y |x)]] + H(y) ≥ E x∼G θ (y) [E y ∼p θ (y|x) [log Q(y |x)]] + H(y)

y∼p data (x,y),z∼E(x),y ∼p(y) [log G θ (y, z) + I(y ; G θ (y , z))] B.3 Reconstruction as a function of number of atoms

Figure 5: Accuracy score as a function of the number of heavy atoms in the molecule(x axis) for different size of the latent code

Molecular graph reconstruction task. JT-VAE result is taken fromJin et al. (2018) and uses a latent code of size 56. We compared the expressivity of our decoder in a molecular reconstruction task with the JT-VAE decoder. JT-VAE model refers to the original VAE framework as described in inJin et al. (2018) whereas JT-AE is the adapted deterministic version of it (we removed the stereochemistry information of the molecules as we believe it is an unnecessary burden for the model). The encoder is updated only during the reconstruction phase where we sample attributes y from the true posterior. The encoder loss is a linear combination of the molecular graph reconstruction (L rec ) and the property reconstruction( L prop ) s.t. L

Table content: <table><tr><td>Reconstruction</td></tr></table>

Constrained penalized LogP maximisation task: each row shows a different level of similarity constraint δ and columns are for improvements (Imp.), similarity to the original query (Sim.), and the success rate (Suc.). Values for other models are taken fromYou et al.

Table content: <table/>

Peer review (iclr2019):
 This paper proposed a variant of the graph variational autoencoder [1] to do generative modeling of graphs. The author introduced an additional conditional variable (e.g., property value) into the decoder. By backpropagating through the discriminator, the model is able to find the graph with desired property value. Overall the paper reads well and is easy to follow. The conditional generation of graphs seems also helpful regarding the empirical performance. However, there are several concerns regarding the paper: 1) The edge factorization-based modeling is not new. In fact [1] already uses the node embeddings to factorize the adjacency matrix. This paper models extra information including node tags and edge types, but these are not fundamental differences compared to [1]. 2) The paper claims the method is ‘cheaper’ and ‘scalable’. Since essentially the computation cost is similar to [1] which requires at least O(n^2) to generate a graph with n nodes, I’m not super confident about the author’s claim. Though this can be parallelized, but the memory cost is still in this order of magnitude, which might be too much for a sparse graph. Also there’s no large graph generative modeling experiments available. 3) Continue with 2), the adjacency matrix of a large graph (e.g., graph with more than 1k nodes) doesn’t have to be low rank. So modeling with factorization (with typically ~256 embedding size) may not be suitable in this case. Some minor comments: 4) Regarding Eq (2), why the lstm is used, instead of some simple order invariant aggregation? 5) the paper needs more refinement. E.g., in the middle of page 2 there is a missing citation. [1] Kipf & Welling, Variational Graph Auto-Encoders, https://arxiv.org/pdf/1611.07308.pdf