Stepping Back to SMILES Transformers for Fast Molecular Representation Inference

1 INTRODUCTION . Recent years have witnessed the remarkable progress in the combination of deep learning and molecular science , with regard to tasks including molecular property prediction ( Xiong et al. , 2019 ; Li et al. , 2021 ) , conformation generation ( Xu et al. , 2021 ) , and virtual screening for drug discovery ( Gentile et al. , 2020 ; Stevenson et al. , 2021 ) . Meanwhile , encoding molecules into numerical vectors , or molecular representation , serves as the very first yet important step . As 2D structures of molecules , i.e . atom-bond connectivity , can be naturally represented as graphs , neural networks for graph data have become one of the predominant choices . Exploiting network architectures from graph message-passing networks ( Gilmer et al. , 2017 ) to Transformers ( Rong et al. , 2020 ; Ying et al. , 2021 ) , these graph-based methods have reached the state-of-the-art performances on standard benchmark datasets . However , in real applications such as virtual screening , operations on very large datasets are required , and high-throughput molecular encoders are indispensable . In most existing chemical libraries such as ZINC ( Irwin & Shoichet , 2005 ) , molecules are stored with SMILES ( Simplified Molecular-Input Line-Entry System , Weininger ( 1988 ) ) , a line notation format with strict generating rules that encodes complete 2D structural information . In order to apply graph-based models to SMILES inputs , an additional SMILES-to-Graph conversion is needed . Operated upon an unparallelizable deterministic algorithm ( Weininger , 1990 ) with high time complexity proved by experiments , this conversion has become a serious efficiency bottleneck that restricts graph-based model from achieving high feature extraction throughput on SMILES datasets . Researchers have also been exploring ways to generate molecular representations ( or fingerprints ) from SMILES strings in an end-to-end fashion , like the topological ECFP ( Rogers & Hahn , 2010 ) and Transformer-based SMILES language models ( Honda et al. , 2019 ; Wang et al. , 2019 ; Chithrananda et al. , 2020 ) . However , without pretraining , their overall benchmark results on standard datasets are outperformed by basic graph-based models like GCN ( Kipf & Welling , 2016 ) . While for competitive pretrained models like MolFormer ( Ross et al. , 2021 ) , the inference speed is largely restricted by huge model size , and pretraining requires massive SMILES data and computational resources . Based on the discussion above , we propose ST-KD , an end-to-end SMILES Transformer architecture incorporated with Knowledge Distillation ( KD ) from graph Transformers . Considering SMILES implicitly encodes the complete 2D molecular structure , we believe the key reason for the low performances of current SMILES-based models is that , without structure-related supervision , the hidden atom-bond connectivity within SMILES is not well captured in training . During KD from graph Transformer ( teacher ) to SMILES Transformer ( student ) , the student is forced to mimic the teacher while learning structural knowledge of molecules , for the distilled hidden representations and attention weights are concentrated knowledge the teacher have learned from molecular graphs . This knowledge transfer process is capable of bringing a magnificent performance boost to SMILES Transformer with little loss of efficiency , as shown in Figure 1 . As SMILES is generated from strict grammar rules and contains different types of tokens ( like tokens for atoms , bonds , ring-breaks , etc . ) , applying standard NLP data preprocess methods will lead to faulty tokenization and ambiguous input embeddings . Hence , we build an efficient module to transform SMILES strings into unified token sequences embedded with clear chemical definition , then attach structure-based positional encodings to them . Besides , graph Transformer and SMILES Transformer have different procedures to update token embeddings and compute pairwise attention interactions , so most standard Transformer distillation methods are infeasible . To solve this problem , in addition to feature distillation , we add parameterized attention biases to ST-KD ’ s self-attention layers , and partially supervise on the biases to transfer knowledge in attention blocks . Ablation studies have proved our redesigned distillation techniques are essential to optimize the teaching results . Directly taking SMILES as input , ST-KD has averted the SMILES-to-graph bottleneck and can run 3-14× faster than graph-based models at inference . Compared to pretrained SMILES models , STKD is much more light-weight and efficient at training , without compromising on performance on most tasks . Boosted by the solid network structure and knowledge distillation from cutting-edge graph Transformer ( Ying et al. , 2021 ) , ST-KD is qualified to deliver competitive , even state-of-theart results on public molecular benchmarks like MoleculeNet ( Wu et al. , 2018 ) and OGB-LSC ( Hu et al. , 2020 ) . Code for reproducing our results is provided in the supplementary material . 2 RELATED WORK . SMILES and SMILES-based Fingerprint Generators SMILES ( simplified molecular-input line-entry system ) , first proposed in ( Weininger , 1988 ; Weininger et al. , 1989 ; Weininger , 1990 ) , is a universal line notation format to represent 2D molecular structure . For example , the SMILES string c1ccccc1C can specify the structure of toluene . SMILES has complex grammars and all symbols have specific chemical definitions , e.g . C and N represent carbon atom and nitrogen atom , - and = represent single bond and double bond , 1 and ( represent ( part of ) ring-break and branching bond . An important fact is that the original SMILES system does not create a bijective mapping between the set of all possible SMILES sequences and molecules . A legal SMILES sequence can be translated into one certain molecule , but it is possible for a molecule to have multiple SMILES representations , e.g . both CCO and OCC specify ethanol . Before the age of deep learning , researchers have developed multiple SMILES-based molecular fingerprint generators , including the hash-based methods ( Glen et al. , 2006 ; Hu et al. , 2009 ; Rogers & Hahn , 2010 ) and task-driven methods ( O ’ Boyle et al. , 2011 ) . Extended-Connectivity FingerPrint ( ECFP ) ( Rogers & Hahn , 2010 ) is a famous method for efficiently generating topological structure fingerprints for SMILES . On the other side , the recent spurt of deep models has inspired researchers to build SMILES fingerprint generators by training from data without expert knowledge ( Xu et al. , 2017 ; Honda et al. , 2019 ; Zhang et al. , 2018 ; Wang et al. , 2019 ; Chithrananda et al. , 2020 ) . Treating SMILES as a formal language , most of these methods take NLP models like Transformer ( Vaswani et al. , 2017 ) and BERT ( Devlin et al. , 2018 ) as backbone , then train the model with pretrain-finetune paradigm for downstream tasks . These methods have shown promising results in specific tasks like molecular property prediction ( Ross et al. , 2021 ) , chemical reaction prediction ( Schwaller et al. , 2019 ) and drug discovery ( Honda et al. , 2019 ) . As for molecular property prediction , current SMILES-based models can only reach competitive performances against graph-based models with extensive pretraining , like ( Ross et al. , 2021 ; Irwin et al. , 2021 ) . Compared with largescale pretraining on massive data , the distillation strategy we employ requires no big data and is much more effective and efficient , which will be discussed in Section 4.2 . Graph-based Molecular Representation Learning Graph-based deep models ( Duvenaud et al. , 2015 ; Gilmer et al. , 2017 ; Xiong et al. , 2019 ; Rong et al. , 2020 ; Ying et al. , 2021 ) have taken the leading role in molecular representation learning . Some studies ( Cho & Choi , 2019 ; Klicpera et al. , 2020 ; Song et al. , 2020 ; Li et al. , 2021 ) also begin to enhance molecular representations with 3D conformation . Graphormer ( Ying et al. , 2021 ) is a newly proposed graph Transformer model for graph representation learning , which attains the state-of-the-art performance on many public datasets and wins the OGB Large-Scale Challenge Hu et al . ( 2021 ) in quantum chemistry regression dataset . In experiments , we pick Graphormer as the teacher model , considering its outstanding performance and Transformer-based architecture . Transformer Distillation and Cross-modal Knowledge Transfer Transformer ( Vaswani et al. , 2017 ) has become one of the most popular building blocks in deep models with its powerful self-attention mechanism , which can capture long-term dependencies within sequences of tokens . Transformer-based pretrain language models like BERT Devlin et al . ( 2018 ) exhibit excellent performances but are also computationally expensive . There have been many works that attempt to compress Transformer-based models with knowledge distillation ( Jiao et al. , 2019 ; Sun et al. , 2019a ; 2020 ; Wang et al. , 2020 ) . The distilled knowledge may be soft target probabilities , embedding outputs , hidden representations or attention weight distributions . Existing works have also explored cross-modal knowledge transfer , including contrastive cross-modal distillation ( Tian et al. , 2019 ) and graph-to-text distillation ( Dong et al. , 2020 ) . ST-KD is built upon the Transformer architecture and borrows some insights in distillation methods above . However , since the difference between SMILES and other types of structured knowledge is essential , the distillation techniques we apply are completely redesigned to bridge the knowledge transfer between graph-based models and SMILES-based models . 3 METHOD . 3.1 PRELIMINARIES . Problem Formulation We formulate molecular representation learning as a supervised learning task , which takes notations of molecules as inputs , properties of molecules as supervisions . A moleculeM can be denoted as an attributed graph GM = ( V , E , n , m , XV , XE ) , or a SMILES line notation SOrgM = ( s1 , s2 , . . . , sl ) , where V is the set of n atoms , E ⊂ V × V is the set of m bonds , XV = [ xV1 , . . . , x V n ] > ∈ Rn×dV is the atom feature matrix , XE = [ xE1 , . . . , xEm ] > ∈ Rm×dE is the bond feature matrix , s1 , s2 , . . . , sl ∈ { C , N , - , =,1 , ( , . . . } are SMILES symbols . Overview Figure 2 is an illustration of the proposed ST-KD architecture . SMILES strings are first transformed to generate initial input embeddings , then sent to the stacked layers . ST-KD is built by standard Transformer layers and biased Transformers layers we introduce in Section 3.3 , which are designed to enable the knowledge transfer from graph Transformer attention block to SMILES Transformer . Transformer and Multi-head Attention A Transformer encoder layer is composed of a multihead attention ( MHA ) block and feed-forward networks ( FFN ) with residue connections . Let H = [ h1 , . . . , hl ] > ∈ Rl×d be the input embeddings , the calculation process of multi-head attention can be formulated as : Qi = HW Q i , Ki = HW K i , Vi = HW V i , ( 1 ) Ai = Softmax ( QiK > i√ dk ) , headi = AiVi , for i = 1 , . . . , h , ( 2 ) MHA ( H ) = Concat ( head1 , . . . , headh ) WO , ( 3 ) where l is length of input sequence , h is number of attention heads , WQi ∈ Rd×dk , WKi ∈ Rd×dk , W Vi ∈ Rd×dv and WQ ∈ Rhdv×d are projection parameter matrices , d , dk , dv are the dimension of hidden layers , keys and values . Graph Transformer Graph Transformers ( Ying et al. , 2021 ; Dwivedi & Bresson , 2020 ) generally take graphs as node sequences , update the node representations using attention mechanism and linear layers , and leverage the structural information into the Transformer architecture by positional encodings and attention biases . Given molecule M , it should be fed into the graph Transformer model as a featurized graph ( V , E , n , m , XV , XE ) with initial token representations HGraph0 = [ x0 , x V 1 , . . . , x V n ] > ∈ RlGraph×dV , lGraph = n+ 1 , ( 4 ) where x0 represents for a virtual token for molecule embeddings . Then H Graph 0 is updated by stacked Graph Transformer layers to produce the final result .

Summary: Authors propose knowledge distillation from graph Transformer (GT) to SMILES Transformer toward learning a molecular representation with similar performance to Graph-based models while having low inference time (by avoiding smiles-to-graph conversion bottleneck).  Main contributions are:  (1)  An end- to-end SMILES Transformer for molecular representation learning is proposed that is boosted by graph transformer to smiles transformer knowledge distillation. (2) Biased transformer layers are designed to achieve knowledge transfer; learnable SMILES token embedding are used to encode structural information.  (3) ST-KD shows competitive results on PCQM4M-LSC and QM9, with fast inference compared to existing graph models.  (4)