DEFactor: Differentiable Edge Factorization-based Probabilistic Graph Generation

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 . 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 : H l = σ ( ∑ e [ D− 1 2 e EeD − 12 e H l−1W le ] +D−1H l−1W ls ) , ( 1 ) where Ee is the e-th frontal slice of the adjacency tensor , De the corresponding degree tensor and W le and W ls 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 : z = gagg ( feLSTM ( { HK } ) ) . ( 2 ) 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 = [ s0 , s1 , ... , sn ] 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 ( si|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 : ht+1 = fdLSTM ( gin ( [ z , st ] ) , ht ) ( 4 ) st+1 = fembed ( [ ht+1 , z ] ) , ( 5 ) where fdLSTM corresponds to the LSTM recurrence operation and gin and fembed 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 ∈ Rn×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 : p ( E : , : ,k|S ) = n∏ i=1 n∏ j=1 p ( Ei , j , k|si , sj ) . ( 6 ) We model this by a set of edge-specific factors U = ( u1 , · · · , ue ) ∈ Re×p such that we can reconstruct the adjacency tensor as : Ẽi , j , k = σ ( sTi Dksj ) = p ( Ei , j , k|si , sj ) , ( 7 ) where σ is the logistic sigmoid function , Dk the corresponding diagonal matrix of the vector uk and the factors ( ui ) ∈ Re×p are learned parameters . We reconstruct the node features ( i.e . the atom type ) with a simple affine transformation such that : Ñi , : = p ( Ni|si ) = softmax ( Wsi ) , ( 8 ) where W ∈ Rp×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 .

Summary: In this paper, authors propose a deep generative model and a variant for graph generation and conditional graph generation respectively. It exploits an encoder which is built based on GCN and GraphSAGE, a autoregressive LSTM decoder which generates the graph embedding, and a factorized edge based probabilistic model for generating edge and node type. For conditional generation, authors also propose a discriminating training scheme based on maximizing the mutual information. Experiments on ZINC dataset show that the proposed method is promising.