Graph Residual Flow for Molecular Graph Generation

1 INTRODUCTION . We propose a deep generative model for molecular graphs based on invertible functions . We especially focus on introducing an invertible function that is tuned for the use in graph structured data , which allows for flexible mappings with less number of parameters than previous invertible models for graphs . Molecular graph generation is one of the hot trends in the graph analysis with a potential for important applications such as in silico new material discovery and drug candidate screening . Previous generative models for molecules deal with string representations called SMILES ( e.g . Kusner et al . ( 2017 ) ; Gómez-Bombarelli et al . ( 2018 ) ) , which does not consider graph topology . Recent models such as ( Jin et al. , 2018 ; You et al. , 2018 ; De Cao & Kipf , 2018 ; Madhawa et al. , 2019 ) are able to directly handle graphs . Several researchers are investigating this topic using sophisticated statistical models such as variational autoencoders ( VAEs ) ( Kingma & Welling , 2014 ) , adversarial loss-based models such as generative adversarial networks ( GANs ) ( Goodfellow et al. , 2014 ; Radford et al. , 2015 ) , and invertible flows ( Kobyzev et al. , 2019 ) and have achieved desirable performances . The decoders of these graph generation models generate a discrete graph-structured data from a ( typically continuous ) representation of a data sample , which is modeled by aforementioned statistical models . In general , it is difficult to design a decoder that balances the efficacy of the graph generation and the simplicity of the implementation and training . For example , MolGAN ( De Cao & Kipf , 2018 ) has a relatively simple decoder but suffers from generating numerous duplicated graph samples . The state-of-the-art VAE-based models such as ( Jin et al. , 2018 ; Liu et al. , 2018 ) have good generation performance but their decoding scheme is highly complicated and requires careful training . On the contrary , invertible flow-based statistical models ( Dinh et al. , 2015 ; Kobyzev et al. , 2019 ) do not require training for their decoders because the decoders are simply the inverse mapping of the encoders and are known for good generation performances in image generation ( Dinh et al. , 2017 ; Kingma & Dhariwal , 2018 ) . Liu et al . ( 2019 ) proposes an invertible-flow based graph generation model . However , their generative model is not invertible because its decoder for graph structure is not built upon invertible flows . The GraphNVP by Madhawa et al . ( 2019 ) is the seminal fully invertible-flow approach for graph generation , which successfully combines the invertible maps with the generic graph convolutional networks ( GCNs , e.g Kipf & Welling ( 2017 ) ; Schlichtkrull et al . ( 2017 ) ) . However , the coupling flow ( Kobyzev et al. , 2019 ) used in the GraphNVP has a serious drawback when applied to sparse graphs such as molecular graphs we are interested in . The coupling flow requires a disjoint partitioning of the latent representation of the data ( graph ) in each layer . We need to design this partitioning carefully so that all the attributes of a latent representation are well mixed through stacks of mapping layers . However , molecular graphs are highly sparse in general : degree of each node atom is at most four ( valency ) , and only few kind of atoms comprise the majority of the molecules ( less diversity ) . Madhawa et al . ( 2019 ) argued that only a specific form of partitioning can lead to a desirable performance owing to sparsity : for each mapping layer , the representation of only one node is subject to update and all the other nodes are kept intact . In other words , a graph with 100 nodes requires at least 100 layers . But with the 100 layers , only one affine mapping is executed for each attribute of the latent representation . Therefore , the complexity of the mappings of GraphNVP is extremely low in contrast to the number of layer stacks . We assume that this is why the generation performance of GraphNVP is less impressive than other state-of-the-art models ( Jin et al. , 2018 ; Liu et al. , 2018 ) in the paper . In this paper we propose a new graph flow , called graph residual flow ( GRF ) : a novel combination of a generic GCN and recently proposed residual flows ( Behrmann et al. , 2019 ; Song et al. , 2019 ; Chen et al. , 2019 ) . The GRF does not require partitioning of a latent vector and can update all the node attributes in each layer . Thus , a 100 layer-stacked flow model can apply the ( non-linear ) mappings 100 times for each attribute of the latent vector of the 100-node graph . We derive a theoretical guarantee of the invertibility of the GRF and introduce constraints on the GRF parameters , based on rigorous mathematical calculations . Through experiments with most popular graph generation datasets , we observe that a generative model based on the proposed GRF can achieve a generation performance comparable to the GraphNVP Madhawa et al . ( 2019 ) , but with much fewer trainable parameters . To summarize , our contributions in this paper are as follows : • propose the graph residual flow ( GRF ) : a novel residual flow model for graph generation that is compatible with a generic GCNs . • prove conditions such that the GRFs are invertible and present how to keep the entire network invertible throughout the training and sampling . • demonstrate the efficacy of the GRF-based models in generating molecular graphs ; in other words , show that a generative model based on the GRF has much fewer trainable parameters compared to the GraphNVP , while still maintaining a comparable generation performance . 2 BACKGROUND . 2.1 GRAPHNVP . We first describe the GraphNVP ( Madhawa et al. , 2019 ) , the first fully invertible model for chemical graph generation , as a baseline . We simultaneously introduce the necessary notations for graph generative models . We use the notation G = ( A , X ) to represent a graph G comprising an adjacency tensor A and a feature matrix X . Let N be the number of nodes in the graph , M be the number of the types of nodes , and R be the number of the types of edges . Then , A ∈ { 0 , 1 } N×N×R and X ∈ { 0 , 1 } N×M . In the case of molecular graphs , G = ( A , X ) represents a molecule with R types of bonds ( single , double , etc . ) and M the types of atoms ( e.g. , oxygen , carbon , etc. ) . Our objective is to train an invertible model fθ with parameters θ that maps G into a latent point z = fθ ( G ) ∈ RD= ( N×N×R ) + ( N×M ) . We describe fθ as a normalizing flow composed of multiple invertible functions . Let z be a latent vector drawn from a known prior distribution pz ( z ) ( e.g. , Gaussian ) : z ∼ pz ( z ) . After applying a variable transformation , the log probability of a given graph G can be calculated as : log ( pG ( G ) ) = log ( pz ( z ) ) + log ( ∣∣∣∣det ( ∂z∂G ) ∣∣∣∣ ) , ( 1 ) where ∂z∂G is the Jacobian of fθ at G. In ( Madhawa et al. , 2019 ) fθ is modeled by two types of invertible non-volume preserving ( NVP ) mappings ( Dinh et al. , 2017 ) . The first type of mapping is the one that transforms the adjacency tensor , and the second type is the one that transforms the node attribute X . Let us divide the hidden variable z into two parts z = [ zX , zA ] ; the former zX is derived from invertible mappings of X and the latter zA is derived from invertible mappings of A . For the mapping of the feature matrix X , the GraphNVP provides a node feature coupling : z ( ` ) X [ ` , : ] ← z ( ` −1 ) X [ ` , : ] ◦ exp ( s ( z ( ` −1 ) X [ ` − , : ] , A ) ) + t ( z ( ` −1 ) X [ ` − , : ] , A ) , ( 2 ) where ` indicates the layer of the coupling , functions s and t stand for scale and translation operations , respectively , and ◦ denotes element-wise multiplication . We use zX [ ` − , : ] to denote a latent representation matrix of X ′ excluding the ` th row ( node ) . The rest of the rows of the feature matrix remains the same as follows . z ( ` ) X [ ` − , : ] ← z ( ` −1 ) X [ ` − , : ] . ( 3 ) s and t are modeled by a generic GCN , requiring the adjacency information of nodes , A , for better interactions between the nodes . For the mapping of the adjacency tensor , the GraphNVP provides an adjacency coupling : z ( ` ) A [ ` , : , : ] ← z ( ` −1 ) A [ ` , : , : ] ◦ exp ( s ( z ( ` −1 ) A [ ` − , : , : ] ) ) + t ( z ( ` −1 ) A [ ` − , : , : ] ) . ( 4 ) The rest of the rows remain as they are , as follows : z ( ` ) A [ ` − , : , : ] ← z ( ` −1 ) A [ ` − , : , : ] . ( 5 ) For adjacency coupling , we employ simple multi-layer perceptrons ( MLPs ) for s and t. The abovementioned formulations map only those variables that are related to a node ` in each ` -th layer ( Eqs . ( 2,4 ) ) , and the remaining nodes ` − are kept intact ( Eqs . ( 3,5 ) ) ; i.e . the partitioning of the variables always occurs in the first axis of tensors . This limits the parameterization of scaling and translation operations , resulting in reduced representation power of the model . In the original paper , the authors mention : “ masking ( switching ) ... w.r.t the node axis performs the best . ... We can easily formulate ... the slice indexing based on the non-node axis ... results in dramatically worse performance due to the sparsity of molecular graph. ” Here , sparsity can be described in two ways : one is the sparsity of non-carbon atoms in organic chemicals , and the other is the low degrees of atom nodes ( because of valency ) . 2.2 INVERTIBLE RESIDUAL BLOCKS . One of the major drawbacks of the partition-based coupling flow is that it covers a fairly limited family of mappings . Instead , the coupling flow offers computational cheap and analytic form of inversions . A series of recent invertible models ( Behrmann et al. , 2019 ; Song et al. , 2019 ; Chen et al. , 2019 ) propose a different approach for invertible mappings , called residual flow ( Kobyzev et al. , 2019 ) . They formulate ResNets ( He et al. , 2016 ) , which have been successful in image recognintion , as invertible mappings . The general idea is described as follows . Our objective is to develop an invertible residual layer for a vector z : z ( ` +1 ) = z ( ` ) + R ( z ( ` ) ) , ( 6 ) where z ( ` ) is the representation vector at the ` th layer , and R is a residual block . If we correctly constrain R , then we can assure the invertibility of the above-mentioned residual layer . i-ResNet ( Behrmann et al. , 2019 ) presents a constraint regarding the Lipschitz constant of R. MintNet ( Song et al. , 2019 ) limits the shape of the residual block R and derives the non-singularity requirements of the Jacobian of the ( limited ) residual block . Notably , the ( invertible ) residual connection ( Eq . ( 6 ) ) does not assume the partition of variables into “ intact ” and “ afine-map ” parts . This means that each layer of invertible residual connection updates all the variables at once . In both the aforementioned papers , local convolutional network architecture ( He et al. , 2016 ) of the residual block R is proposed for image tensor inputs , which can be applied for image generation/reconstructions for experimental validations . For example , in i-ResNet , the residual block is defined as : R ( x ) =W3 ◦ φ ◦W2 ◦ φ ◦W1 ( x ) , ( 7 ) where φ denotes a contractive nonlinear function such as ReLU and ELU , W· are ( spatially ) local convolutional layers ( i.e . aggregating the neighboring pixels ) . In this case , we put a constraint that the spectral norms of all W s are less than unity for the Lipschitz condition .

Summary: GraphNVP is the first paper to introduce the concept of "invertible flow", that is to construct the invertible mapping from latent vector z to the graph G. By constructing the mapping from G to z, GraphNVP first changes the discrete feature vector into continuous variables, then update this matrix representation by scaling and transforming functions (Eq. (2)-(5) in this GRF paper). In each iteration the matrix is only updated by one row (one slice for the tensor), while keep other rows intact. Then for constructing the inverse mapping, we can first sample a random vector and then apply the “inverse” of the update rule to recover the edge matrix and node matrix respectively.