Year: 2021

Abstract:
In the intersection of molecular science and deep learning, tasks like virtual screening have driven the need for a high-throughput molecular representation generator on large chemical databases. However, as SMILES strings are the most common storage format for molecules, using deep graph models to extract molecular feature from raw SMILES data requires an SMILES-to-graph conversion, which significantly decelerates the whole process. Directly deriving molecular representations from SMILES is feasible, yet there exists a performance gap between the existing unpretrained SMILES-based models and graph-based models at large-scale benchmark results, while pretrain models are resource-demanding at training. To address this issue, we propose ST-KD, an end-to-end SMILES Transformer for molecular representation learning boosted by Knowledge Distillation. In order to conduct knowledge transfer from graph Transformers to ST-KD, we have redesigned the attention layers and introduced a pre-transformation step to tokenize the SMILES strings and inject structure-based positional embeddings. Without expensive pretraining, ST-KD shows competitive results on latest standard molecular datasets PCQM4M-LSC and QM9, with 3-14× inference speed compared with existing graph models.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHOD

heading: 3.1 PRELIMINARIES
text: 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.

heading: 3.2 SMILES PRE-TRANSFORMATION AND INPUT EMBEDDINGS
text: In this section, we transform irregular SMILES sequences into a unified representation form, in which all atom and bond tokens (maybe omitted in the original SMILES) are parsed out and embedded with structural information.
SMILES Pre-Transformation At SMILES pre-tranformation, we convert a general SMILES line notation SOrgM = (s1, s2, . . . , sl) into a sequence SM = (sV1 , . . . , sVn , sE1 , . . . , sEm) that consists of SMILES symbols, where sV1 , . . . , s V n ∈ {C,N,O,c, . . .} are n atom symbols, sE1 , . . . , sEm ∈ {-,=,:, . . .} are m bond symbols. Atom-bond connectivity is also decoded in this process for computing positional embeddings. We build a SMILES parsing module to perform this step, and we omit its detailed description for simplicity due to the extremely complex SMILES grammar. One can refer to the code for our implementation details.
Input Representations A virtual token [V] is first added to the processed sequence SM , which represents the molecule embedding and its representation will be updated as normal tokens, like the [CLS] token in BERT. Then for a given token, its input representation is constructed by summing up its corresponding token embedding and positional embedding. We use a SMILES symbol vocabulary to generate learnable token embeddings, and positional embeddings are calculated from the decoded molecular structure in the previous step. We first initialize a learnable parameter tensor EPE ∈ Rnmax×d, where nmax is maximum number of atoms of molecules in the dataset, d is the hidden size. Then in the processed SMILES token sequence with special token ([V], sV1 , . . . , s V n , s E 1 , . . . , s E m), for atom token s V i with token embedding h TE,V i ∈ Rd, its input embedding is computed by hVi = h TE,V i +E PE i,: , (5)
for bond token sEj with token embedding h TE,E j ∈ Rd and links atom sVu and sVv , its input embedding is generated using hEj = h TE,E j +E PE u,: +E PE v,:, (6)
as we assign no position embeddings to the special token [V], its input embedding h0 equals its token embedding. The positional embeddings we inject into the input representations are flexible and able to encode complete structural information. Finally, the input embeddings can be formulated as
HSMILES0 = [h0,h V 1 , . . . ,h V n ,h E 1 , . . . ,h E m] > ∈ RlSMILES×d, lSMILES = n+m+ 1. (7)
From here we assume for j = 1, . . . , n, xVj in H Graph 0 and h V j in H SMILES 0 represent the same atom j in M , with such assumption for every 1 ≤ r, s ≤ n, attention interaction from the r-th token to s-th token stands for the same atom-atom correlation in SMILES Transformer and Graph Transformer.

heading: 3.3 SMILES TRANSFORMER WITH KNOWLEDGE DISTILLATION
text: In this section we discuss the knowledge distillation from graph Transformer (GT) to SMILES Transformer (ST), a knowledge transfer process that bridges models with distinct input representations.
Challenges For distillation between Transformer layers in language models, features like intermediate states and self-attention distributions can be used as transferred knowledge (Sun et al., 2019b; Jiao et al., 2019; Wang et al., 2020). However, we can not bring most of existing methods to distillation from GT to ST due to unmatched token sequences, for an input token sequence of ST has lSMILES = n+m+ 1 tokens, while for GT it has lGraph = n+ 1 tokens. For instance, if we simply distill the GT atom token embeddings or transfer the atom-atom attention weights, experiments have shown this unbalanced teaching (no bond token representations or atom-bond token interactions in ST are supervised) only results in subtle improvements.
Global View of the Distillation Method Illustrated by Figure 2, our proposed method for distilling knowledge from a GT layer to a ST layer has two parts. First, both models use a virtual token for updating molecular embeddings, so we perform a straightforward feature distillation on the virtual token embeddings to transfer the teacher’s learned molecular feature while avoiding the unmatching issue when distilling on other tokens. Second, to transfer the GT attention weights without disrupting the attention interactions in ST, we introduce a supportive attention bias module to ST layers. This attention bias module learns the atom-atom attention distributions in the teacher model, and highlights the learned pairwise interactions by adding scalar biases in the calculation of attention weights. We describe the detailed distillation scheme in the following paragraphs.
Molecular Feature Distillation Consider the distillation from layer p of GT to layer q of ST with corresponding output molecule emebddings xp0 ∈ RdV and h q 0 ∈ Rd. Then the feature distillation loss can be defined as L(p,q)feat = MSE(x p 0,h q 0Wfeat), (8) here MSE() specifies mean square loss function, Wfeat ∈ Rd×dV is a learnable linear transformation matrix to transform features of the student into the same space as the teacher’s features, which can also be removed if d = dV .
Biased Multi-head Attention We have introduced the multi-head attention mechanism (MHA) in Section 3.1. Here define biased Transformer layer by adding learnable attention biases to MHA, which serve as a side module to learn and emphasize knowledge from teacher attention weights. Following the notations, for every attention head i, the attention bias matrix is calculated by
Bi = (HW Q,bias i )(HW K,bias i ) >, (9)
where WQ,biasi ,W K,bias i ∈ Rd×dk are learnable attention bias projections not shared among heads. With an associated attention bias mask M ∈ {0, 1}l×l, the computing step of biased multi-head attention block can be defined as
Ai = Softmax( QiK > i√
dk +M Bi), headi = AiVi, for i = 1, . . . , h, (10)
Biased-MHA(H,M) = Concat(head1, . . . , headh)WO. (11)
Attention Weight Distillation via Biases The proposed ST-KD consists of standard Transformer layers and biased Transformer layers. We perform the distillation from layer p of GT to biased Transformer layer q of ST-KD. Suppose layer p outputs the attention weight tensor Ap ∈ [0, 1]h′×lGraph×lGraph , where h′ is the number of attetion heads and lGraph = n + 1. Our goal is to transfer knowledge in Ap to layer q via attention biases. We first stack bias matrices of each attention head in MHA of layer q into the attention bias tensor
Bq = [B1,B2, . . . ,Bh] ∈ Rh×lSMILES×lSMILES , (12) and we introduce a attention mask matrix M ∈ {0, 1}lSMILES×lSMILES to mask out pairwise attention interactions except atom-atom or atom-vn (virtual node) ones. Specifically, the definition is
Mi,j = { 1, if 1 ≤ i ≤ n and 0 ≤ j ≤ n, 0, else.
(13)
Notably, M filters out attention interactions that won’t match Ap, and also serves as attention bias mask in the following attention calculations. Then we multiply Bq by M , and reduce its size to fit Ap (removed elements are all zeros, according to definition of M ) in
B̃ q = (M Bq):,0:lGraph,0:lGraph ∈ Rh×lGraph×lGraph , (14)
where M is broadcasted at the element-wise product. Finally, by applying the softmax function on the last dimension of B̃ q , a weight distribution that matches Ap is predicted from the attention biases. By penalizing the distance between the computed distribution and masked Ap, the attention bias module is able to learn the knowledge inside teacher attention weights. We formulate the attention weight distillation loss as
L(p,q)attn = MSE(M0:lGraph,0:lGraph A p,Softmax(WattnB̃ q )), (15)
here Wattn ∈ Rh ′×h is an optional learnable linear projection used when the teacher and student have different numbers of attention heads. Unexpectedly, in experiments we find the simple mean square loss function is the best choice for model convergence, rather than cross entropy or KL-divergence.
Final Loss Function A task-related loss Ltask is also need during supervised training. Using the above objectives, we can unify the loss function of ST-KD to
L = Ltask + α ∑
(p,q)∈Pfeat
L(p,q)feat + β ∑
(p,q)∈Pattn
L(p,q)attn , (16)
where Pfeat and Pattn are the set of all GT-ST layer pairs for feature and attention weight distillation, α and β stands for the weight for feature distillation loss and attention weight distillation loss.

heading: 4 EXPERIMENTS
text: We first conduct knowledge distillation experiments on the recent PCQM4M-LSC dataset (Hu et al., 2021) for quantum chemistry regression, which is the latest large-scale standard dataset for molecular property prediction and contains more than 3.8M molecules. Then we investigate the transfer learning capacity of ST-KD on several molecular property prediction tasks, with models pretrained by knowledge distillation in the previous step. Finally, we perform tests to support our claim on the efficiency superioiriy of SMILES-based models over graph-based models, and run ablation studies to further inspect each component of ST-KD. A detailed description of datasets can be found in Appendix A.1.

heading: 4.1 KNOWLEDGE DISTILLATION ON PCQM4M-LSC
text: Baselines For graph-based methods, we compare ST-KD with GCN (Kipf & Welling, 2016) and GIN (Xu et al., 2018) with their variants with virtual node attached to improve graph property prediction performance (Gilmer et al., 2017). They are standard baselines and achieve good performances on official leaderboard 1. In terms of SMILES-based models, we test the MLP models over the ECFP fingerpint which is derived using a variant of Morgan algorithm (Morgan, 1965), and a SMILES Transformer ST-BASE that uses a standard SMILES tokenizer and sinusoidal positional encodings. We also apply data augmentation on ST-BASE by adding 5 random equivalent SMILES for every data in the training set.
ST-KD Settings We pick Graphormer (Ying et al., 2021) as the teacher model since Graphormer is built upon the Transformer architecture and achieves current state-of-the-art performance on the PCQM4M dataset. We have reproduced the training process of Graphormer on PCQM4M dataset with official code2 and reported our results3. Graphormer has 12 Transformer layers with hidden dimension 768 and 32 attention heads.
We build the SMILES Transformer with 3 Transformer layers followed by 3 biased Transformer layers. We set the hidden dimension d to 512, number of attention heads h to 16, dimension of FFN layers to 2048. For distillation layer mappings, we perform feature distillation and attention weight distillation simultaneously from the last three layers of Graphormer to the three biased Transformer layers, i.e. the set of all GT-ST distillation layer pairs is Pfeat = Pattn = {(9, 3), (10, 4), (11, 5)}. For loss weights, we set α = 0.5 and β = 2.0. All models are trained on 1 or 2 NVIDIA RTX 3090 GPU for up to 2 days. A complete description of training details is available in Appendix A.2.1.
Results Table 1 summarizes model performances on PCQM4M-LSC dataset. ST-KD surpasses all previous SMILES-based models by a very large margin and performs better than the best GNN baseline GIN-VIRTUAL, which proves that the performance gap between SMILES-based models
1https://github.com/snap-stanford/ogb/tree/master/examples/lsc/pcqm4m# performance
2https://github.com/microsoft/Graphormer 3Training is performed on 2 NVIDIA RTX 3090 GPUs for 2 days with batchsize 1024.
and graph-based models can be overcome with the proposed strategies. Moreover, performance of SMILES Transformer attains a magnificent boost from 0.2062 to 0.1379, augmented with knowledge distillation and SMILES preprocess techniques we propose. Ablation studies on it can be found in Section 4.4. Still, ST-KD has a long way to go before catching up with its Graphormer teacher, but the progress we make so far has shown the extensive potentials of SMILES-based models.

heading: 4.2 MOLECULAR PROPERTY PREDICTION
text: In this section, we evaluate performances of ST-KD on popular molecular property prediction datasets QM9, QM8, QM7, FreeSolv and HIV in MoleculeNet (Wu et al., 2018). We mainly explore the transfer learning capacity of ST-KD by fine-tuning the model trained by knowledge distillation on PCQM4M-LSC in the previous step.
Baselines In this section, we benchmark ST-KD with GCN-VIRTUAL, GIN-VIRTUAL, Graphormer, MLP-ECFP and ST-BASE mentioned above in Section 4.1, along with MoleculeNet (best performances collected in (Wu et al., 2018)), MoLFormer-XL (Ross et al., 2021) and AttentiveFP (Xiong et al., 2019), which is a graph-based deep model that generates molecular fingerprints with local and global attentive layers. For Graphormer, we initialize its weights with the final checkpoint from PCQM4M-LSC training in Section 4.1.
Settings For dataset split, we split the dataset randomly into 8:1:1 as training, validation and test sets if no predefined split is available. For all models, we train them for 3 times with different random seeds, and report the means and standard deviations of performances. For ST-KD, in all datasets, we use the same model hyperparameters as section 4.1 and model weights are initialized with the checkpoint of best validation set performance saved in distillation on PCQM4M-LSC. Appendix A.2.2 presents detailed training procedures. Results are reported in Table 2.
Results ST-KD shows outstanding performance on the QM datasets, being state-of-the-art and outperforms both Graphormer and pretrained MoLFormer-XL. On FreeSolv and HIV, MolFormerXL shows leading performance due to large-scale pretraining, while ST-KD still delivers competitive results with performable graph-based baselines. Compared to ST-BASE, ST-KD gains significant performance improvement on all downstream molecular property prediction tasks, which confirms the effectiveness of our proposed distillation method. Despite being outperformed by the Graphormer teacher on PCQM4M-LSC, on downstream tasks, ST-KD performs better on QM datasets and slightly worse on FreeSolv and HIV, indicating that SMILES-based models can fit better on certain tasks than graph-based models. Compared to MoLFormer-XL (Ross et al., 2021), which requires pretraining on 16 NVIDIA V100 GPUs for 208 hours with at least 81M parameter size 4, ST-KD is light-weight (20.4M, 32-bit float), efficient at training (2 NVIDIA RTX 3090 GPUs
4No MolFormer-XL code is available, we estimate this given that MoLFormer-XL has 12 layers with hidden dimension 768, and we assume parameters are stored with 32-bit float numbers.
for 80 hours, including training Graphormer), and achieves better results on QM8 and QM9. This demonstrates that knowledge distillation is an effective and efficient way to raise the performance of SMILES-based models to state-of-the-art level, with contrast to large-scale pretraining.

heading: 4.3 EFFICIENCY TESTS
text: Settings We follow the OGB-LSC official guide for measuring total SMILES-to-target inference time in 5, and apply all related parameters. To replicate a real-world scenario, we run tests on PCQM4M-LSC and a ZINC subset containing 1M molecules sampled from the public ZINC15 database (Sterling & Irwin, 2015), which represents for real molecular data in drug discovery. All models take raw SMILES strings as input, with their inference time being evaluated on a single NVIDIA GeForce RTX 3090 GPU and an Intel(R) Xeon(R) Platinum 8260C CPU @ 2.30GHz, 256G RAM with warm up. All models are directly taken from Section 4.1. The inference time is measured by millisecond (ms) per molecule. We report the results by the mean of 3 runs in Table 3.
Dataset PCQM4M-LSC ZINC Subset Input Graph SMILES Graph SMILES
Results Table 3 summarizes the inference time of all models. It can be observed that ST-KD reaches a 3-14× higher inference speed from raw SMILES strings than graph-based models, which proves our statements. Compared to ST-BASE, ST-KD runs relatively slower with the additional SMILES preprocessing step and attention bias computations, but has significantly better performance. We can infer from the Table that the SMILES-to-Graph conversion is a serious efficiency bottleneck for graph-based models, as we have mentioned in Section 1.
Comparison with SMILES Pretrain Models We are unable to run inference time tests of most SMILES pretrain models since code is not publicly available, but we can make reliable conjectures about their inference speed with model hyperparameters. Considering MoLFormer-XL, which has 12 attention layers with hidden dimension 768, we can reasonably expect it to run at least 2-3 times slower than ST-KD (6 layers with hidden dimension 512) at inference. This also confirms the efficiency superiority of ST-KD over SMILES pretrain models.

heading: 4.4 ABLATION STUDIES
text: We run a series of ablation studies to further look into the importance of each component in the designed ST-KD architecture on the knowledge distillation experiments with PCQM4M-LSC dataset. Other model hyperparameter settings and training procedures stay the same as in Section 4.1. Table 4 demonstrates that all techniques we have plugged into ST-KD are necessary to raise the perfomance of SMILES Transformer to its best level. With the SMILES pre-transformation, the Transformer gains increased performance, since the input sequences now consist of well-structured tokens and meaningful input emebddings. As for knowledge distillation, both type of the distilled knowledge contribute to the teaching process. The performance of ST-KD can be optimized only with the combination of feature distillation and attention weight distillation.
5In Measuring the Test Inference Time of https://github.com/snap-stanford/ogb/tree/ master/examples/lsc/pcqm4m#performance and code test inference gnn.py in this directory.

heading: 5 CONCLUSION
text: We have proposed ST-KD, a novel architecture for SMILES Transformer empowered by knowledge distillation and shown its competitive results on benchmarks. As high-throughput and highperformance molecular feature extractors, SMILES-based models exhibit significant practical potentials in handling large-scale molecular data. With initial results being encouraging, challenges still exist. Our distillation approach cannot be applied without a well-learned graph model as teacher. In distillation experiments, Graphormer outperforms its student well, so results of ST-KD can be further improved with a more effective knowledge transfer strategy. We leave them as future works.

heading: A APPENDIX
text: A.1 DATASETS
Statistics of datasets we use in this paper are shown in Table 5. PCQM4M-LSC is a large-scale graph property prediction dataset in the recent OGB-LSC Large-Scale Challenge (Hu et al., 2021), originally curated under the PubChemQC project (Nakata & Shimazaki, 2017). The task of PCQM4MLSC is to predict HOMO-LUMO gap of molecules calculated by DFT (Density Functional Theory) with 2D structure of mulecules. Test labels of PCQM4M-LSC is now unavailable, and we use validation set to evaluate model performance. QM9, QM8 and QM7 are datasets containing organic molecules with up to 9, 8 and 7 heavy atoms (any atom that is not hydrogen) and their quantitative quantum-chemical properties. For QM7, we use a sanitized version by DeepChem (Ramsundar et al., 2019) available in https://deepchemdata.s3-us-west-1.amazonaws. com/datasets/qm7.csv. FreeSolv contains molecules with their corresponding quantitative hydration free energy. HIV contains molecules labeled by their HIV active labels, with a predefined scaffold data split.
A.2 EXPERIMENT DETAILS
A.2.1 KNOWLEDGE DISTILLATION ON PCQM4M-LSC
All parameters involved are listed in Table 6. During experiments we find by adding an initial stage when there is no task-related supervision, the model can have a faster convergence. Parameters of this initial stage is also shown.
A.2.2 MOLECULAR PROPERTY PREDICTION EXPERIMENTS
For ST-KD, we load the model checkpoint with best performance on validation set saved during distillation as initial weights for molecular property prediction experiments in this section. In experiments we discover that results can benefit from changing the model dropout value at this finetune training process. All related parameters are listed in Table 7 and 8.
A.2.3 ADDITIONAL INFORMATION ON PERFORMANCES OF ST-KD ON QM9 DATASET
QM9 is a dataset with 12 regression targets. We follow the conventions to standardize the targets and fit them simultaneously. Here we report ST-KD’s test set performance on 12 tasks seperately.

Review summaries:
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)

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.

The paper targets speeding up the inference process (specifically, SMILES-to-graph preprocessing) in virtual screening by adopting SMILES-based transformer for molecule representation learning. To ensure the performance of SMILES-based transformer, authors propose to use graph-based transformer as teacher network to conduct knowledge distillation (ST-KD). Numerical results claim for the competitive performance and faster inference time.