Year: 2018

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.

heading: 1 INTRODUCTION
text: 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.

heading: 2 METHOD

heading: 2.1 CHEMICAL REACTION AS MARKOV DECISION PROCESS OF GRAPH TRANSFORMATIONS
text: 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.

heading: 2.2 GRAPH TRANSFORMATION POLICY NETWORK
text: 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.

heading: 2.2.1 GRAPH NEURAL NETWORK
text: 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.

heading: 2.2.2 NODE PAIR PREDICTION NETWORK
text: 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.

heading: 2.2.3 POLICY NETWORK
text: 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: p (ξτ = 1) = sigmoid ( f signal ([ hτ−1, g (ZτK) ])) (6) 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τ−1K = g (Z τ K) =
1
K K∑ k=1 Wzτ−1ukvk
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 ((uk, vk) τ ) = softmaxK ( sτukvk ) (7)
where sτukvk is the score of the atom pair (uk, vk) τ 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:
p (bτ |(u, v)τ ) = softmaxB ( f bond ([ hτ−1, zτuv, (eb − ebold) ])) (8)
where B is the total number of bond types; zτuv is the representation vector of (u, v) τ computed in Eq. (4); bold is the old bond of (u, v); ebold and eb are the embedding vectors corresponding to the bond type bold and b, respectively; and f bond is a neural network.

heading: 2.3 UPDATING STATES
text: 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τ+1i of the new intermediate graph Gτ+1 for i ∈ V . These updates are presented in Appendix A.2.

heading: 2.4 TRAINING
text: Loss function plays a central role in achieving fast training and high performance. We design the following loss:
L = λ1LA2C + λ2Lvalue + λ3Latom pair + λ4Lover length + λ5Lin topK
where LA2C is the Advantage Actor-Critic (A2C) loss (Mnih et al., 2016) to account for the correct sequence of reaction triples; Lvalue is the loss for estimating the value function used in A2C; Latom pair accounts for binary change in the bond of an atom pair; Lover length penalizes long predicted sequences; and Lin topK 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.

heading: 2.4.1 REACTION TRIPLE LOSS
text: The loss follows a policy gradient method known as Advantage Actor-Critic (A2C):
LA2C = − Tend−1∑ τ=0 ( Aτsignal log p (ξ τ ) +Aτatom pair log p ((u, v) τ ) +Aτbond log p (b τ ) )
−ATendsignal log π ( ξTend ) (9)
where Tend is the first step that ξ = 0; Asignal, Aatom pair and Abond are called advantages. To compute these advantages, we use the unbiased estimations called Temporal Different errors, defined as:
Aτsignal = r τ signal + γVφ ( Zτ+1K ) − Vφ (ZτK) (10)
Aτatom pair = r τ atom pair + γVφ ( Zτ+1K ) − Vφ (ZτK) (11)
Aτbond = r τ bond + γVφ ( Zτ+1K ) − Vφ (ZτK) (12)
where rτsignal, r τ atom pair, r τ bond are immediate rewards at step τ ; at the final step τ = Tend, 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:
Lvalue = 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 Twrong (Twrong < Tend), the whole predicted sequence will be incorrect regardless of what will be predicted from Twrong + 1 to Tend. Therefore, computing the loss for actions from Twrong + 1 to Tend is redundant. ii) More importantly, the incorrect updates of the graph structure at subsequent steps from Twrong + 1 to Tend 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 ≤ tfirst wrong 0 if t > tfirst wrong where tfirst 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:
LA2C = − 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:
Lvalue = T∑ τ=0 ζτ ‖Vφ (ZτK)−Rτ‖ 2

heading: 2.4.2 REACTION ATOM PAIR LOSS
text: 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:
Latom pair = − Tfirst wrong∑ τ=0 ∑ i∈V ∑ j∈V,j 6=i ηijτ (yij log pij + (1− yij) log(1− pij)) (15)
where Tfirst wrong = ⌊ tfirst wrong
3
⌋ ; ηijt ∈ {0, 1} is a mask of the atom pair (i, j) at step τ ; yij ∈ {0, 1} is
the label indicating whether the atom pair (i, j) is a reaction atom pair or not; pij = sigmoid(sij) (see Eq. (5)).

heading: 2.4.3 CONSTRAINT ON THE SEQUENCE LENGTH
text: One major difficulty of the chemical reaction prediction problem is to know exactly when to stop prediction so we can make accurate inference. By forcing the model to stop immediately when making wrong prediction, we can prevent cumulative error and significantly reduce variance during training. But it also comes with a cost: The model cannot learn (because it does not have to learn) when to stop. This phenomenon can be visualized easily as the model predicts 1 for the signal at
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:
Lover length = − ∑
T gtend≤τ<Tend
log p (ξτ = 0) (16)
where T gtend is the end step of the ground-truth sequence. Note that the loss in Eq. (16) is not applied when Tend ≤ T gtend. The reason is that forcing ξτ = 1 with Tend ≤ τ < T gt end is not theoretically correct because all the signals after Tend are assumed to be 0. The incentive to force Tend close to T gt end when it is smaller than T gtend 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 (ũ, ṽ)τ into the top-K atom pairs at each step τ < T gtend. However, this is not guaranteed, especially when τ comes close to T gt end. To encourage the ground-truth reaction atom pair (ũ, ṽ)τ with the highest score to appear in the topK, we introduce an additional rank-based loss:
Lin topK = − Tfirst wrong∑ τ=0 log p ((ũ, ṽ)τ in top K)
where p ((ũ, ṽ)τ in top K) is computed as:
p ((ũ, ṽ)τ in top K) = exp (sτũṽ) exp (sτũṽ) + ∑K k=1 exp ( sτukvk ) (17)

heading: 3 EXPERIMENTS

heading: 3.1 DATASET
text: 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).
1https://www.rdkit.org/

heading: 3.2 REACTION ATOM PAIR PREDICTION
text: 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 xti (different from the node feature vector vi) 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.

heading: 3.4 REACTION PRODUCT PREDICTION
text: 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).

heading: 4 RELATED WORK

heading: 4.1 LEARNING TO PREDICT CHEMICAL REACTION
text: 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 & 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.
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.

heading: 4.2 GRAPH NEURAL NETWORKS FOR MODELING MOLECULES
text: 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).

heading: 4.3 REINFORCEMENT LEARNING FOR STRUCTURAL REASONING
text: 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).

heading: 5 DISCUSSION
text: 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.
2Structural reasoning is a problem of inferring or generating new structure (e.g. objects with relations)

heading: A APPENDIX

heading: A.1 GRAPH NEURAL NETWORK
text: 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.

heading: GRAPH NOTATIONS
text: Input to our GNN is a graph G = (V, E) in which each node i ∈ V is represented by a node feature vector vi and each edge (i, j) ∈ E is represented by an edge feature vector eij . For example of molecular graph, the node feature vector vi may include chemical information about the atom i such as its type, charge and degree. Similarly, eij 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, eij) | j is a neighbor node of i} If we only care about the neighbor nodes of i not their links, we use the notation Nn(i) defined as:
Nn(i) ≡ {j | j is a neighbor node of i} In addition to vi, node i also has a state vector xi 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 x0i is the nonlinear mapping of vi:
x0i = σ (Wvi + b) (18)

heading: COMPUTING NEIGHBOR MESSAGES
text: At the message passing step t, we compute the messagemtij from every neighbor node j ∈ Nn(i) to node i as:
mtij = f ( xti,x t j , eij ) = σ ( W [ xti,x t j , eij ] + b )
(19) where [·] denotes concatenation; and σ is a nonlinear function.

heading: AGGREGATING NEIGHBOR MESSAGES
text: Then, we aggregate all the messages sent to node i into a single message vector by averaging:
mti = 1 |Nn(i)| ∑
j∈Nn(i)
mtij (20)
where |Nn(i)| is the number of neighbor nodes of node i.

heading: UPDATING NODE STATE
text: Finally, we update the state of node i as follows:
xt+1i = g ( xti,m t i,vi ) (21)
where g(.) is a Highway Network (Srivastava et al., 2015):
xt+1i = Highway ( xti,m t i,vi ) (22)
= α ∗ x̃t+1i + (1−α) ∗ x t i (23)
where x̃t+1i is the nonlinear part which is computed as: x̄t+1i = σ ( W1 [ xti,m t i,v t i ] + b1 ) and α is the gate controlling the flow of information:
α = sigmoid(W2 [ xti,m t i,v t i ] + b2)
By combining Eqs. (19,20,22) together, one step of message passing update for node i can be written in a generic way as follows:
xt+1i = MessagePassing ( xti,vi,N (i) ) (24)

heading: A.2 UPDATING STATES

heading: UPDATING RNN STATE
text: We keep the old representation of the edge that have been modified in the hidden memory of the RNN as follows:
hτ = GRU ( hτ−1, zτuv ) (25)
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.

heading: UPDATING GRAPH STRUCTURE AND NODE STATES
text: 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:
N τ (u) = ( N τ−1(u)\ ( v, bold )) ∪ (v, b) (26)
N τ (v) = ( N τ−1(v)\ ( u, bold )) ∪ (u, b) (27)
Next, to update the node states, our model performs one step of message passing for u and v with their new neighbor sets:
xτu = MessagePassing ( xτ−1u ,vu,N τ (u) ) (28)
xτv = MessagePassing ( xτ−1v ,vv,N τ (v) ) (29)
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.

heading: A.3 MODEL CONFIGURATIONS
text: 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).
Graph Neural Network (GNN) There are 72 different types of atom depending on their atomic numbers and 5 different types of bond including NULL, SINGLE, DOUBLE, TRIPLE and AROMATIC. 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 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 zij of an atom pair (i, j) using a neural network with hidden size of 71. The second part maps zij to an unnormalized score sij 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 = LA2C + 0.5× Lvalue + Latom pair + 0.2× Lover length + 0.2× Lin topK
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 106 and the batch size is 20.

heading: A.4 DECODING WITH BEAM SEARCH
text: 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 τ :
log p ( (ξ, u, v, b)0:τ |G ) = log p ( (ξ, u, v, b)τ |G, (ξ, u, v, b)0:τ−1 ) +
log p ( (ξ, u, v, b)0:τ−1|G ) (30)
Computing all configurations of (ξ, u, v, b)τ jointly is very memory demanding, however. Thus, we decompose the first term as follows:
log p ( (ξ, u, v, b)τ |G, (ξ, u, v, b)0:τ−1 ) = log p ( ξτ |G, (ξ, u, v, b)0:τ−1 ) +
log p ( (u, v)τ |ξτ ,G, (ξ, u, v, b)0:τ−1 ) +
log ( bτ |(ξ, u, v)τ ,G, (ξ, u, v, b)0:τ−1 ) At step τ , we do beam search for the signal ξτ , then the atom pair (u, v)τ and finally the bond type bτ . Algorithm 1 describes beam search in detail. Some notable technicalities are:
• We only do beam search for (u, v) and b if the prediction is ongoing, i.e., when ξτ = 1. To keep track of this, we use a boolean vector C of length N with C0 is initialized to be all true. • To avoid beam search favoring short sequences, we normalize the log probability scores
over sequence lengths. This is shown in lines 10, 17, 32 and 47

heading: BEAM WIDTH ANALYSIS
text: Table 5 reports how beam width affects the decoding performance on the USPTO-15k dataset. Surprisingly, the top-1 accuracy in case of beam width3 of 1 is higher than the those when beam widths range from 2 to 15. It means that large beam width is not always good in our situation. However, at beam width of 20, our beam search achieves the best results for different values of k. Thus, we set the beam width to 20 in subsequent experiments.

heading: A.5 USING REAGENT INFORMATION EXPLICITLY
text: As can be seen from Table 6, reagent molecules account for about a half of the input molecules on average and 60-80% of all reactions containing reagents. It suggests that the proper use of reagent information will lead to better prediction. In our model, before computing the scores for all atom pairs, we append to the representation vector of every atom a binary scalar indicating whether this atom comes from a reagent molecule or not. Then, at the top-K atom pair selection step, we also exclude all atom pairs that have either atoms belong to a reagent molecule. The improvement in prediction accuracy on the validation set of USPTO-15k is shown in Fig. 4.

heading: A.6 COMPARISON WITH ELECTRO
text: In method Both GTPN and ELECTRO (Bradshaw et al., 2018) are able to explain the mechanism behind a reaction. ELECTRO regards a reaction as an ordered sequence that alternates between removing and adding a single bond. Our model, on the other hand, assumes no specific order of transformations as well as the amount of valences that a bond can change. Thus, our model is more generic than ELECTRO and can cover a much larger set of reactions.
3Note that beam search with beam width = 1 is different from greedy search as in beam search, as we model the whole sequence probability.
Algorithm 1 Reaction triple prediction using beam search. Input: A multi-graph G consisting of reactant and reagent molecules, number of bond types E, max
prediction steps T , beam width N 1: P 0 = [(−1,−1,−1,−1), ...] .The best N subsequences of (ξ, u, v, b) 2: S0 = [0, ...] .The length-normalized log joint probabilities of the best N subsequences 3: C0 = [True, ...] .The continuation indicator of the best N subsequences
4: Perform L steps of message passing for all nodes using Eq. (1) 5: x0i = xi ∀i ∈ V .The initial states of all nodes before decoding 6: N 0(i) = N (i) ∀i ∈ V .The initial neighbor set of all nodes before decoding 7: h0 is loaded from the saved model .The initial RNN hidden state before decoding
8: for τ from 1 to T do 9: Find the top K atom pairs { (uk, vk) τ | k = 1,K } using Eqs. (4,5)
10: Sτ−1;0 = Sτ−1 × τ−1τ .Superscript 0 denotes the sub-step 0 11: P τ−1;0 = P τ−1; Cτ−1;0 = Cτ−1
12: Beam search for continuation signals 13: 14: Rsignal = ∅ .Stores the log joint probabilities for N × 2 possible signals 15: for n from 1 to N do 16: Compute p ( ξτ | P τ−1;0n ) using Eq. (6)
17: Add Cτ−1;0 × 1τ log p ( ξτ = δ | P τ−1;0n ) + Sτ−1;0n to R
signal for δ ∈ {True,False} 18: end for 19: Sort Rsignal in descending order 20: Sτ−1;1 = Rsignal0:N 21: ξ̄τ ≡ output signal of N beams in Rsignal0:N 22: Iτ−1;1 ≡ indices of N beams in Rsignal0:N 23: P τ−1;1 = extract ( P τ−1;0, Iτ−1;1
) 24: Cτ−1;1 = extract ( Cτ−1;0, Iτ−1;1
) 25: Cτ−1;1n = C τ−1;1 n ∧ ξτn ∀n ∈ 1, N 26:
27: Beam search for atom pairs 28: 29: Ratom pair = ∅ .Stores the log joint probabilities for N ×K possible atom pairs 30: for n from 1 to N do 31: Compute p ( (u, v)τ |ξ̄τn, P τ−1;1n ) using Eq. (7)
32: Add Cτ−1;1 × 1τ log p ( (u, v)τk|ξτn, P τ−1;1n ) + Sτ−1;1n to R
atom pair ∀k ∈ 1,K 33: end for 34: Sort Ratom pair in descending order 35: Sτ−1;2 = Ratom pair0:N 36: (ū, v̄)τ ≡ output atom pair of N beams in Ratom pair0:N 37: Iτ−1;2 ≡ indices of N beams in Ratom pair0:N 38: P τ−1;2 = extract ( P τ−1;1, Iτ−1;2
) 39: Cτ−1;2 = extract ( Cτ−1;1, Iτ−1;2
) 40: ξ̄τ = extract ( ξ̄τ , Iτ−1;2
) 41:
42: Beam search for bonds 43: 44: Rbond = ∅ .Stores the log joint probabilities for N ×B possible bonds
Algorithm 2 Reaction triple prediction using beam search (cont.) 45: for n from 1 to N do 46: Compute p ( bτ | (ξ̄, ū, v̄)τn, P τ−1n ) using Eq. (8)
47: Add Cτ−1;2 × 1τ log p ( bτ = β | (ξ̄, ū, v̄), P τ−1n ) + Sτ−1b to R
bond ∀β ∈ 1, B 48: end for 49: Sort Rbond in descending order 50: Sτ−1;3 = Rbond0:N 51: b̄τ ≡ output bond of N beams in Rbond0:N 52: Iτ−1;3 ≡ indices of N beams in Rbond0:N 53: P τ−1;3 = extract ( P τ−1;2, Iτ−1;3
) 54: Cτ−1;3 = extract ( Cτ−1;2, Iτ−1;3
) 55: ξ̄τ = extract ( ξ̄τ , Iτ−1;3
) 56: (ū, v̄)τ = extract ( (ū, v̄)τ , Iτ−1;3
) 57:
58: Sτ = Sτ−1;3; Cτ = Cτ−1;3 59: P τn = append ( P τ−1;3n , (ξ̄, ū, v̄, b̄) τ n ) 60: for n from 1 to N do 61: Update the N τ (ūn) and N τ (v̄n) for all n = 1, N using Eqs. (26,27) 62: Update xτūn and x τ v̄n using Eq. (1) 63: Perform m steps of message passing for all nodes in the graph 64: Update hτ using Eq. (25) 65: end for 66: end for Output: PT , ST
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.

heading: A.7 ERROR ANALYSIS
text: 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.

Review summaries:
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.