A Generative Model For Electron Paths

1 INTRODUCTION . The ability to reliably predict the products of chemical reactions is of central importance to the manufacture of medicines and materials , and to understand many processes in molecular biology . Theoretically , all chemical reactions can be described by the stepwise rearrangement of electrons in molecules ( Herges , 1994b ) . This sequence of bond-making and breaking is known as the reaction mechanism . Understanding the reaction mechanism is crucial because it not only determines the products ( formed at the last step of the mechanism ) , but it also provides insight into why the products are formed on an atomistic level . Mechanisms can be treated at different levels of abstraction . On the lowest level , quantum-mechanical simulations of the electronic structure can be performed , which are prohibitively computationally expensive for most systems of interest . On the other end , chemical reactions can be treated as rules that “ rewrite ” reactant molecules to products , which abstracts away the individual electron redistribution steps into a single , global transformation step . To combine the advantages of both approaches , chemists use a powerful qualitative model of quantum chemistry colloquially called “ arrow pushing ” , which simplifies the stepwise electron shifts using sequences of arrows which indicate the path of electrons throughout molecular graphs ( Herges , 1994b ) . Recently , there have been a number of machine learning models proposed for directly predicting the products of chemical reactions ( Coley et al. , 2017 ; Jin et al. , 2017 ; Schwaller et al. , 2018 ; Segler and Waller , 2017a ; Segler et al. , 2018 ; Wei et al. , 2016 ) , largely using graph-based or machine translation models . The task of reaction product prediction is shown on the left-hand side of Figure 1 . In this paper we propose a machine learning model to predict the reaction mechanism , as shown on the right-hand side of Figure 1 , for a particularly important subset of organic reactions . We argue that our model is not only more interpretable than product prediction models , but also allows easier encoding of constraints imposed by chemistry . Proposed approaches to predicting reaction mechanisms have often been based on combining hand-coded heuristics and quantum mechanics ( Bergeler et al. , 2015 ; Kim et al. , 2018 ; Nandi et al. , 2017 ; Segler and Waller , 2017b ; Rappoport et al. , 2014 ; Simm and Reiher , 2017 ; Zimmerman , 2013 ) , rather than using machine learning . We call our model ELECTRO , as it directly predicts the path of electrons through molecules ( i.e. , the reaction mechanism ) . To train the model we devise a general technique to obtain approximate reaction mechanisms purely from data about the reactants and products . This allows one to train our a model on large , unannotated reaction datasets such as USPTO ( Lowe , 2012 ) . We demonstrate that not only does our model achieve impressive results , surprisingly it also learns chemical properties it was not explicitly trained on . 2 BACKGROUND . We begin with a brief background from chemistry on molecules and chemical reactions , and then review related work in machine learning on predicting reaction outcomes . We then describe a particularly important subclass of chemical reactions , called linear electron flow ( LEF ) reactions , and summarize the contributions of this work . 2.1 MOLECULES AND CHEMICAL REACTIONS . Organic ( carbon-based ) molecules can be represented via a graph structure , where each node is an atom and each edge is a covalent bond ( see example molecules in Figure 1 ) . Each edge ( bond ) represents two electrons that are shared between the atoms that the bond connects . Electrons are particularly important for describing how molecules react with other molecules to produce new ones . All chemical reactions involve the stepwise movement of electrons along the atoms in a set of reactant molecules . This movement causes the formation and breaking of chemical bonds that changes the reactants into a new set of product molecules ( Herges , 1994a ) . For example , Figure 1 ( Right ) shows how electron movement can break bonds ( red arrows ) and make new bonds ( green arrows ) to produce a new set of product molecules . 2.2 RELATED WORK . In general , work in machine learning on reaction prediction can be divided into two categories : ( 1 ) Product prediction , where the goal is to predict the reaction products , given a set of reactants and reagents , shown in the left half of Figure 1 ; and ( 2 ) Mechanism prediction , where the goal is to determine how the reactants react , i.e. , the movement of electrons , shown in the right of Figure 1 . Product prediction . Recently , methods combining machine learning and template-based molecular rewriting rules have been proposed ( Coley et al. , 2017 ; Segler and Waller , 2017a ; Segler et al. , 2018 ; Wei et al. , 2016 ; Zhang and Aires-de Sousa , 2005 ) . Here , a learned model is used to predict which rewrite rule to apply to convert one molecule into another . While these models are readily interpretable , they tend be brittle . Another approach , introduced by Jin et al . ( 2017 ) , constructs a neural network based on the Weisfeiler-Lehman algorithm for testing graph isomorphism . They use this algorithm ( called WLDN ) to select atoms that will be involved in a reaction . They then enumerate all chemically-valid bond changes involving these atoms and learn a separate network to rank the resulting potential products . This method , while leveraging new techniques for deep learning on graphs , can not be trained end-to-end because of the enumeration steps for ensuring chemical validity . Schwaller et al . ( 2018 ) represents reactants as SMILES ( Weininger , 1988 ) strings and then uses a sequence to sequence network ( specifically , the work of Zhao et al . ( 2017 ) ) to predict product SMILES . While this method ( called Seq2Seq ) is end-to-end trainable , the SMILES representation is quite brittle as often single character changes will not correspond to a valid molecule . These latter two methods , WLDN and Seq2Seq , are state-of-the-art on product prediction and have been shown to outperform the above template-based techniques ( Jin et al. , 2017 ) . Thus we compare directly with these two methods in this work . Mechanism prediction . The only other work we are aware of to use machine learning to predict reaction mechanisms are Fooshee et al . ( 2018 ) ; Kayala and Baldi ( 2011 ; 2012 ) ; Kayala et al . ( 2011 ) . All of these model a chemical reaction as an interaction between atoms as electron donors and as electron acceptors . They predict the reaction mechanisms via two independent models : one that identifies these likely electron sources and sinks , and another that ranks all combinations of them . These methods have been run on small expert-curated private datasets , which contain information about the reaction conditions such as the temperature and anion/cation solvation potential ( Kayala and Baldi , 2011 , §2 ) . In contrast , in this work , we aim to learn reactions from noisy large-scale public reaction datasets , which are missing the required reaction condition information required by these previous works . As we can not yet apply the above methods on the datasets we use , nor test our models on the datasets they use ( as the data are not yet publicly released ) , we can not compare directly against them ; therefore , we leave a detailed investigation of the pros and cons of each method for future work . As a whole , this related work points to at least two main desirable characteristics for reaction prediction models : 1 . End-to-End : There are many complex chemical constraints that limit the space of all possible reactions . How can we differentiate through a model subject to these constraints ? 2 . Mechanistic : Learning the mechanism offers a number of benefits over learning the products directly including : interpretability ( if the reaction failed , what electron step went wrong ) , sparsity ( electron steps only involve a handful of atoms ) , and generalization ( unseen reactions also follow a set of electron steps ) . Table 1 describes how the current work on reaction prediction satisfies these characteristics . In this work we propose to model a subset of mechanisms with linear electron flow , described below . 2.3 LINEAR ELECTRON FLOW REACTIONS . Reaction mechanisms can be classified by the topology of their “ electron-pushing arrows ” ( the red and green arrows in Figure 1 ) . Here , the class of reactions with linear electron flow ( LEF ) topology is by far the most common and fundamental , followed by those with cyclic topology ( Herges , 1994a ) . In this work , we will only consider LEF reactions that are heterolytic , i.e. , they involve pairs of electrons.1 1The treatment of radical reactions , which involve the movement of single electrons , will be the topic of future work . If reactions fall into this class , then a chemical reaction can be modelled as pairs of electrons moving in a single path through the reactant atoms . In arrow pushing diagrams representing LEF reactions , this electron path can be represented by arrows that line up in sequence , differing from for example pericyclic reactions in which the arrows would form a loop ( Herges , 1994a ) . Further for LEF reactions , the movement of the electrons along the linear path will alternately remove existing bonds and form new ones . We show this alternating structure in the right of Figure 1 . The reaction formally starts by ( step 1 ) taking the pair of electrons between the Li and C atoms and moving them to the C atom ; this is a remove bond step . Next ( step 2 ) a bond is added when electrons are moved from the C atom in reactant 1 to a C atom in reactant 2 . Then ( step 3 ) a pair of electrons are removed between the C and O atoms and moved to the O atom , giving rise to the products . Predicting the final product is thus a byproduct of predicting this series of electron steps . Contributions . We propose a novel generative model for modeling the reaction mechanism of LEF reactions . Our contributions are as follows : • We propose an end-to-end generative model for predicting reaction mechanisms , ELECTRO , that is fully differentiable . It can be used with any deep learning architecture on graphs . • We design a technique to identify LEF reactions and mechanisms from purely atom-mapped reactants and products , the primary format of large-scale reaction datasets . • We show that ELECTRO learns chemical knowledge such as functional group selectivity without explicit training . 3 THE GENERATIVE MODEL . In this section we define a probabilistic model for electron movement in linear electron flow ( LEF ) reactions . As described above ( §2.1 ) all molecules can be thought of as graphs where nodes correspond to atoms and edges to bonds . All LEF reactions transform a set of reactant graphs , M0 into a set of product graphsMT+1 via a series of electron actions P0 : T = ( a0 , . . . , aT ) . As described , these electron actions will alternately remove and add bonds ( as shown in the right of Figure 1 ) . This reaction sometimes includes additional reagent graphs , Me , which help the reaction proceed , but do not change themselves . We propose to learn a distribution pθ ( P0 : T | M0 , Me ) over these electron movements . We first detail the generative process that specifies pθ , before describing how to train the model parameters θ . To define our generative model , we describe a factorization of pθ ( P0 : T | M0 , Me ) into three components : 1. the starting location distribution pstartθ ( a0 | M0 , Me ) ; 2. the electron movement distribution pθ ( at | Mt , at−1 , t ) ; and 3. the reaction continuation distribution pcontθ ( ct | Mt ) . We define each of these in turn and then describe the factorization ( we leave all architectural details of the functions introduced to the appendix ) . Starting Location . At the beginning the model needs to decide on which atom a0 starts the path . As this is based on ( i ) the initial set of reactantsM0 and possibly ( ii ) a set of reagentsMe , we propose to learn a distribution pstartθ ( a0 | M0 , Me ) . To parameterize this distribution we propose to use any deep graph neural network , denoted hA ( · ) , to learn graph-isomorphic node features from the initial atom and bond features2 ( Duvenaud et al. , 2015 ; Kipf and Welling , 2017 ; Li et al. , 2016 ; Gilmer et al. , 2017 ) . We choose to use a 4 layer Gated Graph Neural Network ( GGNN ) ( Li et al. , 2016 ) , for which we include a short review in the appendix . Given these atom embeddings we also compute graph embeddings ( Li et al. , 2018 , §B.1 ) ( also called an aggregation graph transformation ( Johnson , 2017 , §3 ) ) , which is a vector that represents the entire molecule setM that is invariant to any particular node ordering . Any such function g ( · ) that computes this mapping can be used here , but the particular graph embedding function we use is inspired by Li et al . ( 2018 ) , and described in detail in Appendix B . We can now parameterize pstartθ ( a0 | M0 , Me ) as pstartθ ( a0 | M0 , Me ) = softmax [ f start ( hA ( M0 ) , greagent ( Me ) ) ] , ( 1 ) 2The molecular features we use are described in Table 4 in Appendix B. where f start is a feedforward neural network which computes logits x ; the logits are then normalized into probabilities by the softmax function , defined as softmax [ x ] = ex/ ∑ i e xi . Electron Movement . Observe that since LEF reactions are a single path of electrons ( §2.3 ) , at any step t , the next step at in the path depends only on ( i ) the intermediate molecules formed by the action path up to that pointMt , ( ii ) the previous action taken at−1 ( indicating where the free pair of electrons are ) , and ( iii ) the point of time t through the path , indicating whether we are on an add or remove bond step . Thus we will also learn the electron movement distribution pθ ( at | Mt , at−1 , t ) . Similar to the starting location distribution we again make use of a graph-isomorphic node embedding function hA ( M ) . In contrast , the above distribution can be split into two distributions depending on the parity of t : the remove bond step distribution premoveθ ( at | Mt , at−1 ) when t is odd , and the add bond step distribution paddθ ( at | Mt , at−1 ) when t is even . We parameterize the distributions as premoveθ ( at | Mt , at−1 ) ∝ βremove softmax [ f remove ( hA ( Mt ) , at−1 ) ] , ( 2 ) paddθ ( at | Mt , at−1 ) ∝ βadd softmax [ f add ( hA ( Mt ) , at−1 ) ] ( 3 ) pθ ( at | Mt , at−1 , t ) = { premoveθ ( at | Mt , at−1 ) if t is odd paddθ ( at | Mt , at−1 ) otherwise ( 4 ) The vectors βremove , βadd are masks that zero-out the probability of certain atoms being selected . Specifically , βremove sets the probability of any atoms at to 0 if there is not a bond between it and the previous atom at−13 . The other mask vector βadd masks out the previous action , preventing the model from stalling in the same state for multiple time-steps . The feedforward networks fadd ( · ) , f remove ( · ) and other architectural details are described in Appendix C. Reaction Continuation / Termination . Additionally , as we do not know the length of the reaction T , we introduce a latent variable ct ∈ { 0 , 1 } at each step t , which describes whether the reaction continues ( ct=1 ) or terminates ( ct=0 ) 4 . We also define an upper bound Tmax on the number of reaction steps . 3One subtle point is if a reaction begins with a lone-pair of electrons then we say that this reaction starts by removing a self-bond . Thus , in the first remove step βremove it is possible to select a1=a0 . But this is not allowed via the mask vector in later steps . 4An additional subtle point is that we do not allow the reaction to stop until until it has picked up an entire pair ( ie c1 = 1 ) . Algorithm 1 The generative steps of ELECTRO ( given that the model chooses to react , ie c0 = 1 ) . Input : Reactant molecules M0 ( consisting of atoms A ) , reagents Me , atom embedding function hA ( · ) , graph embedding functions greagent ( · ) and gcont ( · ) , additional logit functions f start ( · ) , f remove ( · ) , fadd ( · ) , time steps Tmax 1 : pstartθ ( a | M0 , Me ) , softmax [ f start ( hA ( M0 ) , greagent ( Me ) ) ] .a starts reaction 2 : a0 ∼ pstartθ ( a | M0 , Me ) 3 : M1 ←M0 .The molecule does not change until complete pair picked up 4 : c1 , 1 .You can not stop until picked up complete pair 5 : for t = 1 , . . . , Tmax do 6 : if t is odd then 7 : premoveθ ( at|Mt , at−1 ) ∝ βremovesoftmax [ f remove ( hA ( Mt ) , at−1 ) ] 8 : at ∼ premoveθ ( at|Mt , at−1 ) .electrons remove bond between at and at−1 9 : else 10 : paddθ ( at|Mt , at−1 ) ∝ βaddsoftmax [ fadd ( hA ( Mt ) , at−1 ) ] 11 : at ∼ paddθ ( at|Mt , at−1 ) .electrons add bond between at and at−1 12 : end if 13 : Pt = P0 : t−1 ∪ at 14 : Mt+1 ←Mt , at .modify molecules based on previous molecule and action 15 : pcontθ ( ct+1 | Mt+1 ) , σ ( gcont ( Mt+1 ) ) 16 : ct+1 ∼ pcontθ ( ct+1 | Mt+1 ) .whether to continue reaction 17 : if ct+1 = 0 then 18 : break 19 : end if 20 : end for Output : Electron path P0 : t The final distribution we learn is the continuation distribution pcontθ ( ct | Mt ) . For this distribution we learn a different graph embedding function gcont ( · ) to decide whether to continue or not : pcontθ ( ct | Mt ) = σ ( gcont ( Mt ) ) . ( 5 ) where σ is the sigmoid function σ ( a ) = 1/ ( 1 + e−a ) . Path Distribution Factorization . Given these distributions we can define the probability of a path P0 : T with the distribution pθ ( P0 : T | M0 , Me ) , which factorizes as pθ ( P0 : T | M0 , Me ) = pcontθ ( c0 | M0 ) pstartθ ( a0 | M0 , Me ) ( 6 ) × [ T∏ t=1 pcontθ ( ct | Mt ) pθ ( at | Mt , at−1 , t ) ] ( 1− pcontθ ( cT+1 | MT+1 ) ) , Figure 2 gives a graphical depiction of the generative process on a simple example reaction . Algorithm 1 gives a more detailed description . Training We can learn the parameters θ of all the parameterized functions , including those producing node embeddings , by maximizing the log likelihood of a full path log pθ ( P0 : T | M0 , Me ) . This is evaluated by using a known electron path a ? t and intermediate productsM ? t extracted from training data , rather than on simulated values . This allows us to train on all stages of the reaction at once , given electron path data . We train our models using Adam ( Kingma and Ba , 2015 ) and an initial learning rate of 10−4 , with minibatches consisting of a single reaction , where each reaction often consists of multiple intermediate graphs . Prediction Once trained , we can use our model to sample chemically-valid paths given an input set of reactantsM0 and reagentsMe , simply by simulating from the conditional distributions until sampling a continue value equal to zero . We instead would like to find a ranked list of the top-K predicted paths , and do so using a modified beam search , in which we roll out a beam of width K until a maximum path length Tmax , while recording all paths which have terminated . This search procedure is described in detail in Algorithm 2 in the appendix .

Summary: The paper presents a novel method for predicting organic chemical reactions, in particular, for learning (Robinson-Ingold's) ''arrow pushing" mechanisms in an end-to-end manner. Organic molecules consist of covalent bonds (that's why we can model them as molecular graphs), and organic reactions are recombinations of these bonds. As seen in organic chemistry textbooks, traditional chemists would qualitatively understand organic reactions as an alternating series of electron movements by bond breaking (bond cleavage) and bond forming (bond formation). Though now quantum chemistry calculations can give accurate quantitative predictions, these qualitative understanding of organic reactions still also gives strong foundations to consider and develop organic reactions. The proposed method tries to learn these series of bond changes directly through differentiable architectures consisting of three graph neural networks: 1) the one for determining the initial atom where electron movements start, 2) the one for representing state transitions from  the previous bond change to the next, and 3) the one for determining when the electron movements end. Experimental evaluations illustrate the quantitative improvement in final product prediction against existing methods, as well as give chemical intuitions that the proposed method can detect a class of LEFs (linear electron flows).