Molecular Graph Enhanced Transformer for Retrosynthesis Prediction

1 INTRODUCTION . Retrosynthesis prediction aims to predict a set of suitable reactants that can synthesize the desired molecule via a series of reactions . It pushes forward an immense influence in agriculture , medical treatment , drug discovery and so on . However , the retrosynthesis prediction is challenging since there are massive possible synthetic routes available and it is often difficult to navigate the direction of retrosynthesis process . Indeed , each bond in the target molecule may represent a possible retrosynthetic disconnection , leading to a vast space of possible starting materials . Besides , the difference between two synthetic routes may be subtle , which usually depends on the global structures . Actually , planning a proper retrosynthetic route for a complex molecule is also a tough work even for the professional chemists . One of the prevailing methods is to deem the retrosynthesis prediction as a machine translation task . This analogy is comprehensible since every molecule has a unique text representation named SMILES ( Weininger , 1988 ) . In this case , given a target molecule written in SMILES notation , the retrosynthesis prediction is just to predict a string of SMILES which represents the reactants . Based on this idea , Liu et al . ( 2017 ) first applied the LSTM with attention mechanism in retrosynthesis prediction and achieved comparable performance compared with previous traditional methods . Whereafter , many works ( Karpov et al. , 2019 ; Zheng et al. , 2019 ; Lin et al. , 2019 ; Lee et al. , 2019 ) tried to employ Transformer ( Vaswani et al. , 2017 ) , which is a more powerful Sequence-toSequence ( Seq2Seq ) model , to improve prediction accuracy in retrosynthesis . However , these methods just utilize the sequential representations of the molecule , while ignoring the natural topological connections between atoms within the molecule . These atomic connections can provide more flexible and accurate chemical information , which is critical in many related chemical tasks like molecular representation ( Duvenaud et al. , 2015 ; Gilmer et al. , 2017 ) and chemical reaction prediction ( Jin et al. , 2017 ; Do et al. , 2019 ) . We believe that the absence of this molecular graph information hinders the further improvement of the present methods for retrosynthesis . How to effectively make use of this natural graphical information of the molecular structure , therefore , becomes a vital problem . To tackle this problem , we propose Graph Enhanced Transformer ( GET ) framework that can enjoy the advantage of both graph-level representations and sequence-level representations . Specifically , to solve the retrosynthesis problem , we design an improved Graph Neural Network ( GNN ) called Graph Attention with Edge and Skip-connection ( GAES ) to learn each atom ’ s representation , and try four strategies to incorporate it with the original SMILES representation in the encoder . The main contributions of this paper are as follows : • We propose a new framework called GET that fuses graphical representations with sequential representations of the target molecule to solve retrosynthesis prediction task . • We design a powerful GNN called GAES that learns high-quality representations of atom nodes in a self-attention manner with bond features , and it is less affected by the side-effect of stacking more layers . GET is evaluated on USPTO-50K , a common benchmark dataset for retrosynthesis . Experimental results show that our model achieves new records for top-1 prediction accuracy in the state-of-the-art methods and outperforms Seq2Seq-based methods in all tested top-n accuracy , demonstrating the effectiveness of fusing the molecular graph information with the SMILES sequence information . 2 RELATED WORK . Prior work on retrosynthesis can be mainly summarized into two categories : template-based methods and template-free methods . Template-based Methods The majority of computer-aided retrosynthetic methods in the early period were relied on encoding reaction templates or generalized subgraph matching rules . LHASA ( Corey & Wipke , 1969 ) was the first software for retrosynthetic analysis . Recently , one of the most well-known retrosynthesis analysis tool is Synthia ( Szymkuć et al. , 2016 ) that integrated about 70,000 hand-encoded reaction rules collected by manual . Based on the 60,000 reaction templates derived from 12 million single-step reaction examples , Schreck et al . ( 2019 ) introduced Reinforcement Learning ( RL ) into this area by treating retrosynthesis as a game whose goal is to identify policies that make ( near ) optimal reaction choices during each step of retrosynthetic planning . Segler et al . ( 2018 ) extracted two sets of transformation rules and combined Monte Carlo tree search with symbolic AI to discover possible retrosynthetic routes . Besides manual extracted rules , some works ( Law et al. , 2009 ; Segler & Waller , 2017 ; Coley et al. , 2017 ) tried to collect reaction templates automatically and perform retrosynthesis based on these automated templates . Although template-based methods work well in many cases , they still face a serious drawback that they generally can not achieve accurate prediction accuracy outside of their known templates . Template-free Methods Our model just belongs to this category . Emerging template-free methods are to treat retrosynthesis as a machine translation task as introduced in section 1 . Since these methods do not need any reaction templates and prior chemistry knowledge , they are attracting more and more attention from academia . Moreover , without the constraint of fixed templates , they have the potential of discovering novel synthetic routes . The most related work to ours are Zheng et al . ( 2019 ) Lin et al . ( 2019 ) ; Karpov et al . ( 2019 ) ; Lee et al . ( 2019 ) that apply Transformer in retrosynthesis prediction . 3 BACKGROUND . In this section , we first introduce how GNN is used in learning the molecule ( or atoms ) representation , and then describe how the Transformer model is previously applied in retrosynthesis prediction . 3.1 GNN FOR MOLECULE REPRESENTATION LEARNING . GNN has been widely used in learning the representation of the molecule and its atoms . Naturally , molecules can be represented as graph structure with atoms as nodes and bonds as edges . Suppose that a molecular graph G has initial node representations hv and edge representations evw , a typical one-layer GNN can learn new and more powerful node representations from G by the following message passing process described in Gilmer et al . ( 2017 ) : mv = ∑ w∈N ( v ) M ( hv , hw , evw ) , ( 1 ) hnewv = U ( hv , mv ) , ( 2 ) where N ( v ) denotes the neighbors of node v in graph G , M is the message function that is responsible for collecting information from neighbors , and U is the update function for fusing collected information mv with old node representation hv to obtain the new node representation hnewv . Further , we can stack several these GNN layers to capture higher-order neighbors ’ information . Then a readout function R can be used to integrate all node representations into a whole graph representation g : g = R ( { hv|v ∈ G } ) . ( 3 ) The hv and g , which represent the atoms and the whole molecule , are often trained in an end-to-end way for a specific chemical task , such as chemical properties prediction , reaction prediction and molecule optimization . 3.2 TRANSFORMER FOR RETROSYNTHESIS PREDICTION . Transformer ( Vaswani et al. , 2017 ) is a Seq2Seq model that has shown excellent performance in machine translation task . Also , it has been applied in chemical reaction prediction and retrosynthesis prediction before . Given an input SMILES that represents the target molecule and a specified reaction type ( optional ) , retrosynthesis prediction is to predict the output SMILES which represents the possible reactants that can synthesis the target molecule in the specified reaction type . Thus , retrosynthesis prediction can be deemed as a machine translation task whose source language is target molecule SMILES and the target language is reactants SMILES . In this view , Transformer can be applied to retrosynthesis prediction as the same as to machine translation . Since Transformer is a mature model that has been widely used in natural language processing ( NLP ) , we just give a simple introduction here . Specifically , Transformer follows an encoder-decoder structure and is composed of several combinations of multi-head attention layers and position-wise feed forward layers . The encoder consists of a stack of N = 6 identical layers . Each layer includes two main components : ( multi-head ) self-attention layer and feed-forward network . Given an input vectors ( p1 , ... , pn ) , p ∈ Rd , the t-th output st of the self-attention layer is calculated by : qt = Wqpt , km = Wmpm , vt = Wvpt , st = n∑ m=1 softmax ( < qt , km > √ dk ) vt , ( 4 ) where dk is the dimension of q and k , Wq , Wm , Wv are weight matrices . One such operation is called one head , and we can concatenate several heads to change to multi-head self attention . The feed-forward network is composed of two linear transformations with a ReLU activation : FFN ( x ) = max ( 0 , xW1 + b1 ) W2 + b2 , ( 5 ) where W1 , W2 are weight matrices and b1 , b2 are biases . Similarly , the decoder is also mainly composed of multi-attention layers and position-wise feed forward layers . It will generate the output SMILES step by step . At step t , it utilizes the encoder ’ s output ( p′1 , ... , p ′ n ) and all previous steps ’ output ( x ′ 1 , ... , x ′ t−1 ) to generate the next SMILES character x′t . This process repeated until generating a specific termination character , i.e. , x ′ t = < EOS > . 4 GRAPH ENHANCED TRANSFORMER . In this section , we provide the details about our Graph Enhanced Transformer ( GET ) framework for retrosynthesis prediction . Figure 1 shows an overview of GET . On the whole , GET is accord with typical encoder-decoder structure , of which the integral encoder is composed of the graph encoder and transformer encoder for learning the representation in graph-level and sequence-level respectively . Given a target molecule ’ s SMILES , it will first pass through the two encoders somehow to get the hidden representation of each character , and then the decoder will utilize these hidden representations to generate the output SMILES . 4.1 GRAPH ENCODER . We design a new powerful GNN called Graph Attention with Edge and Skip-connection ( GAES ) as the graph encoder , which can learn the representation of each atom in a molecule . This graphical representation reflects the connection of atoms within a molecule and may play a significant role in further alleviating the long-term dependency problem to avoid generating chemically invalid output . We use RDKit ( Landrum , 2016 ) to transform a SMILES into the molecular graph , whose nodes are atoms and edges are chemical bonds . The input representation of the atom node is a 21-dimensional vector that contains some chemical information about the atom , and the detail can be found in Table 1 ( nearly consistent with Gilmer et al . ( 2017 ) ) . The input representation of the edge is a 4-dimensional one-hot vector that encodes the bond types including single , double , triple and aromatic . Since the SMILES sequence ( x1 , ... , xn ) has been transformed into a graph G with the input representations ( h1 , ... , hN ) for nodes and { eij } for edges that exist between node i and node j , our GNN will produce new representation h′i for each node i by the following message passing operations : αij = exp ( LeakyReLU ( aT [ Whi‖Whj‖eij ] ) ) ∑ k∈Ni exp ( LeakyReLU ( a T [ Whi‖Whk‖eik ] ) ) , ( 6 ) ĥi = σ ∑ j∈Ni αijWhj , ( 7 ) where a ∈ R ( 2F ′+E ) is a weight vector for attention mechanism and W ∈ RF ′×F is a weight matrix for transforming the node features , so F is the input dimension of nodes , F ′ is the output dimension of nodes and E is the dimension of edges . Ni is the set of first-order neighbors of node i ( including itself ) . σ is an activation function , e.g. , ReLU . In practice , we performK multi-head attention ( Veličković et al. , 2017 ) to enrich the model capacity and to stabilize the learning process . Each attention head has its own parameters and we average their outputs to get better representation : ĥi = σ 1 K K∑ k=1 ∑ j∈Ni αkijW khj . ( 8 ) The above operations can be seen as GAT ( Veličković et al. , 2017 ) extended to include edge features . Then , to mitigate the accuracy reduction issue ( Kipf & Welling , 2016 ) caused by stacked graph convolution layers , we adopt the gated skip-connection mechanism ( Ryu et al. , 2018 ) to get the final representation : zi = sigmod ( U1ĥi +U2hi + b ) , ( 9 ) h′i = zi ĥi + ( 1− zi ) hi , ( 10 ) where U1 , U2 and b are trainable parameters . Note that the above operations are just in one layer of our GAES , and we can stack several layers to capture the information about higher-order atom neighbors so that to obtain more comprehensive representations .

Summary: This paper focuses on the retrosynthesis prediction problem which to my understanding is the factorization of a target molecule into simpler structures. Previously, retrosynthesis prediction has been tackled as a language translation problem by using a Seq2Seq algorithm called "Transformer". This sequential approach was possible because molecules can be expressed as strings using the following format: Simplified Molecular-Input Line-Entry System (SMILES).