Permutation invariant graph-to-sequence model for template-free retrosynthesis and reaction prediction

1 INTRODUCTION . Retrosynthetic analysis ( Corey , 1988 ; Corey & Cheng , 1989 ) and its reverse problem , reaction outcome prediction ( Corey & Wipke , 1969 ) , are two fundamental problems in computer-aided organic synthesis . The former tries to propose possible reaction precursors given the desirable product , whereas the latter aims to predict the major products given reactants . Historically , they were tackled using rule-based expert systems such as LHASA ( Corey et al. , 1972 ) . Recent developments in machine learning have led to a number of new template-based , graph edit-based , and translation-based methods , for which we give a detailed review in Section 3.1 . For both tasks , translation-based approaches have grown popular , possibly because the end-to-end formulation is procedurally simple . Since most organic molecules can be represented as SMILES strings ( Neglur et al. , 2005 ) , retrosynthesis can be cast as a translation from product SMILES to reactant SMILES ( Liu et al. , 2017 ) , and so can reaction outcome prediction ( Nam & Kim , 2016 ) . Modeling these tasks as machine translation problems enables the use of neural architectures that are well-studied and well-optimized in the field of Natural Language Processing ( NLP ) . Several of the best performing models across multiple benchmark datasets ( Tetko et al. , 2020 ; Irwin et al. , 2021 ; Sun et al. , 2020 ; Seo et al. , 2021 ; Wang et al. , 2021b ) have used the Transformer architecture ( Vaswani et al. , 2017 ) as the backbone on SMILES representations , showing the effectiveness of translation-based formulation . However , SMILES representations do not provide bijective mappings to molecular structure . As a result , data augmentation with chemically-equivalent SMILES ( Bjerrum , 2017 ) has become a common practice to improve empirical performance . Augmenting the training data with just 1 equivalent reaction SMILES by permuting both the inputs and the outputs can already produce noticeable improvements of 0.8 % to 4.3 % ( Schwaller et al. , 2019 ; Seo et al. , 2021 ) . Incorporating 9 ( Wang et al. , 2021b ) , or up to 100 ( Tetko et al. , 2020 ) equivalent SMILES can provide additional gains . While these efforts demonstrate the effectiveness of SMILES augmentation , this introduces a non-trivial choice as to how much augmentation should be done . The performance has not saturated even with 100 augmented SMILES ( Tetko et al. , 2020 ) , but this may significantly complicate the model pipeline , especially during test time which typically necessitates de-duplication , ensembling , and heuristic-based scoring ( Wang et al. , 2021b ; Tetko et al. , 2020 ) . In this paper , we propose a novel graph-to-sequence architecture called Graph2SMILES to solve the tasks of retrosynthesis and reaction prediction . We first design a sequential graph encoder with an attention-augmented directed message passing neural network ( D-MPNN ) based on Yang et al . ( 2019 ) and optionally augmented by attention-based message updates , followed by a Transformerbased global attention encoder with graph-aware positional embedding . We then pair the graph encoder with a Transformer decoder to transform molecular graph inputs into SMILES outputs , without using sequence representations of input SMILES at all . As such , we guarantee the permutation invariance of Graph2SMILES to the input , eliminating the need for input-side augmentation altogether . Our main contributions can be summarized as follows : 1 . We propose a Graph2SMILES architecture with two encoder components modeling local and global atomic interactions respectively , to address forward prediction and one-step retrosynthesis as graph-to-sequence tasks . 2 . We design a graph-aware positional embedding to further enhance performance . It is easily generalizable to graphs containing two or more molecules , while not requiring any pretraining or joint training with auxiliary tasks . 3 . We demonstrate the adequacy of graph representations alone by showing that Graph2SMILES outperforms Transformer baselines on predictive chemistry tasks without needing any input-side SMILES augmentation . Our Graph2SMILES architecture achieves state-of-the-art top-1 accuracy on common benchmarks among methods that do not make use of reaction templates , atom mapping , pretraining , or data augmentation strategies . We emphasize that Graph2SMILES is a backbone architecture , and hence a drop-in replacement for the Transformer model . As such , Graph2SMILES can be plugged into any method for molecular transformation tasks that uses the Transformer model , while retraining the benefits from techniques or features orthogonal to the architecture itself . 2 METHODS . 2.1 GRAPH AND SEQUENCE REPRESENTATIONS OF MOLECULES . There are multiple ways of representing molecular structures , such as by molecular fingerprints ( Rogers & Hahn , 2010 ) , by SMILES strings ( Weininger , 1988 ) , or as molecular graphs with atoms as nodes and bonds as edges . We represent the input molecule ( s ) as graphs , and the output molecule ( s ) as SMILES strings , thus modeling both reaction prediction and retrosynthesis as graph-to-sequence transformations . Formally , let Gin denotes the molecular graph input , which can contain multiple subgraphs for different molecules , with a total of N atoms . Following the convention in Somnath et al . ( 2020 ) , we describe Gin = ( V , E ) with atoms V and bonds E . Each atom u ∈ V has a feature vector xu ∈ Ra , and each directed bond ( u , v ) ∈ E from atom u to v has its feature vector as xuv ∈ Rb . The details of the atom and bond features used can be found in Appendix A . We build the input molecular graphs from their SMILES strings with RDKit ( Landrum , 2016 ) . Note that all feature vectors are invariant to the order of atoms and bonds , as well as to how the original SMILES strings are written . We represent the output as a sequence of SMILES tokens Sout = { s1 , s2 , . . . , sn } , where the tokens { si } are obtained from the canonical SMILES using the regex tokenizer in Schwaller et al . ( 2019 ) . 2.2 GRAPH2SMILES . The Graph2SMILES model is a variant of the encoder-decoder model ( Cho et al. , 2014 ) commonly used for machine translation . Figure 1 displays the architecture of Graph2SMILES with the permutation invariant graph encoding process shown within the blue dashed box . We replace the encoder part of the standard Transformer model ( Vaswani et al. , 2017 ) used in Molecular Transformer ( Schwaller et al. , 2019 ) with a novel attention-augmented directed message passing encoder , followed by a global attention encoder with carefully designed graph-aware positional embedding . Each module has its intuitive function : the D-MPNN captures the local chemical context , the global attention encoder allows for global-level information exchange , and the graph-aware positional embedding enables the attention encoder to make use of topological information more explicitly . The permutation invariant encoding process eliminates the need for SMILES augmentation for the input side altogether , simplifying data preprocessing and potentially saving training time . 2.2.1 ATTENTION AUGMENTED DIRECTED MESSAGE PASSING ENCODER . The first module of the graph encoder is a D-MPNN ( Yang et al. , 2019 ) with Gated Recurrent Units ( GRUs ) ( Cho et al. , 2014 ) used for message updates ( Jin et al. , 2018 ; Somnath et al. , 2020 ) . Unlike atom-oriented message updates in edge-aware MPNNs ( Hu et al. , 2020 ; Yan et al. , 2020 ; Mao et al. , 2021 ; Wang et al. , 2021a ) , updates in D-MPNN are oriented towards directed bonds to prevent totters , or messages being passed back-and-forth between neighbors ( Mahé et al. , 2004 ; Yang et al. , 2019 ) . We optionally augment the D-MPNN with attention-based message updates inspired from Graph Attention Network ( Veličković et al. , 2018 ; Brody et al. , 2021 ) . We term this variant as Directed Graph Attention Network ( D-GAT ) and refer to the original D-MPNN variant used in Somnath et al . ( 2020 ) as Directed Graph Convolutional Network ( D-GCN ) , while keeping both D-GAT and D-GCN as possible design choices for the D-MPNN . For D-GAT , at each message passing step t , the message mt+1uv associated with each directed bond ( u , v ) ∈ E is updated using suv = AttnSum ( xu , xuv , { mtwu } w∈N ( u ) \v ) ( 1 ) zuv = σ ( Wz [ xu ; xuv ; suv ] + bz ) ( 2 ) ruv = σ ( Wr [ xu ; xuv ; suv ] + br ) ( 3 ) m̃uv = tanh ( W [ xu ; xuv ] +Uruv + b ) ( 4 ) mt+1uv = ( 1− zuv ) suv + zuv m̃uv ( 5 ) where Wz , Wr , W , U and bz , br , b are the learnable weights and biases respectively . σ is the sigmoid function , `` ; '' indicates concatenation , and is the element-wise product . AttnSum is the attention-based message aggregation defined as ewu = a TLeakyReLU ( Wqk [ xu ; xuv ; m t wu ] + bqk ) ( 6 ) awu = exp ( ewu ) ∑ w∈N ( u ) \v exp ( ewu ) ( 7 ) AttnSum ( xu , xuv , { mtwu } w∈N ( u ) \v ) = ∑ w∈N ( u ) \v awu ( Wvm t wu + bv ) ( 8 ) where Wqk , bqk , a are the learnable parameters for the attention scores , and Wv , bv are the parameters for the value vectors . We have omitted a self-loop in message aggregation , as it did not have any noticeable effect on performance from our experiments . In Eqn ( 3 ) we simplify the reset gate to be shared for all incoming edges , in contrast to Somnath et al . ( 2020 ) where a separate reset gate is defined for each edge . Lastly , after T iterations , we obtain the atom representations hu with similar attention-based aggregation ( with different parameters ) over the bond messages coming into each atom u , followed by a single output layer with weight Wo and GELU activation ( Hendrycks & Gimpel , 2016 ) . mu = AttnSum ′ ( xu , { m ( T ) wu } w∈N ( u ) ) ( 9 ) hu = GELU ( Wo [ xu ; mu ] ) ( 10 ) We use multi-headed attention in our formulation similar to Brody et al . ( 2021 ) . 2.2.2 GLOBAL ATTENTION ENCODER WITH GRAPH-AWARE POSITIONAL EMBEDDING . To capture global interactions , the atom representations coming out of the D-MPNN are fed into a global attention encoder , which is a variant of the Transformer encoder . We incorporate graphaware positional embedding , adapted from the relative positional embedding used in TransformerXL ( Dai et al. , 2019b ) as follows . Firstly , in the standard Transformer either with sinusoidal encoding ( Vaswani et al. , 2017 ; Schwaller et al. , 2019 ) or learnable ( Devlin et al. , 2019 ) absolute positional embedding , the attention score between atoms u and v can be decomposed as eabsu , v = h T uW̃ T q W̃khv + h T uW̃ T q W̃kpv + p T uW̃ T q W̃khv + p T uW̃ T q W̃kpv ( 11 ) where W̃q , W̃k are weights for the keys and queries , and pu , pv are the absolute positional encoding or embedding corresponding to atoms u and v. Similar to Transformer-XL , we reparameterize the four terms . Instead of using sequence-based relative positional embedding ru−v , we use a learnable embedding term ru , v that is dependent on the shortest path length between u and v erelu , v = ( hTuW̃ T q + c T ) W̃khv + ( hTuW̃ T q + d T ) W̃k , Rru , v = ( hTuW̃ T q + c T ) W̃khv + ( hTuW̃ T q + d T ) r̃u , v ( 12 ) The trainable biases are renamed as c and d to avoid confusion and shared across all layers . Intuitively , the two terms in Eqn ( 12 ) model the interactions between inputs ( hu and hv ) , and between input and relative graph position ( hu and r̃u , v ) respectively . Unlike Transformer-XL , we forgo the inductive bias built into the sinusoidal encoding , merging W̃k , Rru , v into a single learnable r̃u , v . This makes the relative positional embedding easily generalizable to atoms not within the same molecule , which is particular useful for reaction outcome prediction and , more generally , tasks with more than two input molecules . We also bucket the distances similar to Raffel et al . ( 2020 ) such that Bu , v = distance ( u , v ) , if distance ( u , v ) < 8 8 , if 8 ≤ distance ( u , v ) < 15 9 , if 15 ≤ distance ( u , v ) and u , v are in the samemolecule 10 , if u , v are not in the samemolecule Bu , v is then used to look up r̃u , v in the trainable positional embedding matrix . The rest of the global attention encoder mostly follows a standard Transformer with multi-headed self-attention ( Vaswani et al. , 2017 ) , layer normalization ( Ba et al. , 2016 ) , and position-wise feed forward layers , except no positional information is added to the value vectors within each Transformer layer .

Summary: This paper proposes a graph-to-sequence architecture called Graph2SMILES for the retrosynthesis and the reaction outcome prediction. Graph2SMILES uses an attention-augmented D-MPNN encoder to capture the local information and a global attention encoder with graph-aware positional embeddings to capture the global information. Experiments show that Graph2SMILES is competitive with Transformer baselines but does not outperform state-of-the-art methods on tasks of the one-step retrosynthesis and the reaction outcome prediction.