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: The paper proposes a GNN-based extension of transformers, which have been shown to be effective for reaction prediction etc. before. In particular, the GNN-based embedding of molecules in the reaction embeddings overcomes the artificial bias inherent in the often applied sequence embeddings. The experiments show that there are sometimes increases of performance in reaction and retrosynthesis prediction - and the approach could be applied to similar problems.