Learning Multimodal Graph-to-Graph Translation for Molecule Optimization

1 INTRODUCTION . The goal of drug discovery is to design molecules with desirable chemical properties . The task is challenging since the chemical space is vast and often difficult to navigate . One of the prevailing approaches , known as matched molecular pair analysis ( MMPA ) ( Griffen et al. , 2011 ; Dossetter et al. , 2013 ) , learns rules for generating “ molecular paraphrases ” that are likely to improve target chemical properties . The setup is analogous to machine translation : MMPA takes as input molecular pairs { ( X , Y ) } , where Y is a paraphrase of X with better chemical properties . However , current MMPA methods distill the matched pairs into graph transformation rules rather than treating it as a general translation problem over graphs based on parallel data . In this paper , we formulate molecular optimization as graph-to-graph translation . Given a corpus of molecular pairs , our goal is to learn to translate input molecular graphs into better graphs . The proposed translation task involves many challenges . While several methods are available to encode graphs ( Duvenaud et al. , 2015 ; Li et al. , 2015 ; Lei et al. , 2017 ) , generating graphs as output is more challenging without resorting to a domain-specific graph linearization . In addition , the target molecular paraphrases are diverse since multiple strategies can be applied to improve a molecule . Therefore , our goal is to learn multimodal output distributions over graphs . To this end , we propose junction tree encoder-decoder , a refined graph-to-graph neural architecture that decodes molecular graphs with neural attention . To capture diverse outputs , we introduce stochastic latent codes into the decoding process and guide these codes to capture meaningful molecular variations . The basic learning problem can be cast as a variational autoencoder , where the posterior over the latent codes is inferred from input molecular pair ( X , Y ) . Further , to avoid invalid translations , we propose a novel adversarial training method to align the distribution of graphs generated from the model using randomly selected latent codes with the observed distribution of valid targets . Specifically , we perform adversarial regularization on the level of the hidden states created as part of the graph generation . We evaluate our model on three molecular optimization tasks , with target properties ranging from drug likeness to biological activity.1 As baselines , we utilize state-of-the-art graph generation methods ( Jin et al. , 2018 ; You et al. , 2018a ) and MMPA ( Dalke et al. , 2018 ) . We demonstrate that our model excels in discovering molecules with desired properties , outperforming the baselines across 1Code and data are available at https : //github.com/wengong-jin/iclr19-graph2graph different tasks . Meanwhile , our model can translate a given molecule into a diverse set of compounds , demonstrating the diversity of learned output distributions . 2 RELATED WORK . Molecular Generation/Optimization Prior work on molecular optimization approached the graph translation task through generative modeling ( Gómez-Bombarelli et al. , 2016 ; Segler et al. , 2017 ; Kusner et al. , 2017 ; Dai et al. , 2018 ; Jin et al. , 2018 ; Samanta et al. , 2018 ; Li et al. , 2018a ) and reinforcement learning ( Guimaraes et al. , 2017 ; Olivecrona et al. , 2017 ; Popova et al. , 2018 ; You et al. , 2018a ) . Earlier approaches represented molecules as SMILES strings ( Weininger , 1988 ) , while more recent methods represented them as graphs . Most of these methods coupled a molecule generator with a property predictor and solved the optimization problem through Bayesian optimization or reinforcement learning . In contrast , our model is trained to translate a molecular graph into a better graph through supervised learning , which is more sample efficient . Our approach is closely related to matched molecular pair analysis ( MMPA ) ( Griffen et al. , 2011 ; Dossetter et al. , 2013 ) in drug de novo design , where the matched pairs are hard-coded into graph transformation rules . MMPA ’ s main drawback is that large numbers of rules have to be realized ( e.g . millions ) to cover all the complex transformation patterns . In contrast , our approach uses neural networks to learn such transformations , which does not require the rules to be explicitly realized . Graph Neural Networks Our work is related to graph encoders and decoders . Previous work on graph encoders includes convolutional ( Scarselli et al. , 2009 ; Bruna et al. , 2013 ; Henaff et al. , 2015 ; Duvenaud et al. , 2015 ; Niepert et al. , 2016 ; Defferrard et al. , 2016 ; Kondor et al. , 2018 ) and recurrent architectures ( Li et al. , 2015 ; Dai et al. , 2016 ; Lei et al. , 2017 ) . Graph encoders have been applied to social network analysis ( Kipf & Welling , 2016 ; Hamilton et al. , 2017 ) and chemistry ( Kearnes et al. , 2016 ; Gilmer et al. , 2017 ; Schütt et al. , 2017 ; Jin et al. , 2017 ) . Recently proposed graph decoders ( Simonovsky & Komodakis , 2018 ; Li et al. , 2018b ; Jin et al. , 2018 ; You et al. , 2018b ; Liu et al. , 2018 ) focus on learning generative models of graphs . While our model builds on Jin et al . ( 2018 ) to generate graphs , we contribute new techniques to learn multimodal graph-to-graph mappings . Image/Text Style Translation Our work is closely related to image-to-image translation ( Isola et al. , 2017 ) , which was later extended by Zhu et al . ( 2017 ) to learn multimodal mappings . Our adversarial training technique is inspired by recent text style transfer methods ( Shen et al. , 2017 ; Zhao et al. , 2018 ) that adversarially regularize the continuous representation of discrete structures to enable end-to-end training . Our technical contribution is a novel adversarial regularization over graphs that constrains their scaffold structures in a continuous manner . 3 JUNCTION TREE ENCODER-DECODER . Our translation model extends the junction tree variational autoencoder ( Jin et al. , 2018 ) to an encoder-decoder architecture for learning graph-to-graph mappings . Following their work , we interpret each molecule as having been built from subgraphs ( clusters of atoms ) chosen from a vocabulary of valid chemical substructures . The clusters form a junction tree representing the scaffold structure of molecules ( Figure 1 ) , which is an important factor in drug design . Molecules are decoded hierarchically by first generating the junction trees and then combining the nodes of the tree into a molecule . This coarse-to-fine approach allows us to easily enforce the chemical validity of generated graphs , and provides an enriched representation that encodes molecules at different scales . In terms of model architecture , the encoder is a graph message passing network that embeds both nodes in the tree and graph into continuous vectors . The decoder consists of a tree-structured decoder for predicting junction trees , and a graph decoder that learns to combine clusters in the predicted junction tree into a molecule . Our key departures from Jin et al . ( 2018 ) include a unified encoder architecture for trees and graphs , along with an attention mechanism in the tree decoding process . 3.1 TREE AND GRAPH ENCODER . Viewing trees as graphs , we encode both junction trees and graphs using graph message passing networks . Specifically , a graph is defined as G = ( V , E ) where V is the vertex set and E the edge set . Each node v has a feature vector fv . For atoms , it includes the atom type , valence , and other atomic properties . For clusters in the junction tree , fv is a one-hot vector indicating its cluster label . Similarly , each edge ( u , v ) ∈ E has a feature vector fuv . Let N ( v ) be the set of neighbor nodes of v. There are two hidden vectors νuv and νvu for each edge ( u , v ) representing the message from u to v and vice versa . These messages are updated iteratively via neural network g1 ( · ) : ν ( t ) uv = g1 ( fu , fuv , ∑ w∈N ( u ) \v ν ( t−1 ) wu ) ( 1 ) where ν ( t ) uv is the message computed in the t-th iteration , initialized with ν ( 0 ) uv = 0 . In each iteration , all messages are updated asynchronously , as there is no natural order among the nodes . This is different from the tree encoding algorithm in Jin et al . ( 2018 ) , where a root node was specified and an artificial order was imposed on the message updates . Removing this artifact is necessary as the learned embeddings will be biased by the artificial order . After T steps of iteration , we aggregate messages via another neural network g2 ( · ) to derive the latent vector of each vertex , which captures its local graph ( or tree ) structure : xu = g2 ( fu , ∑ v∈N ( u ) ν ( T ) vu ) ( 2 ) Applying the above message passing network to junction tree T and graph G yields two sets of vectors { xT1 , · · · , xTn } and { xG1 , · · · , xGn } . The tree vector xTi is the embedding of tree node i , and the graph vector xGj is the embedding of graph node j . 3.2 JUNCTION TREE DECODER . We generate a junction tree T = ( V , E ) with a tree recurrent neural network with an attention mechanism . The tree is constructed in a top-down fashion by expanding the tree one node at a time . Formally , let Ẽ = { ( i1 , j1 ) , · · · , ( im , jm ) } be the edges traversed in a depth first traversal over tree T , where m = 2|E| as each edge is traversed in both directions . Let Ẽt be the first t edges in Ẽ . At the t-th decoding step , the model visits node it and receives message vectors hij from its neighbors . The message hit , jt is updated through a tree Gated Recurrent Unit ( Jin et al. , 2018 ) : hit , jt = GRU ( fit , { hk , it } ( k , it ) ∈Ẽt , k 6=jt ) ( 3 ) Topological Prediction When the model visits node it , it first computes a predictive hidden state ht by combining node features fit and inward messages { hk , it } via a one hidden layer network . The model then makes a binary prediction on whether to expand a new node or backtrack to the parent of it . This probability is computed by aggregating the source encodings { xT∗ } and { xG∗ } through an attention layer , followed by a feed-forward network ( τ ( · ) stands for ReLU and σ ( · ) for sigmoid ) : ht = τ ( W d 1 fit +W d 2 ∑ ( k , it ) ∈Ẽt hk , it ) ( 4 ) cdt = attention ( ht , { xT∗ } , { xG∗ } ; Udatt ) ( 5 ) pt = σ ( ud · τ ( W d3 ht +W d4 cdt ) ) ( 6 ) Here we use attention ( · ; Udatt ) to mean the attention mechanism with parameters Udatt . It computes two set of attention scores { αT∗ } , { αG∗ } ( normalized by softmax ) over source tree and graph vectors respectively . The output cdt is a concatenation of tree and graph attention vectors : cdt = [ ∑ i αTi , tx T i , ∑ i αGi , tx G i ] ( 7 ) Label Prediction If node jt is a new child to be generated from parent it , we predict its label by clt = attention ( hit , jt , { xT∗ } , { xG∗ } ; U latt ) ( 8 ) qt = softmax ( U lτ ( W l1hit , jt +W l 2c l t ) ) ( 9 ) where qt is a distribution over the label vocabulary and U latt is another set of attention parameters .

Summary: This paper proposed an extension of JT-VAE [1] into the graph to graph translation scenario. To help make the translation model predicting diverse and valid outcomes, the author added the latent variable to capture the multi-modality, and an adversarial regularization in the latent space. Experiment on molecule translation tasks show significant improvement over existing methods.