Title: GRAPH TRANSFORMATION POLICY NETWORK FOR CHEMICAL REACTION PREDICTION

Abstract: We address a fundamental problem in chemistry known as chemical reaction product prediction. Our main insight is that the input reactant and reagent molecules can be jointly represented as a graph, and the process of generating product molecules from reactant molecules can be formulated as a sequence of graph transformations. To this end, we propose Graph Transformation Policy Network (GTPN)a novel generic method that combines the strengths of graph neural networks and reinforcement learning to learn the reactions directly from data with minimal chemical knowledge. Compared to previous methods, GTPN has some appealing properties such as: end-to-end learning, and making no assumption about the length or the order of graph transformations. In order to guide model search through the complex discrete space of sets of bond changes effectively, we extend the standard policy gradient loss by adding useful constraints. Evaluation results show that GTPN improves the top-1 accuracy over the current state-of-the-art method by about 3% on the large USPTO dataset. Our model's performances and prediction errors are also analyzed carefully in the paper.

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 Figure 1 : A sample reaction represented as a set of graph transformations from reactants (leftmost) to products (rightmost) . Atoms are labeled with their type (Carbon, Oxygen,...) and their index (1, 2,...) in the molecular graph. The atom pairs that change connectivity and their new bonds (if existed) are highlighted in green. There are two bond changes in this case: 1) The double bond between O:1 and C:2 becomes single. 2) A new single bond between C:2 and C:10 is added. 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 T end < 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 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 s 0 = G 0 = 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 

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 τ -1 i , v i , N τ (i)

(1) where m is the number of message passing steps; v i is the feature vector of node i; N τ (i) is the set of all neighbor nodes of node i; and x τ -1 i is the state vector of node i at the previous step. When τ = 0, x τ -1 i is initialized from v i 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 = σ W 1 h τ -1 , (x τ i + x τ j ), e ij + b 1 (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; W 1 and b 1 are parameters; h τ -1 is the hidden state of the RNN at the previous step; and e ij 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 z ij 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 = σ V 1 (x τ i + x τ j ), e ij + c 1 a τ ij = softmax V 2 r τ ij + c 2 c τ i = j∈V a ij x τ j z τ ij = σ W 1 h τ -1 , (x τ i + x τ j ), (c τ i + c τ j ), e ij + b 1 (4) s τ ij = f atom pair z τ ij

(5) where a ij is the attention score from node i to every other node j; c i 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

(u k , v k )|k = 1, K , s u k v k |k = 1, K and Z K = z u k v k |k = 1, K , respectively.

2.2.3. POLICY NETWORK:

Predicting continuation signal To account for varying number of transformation steps, PN generates a continuation signal ξ τ ∈ {0, 1} to indicate whether prediction should continue or terminate.

ξ τ is drawn from a Bernoulli distribution:

EQUATION

where h τ -1 is the previous RNN state; Z τ K is the set of representation vectors of the top K atom pairs at the current step; f signal is a neural network; g is a function that maps an unordered set of inputs to an output vector. For simplicity, we use a mean function:

z τ -1 K = g (Z τ K ) = 1 K K k=1 W z τ -1 u k v k

Predicting atom pair At the next sub-step, PN predicts which atom pair changes its bond during the reaction by sampling from the top-K atom pairs with probability:

p ((u k , v k ) τ ) = softmax K s τ u k v k (7)

where s τ u k v k is the score of the atom pair (u k , v k ) τ computed in Eq. ( 5). After predicting the atom pair (u, v) τ , we will mask it to ensure that it could not be in the top K again at future steps.

Predicting bond type Given an atom pair (u, v) τ sampled from the previous sub-step, we predict a new bond type b τ between u and v to get a complete reaction triple (u, v, b) τ using the probability:

EQUATION

where B is the total number of bond types; z τ uv is the representation vector of (u, v) τ computed in Eq. ( 4); b old is the old bond of (u, v); e b old and e b are the embedding vectors corresponding to the bond type b old and b, respectively; and f bond is a neural network.

2.3. UPDATING STATES:

After predicting a complete reaction triple (u, v, b) τ , our model updates: i) the new recurrent hidden state h τ , and ii) the new node representation vectors x τ +1 i of the new intermediate graph G τ +1 for i ∈ V. These updates are presented in Appendix A.2.

2.4. TRAINING:

Loss function plays a central role in achieving fast training and high performance. We design the following loss:

L = λ 1 L A2C + λ 2 L value + λ 3 L atom pair + λ 4 L over length + λ 5 L in top K

where L A2C is the Advantage Actor-Critic (A2C) loss (Mnih et al., 2016) to account for the correct sequence of reaction triples; L value is the loss for estimating the value function used in A2C; L atom pair accounts for binary change in the bond of an atom pair; L over length penalizes long predicted sequences; and L in top K is the rank loss to force a ground-truth reaction atom pair to appear in the top-K; and λ 1 , ..., λ 5 > 0 are tunable coefficients. The component losses are explained in the following.

2.4.1. REACTION TRIPLE LOSS:

The loss follows a policy gradient method known as Advantage Actor-Critic (A2C):

L A2C = - Tend-1 τ =0 A τ signal log p (ξ τ ) + A τ atom pair log p ((u, v) τ ) + A τ bond log p (b τ ) -A Tend signal log π ξ Tend (9)

where T end is the first step that ξ = 0; A signal , A atom pair and A bond are called advantages. To compute these advantages, we use the unbiased estimations called Temporal Different errors, defined as:

EQUATION

) where r τ signal , r τ atom pair , r τ bond are immediate rewards at step τ ; at the final step τ = T end , the model receives additional delayed rewards; γ is the discount factor; and V φ is the parametric value function. We train V φ using the following mean square error loss:

L value = Tend τ =0 V φ (Z τ K ) -R τ 2 (13)

where R τ is the return at step τ .

Episode termination during training Although the loss defined in Eq. ( 9) is correct, it is not good to use in practice because: i) If our model selects a wrong sub-action at any sub-step of the step T wrong (T wrong < T end ), the whole predicted sequence will be incorrect regardless of what will be predicted from T wrong + 1 to T end . Therefore, computing the loss for actions from T wrong + 1 to T end is redundant. ii) More importantly, the incorrect updates of the graph structure at subsequent steps from T wrong + 1 to T end will lead to cumulative prediction errors which make the training of our model much more difficult.

To resolve this issue, during training, we use a binary vector ζ ∈ {0, 1} 3T to keep track of the first wrong sub-action:

ζ t = 1 if t ≤ t first wrong 0 if t > t first wrong

where t first wrong denotes the sub-step at which our model chooses a wrong sub-action the first time. The actor-critic loss in Eq. ( 9) now becomes:

L A2C = - T τ =0 ζ τ A τ signal log p (ξ τ ) + ζ (τ +1) A τ atom pair log p ((u, v) τ ) + ζ (τ +2) A τ bond log p (b τ ) (14)

where T is the maximum number of steps. Similarly, we change the value loss into:

L value = T τ =0 ζ τ V φ (Z τ K ) -R τ 2

2.4.2. REACTION ATOM PAIR LOSS:

To train our model to assign higher scores to reaction atom pairs and lower to non-reaction atom pairs, we use the following cross-entropy loss function:

EQUATION

where T first wrong = tfirst wrong 3

; η ijt ∈ {0, 1} is a mask of the atom pair (i, j) at step τ ; y ij ∈ {0, 1} is the label indicating whether the atom pair (i, j) is a reaction atom pair or not; p ij = sigmoid(s ij ) (see Eq. ( 5)).

2.4.3. CONSTRAINT ON THE SEQUENCE LENGTH:

One major difficulty of the chemical reaction prediction problem is to know exactly when to stop prediction so we can make accurate inference. 

EQUATION

where T gt end is the end step of the ground-truth sequence. Note that the loss in Eq. ( 16) is not applied when T end ≤ T gt end . The reason is that forcing ξ τ = 1 with T end ≤ τ < T gt end is not theoretically correct because all the signals after T end are assumed to be 0. The incentive to force T end close to T gt end when it is smaller than T gt end has already been included in the advantages in Eq. ( 14).

2.4.4. CONSTRAINT ON THE TOP-K ATOM PAIRS:

Ideally, the loss from Eq. ( 15) pushes a reaction atom pair (ũ, ṽ) τ into the top-K atom pairs at each step τ < T gt end . However, this is not guaranteed, especially when τ comes close to T gt end . To encourage the ground-truth reaction atom pair (ũ, ṽ) τ with the highest score to appear in the top K, we introduce an additional rank-based loss:

L in top K = - Tfirst wrong τ =0 log p ((ũ, ṽ) τ in top K)

where p ((ũ, ṽ) τ in top K) is computed as:

p ((ũ, ṽ) τ in top K) = exp (s τ ũṽ ) exp (s τ ũṽ ) + K k=1 exp s τ u k v k (17) 3 EXPERIMENTS 3.1 DATASET

We evaluate our model on two standard datasets USPTO-15k (15K reactions) and USPTO (480K reactions) which have been used in previous works (Jin et al., 2017; Schwaller et al., 2018; Bradshaw et al., 2018) . Details about these datasets are given in Table 1 . The USPTO dataset contains reactant, reagent and product molecules represented as SMILES strings. Using RDKit1 , we convert the SMILES strings into molecule objects and store them as graphs. For each reaction, every atom in the reactant and reagent molecules is identified with a unique "atom map number". This identity is the same in the products. Using this knowledge, we compare every atom pair in the input molecules with the correspondent in the product molecules to obtain a ground-truth set of reaction triples for training.

In USPTO-15k, the ground-truth sets of reaction triples was precomputed by (Jin et al., 2017) .

Model USPTO-15k USPTO C@6 C@8 C@10 C@6 C@8 C@10 WLN (Jin et al., 2017) 81 

3.2. REACTION ATOM PAIR PREDICTION:

In this section, we test our model's ability to identify reaction atom pairs by formulating it as a ranking problem with the scores computed in Eq. ( 5). Similar to (Jin et al., 2017) , we use Coverage@k as the evaluation metric, which is the proportion of reactions that have all groundtruth reaction atom pairs appear in the top k predicted atom pairs. We compare our proposed graph neural network (GNN) with Weisfeiler-Lehman Network (WLN) (Jin et al., 2017) and Column Network (CLN) (Pham et al., 2017) . Since our GNN explicitly uses reagent information to compute the scores of atom pairs, we modify the implementation of WLN and CLN accordingly for fair comparison. From Table 2 , we observe that our GNN clearly outperforms WLN and CLN in all cases. We attribute this improvement to the use of a separate node state vector x t i (different from the node feature vector v i ) for updating the structural information of a node (see Eq. ( 21)). The other two models, on the other hand, only use a single vector to store both the node features and structure, hence, some information may be lost. In addition, using explicit reagent information boosts the prediction accuracy, which improves the WLN by 1-7% depending on the metrics. The presence of reagent information reduces the number of atom pairs to be searched on and contributes to the likelihood of reaction atom pairs. Further results are presented in Appendix A.5.

3.3. TOP-K ATOM PAIR EXTRACTION:

The performance of our model depends on the number of selected top atom pairs K. The value of K presents a trade-off between coverage and efficiency. In addition to the metric Coverage@k in Sec. 3.2, we use Recall@k which is the proportion of correct atom pairs that appear in top k to find the good K. Fig. 3 shows Coverage@k and Recall@k for the USPTO dataset with respect to k. We see that both curves increase rapidly when k < 10 and stablize when k > 10. We also ran experiments with k = 10, 15, 20 and observed that their prediction results are quite similar. Hence, in what follows we select K = 10 for efficiency.

Model:

USPTO-15k USPTO P@1 P@3 P@5 P@1 P@3 P@5 WLDN (Jin et al., 2017) 76 

3.4. REACTION PRODUCT PREDICTION:

This experiment validates GTPN on full reaction product prediction against the recent state-of-the-art methods (Jin et al., 2017; Schwaller et al., 2018) using the accuracy metric. The recent method ELECTRO (Bradshaw et al., 2018) is not compatible here because it was only evaluated on a subset of USPTO limited to linear chain topology. Comparison against ELECTRO is reported separately in Appendix A.6. Table 3 shows the prediction results. We produce multiple reaction product candidates by using beam search decoding with beam width N = 20. Details about beam search and its behaviors are presented in Appendix A.4.

In brief, we compute the normalized-over-length log probabilities of N predicted sequences of reaction triples and sort these values in descending order to get a rank list of N possible reaction outcomes. Given a predicted sequence of reaction triples (u, v, b) 0:T , we can generate reaction products from input reactants simply by replacing the old bond of (u, v) τ with b τ . However, these products are not guaranteed to be valid (e.g., maximum valence constraint violation or aromatic molecules cannot be kekulized) so we post-process the outputs by removing all invalid products. The removal increases the top-1 accuracy by about 8% and 10% on USPTO-15k and USPTO, respectively. Due to the permutation invariance of the predicted sequence of reaction triples, some product candidates are duplicate and will also be removed. This does not lead to any change in P@1 but slightly improves P@3 and P@5 by about 0.5-1% on the two datasets.

Overall, GTPN with beam search and post-processing outperforms both WLDN (Jin et al., 2017) and Seq2Seq (Schwaller et al., 2018) in the top-1 accuracy. For the top-3 and top-5, our model's performance is comparable to WLDN's on USPTO-15k and is worse than WLDN's on USPTO. It is not surprising since our model is trained to accurately predict the top-1 outcomes instead of ranking the candidates directly like WLDN. It is important to emphasize that we did not tune the model hyper-parameters when training on USPTO but reused the optimal settings from USPTO-15k (which is 25 times smaller than USPTO) so the results may not be optimal (see Appendix A.3 for more training detail). Baldi, 2011; Jin et al., 2017; Fooshee et al., 2018) and template-based (Wei et al., 2016; Segler & Waller, 2017; Coley et al., 2017) . Both groups share the same mechanism: running multiple stages with the aid of reaction templates or rules. For example, in (Wei et al., 2016) the authors proposed a two-stage model that first classifies reactions into different types based on the neural fingerprint vectors (Duvenaud et al., 2015) of reactant and reagent molecules. Then, it applies pre-designed SMARTS transformation on the reactants with respect to the most suitable predicted reaction type to generate the reaction products.

4. RELATED WORK:

The work of (Jin et al., 2017) treats a reaction as a set of bond changes so in the first step, they predict which atom pairs are likely to be reactive using a variant of graph neural networks called Weisfeiler-Lehman Networks (WLNs). In the next step, they do almost the same as (Coley et al., 2017) by modifying the bond type between the selected atom pairs (with chemical rules satisfied) to create product candidates and rank them (with reactant molecules as addition input) using another kind of WLNs called Weifeiler-Lehman Different Networks (WLDNs).

To the best of our knowledge, (Jin et al., 2017) is the first work that achieves remarkable results (with the Precision@1 is about 79.6%) on the large USPTO dataset containing more than 480 thousands reactions. Works of (Nam & Kim, 2016) and (Schwaller et al., 2018) avoid multi-stage prediction by building a seq2seq model that generates the (canonical) SMILES string of the single product from the concatenated SMILES strings of the reactants and reagents in an end-to-end manner. However, their methods cannot deal with sets of reactants/reagents/products properly as well as cannot provide concrete reaction mechanism for every reaction.

The most recent work on this topic is (Bradshaw et al., 2018) which solves the reaction prediction problem by predicting a sequence of bond changes given input reactants and reagents represented as graphs. To handle ordering, they only select reactions with predefined topology. Our method, by contrast, is order-free and can be applied to almost any kind of reactions.

4.2. GRAPH NEURAL NETWORKS FOR MODELING MOLECULES:

In recent years, there has been a fast development of graph neural networks (GNNs) for modeling molecules. These models are proposed to solve different problems in chemistry including toxicity prediction (Duvenaud et al., 2015) , drug activity classification (Shervashidze et al., 2011; Dai et al., 2016; Pham et al., 2018) , protein interface prediction (Fout et al., 2017) and drug generation (Simonovsky & Komodakis, 2018; Jin et al., 2018) . Most of them can be regarded as variants of message-passing graph neural networks (MPGNNs) (Gilmer et al., 2017) .

4.3. REINFORCEMENT LEARNING FOR STRUCTURAL REASONING:

Reinforcement learning (RL) has become a standard approach to many structural reasoning problems2 because it allows agents to perform discrete actions. A typical example of using RL for structural reasoning is drug generation (Li et al., 2018; You et al., 2018) . Both (Li et al., 2018) and (You et al., 2018) learn the same generative policy whose action set including: i) adding a new atom or a molecular scaffold to the intermediate graph, ii) connecting existing pair of atoms with bonds, and iii) terminating generation. However, (You et al., 2018) uses an adversarial loss to enforce global chemical constraints on the generated molecules as a whole instead of using the common reconstruction loss as in (Li et al., 2018) . Other examples are path-based relational reasoning in knowledge graphs (Das et al., 2018) and learning combinatorial optimization over graphs (Khalil et al., 2017) .

5. DISCUSSION:

We have introduced a novel method named Graph Transformation Policy Network (GTPN) for predicting products of a chemical reaction. GTPN uses graph neural networks to represent input reactant and reagent molecules, and uses reinforcement learning to find an optimal sequence of bond changes that transforms the reactants into products. We train GTPN using the Advantage Actor-Critic (A2C) method with appropriate constraints to account for notable aspects of chemical reaction. Experiments on real datasets have demonstrated the competitiveness of our model.

Although the GTPN was proposed to solve the chemical reaction problem, it is indeed generic to solve the graph transformation problem, which can be useful in reasoning about relations (e.g., see (Zambaldi et al., 2018) ) and changes in relation. Open rooms include addressing dynamic graphs over time, extending toward full chemical planning and structural reasoning using RL.

A APPENDIX:

A.1 GRAPH NEURAL NETWORK

In this section, we describe our graph neural network (GNN) in detail. Since our GNN does not use the recurrent hidden state h τ , we exclude the time step τ from our notations for clarity. Instead, we use t to denote a message passing step.

GRAPH NOTATIONS:

Input to our GNN is a graph G = (V, E) in which each node i ∈ V is represented by a node feature vector v i and each edge (i, j) ∈ E is represented by an edge feature vector e ij . For example of molecular graph, the node feature vector v i may include chemical information about the atom i such as its type, charge and degree. Similarly, e ij captures the bond type between the two atoms i and j. We denote by N (i) the set of all neighbor nodes of node i together with their links to node i:

N (i) ≡ {(j, e ij )

| j is a neighbor node of i} If we only care about the neighbor nodes of i not their links, we use the notation N n (i) defined as:

N n (i) ≡ {j | j is a neighbor node of i}

In addition to v i , node i also has a state vector x i to store information about itself and the surrounding context. This state vector is updated recursively using the neural message passing method (Battaglia et al., 2016; Pham et al., 2017; Hamilton et al., 2017; Gilmer et al., 2017; Schlichtkrull et al., 2018) . The initial state x 0 i is the nonlinear mapping of v i :

EQUATION

COMPUTING NEIGHBOR MESSAGES At the message passing step t, we compute the message m t ij from every neighbor node j ∈ N n (i) to node i as:

EQUATION

where [•] denotes concatenation; and σ is a nonlinear function.

AGGREGATING NEIGHBOR MESSAGES:

Then, we aggregate all the messages sent to node i into a single message vector by averaging:

m t i = 1 |N n (i)| j∈Nn(i) m t ij (20)

where |N n (i)| is the number of neighbor nodes of node i.

UPDATING NODE STATE:

Finally, we update the state of node i as follows:

EQUATION

) where g(.) is a Highway Network (Srivastava et al., 2015) :

x t+1 i = Highway x t i , m t i , v i (22) = α * xt+1 i + (1 -α) * x t i (23)

where xt+1 i is the nonlinear part which is computed as:

xt+1 i = σ W 1 x t i , m t i , v t i +

b 1 and α is the gate controlling the flow of information: Eqs. (19, 20, 22) together, one step of message passing update for node i can be written in a generic way as follows:

α = sigmoid(W 2 x t i , m t i , v t i + b 2 ) By combining

EQUATION

A.2 UPDATING STATES

UPDATING RNN STATE:

We keep the old representation of the edge that have been modified in the hidden memory of the RNN as follows:

EQUATION

) where GRU stands for Gated Recurrent Units (Cho et al., 2014) ; z τ uv is the representation vector of the atom pair (u, v) τ including its old bond (see Eq. 4). Eq. ( 25) allows the model to keep track of all the changes happening to the graph so far so it can make more accurate prediction later.

UPDATING GRAPH STRUCTURE AND NODE STATES:

After predicting a reaction triple (u, v, b) τ at step τ , we update the graph structure and node states based on the new bond change. First, to update the graph structure, we simply update the neighbor set of u and v with information from the other atom and the new bond type b as follows:

EQUATION

) Next, to update the node states, our model performs one step of message passing for u and v with their new neighbor sets:

EQUATION

) where the MessagePassing(.) function is defined in Eq. ( 24). For other nodes in the graph to be aware of the new structures of u and v, we need to perform several message passing steps for all nodes in the graph after Eqs. (28, 29) . However, it is very costly to run for every prediction step τ . Sometimes it is unnecessary since far-away bonds are less likely to be affected by the current bond change (unless the far-way bonds and the new bond are in an aromatic ring). Therefore, in our model, we limit the number of message passing updates for all nodes at step τ to be 1.

A.3 MODEL CONFIGURATIONS:

We optimize our model's hyper-parameters in two stages: First, we tune the hyper-parameters of the GNN and the NPPN for the reaction atom pair prediction task. Then, we fix the optimal settings of the first two components and optimize the hyper-parameters of the PN for the reaction product prediction task.

We provide details about the settings that give good results on the USPTO-15k dataset below. With these settings, we trained another model on the USPTO dataset from scratch. Because training on the large dataset such as the USPTO takes time, we did not tune hyper-parameters on the USPTO, eventhough it is possible to increase model sizes for better performance.

Unless explicitly stated, all neural networks in our model have 2 layers with the same number of hidden units, ReLU activation and residual connections (He et al., 2016) 4 . These attributes are normalized to the range of [0, 1] and are concatenated to the atom embedding vector to form a final atom feature vector of size 56. The state vector and the neighbor message vector for an atom both have the size of 99.

The number of message passing steps is 6.

Node Pair Prediction Network (NPPN) This component consists of two parts. The first part computes the representation vector z ij of an atom pair (i, j) using a neural network with hidden size of 71. The second part maps z ij to an unnormalized score s ij using the function f atom pair (see Eqs. (3, 5) ). This function is also a neural network with hidden size of 51.

Policy Network (PN) The recurrent network is a GRU (Cho et al., 2014) with 101 hidden units. The value function V φ is a neural network with 99 hidden units. The two functions f signal for computing signal scores (see Eq. ( 6)) and f bond for computing scores over bond types (see Eq. ( 8)) are neural networks with 81 hidden units.

Training At each step, we set the reward to be 1.0 for correct prediction of signal/atom pair/bond type and -1.0 for incorrect prediction. After the prediction sequence is terminated (zero signal was emitted), we check whether the entire set of predicted reaction triples is correct or not. If it is correct, we give the model a reward value of 2.0, otherwise -2.0. From the rewards and estimated values for signal, atom pair and bond type, we define the Advantage Actor Critic loss (A2C) as in Eq. ( 14). The coefficients of components in the final loss L are set empirically as follows:

L = L A2C + 0.5 × L value + L atom pair + 0.2 × L over length + 0.2 × L in top K

We trained our model using Adam (Kingma & Ba, 2015) with the initial learning rate of 0.001 for both USPTO-15k and USPTO. For USPTO-15k, the learning rate will decrease by half if the Precision@1 does not improve on the validation set after 1,000 steps until it reaches the minimum value of 5 × 10 -5 . For USPTO, the decay rate is 0.8 after every 500 steps of no improvement until reaching the minimum learning rate is 2 × 10 -5 . The maximum number of training iterations is 10 6 and the batch size is 20.

A.4 DECODING WITH BEAM SEARCH:

For decoding, our model generates a sequence of reaction triples (including the stop signal) (ξ, u, v, b) by taking the best (u, v) and b at every step until it outputs a zero signal (ξ = 0). In other words, it computes the argmax of p (ξ, u, v, b) τ | G, (ξ, u, v, b) 0:τ -1 at every step τ . However, this algorithm is not robust for the sequence generation task because just a single error at a step may destroy the entire sequence. To overcome this issue, we employ beam search for decoding.

During beam search, we keep track of N > 1 best subsequences at every step τ . N is called beam width. Instead of modeling the conditional distribution of generating an output at the current step τ , we model the joint distribution of the whole subsequence that has been generated from 0 to τ :

EQUATION

Processed USPTO P@1 P@3 P@5 WLDN (Jin et al., 2017) 84.0 91.1 92.3 ELECTRO (Bradshaw et al., 2018) Errors grouped by length. In (a), blue: all reactions having that sequence length; orange: correct predicted reactions. In (b), red: the predicted sequence is shorter (than the groundtruth sequence); green: the predicted and the groundtruth have the same length; blue: the predicted sequence is longer; number indidate the average length.

In performance To do a fair comparison with ELECTRO (Bradshaw et al., 2018) , we follow their procedure described in the paper to prepare a new test set that contains only reactions with linear chain topology and single-valence bond changes. It results in 29,808 reactions, close to the reported number of 29,360 in (Bradshaw et al., 2018) . We reuse our old model (see Section 3.4) trained on the original USPTO dataset. We also use beam search decoding and post-processing as similar to (Bradshaw et al., 2018) . From Table 7 , we see that GTPN achieves the highest top-1 accuracy of 87.35%, outperforming ELECTRO and WLDN by 0.35% and 3%, respectively. For the top-3 and top-5 accuracies, our model, however, does worse than the other two. Especially, while both ELECTRO and WLDN have big jumps from P@1 to P@3 with about 7% improvement, GTPN only has 3% increase. We conjecture that this problem mainly comes from the fact that GTPN was not optimized on the compatible training and validation sets.

A.7 ERROR ANALYSIS:

In this section, we analyze several error types that our model makes during prediction. All the results below are computed on the USPTO-15k dataset by using beam search decoding with the beam width N = 20 and no post-processing.

Errors grouped by number of bond changes Fig. 5 shows the top-1 accuracies for reactions with different number of bond changes. Our model performs poorly on reactions with many bond changes. However, those kinds of reactions only accounts for a small proportion in the dataset. From Fig. 5b , we see that the lengths of the error sequences tends to be shorter than the lengths of the groundtruth sequences.

Errors caused by signal/atom pair/bond type We define a sub-action causing error as the first sub-action that our model makes a wrong decision. In Fig. 6a , we plot the the proportion of errors with respect to the three kinds of sub-actions. Clearly, atom pair prediction causes the most errors (nearly two third). This makes sense because this sub-action is harder than signal prediction and bond type prediction. Therefore, more effort should be put on improving the prediction of atom pairs.

Errors caused by symmetry There exists cases in which different sequences of bond changes can result in the same products due to symmetric graph structures. Errors caused by symmetry account for 5.7% of the top-1 errors on the USPTO-15k dataset as shown in Fig. 6b . For better understanding, we provide a short list of wrong reaction triple predictions caused by symmetry in Fig. 7 . In this list, the top-1 products (along the second column) are incorrect while the top-2 products (along the third column) are correct though both have the same probability.

:

https://www.rdkit.org/

Structural reasoning is a problem of inferring or generating new structure (e.g. objects with relations)

Note that beam search with beam width = 1 is different from greedy search as in beam search, as we model the whole sequence probability.

Figures and Tables:

Figure 2: Workflow of a Graph Transformation Policy Network (GTPN). At every step of the forward pass, our model performs 7 major functions: 1) Computing the atom representation vectors, 2) Computing the most possible K reaction atom pairs, 3) Predicting the continuation signal ξ, 4) Predicting the reaction atom pair (u, v), 5) Predicting a new bond b of this atom pair, 6) Updating the atom representation vectors, and 7) Updating the recurrent state.

Figure 3: Coverage@k and Recall@k with respect to k for the USPTO dataset.

LEARNING TO PREDICT CHEMICAL REACTION In chemical reaction prediction, machine learning has replaced rule-based methods (Chen & Baldi, 2009) for better generalizability and scalability. Existing machine learning-based techiques are either template-free (Kayala &

Figure 5: Performance with respect to different numbers of bond changes. (a) Top-1 accuracy. (b)Errors grouped by length. In (a), blue: all reactions having that sequence length; orange: correct predicted reactions. In (b), red: the predicted sequence is shorter (than the groundtruth sequence); green: the predicted and the groundtruth have the same length; blue: the predicted sequence is longer; number indidate the average length.

Proportion of the incorrect top-1 products that have similar structure to the groundtruth products.

Figure 6: Errors grouped by the first incorrect sub-actions (a), and errors caused by symmetric structures (b).

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 G Tend . • 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 τ

Table content: <table><tr><td>correctly, it will receive a positive reward for each correct sub-action. Otherwise, a negative</td></tr><tr><td>reward is given. After the prediction process has terminated, if the generated products are</td></tr><tr><td>exactly the same as the groundtruth products, we give the model a positive reward, otherwise</td></tr><tr><td>a negative reward. The concrete reward values are provided in Appendix A.3.</td></tr><tr><td>2.2 GRAPH TRANSFORMATION POLICY NETWORK</td></tr><tr><td>In this section, we describe the architecture of our model -a Graph Transformation Policy Network</td></tr><tr><td>(GTPN). GTPN has three main components namely a Graph Neural Network (GNN), a Node Pair</td></tr><tr><td>Prediciton Network (NPPN), and a Policy Network (PN). Each component is responsible for one or</td></tr><tr><td>several key functions shown in Fig. 2: GNN performs functions 1 and 6; NPPN performs function 2;</td></tr><tr><td>and PN performs functions 3, 4 and 5. Apart from these components, GTPN also has a Recurrent</td></tr><tr><td>Neural Network (RNN) to keep track of the past transformations. The hidden state h of this RNN is</td></tr><tr><td>used by NPPN and PN to make accurate prediction.</td></tr></table>

Statistics of USPTO-15k and USPTO datasets. "changes" means bond changes, "molecules" means reactants and reagents in a reaction; "atoms" and "bonds" are defined for a molecule. Apart from "#reactions", other columns are presented in the format "min | max | mean". every step τ during inference. In order to make the model aware of the correct sequence length during training, we define a loss that punishes the model if it produces a longer sequence than the ground truth sequence:

Table content: <table><tr><td>Dataset</td><td/><td colspan="3">#reactions #changes #molecules</td><td>#atoms</td><td>#bonds</td></tr><tr><td/><td>train</td><td>10,500</td><td colspan="4">1 | 11 | 2.3 1 | 20 | 3.6 4 | 100 | 34.9 3 | 110 | 34.7</td></tr><tr><td>USPTO-15k</td><td>valid</td><td>1,500</td><td colspan="2">1 | 11 | 2.3 1 | 20 | 3.6</td><td>7 | 94 | 34.5</td><td>5 | 99 | 34.2</td></tr><tr><td/><td>test</td><td>3,000</td><td colspan="2">1 | 11 | 2.3 1 | 16 | 3.6</td><td colspan="2">7 | 98 | 34.9 5 | 102 | 34.7</td></tr><tr><td/><td>train</td><td>409,035</td><td>1 | 6 | 2.2</td><td colspan="3">2 | 29 | 4.8 9 | 150 | 39.7 6 | 165 | 38.6</td></tr><tr><td>USPTO</td><td>valid</td><td>30,000</td><td>1 | 6 | 2.2</td><td colspan="3">2 | 25 | 4.8 9 | 150 | 39.6 7 | 158 | 38.5</td></tr><tr><td/><td>test</td><td>40,000</td><td>1 | 6 | 2.2</td><td colspan="3">2 | 22 | 4.8 9 | 150 | 39.8 7 | 162 | 38.7</td></tr></table>

Results for reaction atom pair prediction. C@k is coverage at k. Best results are highlighted in bold. WLN is the original model from(Jin et al., 2017) while WLN is our re-implemented version. Except for WLN , other models explicitly use reagent information.

Table content: <table><tr><td/><td>.6</td><td>86.1</td><td>89.1</td><td>89.8</td><td>92.0</td><td>93.3</td></tr><tr><td>WLN (Jin et al., 2017)</td><td colspan="3">88.45 91.65 93.34</td><td colspan="3">90.97 93.98 95.26</td></tr><tr><td colspan="4">CLN (Pham et al., 2017) 88.68 91.63 93.07</td><td colspan="3">90.72 93.57 94.80</td></tr><tr><td>Our GNN</td><td colspan="3">88.92 92.00 93.57</td><td colspan="3">91.24 94.17 95.33</td></tr></table>

Results for reaction prediction. P@k is precision at k. State-of-the-art results from(Jin et al., 2017) are written in italic. Results from(Schwaller et al., 2018) are marked with and they are computed on a slightly different version of USPTO that contains only single-product reactions. Best results are highlighted in bold. ♦ : With beam search (beam width = 20), ♠ : Invalid product removal,

Table content: <table><tr><td/><td>.7</td><td>85.6</td><td>86.8</td><td>79.6</td><td>87.7</td><td>89.2</td></tr><tr><td>Seq2Seq (Schwaller et al., 2018)</td><td>-</td><td>-</td><td>-</td><td>80.3</td><td>86.2</td><td>87.5</td></tr><tr><td>GTPN</td><td>72.31</td><td>-</td><td>-</td><td>71.26</td><td>-</td><td>-</td></tr><tr><td>GTPN ♦</td><td colspan="6">74.56 82.62 84.23 73.25 80.56 83.53</td></tr><tr><td>GTPN ♦♣</td><td colspan="6">74.56 83.19 84.97 73.25 84.31 85.76</td></tr><tr><td>GTPN ♦♠</td><td colspan="6">82.39 85.60 86.68 83.20 84.97 85.90</td></tr><tr><td>GTPN ♦♠♣</td><td colspan="6">82.39 85.73 86.78 83.20 86.03 86.48</td></tr></table>

. Data types of atom attributes. There are 72 different types of atom depending on their atomic numbers and 5 different types of bond including NULL, SINGLE, DOUBLE, TRIPLE and ARO-MATIC. The size of embedding vectors for atom and bond are 51 and 21, respectively. Apart from atom type, each atom has 5 more attributes listed in Table

Table content: <table><tr><td>Atom attribute</td><td>Data type</td></tr><tr><td>Degree</td><td>numeric</td></tr><tr><td>Explicit valence</td><td>numeric</td></tr><tr><td>Explicit number of Hs</td><td>numeric</td></tr><tr><td>Charge</td><td>numeric</td></tr><tr><td>Part of a ring</td><td>boolean</td></tr><tr><td>Graph Neural Network (GNN)</td><td/></tr></table>

Results for the reaction prediction task. P@k is the precision at k. Best results are highlighted in bold. Meanings of markers in our model: ♦ : With beam search (beam width = 20), ♠ : Invalid product removal, ♣ : Duplicate product removal.

Table content: <table><tr><td/><td>87.0</td><td>94.5</td><td>95.9</td></tr><tr><td>GTPN ♦♠♣</td><td colspan="3">87.35 90.22 90.68</td></tr></table>

Peer review (iclr2019):
 Update: Score increased. ___________________________________ Original review: The paper presents an approach to predict the products of chemical reactions, given the reactants and reagents. It works by stepwise predicting the atom pairs that change their bonds in course of reaction, and then adjusting the bonds between them. This can be interpreted as a stepwise graph transformation. I think this is an interesting applied ML paper with fair results. The presentation is clear and understandable. The experimental setup is reasonable. However, the paper is not ready yet to be accepted in my opinion. I think a higher score is justified if the authors address the following points: - Relation to previous work, originality In contrast to what the authors claim, what is predicted here is not exactly the reaction mechanism, but an implementation of the principle of minimal chemical distance, which was already described by Ugi and coworkers in 1980 [see Jochum, Gasteiger, Ugi, The Principle of Minimum Chemical Distance Angew. Chem. Int. Ed. Engl. 1980, 19, p 495-505]. The “insight” the authors have about treating reagents and reactants jointly is Organic Chemistry 101, and that reactions are stepwise graph transformations was also reported by Ugi et al, however, already in 1979! [Ugi, et al. "New applications of computers in chemistry." Angewandte Chemie International Edition in English 18.2 (1979): 111-123. ] I assume the authors were not aware of these papers, but now they are, so this needs to be modified accordingly, and these papers need to be referred to in the introduction. - Questions: The authors suggest that graph neural networks are more generic that so-called heuristic features (fingerprints) – which, as Duvenaud et al have elaborated, can be interpreted just as graph neural networks themselves – with fixed weights. Also, there are results by the Hochreiter group which show that graph neural networks perform worse that classical chemical features under rigorous testing { DOI: 10.1039/C8SC00148K } Do the authors think their models could also improve if they used the classical fingerprints? Is the GRU really needed to encode the past bond changes? What happens if you remove it? The statement that the method has the advantage of not relying on handcrafted reaction templates is somewhat overselling, because instead it uses a handcrafted complex neural network architecture. How complicated is it to train the network? If you remove some of the “tricks” of shaping the loss function, does it still train well? To what degree is the ranking of the different models just a matter of hyperparameter tuning or different architectures? If you used a different graph neural net instead of an MPNN on top of your GTPN method, what would you expect? Are the differences between the models significant? During prediction, you apply a flag to the starting molecules if they are a reagent or reactant. How do you know upfront what a reagent or reactant is during inference? On page 5 and 7, you speak of the correct sequence of reaction triples (which implies an ordering), even though earlier you claim the algorithm is order-invariant? Where do you get the ground truth labels from? I assume these are already annotated in the data. In the appendix, please replace the pie chart with a bar chart. - Language: I would suggest the authors to adapt the language of their paper towards a more academic tone. Science is not a sports competition of getting slightly higher numbers in benchmarks, but rather about providing insights and explanations. Words like “beating” or “record” are locker room talk, and to be avoided.