GRAPH TRANSFORMATION POLICY NETWORK FOR CHEMICAL REACTION PREDICTION

1 INTRODUCTION . Chemical reaction product prediction is a fundamental problem in organic chemistry . It paves the way for planning syntheses of new substances ( Chen & Baldi , 2009 ) . For decades , huge effort has been spent to solve this problem . However , most methods still depend on the handcrafted reaction rules ( Chen & Baldi , 2009 ; Kayala & Baldi , 2011 ; Wei et al. , 2016 ) or heuristically extracted reaction templates ( Segler & Waller , 2017 ; Coley et al. , 2017 ) , thus are not well generalizable to unseen reactions . A reaction can be regarded as a set ( or unordered sequence ) of graph transformations in which reactants represented as molecular graphs are transformed into products by modifying the bonds between some atom pairs ( Jochum et al. , 1980 ; Ugi et al. , 1979 ) . See Fig . 1 for an illustration . We call an atom pair ( u , v ) that changes its connectivity during reaction and its new bond b a reaction triple ( u , v , b ) . The reaction product prediction problem now becomes predicting a set of reaction triples given the input reactants and reagents . We argue that in order to solve this problem well , an intelligent system should have two key capabilities : ( a ) Understanding the molecular graph structure of the input reactants and reagents so that it can identify possible reactivity patterns ( i.e. , atom pairs with changing connectivity ) . ( b ) Knowing how to choose from these reactivity patterns a correct set of reaction triples to generate the desired products . Recent state-of-the-art methods ( Jin et al. , 2017 ; Bradshaw et al. , 2018 ) have built the first capability by leveraging graph neural networks ( Duvenaud et al. , 2015 ; Hamilton et al. , 2017 ; Pham et al. , 2017 ; Gilmer et al. , 2017 ) . However , these methods are either unaware of the valid sets of reaction triples ( Jin et al. , 2017 ) or limited to sequences of reaction triples with a predefined orders ( Bradshaw et al. , 2018 ) . The main challenge is that the space of all possible configurations of reaction triples is extremely large and non-differentiable . Moreover , a small change in the predicted set of reaction triples can lead to very different reaction products and a little mistake can produce invalid prediction . In this paper , we propose a novel method called Graph Transformation Policy Network ( GTPN ) that addresses the aforementioned challenges . Our model consists of three main components : a graph neural network ( GNN ) , a node pair prediction network ( NPPN ) and a policy network ( PN ) . Starting from the initial graph of reactant and reagent molecules , our model iteratively alternates between modeling an input graph using GNN and predicting a reaction triple using NPPN and PN to generate a new intermediate graph as input for the next step until it decides to stop . The final generated graph is considered as the predicted products of the reaction . Importantly , GTPN does not assume any fixed number or any order of bond changes but learn these properties itself . One can view GTPN as a reinforcement learning ( RL ) agent that operates on a complex and non-differentiable space of sets of reaction triples . To guide our model towards learning a diverse yet robust-to-small-changes policy , we customize our loss function by adding some useful constraints to the standard policy gradient loss ( Mnih et al. , 2016 ) . To the best of our knowledge , GTPN is the most generic approach for the reaction product prediction problem so far in the sense that : i ) It combines graph neural networks and reinforcement learning into a unified framework and trains everything end-to-end ; ii ) It does not use any handcrafted or heuristically extracted reaction rules/templates to predict the products . Instead , it automatically learns various types of reactions from the training data and can generalize to unseen reactions ; iii ) It can interpret how the products are formed via the sequence of reaction triples it generates . We evaluate GTPN on two large public datasets named USPTO-15k and USPTO . Our method significantly outperforms all baselines in the top-1 accuracy , achieving new state-of-the-art results of 82.39 % and 83.20 % on USPTO-15k and USPTO , respectively . In addition , we also provide comprehensive analyses about the performance of GTPN and about different types of errors our model could make . 2 METHOD . 2.1 CHEMICAL REACTION AS MARKOV DECISION PROCESS OF GRAPH TRANSFORMATIONS . A reaction occurs when reactant molecules interact with each other in the presence ( or absence ) of reagent molecules to form new product molecules by breaking or adding some of their bonds . Our main insight is that reaction product prediction can be formulated as predicting a sequence of such bond changes given the reactant and reagent molecules as input . A bond change is characterized by the atom pair ( where the change happens ) and the new bond type ( what is the change ) . We call this atom pair a reaction atom pair and call this atom pair with the new bond type a reaction triple . More formally , we represent the entire system of input reactant and reagent molecules as a labeled graph G = ( V , E ) with multiple connected components , each of which corresponds to a molecule . Nodes in V are atoms labeled with their atomic numbers and edges in E are bonds labeled with their bond types . Given G as input , we predict a sequence of reaction triples that transforms G into a graph of product molecules G′ . As reactions vary in number of transformation steps , we represent the sequence of reaction triples as ( ξ , u , v , b ) 0 , ( ξ , u , v , b ) 1 , ... , ( ξ , u , v , b ) T−1 or ( ξ , u , v , b ) 0 : T for short . Here T is the maximum number of steps , ( u , v ) is a pair of nodes , b is the new edge type of ( u , v ) , and ξ is a binary signal that indicates the end of the sequence . If the sequence ends at Tend < T , ξ0 , ... ξTend−1 will be 1 and ξTend , ... , ξT−1 will be 0 . At every step τ , if ξτ = 1 , we apply the predicted edge change ( u , v , b ) τ on the current graph Gτ to create a new intermediate graph Gτ+1 as input for the next step τ + 1 . This iterative process of graph transformation can be formulated as a Markov Decision Process ( MDP ) characterized by a tuple ( S , A , P , R , γ ) , in which S is a set of states , A is a set of actions , P is a state transition function , R is a reward function , and γ is a discount factor . Since the process is finite and contains no loop , we set the discount factor γ to be 1 . The rest of the MDP tuple are defined as follows : • State : A state sτ ∈ S is an intermediate graph Gτ generated at step τ ( 0 ≤ τ < T ) . When τ = 0 , we denote s0 = G0 = G. • Action : An action aτ ∈ A performed at step τ is the tuple ( ξ , u , v , b ) τ . The action is composed of three consecutive sub-actions : ξτ , ( u , v ) τ , and bτ . If ξτ = 0 , our model will ignore the next sub-actions ( u , v ) τ and bτ , and all the future actions ( ξ , u , v , b ) τ+1 : T . Note that setting ξτ to be the first sub-action is useful in case a reaction does not happen , i.e. , ξ0 = 0 • State Transition : If ξτ = 1 , the current graph Gτ is modified based on the reaction triple ( u , v , b ) τ to generate a new intermediate graph Gτ+1 . We do not incorporate chemical rules such as valency check during state transition because the current bond change may result in invalid intermediate molecules Gτ , but later , other bond changes may compensate it to create the valid final products GTend . • Reward : We use both immediate rewards and delayed rewards to encourage our model to learn the optimal policy faster . At every step τ , if the model predicts ξτ , ( u , v ) τ or bτ correctly , it will receive a positive reward for each correct sub-action . Otherwise , a negative reward is given . After the prediction process has terminated , if the generated products are exactly the same as the groundtruth products , we give the model a positive reward , otherwise a negative reward . The concrete reward values are provided in Appendix A.3 . 2.2 GRAPH TRANSFORMATION POLICY NETWORK . In this section , we describe the architecture of our model − a Graph Transformation Policy Network ( GTPN ) . GTPN has three main components namely a Graph Neural Network ( GNN ) , a Node Pair Prediciton Network ( NPPN ) , and a Policy Network ( PN ) . Each component is responsible for one or several key functions shown in Fig . 2 : GNN performs functions 1 and 6 ; NPPN performs function 2 ; and PN performs functions 3 , 4 and 5 . Apart from these components , GTPN also has a Recurrent Neural Network ( RNN ) to keep track of the past transformations . The hidden state h of this RNN is used by NPPN and PN to make accurate prediction . 2.2.1 GRAPH NEURAL NETWORK . To model the intermediate graph Gτ at step τ , we compute the node state vector xτi of every node i in Gτ by using a variant of the Message Passing Neural Networks ( Gilmer et al. , 2017 ) : xτi = MessagePassing m ( xτ−1i , vi , N τ ( i ) ) ( 1 ) where m is the number of message passing steps ; vi is the feature vector of node i ; N τ ( i ) is the set of all neighbor nodes of node i ; and xτ−1i is the state vector of node i at the previous step . When τ = 0 , xτ−1i is initialized from vi using a neural network . Details about the MessagePassing ( . ) function are provided in Appendix A.1 . 2.2.2 NODE PAIR PREDICTION NETWORK . In order to predict how likely an atom pair ( i , j ) of the intermediate graph Gτ will change its bond , we assign ( i , j ) with a score sτij ∈ R. If sτij is high , ( i , j ) is more probably a reaction atom pair , otherwise , less probably . Similar to ( Jin et al. , 2017 ) , we use two different networks called “ local ” network and “ global ” network for this task . In case of the “ local ” network , sτij is computed as : zτij = σ ( W1 [ hτ−1 , ( xτi + x τ j ) , eij ] + b1 ) ( 2 ) sτij = f atom pair ( zτij ) ( 3 ) where f atom pair is a neural network ; σ is a nonlinear activation function ( e.g. , ReLU ) ; [ . ] denotes vector concatenation ; W1 and b1 are parameters ; hτ−1 is the hidden state of the RNN at the previous step ; and eij is the representation vector of the bond between ( i , j ) . If there is no bond between ( i , j ) we assume that its bond type is “ NULL ” . We consider zij as the representation vector for the atom pair ( i , j ) . The “ global ” network leverages self-attention ( Vaswani et al. , 2017 ; Wang et al. , 2018 ) to detect compatibility between atom i and all other atoms before computing the scores : rτij = σ ( V1 [ ( xτi + x τ j ) , eij ] + c1 ) aτij = softmax ( V2r τ ij + c2 ) cτi = ∑ j∈V aijx τ j zτij = σ ( W1 [ hτ−1 , ( xτi + x τ j ) , ( c τ i + c τ j ) , eij ] + b1 ) ( 4 ) sτij = f atom pair ( zτij ) ( 5 ) where aij is the attention score from node i to every other node j ; ci is the context vector of atom i that summarizes the information from all other atoms . During experiments , we tried both options mentioned above and saw that the “ global ” network clearly outperforms the “ local ” network so we set the “ global ” network as a default module in our model . In addition , since reagents never change their form during a reaction , we explicitly exclude all atom pairs that have either atoms belong to the reagents . This leads to better results than not using reagent information . Detailed analyses are provided in Appendix A.5 . Top-K atom pairs Because the number of atom pairs that actually participate in a reaction is very small ( usually smaller than 10 ) compared to the total number of atom pairs of the input molecules ( usually hundreds or thousands ) , it is much more efficient to identify reaction triples from a small subset of highly probable reaction atom pairs . For that reason , we extract K ( K |V|2 ) atom pairs with the highest scores . Later , we will predict reaction triples taken from these K atom pairs only . We denote the set of top-K atom pairs , their corresponding scores , and representation vectors as { ( uk , vk ) |k = 1 , K } , { sukvk |k = 1 , K } and ZK = { zukvk |k = 1 , K } , respectively .

Summary: This paper develops a novel method for predicting organic chemical reactions, in particular, final product prediction from given reactants. Organic molecules consist of covalent bonds, and hence organic reactions can be regarded as an alternating multistep series of bond breaking and bond forming (i.e. qualitative understanding as "reaction mechanisms" against quantum chemical calculations). The developed method aims to predict this series of bond changes through a form of reinforcement learning guided with neural networks. In this sense, the setup and formulation seem largely inherited from cited previous papers by Bradshaw et al, 2018 or Kayala & Baldi, 2011 (though they are not used any RL formulation). Organic compounds at each elementary step are represented as molecular graphs, and reactions are thus a series of graph transformations. Each bond change can be considered as "action" at that state to form the next states to head for the final product. The method itself seems quite natural: States transitions are shared by an RNN and the hidden states and observations (molecular graphs at each step, and bond changes as actions) are used to learn "policy" and "state transition" for RL via graph neural networks. Atom pairs to lead the bond change are detected from such graph embeddings and state observations through self-attentive architectures. In addition, masking by additional indicator variables is introduced to avoid redundant training as well as determine the termination of the reaction. Experimental evaluations on a standard large benchmark dataset of USPTO show improved prediction performance compared to previous methods of Jin et al, 2017 and Schwaller et al, 2018.