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 end-to-end mechanistic generative model of electron flow in a particular type of chemical reaction (“Linear Electron Flow” reactions) . Interestingly, modeling the flow of electrons aids in the prediction of the final product of a chemical reaction over and above problems which attack this “product prediction problem” directly. The method is also shown to generalize well to held-out reactions (e.g. from a chemistry textbook).