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: High-throughput molecular representation for large chemical databases is one of the kernel requirements for many downstream applications, such as drug discovery and virtual screening etc. Based on the consideration that SMILES strings are the most common storage format for molecules, this paper provides a SMILES-based molecular representation learning model, named ST-KD, by using deep graph models to extract molecular feature from raw SMILES data without SMILES-to-graph conversion. Specifically, the proposed model adopts an end-to-end SMILES Transformer framework for molecular representation learning boosted by Knowledge Distillation. Experimental study shows that ST-KD shows competitive results on latest standard molecular datasets PCQM4M-LSC and QM9, with 3-14 inference speed compared with existing graph models.