G RAPH T RANSFORMATION P OLICY N ETWORK FOR C HEMICAL R EACTION P REDICTION
Anonymous authors
Paper under double-blind review
A BSTRACT
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.
I NTRODUCTION
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 ﬁrst 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 predeﬁned orders (Bradshaw et al., 2018).
The main challenge is that the space of all possible conﬁgurations 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
C:10 C:11 N:
O:
C:2 C:
C:4 C:5 C:
C:
Br:8 C:
C:10 C:11 N:
O:
C:2 C:
C:4 C:5 C:
C:
Br:8 C:
C:10 C:11 N:
O:
C:2 C:
C:4 C:5 C:
C:
Br:8 C:
Reactant s
Intermediate Molecules
Product 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 ﬁnal generated graph is considered as the predicted products of the reaction.
Importantly, GTPN does not assume any ﬁxed 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 uniﬁed 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 signiﬁcantly 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.
M ETHOD 2. C HEMICAL R EACTION AS M ARKOV D ECISION P ROCESS OF G RAPH T RANSFORMATIONS
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 (cid:48) .
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 , ...ξ T end − 1 will be 1 and ξ T end , ..., ξ T − 1 will be 0 .
At every step τ , if ξ τ = 1 , we apply the predicted edge change ( u, v, b ) τ on
[Figure]
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 ﬁnite and contains no loop, we set the discount factor γ to be 1 .
The rest of the MDP tuple are deﬁned as follows:
• State :
A state s τ ∈ S is an intermediate graph G τ generated at step τ (0 ≤ τ < T ) .
When τ = 0 , we denote s
= 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 ﬁrst sub-action is useful in case a reaction does not happen, i e., ξ 0 =• State Transition : If ξ τ = 1 , the current graph G τ is modiﬁed 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 ﬁnal products G T end .
• Reward : We use both immediate rewards and delayed rewards to encourage our model to learn the optimal policy faster.
At every step τ , if the model predicts ξ τ , ( u, v ) τ or b τ correctly, it will receive a positive reward for each correct sub-action.
Otherwise, a negative reward is given.
After the prediction process has terminated, if the generated products are exactly the same as the groundtruth products, we give the model a positive reward, otherwise a negative reward.
The concrete reward values are provided in Appendix A.3.
2. G RAPH T RANSFORMATION P OLICY N ETWORK
In this section, we describe the architecture of our model − a Graph Transformation Policy Network ( GTPN ).
GTPN has three main components namely a Graph Neural Network (GNN), a Node Pair
Prediciton Network (NPPN), and a Policy Network (PN).
Each component is responsible for one or several key functions shown in Fig.
2: GNN performs functions 1 and 6; NPPN performs function 2; and PN performs functions 3, 4 and 5.
Apart from these components, GTPN also has a Recurrent
Neural Network (RNN) to keep track of the past transformations.
The hidden state h of this RNN is used by NPPN and PN to make accurate prediction.
2.2. G RAPH N EURAL N ETWORK
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):
[Equation]
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. N ODE P AIR P REDICTION N ETWORK
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:
[Equation] [Equation]
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:
[Equation]
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 efﬁcient to identify reaction triples from a small subset of highly probable reaction atom pairs.
For that reason, we extract K ( K (cid:28) |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 (cid:8) ( u k , v k ) | k
= 1 , K (cid:9) , (cid:8) s u k v k | k
= 1 , K (cid:9) and Z K = (cid:8)
z u k v k | k
= 1 , K (cid:9) , respectively.
2.2. P OLICY N ETWORK
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:
[Equation]
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:
[Equation]
where s τu k v k is the score of the atom pair ( u k , v k )
τ computed in Eq.
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. U PDATING S TATES
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. T RAINING
Loss function plays a central role in achieving fast training and high performance.
We design the following loss:
[Equation]
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 coefﬁcients.
The component losses are explained in the following.
2.4. R EACTION TRIPLE LOSS
The loss follows a policy gradient method known as Advantage Actor-Critic (A2C):
[Equation]
(9) where T end is the ﬁrst 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, deﬁned as:
[Equation]
where r τ signal , r τ atom pair , r τ bond are immediate rewards at step τ ; at the ﬁnal 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:
[Equation]
where R τ is the return at step τ .
Episode termination during training Although the loss deﬁned 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 difﬁcult.
To resolve this issue, during training, we use a binary vector ζ ∈ { 0 , 1 } 3 T to keep track of the ﬁrst wrong sub-action: {+ζ t = (cid:26) 
+}
{+if t ≤ t ﬁrst wrong 
 if t > t+} ﬁrst wrong where t ﬁrst wrong denotes the sub-step at which our model chooses a wrong sub-action the ﬁrst time.
The actor-critic loss in Eq.
(9) now becomes:
[Equation]
where T is the maximum number of steps.
Similarly, we change the value loss into:
[Equation]
2.4. R EACTION 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 ﬁrst wrong =
(cid:106) t ﬁrstwrong(cid:107) ; η 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.
2.4. C ONSTRAINT ON THE SEQUENCE LENGTH
One major difﬁculty 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 signiﬁcantly 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 [Table] step τ during inference.
In to of sequence length training, we deﬁne a loss that punishes the model if it produces a longer sequence than the ground truth sequence:
[Equation]
(16) where T gtend is the end step of the ground-truth sequence.
Note that the loss in Eq.
(16) is not applied when T end ≤ T gtend .
The reason is that forcing ξ τ = 1 with T end ≤ τ < T gtend 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 gtend when it is smaller than T gtend has already been included in the advantages in Eq.
2.4. C ONSTRAINT 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 gtend .
To encourage the ground-truth reaction atom pair (˜ u, ˜ v ) τ with the highest score to appear in the top K , we introduce an additional rank-based loss:
[Equation]
where p ((˜ u, ˜ v ) τ in top K ) is computed as:
[Equation]
E XPERIMENTS 3. D ATASET
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 RDKit 1 , 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 identiﬁed 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
[Table]
[Figure]
3. R EACTION A TOM P AIR P REDICTION
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.
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 ti (different from the node feature vector v i ) for updating the structural information of a node (see
Eq.
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. T OP - K A TOM P AIR E XTRACTION
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 efﬁciency.
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 ﬁnd 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 efﬁciency. [Table]
3. R EACTION P RODUCT P REDICTION
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@ 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).
R ELATED W ORK 4. L EARNING TO P REDICT C HEMICAL R EACTION
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 ﬁrst classiﬁes reactions into different types based on the neural ﬁngerprint 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 ﬁrst 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 satisﬁed) 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 ﬁrst 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 predeﬁned topology.
Our method, by contrast, is order-free and can be applied to almost any kind of reactions.
4. G RAPH N EURAL N ETWORKS FOR M ODELING M OLECULES
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 classiﬁcation (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. R EINFORCEMENT L EARNING FOR S TRUCTURAL R EASONING
Reinforcement learning (RL) has become a standard approach to many structural reasoning problemsbecause 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).
D ISCUSSION
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 ﬁnd 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 A PPENDIX A. G RAPH N EURAL N ETWORK
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.
G RAPH 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 ) deﬁned 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]
C OMPUTING NEIGHBOR MESSAGES
At the message passing step t , we compute the message m tij from every neighbor node j ∈ N n ( i ) to node i as:
[Equation]
19) where [ · ] denotes concatenation; and σ is a nonlinear function.
A GGREGATING NEIGHBOR MESSAGES
Then, we aggregate all the messages sent to node i into a single message vector by averaging:
[Equation]
20) where |N n ( i ) | is the number of neighbor nodes of node i .
U PDATING NODE STATE
Finally, we update the state of node i as follows:
[Equation]
21) where g ( . ) is a Highway Network (Srivastava et al., 2015):
[Equation] [Equation]
[Equation]
where ˜ x t +1 i is the nonlinear part which is computed as:
[Equation]
and α is the gate controlling the ﬂow of information:
[Equation]
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:
[Equation]
A. U PDATING S TATES U PDATING RNN STATE
We keep the old representation of the edge that have been modiﬁed 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.
U PDATING 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] [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 deﬁned in Eq.
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.
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. M ODEL C ONFIGURATIONS
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 ﬁx the optimal settings of the ﬁrst 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).
[Table]
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 ﬁnal 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 ﬁrst 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.
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.
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 deﬁne the Advantage Actor Critic loss (A2C) as in Eq.
The coefﬁcients of components in the ﬁnal loss L are set empirically as follows:
[Equation]
We trained our model using Adam (Kingma & Ba, 2015) with the initial learning rate of 0.
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 10and the batch size is 20.
A. D ECODING WITH B EAM S EARCH
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 (cid:0) ( ξ, u, v, b ) τ | G , ( ξ, u, v, b ) 0:
τ − 1 (cid: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]
log
[Equation] [Equation]
[Table]
Computing all conﬁgurations of ( ξ, u, v, b ) τ jointly is very memory demanding, however.
Thus, we decompose the ﬁrst term as follows:
[Equation]
At step τ , we do beam search for the signal ξ τ , then the atom pair ( u, v ) τ and ﬁnally 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 C 0 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
B EAM WIDTH ANALYSIS
Table 5 reports how beam width affects the decoding performance on the USPTO-15k dataset.
Surprisingly, the top-1 accuracy in case of beam width 3 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.
A. U SING R EAGENT I NFORMATION E XPLICITLY
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.
A. C OMPARISON WITH ELECTRO
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 speciﬁc 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.
(cid:1) + S τ − 1;0 n to R signal for δ ∈ { True , False } 18: end for 19:
Sort R signal in descending order 20:
S τ − 1;1 =
R signal 0: N 21: ¯ ξ τ ≡ output signal of N beams in R signal 0: N 22:
I τ − 1;1 ≡ indices of N beams in R signal 0: N 23: P τ
[Equation] 1;1 = extract (cid:0) P [Equation] [Equation]
[Equation]
I
[Equation]
1;1 (cid:1) 24: C
[Equation] [Equation]
= extract (cid:0)
[Equation]
(cid:1) 25:
C
[Equation] [Equation]
n =
[Equation]
, N 26: 27:
Beam search for atom pairs 28: 29:
R atom pair = ∅ (cid:46) Stores the log joint probabilities for N × K possible atom pairs 30: for n from 1 to N do 31:
Compute p (cid:0) ( u, v ) τ | ¯ ξ τn , P τ − 1;1 n (cid:1) using Eq.
Add C τ − 1;1 × 1 τ log p (cid:0) ( u, v ) τk | ξ τn , P τ − 1;1 n (cid:1) + S τ − 1;1 n to R atom pair ∀ k ∈ 1 , K 33: end for 34:
Sort R atom pair in descending order 35:
S τ − 1;2 =
R atom pair 0: N 36: (¯ u, ¯ v ) τ ≡ output atom pair of N beams in R atom pair 0: N 37:
I τ − 1;2 ≡ indices of N beams in R atom pair 0: N 38:
P τ − 1;2 = extract (cid:0) P τ − 1;1 , I τ − 1;2 (cid:1) 39:
C τ − 1;2 = extract (cid:0) C τ − 1;1 , I τ − 1;2 (cid:1) 40: ¯ ξ τ = extract (cid:0) ¯ ξ τ , I τ − 1;2 (cid:1) 41: 42:
Beam search for bonds 43: 44:
R bond = ∅ (cid:46) Stores the log joint probabilities for N × B possible bonds
45: for n from 1 to N do 46:
Compute p (cid:0) b τ | (¯ ξ, ¯ u, ¯ v ) τn , P τ − 1 n (cid:1) using Eq.
Add C τ − 1;2 × 1 τ log p (cid:0) b τ = β | (¯ ξ, ¯ u, ¯ v ) , P τ − 1 n (cid:1)
+ S τ − 1 b to R bond ∀ β ∈ 1 , B 48: end for 49:
Sort R bond in descending order 50:
S τ − 1;3 =
R bond 0: N 51: ¯ b τ ≡ output bond of N beams in R bond 0: N 52:
I τ − 1;3 ≡ indices of N beams in R bond 0: N 53:
P τ − 1;3 = extract (cid:0) P τ − 1;2 , I τ − 1;3 (cid:1) 54:
C τ − 1;3 = extract (cid:0) C τ − 1;2 , I τ − 1;3 (cid:1) 55: ¯ ξ τ = extract (cid:0) ¯ ξ τ , I τ − 1;3 (cid:1) 56: (¯ u, ¯ v ) τ
= extract (cid:0) (¯ u, ¯ v ) τ , I τ − 1;3 (cid:1) 57: 58:
S τ = S τ − 1;3 ; C τ = C τ − 1; 59:
P τ n = append (cid:0) P τ − 1;3 n , (¯ ξ, ¯ u, ¯ v, ¯ b ) τn (cid:1) 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.
Update x τ ¯ u n and x τ ¯ v n using Eq.
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: P T , S T
Dataset %reactions containing reagents %reagents over input molecules
USPTO-15k train 63.1% 41.3% valid 65.3% 42.3% test 63.6% 40.9%
USPTO train 79.7% 54.0% valid 80.0% 54.4% test 79.9% 54.2%
[Figure]
[Table]
al
[Table] [Table]
(Bradshaw et al., 2018)
[Table]
[Figure]
In performance To do a fair comparison with ELECTRO (Bradshaw et al., 2018)
[Figure]
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 topand 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. E RROR A NALYSIS
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 deﬁne a sub-action causing error as the ﬁrst 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
[Figure]
(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.
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.
[Table]
Cl:
C:18 O:
Cl: prob: 0.
C:1 C:
C:3 C:4 C:5 C:
O:
H 2 O:
C:9 C:10 F:11 F:12 F:13 C:
C:21N:
C:
C:24O:
O:15C:
Cl:
C:18O:
Cl: prob: 0.
O
Cl
F F F