GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation

1 INTRODUCTION . Designing novel molecular structures with desired properties is a fundamental problem in a variety of applications such as drug discovery and material science . The problem is very challenging , since the chemical space is discrete by nature , and the entire search space is huge , which is believed to be as large as 1033 ( Polishchuk et al. , 2013 ) . Machine learning techniques have seen a big opportunity in molecular design thanks to the large amount of data in these domains . Recently , there are increasing efforts in developing machine learning algorithms that can automatically generate chemically valid molecular structures and meanwhile optimize their properties . Specifically , significant progress has been achieved by representing molecular structures as graphs and generating graph structures with deep generative models , e.g. , Variational Autoencoders ( VAEs ) ( Kingma & Welling , 2013 ) , Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) and Autoregressive Models ( Van Oord et al. , 2016 ) . For example , Jin et al . ( 2018 ) proposed a Junction Tree VAE ( JT-VAE ) for molecular structure encoding and decoding . De Cao & Kipf ( 2018 ) studied how to use GANs for molecular graph generation . You et al . ( 2018a ) proposed an approach called Graph Convolutional Policy Network ( GCPN ) , which formulated molecular graph generation as a sequential decision process and dynamically generates the nodes and edges based on the ∗Equal contribution , with order determined by flipping a coin . Work was done during internship at Mila . 1Code is available at https : //github.com/DeepGraphLearning/GraphAF existing graph substructures . They used reinforcement learning to optimize the properties of generated graph structures . Recently , another very related work called MolecularRNN ( MRNN ) ( Popova et al. , 2019 ) proposed to use an autoregressive model for molecular graph generation . The autoregressive based approaches including both GCPN and MRNN have demonstrated very competitive performance in a variety of tasks on molecular graph generation . Recently , besides the aforementioned three types of generative models , normalizing flows have made significant progress and have been successfully applied to a variety of tasks including density estimation ( Dinh et al. , 2016 ; Papamakarios et al. , 2017 ) , variational inference ( Kingma et al. , 2016 ; Louizos & Welling , 2017 ; Rezende & Mohamed , 2015 ) , and image generation ( Kingma & Dhariwal , 2018 ) . Flow-based approaches define invertible transformations between a latent base distribution ( e.g . Gaussian distribution ) and real-world high-dimensional data ( e.g . images and speech ) . Such an invertible mapping allows the calculation of the exact data likelihood . Meanwhile , by using multiple layers of non-linear transformation between the hidden space and observation space , flows have a high capacity to model the data density . Moreover , different architectures can be designed to promote fast training ( Papamakarios et al. , 2017 ) or fast sampling ( Kingma et al. , 2016 ) depending on the requirement of different applications . Inspired by existing work on autoregressive models and recent progress of deep generative models with normalizing flows , we propose a flow-based autoregressive model called GraphAF for molecular graph generation . GraphAF effectively combines the advantages of autoregressive and flow-based approaches . It has a high model capacity and hence is capable of modeling the density of real-world molecule data . The sampling process of GraphAF is designed as an autoregressive model , which dynamically generates the nodes and edges based on existing sub-graph structures . Similar to existing models such as GCPN and MRNN , such a sequential generation process allows leveraging chemical domain knowledge and valency checking in each generation step , which guarantees the validity of generated molecular structures . Meanwhile , different from GCPN and MRNN as an autoregressive model during training , GraphAF defines a feedforward neural network from molecular graph structures to the base distribution and is therefore able to compute the exact data likelihood in parallel . As a result , the training process of GraphAF is very efficient . We conduct extensive experiments on the standard ZINC ( Irwin et al. , 2012 ) dataset . Results show that the training of GraphAF is significantly efficient , which is two times faster than the state-of-theart model GCPN . The generated molecules are 100 % valid by incorporating the chemical rules during generation . We are also surprised to find that even without using the chemical rules for valency checking during generation , the percentage of valid molecules generated by GraphAF can be still as high as 68 % , which is significantly higher than existing state-of-the-art GCPN . This shows that GraphAF indeed has the high model capability to learn the data distribution of molecule structures . We further fine-tune the generation process with reinforcement learning to optimize the chemical properties of generated molecules . Results show that GraphAF significantly outperforms previous state-of-the-art GCPN on both property optimization and constrained property optimization tasks . 2 RELATED WORK . A variety of deep generative models have been proposed for molecular graph generation recently ( Segler et al. , 2017 ; Olivecrona et al. , 2017 ; Samanta et al. , 2018 ; Neil et al. , 2018 ) . The RVAE model ( Ma et al. , 2018 ) used a variational autoencoder for molecule generation , and proposed a novel regularization framework to ensure semantic validity . Jin et al . ( 2018 ) proposed to represent a molecule as a junction tree of chemical scaffolds and proposed the JT-VAE model for molecule generation . For the VAE-based approaches , the optimization of chemical properties is usually done by searching in the latent space with Bayesian Optimization ( Jin et al. , 2018 ) . De Cao & Kipf ( 2018 ) used Generative Adversarial Networks for molecule generation . The state-of-the-art models are built on autoregressive based approaches ( You et al. , 2018a ; Popova et al. , 2019 ) . ( You et al. , 2018a ) formulated the problem as a sequential decision process by dynamically adding new nodes and edges based on current sub-graph structures , and the generation policy network is trained by a reinforcement learning framework . Recently , Popova et al . ( 2019 ) proposed an autoregressive model called MolecularRNN to generate new nodes and edges based on the generated nodes and edge sequences . The iterative nature of autoregressive model allows effectively leveraging chemical rules for valency checking during generation and hence the proportion of valid molecules generated by these models is very high . However , due to the sequential generation nature , the training process is usually slow . Our GraphAF approach enjoys the advantage of iterative generation process like autoregressive models ( the mapping from latent space to observation space ) and meanwhile calculates the exact likelihood corresponding to a feedforward neural network ( the mapping from observation space to latent space ) , which can be implemented efficiently through parallel computation . Two recent work—Graph Normalizing Flows ( GNF ) ( Liu et al. , 2019 ) and GraphNVP ( Madhawa et al. , 2019 ) —are also flow-based approaches for graph generation . However , our work is fundamentally different from their work . GNF defines a normalizing flow from a base distribution to the hidden node representations of a pretrained Graph Autoencoders . The generation scheme is done through two separate stages by first generating the node embeddings with the normalizing flow and then generate the graphs based on the generated node embeddings in the first stage . By contrast , in GraphAF , we define an autoregressive flow from a base distribution directly to the molecular graph structures , which can be trained end-to-end . GraphNVP also defines a normalizing flow from a base distribution to the molecular graph structures . However , the generation process of GraphNVP is one-shot , which can not effectively capture graph structures and also can not guarantee the validity of generated molecules . In our GraphAF , we formulate the generation process as a sequential decision process and effectively capture the sub-graph structures via graph neural networks , based on which we define a policy function to generate the nodes and edges . The sequential generation process also allows incorporating the chemical rules . As a result , the validity of the generated molecules can be guaranteed . We summarize existing approaches in Table 1 . 3 PRELIMINARIES . 3.1 AUTOREGRESSIVE FLOW . A normalizing flow ( Kobyzev et al. , 2019 ) defines a parameterized invertible deterministic transformation from a base distribution E ( the latent space , e.g. , Gaussian distribution ) to real-world observational space Z ( e.g . images and speech ) . Let f : E → Z be an invertible transformation where ∼ pE ( ) is the base distribution , then we can compute the density function of real-world data z , i.e. , pZ ( z ) , via the change-of-variables formula : pZ ( z ) = pE ( f−1θ ( z ) ) ∣∣∣∣det ∂f−1θ ( z ) ∂z ∣∣∣∣ . ( 1 ) Now considering two key processes of normalizing flows as a generative model : ( 1 ) Calculating Data Likelihood : given a datapoint z , the exact density pZ ( z ) can be calculated by inverting the transformation f , = f−1θ ( z ) ; ( 2 ) Sampling : z can be sampled from the distribution pZ ( z ) by first sample ∼ pE ( ) and then perform the feedforward transformation z = fθ ( ) . To efficiently perform the above mentioned operations , fθ is required to be invertible with an easily computable Jacobian determinant . Autoregressive flows ( AF ) , originally proposed in Papamakarios et al . ( 2017 ) , is a variant that satisfies these criteria , which holds a triangular Jacobian matrix , and the determinant can be computed linearly . Formally , given z ∈ RD ( D is the dimension of observation data ) , the autoregressive conditional probabilities can be parameterized as Gaussian distributions : p ( zd|z1 : d−1 ) = N ( zd|µd , ( αd ) 2 ) , where µd = gµ ( z1 : d−1 ; θ ) , αd = gα ( z1 : d−1 ; θ ) , ( 2 ) where gµ and gα are unconstrained and positive scalar functions of z1 : d−1 respectively to compute the mean and deviation . In practice , these functions can be implemented as neural networks . The affine transformation of AF can be written as : fθ ( d ) = zd = µd + αd · d ; f−1θ ( zd ) = d = zd − µd αd . ( 3 ) The Jacobian matrix in AF is triangular , since ∂zi∂ j is non-zero only for j ≤ i . Therefore , the determinant can be efficiently computed through ∏D d=1 αd . Specifically , to perform density estimation , we can apply all individual scalar affine transformations in parallel to compute the base density , each of which depends on previous variables z1 : d−1 ; to sample z , we can first sample ∈ RD and compute z1 through the affine transformation , and then each subsequent zd can be computed sequentially based on previously observed z1 : d−1 . 3.2 GRAPH REPRESENTATION LEARNING . Following existing work , we also represent a molecule as a graph G = ( A , X ) , where A is the adjacency tensor and X is the node feature matrix . Assuming there are n nodes in the graph , d and b are the number of different types of nodes and edges respectively , then A ∈ { 0 , 1 } n×n×b and X ∈ { 0 , 1 } n×d . Aijk = 1 if there exists a bond with type k between ith and jth nodes . Graph Convolutional Networks ( GCN ) ( Duvenaud et al. , 2015 ; Gilmer et al. , 2017 ; Kearnes et al. , 2016 ; Kipf & Welling , 2016 ; Schütt et al. , 2017 ) are a family of neural network architectures for learning representations of graphs . In this paper , we use a variant of Relational GCN ( RGCN ) ( Schlichtkrull et al. , 2018 ) to learn the node representations ( i.e. , atoms ) of graphs with categorical edge types . Let k denote the embedding dimension . We compute the node embeddings H l ∈ Rn×k at the lth layer of R-GCN by aggregating messages from different edge types : H l = Agg ( ReLU ( { D̃− 1 2 i ẼiD̃ − 12 i H l−1W li } ∣∣i ∈ ( 1 , . . . , b ) ) ) , ( 4 ) where Ei = A [ : , : ,i ] denotes the ith slice of edge-conditioned adjacency tensor , Ẽi = Ei + I , and D̃i = ∑ k Ẽi [ j , k ] . W ( l ) i is a trainable weight matrix for the i th edge type . Agg ( · ) denotes an aggregation function such as mean pooling or summation . The initial hidden node representation H0 is set as the original node feature matrix X . After L message passing layers , we use the the final hidden representation HL as the node representations . Meanwhile , the whole graph representations can be defined by aggregating the whole node representations using a readout function ( Hamilton et al. , 2017 ) , e.g. , summation .

Summary: This paper proposes a generative model architecture for molecular graph generation based on autoregressive flows. The main contribution of this paper is to combine existing techniques (auto-regressive BFS-ordered generation of graphs, normalizing flows, dequantization by Gaussian noise, fine-tuning based on reinforcement learning for molecular property optimization, and validity constrained sampling). Most of these techniques are well-established either for data generation with normalizing flows or for molecular graph generation and the novelty lies in the combination of these building blocks into a framework. Training can be carried out in parallel over the sequential generation process, as no hidden states with sequential dependency are assumed (unlike a regular RNN). Experimental validation is carried out on a standard ZINC molecule generation benchmark (graphs with up to 48 nodes) and the reported metrics are competitive with recent related work.