Stepping Back to SMILES Transformers for Fast Molecular Representation Inference
Authors: Wenhao Zhu
Ziyao Li
Lingsheng Cai
Guojie Song
Venue: ICLR 2022 Submitted (2022)
Keywords: molecular representation learning
knowledge distillation

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 large performance gap between the existing SMILES-based models and graph-based models at benchmark results. 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. ST-KD shows competitive results on latest standard molecular datasets PCQM4M-LSC and QM9, with $3\text{-}14\times$ inference speed compared with existing graph models.

Under review as a conference paper at ICLR 2022
STEPPING BACK TO SMILES TRANSFORMERS FOR FAST MOLECULAR REPRESENTATION INFERENCE
Anonymous authors Paper under double-blind review
ABSTRACT
In the intersection of molecular science and deep learning, tasks like virtual screening have driven the need for 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 ex- tract molecular feature from raw SMILES data requires an SMILEEStto-grah conversion, which significantly decelerates the whole process. Directly deriv- ing molecular representations from SMILES is feasible, yet there exists a per- formance gap between the existing unpretrained AMILES-based models and graph-based models at large-scale benchmark results, while pretrain models are demaadigd at training. To address this issue, we propose ST-KD, end-to-end SMILES Transformer for molecular representation learning boosted by Knowledge Distillation. In order conduct knowledge transfer from graph Transformers to ST-KD, we have redesigned the attention layers and introduced a pre- trarsffofoomnn step to tokenize the SMILES strings and inject structure- -asee positional embeddings. Without expensive pretraining, ST-KD shows competi- tive results on latest standard molecular datasets PCQM4M-LSC and QM9, with 3-14x inference speed compared with existing graph models.
INTRODUCTION
Recent years have witnessed the remarkable progress in the combination of deep learning and molec- ular science, with regard to tasks including molecular property prediction (Xiong et al. 2019 Li iet al. (021), conformation generation (uu et al. 2021). and virtual screening for drug discovery (Gentile et al. 2020 Stevenson al. 2021). Meanwhile, encoding molecules into numerical vec- tors, 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. 202I), these graph-based methods have reached the state-of-ther toffthh-art performances on standard benchmark datasets.
However, in real applications such virtual screening, operations very large datasets are quired, and high-throughput molecular encoders are indispensable. In most existing chemical li- braries 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 generat- ing rules that encodes complete 2D structural information. In order apply graph-based models to SMILES inputs, an additional SMILES-tO-Gree conversion is needed. Operated upon an unpar- allelizable deterministic algorithm (Weininger] 1990) with high time complexity proved by experi- ments, this conversion has become 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). How- ever, 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 mod- els 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.

---

Under review as a conference paper at ICLR 2022
Greshormer (2.790mH (.1274)
(0.152ma, (3777 1379)
(8 5.577mm0..7390
Validation (Inversely ordered)
BCl-umua (0.627ms,
0.20
(A 0.2006
MUP-ECFP (0.950ma 0.2113)
0.1
0.5 Inference time (ms/molecule)
2.0
Figure 1: Inference time accuracy (validation MAE, inversely ordered) PCQM4M-LSC. ST-KD speeds up 3-14x compared with existing graph-based models with competitive accuracy. Explanations of the dataset and baseline models are in Section 4
Based on the discussion above, we propose ST-KD, an end-to-end SMILES Transformer archi- tecture 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 supervi- sion, the hidden atom-bong 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 magnificent performance boost to SMILES Transformer with little loss of efficiency, as shown Figure I
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 snacture-based positional encodings them. Besides, graph Transformer and SMILES Transformer have different procedures to update token embeddings and compute pairwise attention interactions, 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 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-gras -to--rap- bottleneck and can run 3-14x faster than graph-base models inference. Compared to pretrained SMILES models, ST- KD is much more light and efficient at training, without compromising 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-the- art results on public molecular benchmarks like MoleculeNet (Wu et al.) 2018) and OGB-LSC (Hu
al. 2020). Code for reproducing our results is provided in the supplementary material.
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). a universal line notation format to represent 2D molecular structure. For example, the SMILES string eleccccCC C 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
2

---

Under review as a conference paper at ICLR 2022
bond. An important fact is that the original SMILES system does not create 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 molecule to have multiple SMILES representations, e.g. both CCO and OCC specify ethanol.
Before the age of deep learning, researchers have developed multiple AMILES-based molecular fin- gerprint 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. [011]. Extended-Connecttvvve Finger-
Print (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 formal language, most of these methods take NLP models like Trans- former (Vaswani et al. 2017) and BERT (Devlin et al. 2018) as backbone, then train the model with pretrain-finctune paradigm for downstream tasks. These methods have shown promising results in specific tasks like molecular property prediction (Ross et al. 202T), chemical reaction prediction (Schwaller et al. 2019) and drug discovery (Honda et al. 2019)) As for molecular property predic- tion, current MMsLES-based models can only reach competitive performances against graph-bass based models with extensive pretraining, like (Ross al., 2021 Irwin et al. 2021). Compared with large- scale 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 al., 2015 Gilmer et al.] 2017 Xiong et al.] 2019 Rong et al.] [020] (Ying et al.] 2) have taken the leading role in molecular representation learning. Some studies (Cho & Choi, 2019; Klicpera eet al. 2020 Song et al. 2020 ei al. (021) also begin to enhance molecular representations with 3D conformation. Graphormer (Ying et al. 202T) is newly proposed graph Transformer model for graph representation learning, which attains the state-of-the-uu performance on many public datasets and wins the OGB Large-Scale Challenge Hu et al. (202T) in quantum chemistry regression dataset. In experiments, we pick Graphormer as the teacher model, considering its outstanding performance and Transformer-base architecture.
Transformer Distillation and Cross-modal Knowledge Transfer Transformer (Vaswani 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-basee pretrain language models like BERT Devlin et al. (2018) exhibit excellent per- formances but are also computationally expensive. There have been many works that attempt to compress Transformer-based models with knowledge distillation (Tiao et al. (019 SSnn et al. (201aa] 2020 Wang et al. 2020). The distilled knowledge may be soft target probabilities, embedding out- puts, 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 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 ap- ply are completely redesigned to bridge the knowledge transfer between graph-based models and SMILES-baid based models.
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 molecule M can be denoted as an attributed graph Gas = (V, E, n, m, X .x F), or a SMILES line
notation 50rg A = ($1, .2,. st), where V is the set of n atoms, E C V x V is the set of m bonds, X [x) INTTT E =a is the atom feature matrix, X F = [xf [......, x E R gdxdg is
the bond feature matrix, 81, 81 E {c, N, ,,,,, (,...) are SMILES symbols.
3

---

Under review as a conference paper at ICLR 2022
Eroeedrng
Transfer
Graph Transformer Layer
Blased Transtormer Layer
Querien
Keys
Wiless
Graph Transformer
Transformer Layer
Querics
Graph Transformer
Attertion Weights
Transformer Layer
Attention Transter
Graph Transtormer Layer
Kayp
Transformer Layer
Aftertion Biss
C=CCO
SMAL Pre Transformtion
Embeddinga
Converalon
Molecule, SMILES
Token Embeddings
Chemical Dulabase
SMILES Transformer, studintt
Graph Transformer, the teacher
Figure 2: Overview of the proposed ST-KD architecture.
Overview Figure is an illustration of the proposed ST-KD architecture. SMILES strings first transformed to generate initial input embeddings, then sent to the stacked layers. ST-KD built by standard Transformer layers and biased Transformers layers 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 multi- head attention (MHA) block and feed-forward networks (FFN) with residue connections. Let H [hi.... h Rlxd be the input embeddings, the calculation process of multi-head attention can be formulated as:
Q: HW? K. = HHW V, = HWY,
A, Softmax Q,K, ), head, A V. for = h,
(2)
ddz
MHA(H) Concat(headd, head, wo
where is length of input sequence, h number of attention heads, w E Radd W Rdxd W.Y. Rdxd, and w Rd are projection parameter matrices, d, dr, dg are the dimension of hidden layers, keys and values.
Graph Transformer Graph Transformers (Ying et al. 2021 Dwivedi Bresson 2020) generally
take graphs 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 featurized graph (V, E, n, m, XV X F) with initial token representations
Graph = [mo, a
x KTT E laga Xdy AGrraph = 72 + 1,
(4)
where represents for a virtual token for molecule embeddings. Then H is updated by stacked Graph Transformer layers to produce the final result.
3.2 SMILES TRE-TRANSFORMATTION AND INPUT EMBEDDINGS
In this section, transform irregular SMILES sequences into unified representation form, in which all atom and bond tokens (maybe omitted in the original SMILES) are parsed out and embed- ded with structural information.
4

---

Under review as a conference paper at ICLR 2022
SMILES TrenTraasformttion
At SMILES pre-tranformation, convert general SMILES line notation SM (81,82)...,58) into sequence SM ........,88 of SMILES symbols, where .p ,,..,,,,,
that consists O,, c,...) are n atom symbols, sf .f......s. (-,-,;,.... are m bond symbols. Atom-bond connectivity is also decoded in this process for computing positional embeddings. We build SMILES parsing module to perform this step, and omit its detailed description for simplicity due to the extremely complex SMILES grammar. One can refer to the code for our implementation details.
virtual token [V] is first added the processed sequence SM. which represents the molecule embedding and its representation will be updated as normal tokens, like the [CTS] token in BERT. Then for a given token, its input representation 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 cal- culated from the decoded molecular structure in the previous step. We first initialize a learnable parameter tensor E RTm xd, where maxx 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 ; ,,)), for atom token sy with token emhedding h; V.V.V. RRd its input
Input Representations
[[V],)] ,...,... : emhedding is computed by
8% &f. ........................ .....
hy h T. E.
for bond token sf with token embedding is generated using
h TE E Rd and links atom swx and sy its input embedding
h = h T PE
E.E
(6)
as we assign no position embeddings the special token [V], its input embedding ho 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
HOMILES = [ho. bY
h h
h E
E PRISES xd ISMILES + m + 1.
(7)
From here we assume for j = 1, n. x in H Graph and R. in H.MILLSS represent the same atom in M, with such assumption for every r,s 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.
3.3 SMILES TRANSFORMER WITH KNOWLEDGE DISTILLATION
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 represen- tations.
Challenges For distillation between Transformer layers in language models, features like interme- diate states and self self-attention distributions can used as transferred knowledge (Sun et al., 2019b] Jiao et al.) 2019 Wang et al., 2020). However, we can bring most of existing methods dis- tillation from GT to ST due to unmatched token sequences, for an input token sequence of ST has ISMILES = n + m + tokens, while for GT it has GGaph = 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 distill- ing knowledge from a GT layer 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 's learned molecular feature while avoiding the unmatching issue when distilling other tokens. Second, to transfer the GT attention weights without disrupt- ing 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.
5

---

Under review as a conference paper at ICLR 2022
Molecular Feature Distillation Consider the distillation from layer p of GT to layer g of ST with corresponding output molecule emebddings zy E Rdy and hg E R Then the feature distillation loss can be defined as
((P.)) Teal
MSE(a Wi
(8)
here MSB() specifies mean square loss function, feat E Rdxdv is learnable linear transformation matrix to transform features of the student into the same space as the teacher's features, which can also be removed d = dy.
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 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
B. ( Q bins
; "
(HW; hiirs)
where W 2 bias b.i.s E dd are learnable attention bias projections not shared among heads. With an associated attention bias mask M E {0,1)(x, the computing step of biased multi-head attention block can be defined as
Q:K. dr
A; = Softmax
M @ B.), head, = A.V, for i = 1,.......
(10)
Biased-MHA((E H, M) = Concat(headı.... head,) w
(11)
Attention Weight Distillation via Biases The proposed ST-KD consists of standard Trans- former layers and biased Transformer layers. We perform the distillation from layer p of GT to biased Transformer layer 9 of ST-KD. Suppose layer p outputs the attention weight tensor A" E [0, 1] Xlamp Xlamp, where M is the number of attetion heads and 4Gmph = n + 1. Our goal to transfer knowledge in A" layer q via attention biases. We first stack bias matrices of each attention head in MHA of layer into the attention bias tensor
B B2..... BA] E g² xlsxxEs X/SMILES
(12)
and we introduce a attention mask matrix M E {0, 13 iSMILE xISMMLES to mask out pairwise attention interactions except atom-atom or atom-vn (virtual node) ones. Specifically, the definition is
1, i n and 0 inin
Mij.
(13)
0, else.
Notably, M filters out attention interactions that won't match A and also serves as attention bias mask in the following attention calculations. Then we multiply B by M, and reduce its size to fit A" (removed elements are all zeros, according to definition of M) in
(14) where M is broadcasted at the element-wise product. Finally, by applying the softmax function on the last dimension of B" weight distribution that matches A² is predicted from the attention biases. By penalizing the distance between the computed distribution and masked A", the attention bias module is able to learn the knowledge inside teacher attention weights. We formulate the attention weight distillation loss as
B (M c 100a:aam,, E ER A Xlomp,
c(2-a)
AISE(Mo:lces..o Oslengh A", Softmax W sttn B)))
(15)
R& xh 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 KL-divergence.
here W.In
Final Loss Function task-related loss Liss also need during supervised training. Using the above objectives, we can unify the loss function of ST-KD
L = Lisk +
(p.a) Feat
+ 3
(p.9) aatt
(16)
(p-7)E Picn
(p-q)
where Pte. and Pau are the set of all GT-ST layer pairs for feature and attention weight distillation, and 33 stands for the weight for feature distillation loss and attention weight distillation loss.

---

Under review as a conference paper at ICLR 2022
4 EXPERIMENTS
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 molec- ular 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.
4.1 KNOWLEDGE DISTILLATION ON PCQM4M-LSC
Baselines For graph-based methods, we compare ST-KK 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 al. 2017). They standard baselines and achieve good per- formances on official leaderboard | In terms of SMILES-based models, we test the MLP models over the ECFP fingerpint which is derived using variant of Morgan algorithm (Morgan 1965), and a SMILES Transformer ST-BASE that uses standard SMILES tokenizer and sinusoidal positional encodings. We also apply data augmentation on ST-BASE by adding 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 built upon the Transformer architecture and achieves current state-of-the-ee performance the PCQM4M dataset. We have reproduced the training process of Graphormer on PCQM4M dataset with official code and reported our results" 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 to 512, number of attention heads h to 16, dimension of FFN layers 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 Pien Paun {9, 3), (10, 4), (11, 5)). For loss weights, we set a 0.5 and 3 = 2.0. All models are trained on I or 2 NVIDIA RTX 3090 GPU for up days. A complete description of training details is available Appendix A.2.
Model
#Params
Train
Validation
GCN
2.0M
0.1522
0.1702(0.16844)
GIN
3.8M
0.1379
0.1533(0.15104
Graph-based
GCN-VIRTUAL
4.9M
0.1317
0.1542(0.1536%
GIN-VIRTUAL
6.7M
0.1206
2.1421(0.1396'
Graphormer
47.1M
0.0590
(0.0582*)
0.1274 (0. 1234*)
MLP-ECFP
15.8M
0.1032
0.2113
HiILLES-based
ST-BASE
19.1M
0.1483
0.2031
ST-KD
20.4M
0.0877
0.1379
Table 1: Results on PCQM4M-LSC measured by MAE, indicates results cited from OGB-LSC paper (Hu al. 202I) ) and Graphormer paper (Ying al. (02I).
Results Table summarizes model performances on PCQM4M-LSC dataset. ST-KD surpasses all previous EMiLES-based models by very large margin and performs better than the best GNN baseline GIN-VIRTUAL, which proves that the performance gap between SMILES-based models
performance
https:: http:::///ittttt wttp:/cccommccccocmmo http://hh....mmm..oooooo///g///p
TTaining is performed on NVIDIA RTX 3090 GPUs for 2 days with batchsize 1024.
7

---

Under review as a conference paper at ICLR 2022
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 knowl- edge distillation and SMILES preprocess techniques
propose. Ablation studies on 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 make far has shown the extensive potentials of SMILES-based models.
4.2 MOLECULAR PROPERTY PREDICTION
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. (018)), Mol.Forme--XI (Ross et al.) (021) and Atten- tiveFP (Xiong al. 2019), which a graph-based deep model that generates molecular fingerprints with local and global attentive layers. For Graphormer, we initialize its weights with the final check- point 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 no predefined split is available. For all models, we train them for 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 PCOM4M-LSC. Appendix A.2.2 presents detailed training procedures, Results are reported in Table 2
Dataset
QM9
QM8
Q.7
FreeSoly RMSEL
HIV
Metric
Multi-MAE|
Avg-MAEJ
MAEJ
ROC-AUCT
GCN-VIRTUAL.
0.0482
0.0123+000004
74.081 +2
1.013++0.09
75.99+
GIN-VIRTUAL
1.2248_0.0554
0.0111
69.682+1..718
0.852:00.055
77.07
MoleculeNet
2.350*
0.0150+:.0000
94.7+22.7
1.150*
79.20
AttentiveFP
1.292
0.0130=t0.0006
66.2+:2713
0.962:00197
19.10
Graphormer
0.9168=000301
0.0091
44.342=1.419
0.860+0.082
80.51 *
MLP-ECFP
1.2021=0.0573
0.0129+000004
65.845+1.321
0.995+0.007
7..15
MoLFormer-Xi XL
0.6804
0.0102#
0.2308
82.2*
ST-BASE
1.1562-0.0510 0.0169+0000066 0.6215 1.011-
52.348=3.211
1.071
72.10 1.4
ST-KD
0.0080+
43.373+1.770
0.805
80.32++0.7
Table 2: Results on molecular property prediction datasets. indicates results cited from (Wu et al. 2018 Ross al. 2021 Ying al. 2021), for higher is better. 4 contrarily. Additional information on ST-KD's performance on QM9 can be found in Appendix A.2.3
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, MolFormer- XL shows leading performance due to large-scale pretraining, while ST-KD still delivers compet- itive results with performable graph-based baselines. Compared to ST-BASE, ST-KD gains sig- nificant 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, datasets and slightly worse on FreeSolv and HIV, indicating that SMILES-basee models can fit bet- ter certain tasks than graph-based models. Compared which requires pretraining on 16 NVIDIA V100 GPUs for 208 hours with at least 8IM parameter size ST-KD is light-weight (20.4M, 32-bit float), efficient training NVIDIA RTX 3090 GPUs
downstream tasks, ST-KD performs better QM
MoLForrrr-XL (Ross et al. 2021),
No MolFormer-Xt XI. code is available, estimate this given that Mol Former-XI. has 12 layers with hidden dimension 768, and assume parameters stored with 32-hit float numbers.
8

---

Under review as a conference paper at ICLR 2022
for 80 hours, including training Graphormer), and achieves better results QM8 and QM9. This demonstrates that knowledge distillation is an effective and efficient way to raise the performance of SMILES-bai based models to state-of-the-ane level, with contrast to large-scale pretraining.
4.3 EFFICIENCY TESTS
Settings We follow the OGB-LSC official guide for measuring total AMILESStt-targee inference time in [ and apply all related parameters. To replicate a real-world scenario, we run tests on PCQM4M-LSC and ZINC subset containing 1M molecules sampled from the public ZINCI5 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
single NVIDIA GeForce RTX 3090 GPU and an Intel(R) (co 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
Dataset Input
PCQM4M-LSC
ZINC Subset
Graph
SMILES
Graph
SMILES
GCN-VIRTUAL
0.121
0.627
0.204
0,906
Graph-based based
GIN-VIRTUAL
0.118
0.627
0.209
0,913
Graphormer
0.399
2.795
0.573
4.035
MLP-ECFP
0.352
0,467
ST-BASE
0.149
0,230
SMILES-bare
ST-KD
0.192
0.271
Table 3: Results of inference time tests, measured by ms per molecule.
Results Table 3 summarizes the inference time of all models. It can be observed that ST-KD reaches 3-14x higher inference speed from raw SMILES strings than graph-based 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 perfor- mance. We can infer from the Table that the SMILES-to-Gire bottleneck for graph-basee models, as we have mentioned in Section [
conversion is a serious efficiency
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-Xi -LL, 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.
4.4 ABLATION STUDIES
We run a series ablation studies to further look into the importance of each component in the de- signed ST-KD architecture the knowledge distillation experiments with PCQM4M-LSC dataset. Other model hyperparameter settings and training procedures stay the same as in Section Table demonstrates that all techniques we have plugged into ST-KD are necessary raise the perfomance of SMILES Transformer to its best level. With the SMILES pre- 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 ST-KD can be optimized only with the combination of feature distillation and attention weight distillation.
In Measuring the Test Inference Time http:///itto/////ttt..o//////gs/ ma master/examplest 1 pwtpp//mmmmmmmmmmmmm wor and code test wwwt_eeefeeenceeeeeeeeeooo py in this direc- tory.

---

Under review as a conference paper at ICLR 2022
Distillation
Validation Perf.
SMILES Pre-transformation
Feature Transfer
Attention Transfer
MAE
0,2062
0.1808
0,1672
0,1589
0.1379
Table 4: Ablation studies of ST-KD on PCQM4M-LSC dataset with different components removed.
5 CONCLUSION
We have proposed ST-KD, a novel architecture for SMILES Transformer empowered by knowl- edge distillation and shown its competitive results on benchmarks. As high-throughput and high- performance molecular feature extractors, SMILES-based models exhibit significant practical poten- tials in handling large-scale molecular data. With initial results being encouraging, challenges still exist. Our distillation approach cannot be applied without well-learned graph model as teacher. In distillation experiments, Graphormer outperforms its student well, so results of ST-KD can be further improved with more effective knowledge transfer strategy. We leave them as future works.
REFERENCES
Seyone Chithrananda, Gabriel Grand, and Bharath Ramsundar. Chemberta: Large-scale self- supervised pretraining for molecular property prediction. arXiv preprint 1rXXi::0110098855 2020.
Hyeoncheol Cho and Insung S Choi. Enhanced deep-learning prediction of molecular properties via augmentation of bond topology. ChemMedChem, 14(17):1604-169) 2019.
Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint 18..v:8810.048000 2018.
Jin Dong, Marc-Antoine Rondeau, and William L Hamilton. Distilling structured knowledge for text-based relational reasoning. Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP). pp. 6782-6791 2020.
David Duvenaud, Dougal Maclaurin, Jorge Aguilera-Iparraguun othy Hirzel, Alán Aspuru-Guzik, and Ryan Adams. Convolutional networks on graphs for learning molecular fingerprints. arXiv preprint arXiv: 1509.099999 2015.
Rafael Gómez--ombarelli, Tim-
Vijay Prakash Dwivedi and Xavier Bresson. generalization of transformer networks graphs. arXiv preprint arXi::2012.09699 2020.
Francesco Gentile, Vibudh Agrawal, Michael Hsing, Anh-Tien Ton, Fuqiang Ban, Ulf Norinder, Martin E Gleave, and Artem Cherkasov. Deep docking: deep learning platform for augmentation of structure based drug discovery. ACS central science, 6(6):939-949, 2020.
Justin Gilmer, Samuel Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp- 1263-1272. PMLR, 2017.
Robert Glen, Andreas Bender, Catrin H Amby, Lars Carlsson, Scott Boyer, and James Smith. Circular fingerprints: flexible molecular descriptors with applications from physical chemistry to adme. IDrugs, 199) 2006.
Shion Honda, Shoi Shi, and Hiroki R Ueda. Smiles transformer: Pre-trained molecular fingerprint for low data drug discovery. arXiv preprint arXiv:1911.04733 2019.
Weihua Hu, Matthias Fey, Marinka Zimik, Yuxiao Dong, Hongyu Ren, Bowen Liu, Michele Catasta, and Jure Leskovec. Open graph benchmark: Datasets for machine learning on graphs. arXiv preprint W.tv:2005.00688) 2020.
10

---

Under review as a conference paper at ICLR 2022
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-Isc: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.09430, 2021.
Ye Hu, Eugen Lounkine, and Jürgen Bajorath. Improving the search performance of extended con- nectivity fingerprints through activity-oriented feature filtering and application of bit-density- dependent similarity function. ChemMedChem: Chemistry Enabling Drug Discovery, 4(4):540~ 548, 2009.
John J Irwin and Brian K Shoichet. Zinc- free database of commercially available compounds for virtual screening. Journal of chemical information and modeling. 45(1):177-182, 2005.
Ross Irwin, Spyridon Dimitriadis, Jiazhen He, and Esben Bjerrum. Chemformer: A pre-trained transformer for computational chemistry. 2021.
Xiaoqi Jiao, Yichun Yin, Lifeng Shang, Xin Jiang, Xiao Chen, Linlin Li, Fang Wang, and Qun Liu. Tinybert: Distilling bert for natural language understanding. arXiv preprint arXiv: /909. 10351, 2019.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. arXiv preprint arXiv: wet..0909909777 2016.
Johannes Klicpera, Janek GroB, and Stephan Günnemann. Directional message passing for molec- ular graphs. arXiv preprint arXiv:2003.03123: 2020.
Ziyao Li, Shuwen Yang. Guojie Song, and Lingsheng Cai. Hamnet: Conformation-guiided molecular representation with hamiltonian neural networks. arXiv preprins
2021.
Harry I. Morgan. The generation of a unique machine description for chemical structures-a tech-
nique developed at chemical abstracts service. Journal of Chemical Documentation, 5(2):107- 113, 1965.
Maho Nakata and Tomomi Shimazaki. Pubchemqe project: large-scale first-rinciples electronic structure database for data-driven chemistry. Journal of chemical information and modeling, 57 (6):1300-1308, 2017.
Noel M O'Boyle, Casey M Campbell, and Geoffrey R Hutchison. Computational design and selec- tion of optimal organic photovoltaic materials. The Journal of Physical Chemistry C, 115(32):
16200-16210, 2011.
Bharath Ramsundar, Peter Eastman, Patrick Walters, Vijay Pande, Karl Leswing, and Zhenqin Wu. Deep Learning for the Life Sciences. O' O'Reilly Media, 2019. httpt:ttt/////t //tp:p:::/wwww//// amazor com/ Deep- Learnigg Life
David Rogers and Mathew Hahn. Extended-commecttvvin fingerprints. Journal of chemical informa- tion and modeling, 50(5):742-754 2010.
Yu Rong. Yatao Bian, Tingyang Xu, Weiyang Xie, Ying Wei, Wenbing Huang, and Junzhou Huang. Self-supervised graph transformer on large-scale molecular data. arXiv preprint 20Xiv:2007.02835.. 2020.
Jerret Ross, Brian Belgodere, Vijil Chenthamarakshan, Inkit Padhi, Youssef Mroueh, and Payel Das. Do large scale molecular language representations capture important structural information? arXiv preprint arXiv:2106.0955: 2021.
Philippe Schwaller, Teodoro Laino. Théophile Gaudin, Peter Bolgar, Christopher Hunter, Costas Bekas, and Alpha A Lec. Molecular transformer: a model for uncertainty-calibratedd chemical reaction prediction. ACS central science, 5(9):1572- 2019.
Ying Song, Shuangjia Zheng, Zhangming Niu, Zhang- Fu, Yutong Lu, and Yuedong Yang- Communicative representation learning attributed molecular graphs. In IJCAI, pp. 2831-2838, 2020.
Teague Sterling and John J Irwin. Zinc 15-ligand discovery for everyone. Journal of chemical information and modeling, 55(11):233 23337,--333, 2015.

---

Under review as a conference paper at ICLR 2022
Garrett A Stevenson, Derek Jones, Hyojin Kim, WF Bennett, Brian J Bennion, Monica Borucki, Feliza Bourguet, Aidan Epstein, Magdalena Franco, Brooke Harmon, et al. Highhthroughhee vir- tual screening of small molecule inhibitors for sars-cov-2 protein targets with deep fusion models. arXiv preprint arXiv: (0l.:2004.4444,, 2021.
Siqi Sun, Yu Cheng, Zhe Gan, and Jingjing Liu. Patient knowledge distillation for bert model compression. arXiv preprini arXiv: 01199998.99959 2019a.
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobilebert: Task-agnostic compression of bert by progressive knowledge transfer. 2019b.
Zhiqing Sun, Hongkun Yu, Xiaodan Song, Renjie Liu, Yiming Yang, and Denny Zhou. Mobile- bert: a compact tank-agi hert for resource-limited devices. arXiv preprint 20040:00000440484.. 2020.
Yonglong Tian, Dilip Krishnan, and Phillip Isola. Contrastive representation distillation. arXiv
preprint arXiv:1910. 10699, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-5908 6008, 2017.
Sheng Wang, Yuzhi Guo, Yuhong Wang, Hongmao Sun, and Junzhou Huang. Smiles-bert: large scale unsupervised pre-training for molecular property prediction. In Proceedings of the 10th ACM international conference on bioinformatics, computational biology and health informatics, pp. 429-43444 2019.
Wenhui Wang, Furu Wei, Li Dong, Hangbo Bao, Nan Yang, and Ming Zhou. Minilm: Deep self- attention distillation for task-agnostic compression of pre-trained transformers. arXiv preprint 10X1v:2002.10957. 2020.
David Weininger. Smiles, chemical language and information system. introduction method- ology and encoding rules. Journal of chemical information and computer sciences, 28(1):31-36 1988.
David Weininger. Smiles. depict. graphical depiction
depicttion of chemical structures. Journal of chemical information and computer sciences, 30(3):237-243, 1990.
David Weininger, Arthur Weininger, and Joseph L Weininger. Smiles. 2. algorithm for generation of unique smiles notation. Journal of chemical information and computer sciences, 29(2):97-101, 1989.
Zhenqin WU, Bharath Ramsundar, Evan N Feinberg, Joseph Gomes, Caleb Geniesse, Aneesh S Pappu, Karl Leswing, and Vijay Pande. Moleculenet: a benchmark for molecular machine learn- ing. Chemical science, 9():5113-530, 2018.
Zhaoping Xiong, Dingyan Wang, Xiaohong Liu, Feisheng Zhong, Xiaozhe Wan, Xutong Li, Zhao- jun Li, Xiaomin Luo, Kaixian Chen, Hualiang Jiang, et al. Pushing the boundaries of molecular representation for drug discovery with the graph attention mechanism. Journal of medicinal chemistry, 63(16))8749-877 2019.
Keyulu Xu, Weihua Hu, Jure Leskovec, and Stefanie Jegelka. How powerful are graph neural networks? arXiv preprint 1rXiv:1810.00822 2018.
Minkai Xu, Shitong Luo, Yoshua Bengio, Jian Peng, and Jian Tang. Learning neural generative dynamics for molecular conformation generation. arXiv preprint arXiv:2102.10240, 2021.
Zheng Xu, Sheng Wang, Feiyun Zhu, and Junzhou Huang. Seq2seq fingerprint: An unsupervised deep molecular embedding for drug discovery. In Proceedings of the 8th ACM international conference on bioinformatics, computational biology, and health informatics, pp. 285-294, 2017.
Chengxuan Ying, Tianle Cai. Shengjie Luo, Shuxin Zheng, Guolin Ke, Di He, Yanming Shen, and Tie- Yan Liu. Do transformers really perform bad for graph representation? arXiv preprint arXiv:2106. 05234, 2021.
12

---

Under review as a conference paper at ICLR 2022
Xiaoyu Zhang, Sheng Wang, Feiyun Zhu, Zheng Xu, Yuhong Wang. and Junzhou Huang. Seq3seq fingerprint: towards end-to-end semi supervuieeeiee deep drug discovery. In Proceedings of the 2018 ACM International Conference Bioinformatics, Computational Biology. and Health Informat- ics, pp. 404-413, 2018.
A APPENDIX
A.I 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. (02T), origi- nally curated under the PubChemQC project (Nakata & Shimazaki, 2017). The task of PCQM4M- LSC 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 idation 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 quanti- tative guantum-chemice properties. For QM7, we use a sanitized version by DeepChem (Ram- sundar al., 2019) available in https: / /eeppceeedaaateeeeee com/ dat asets qm ccvv FreeSolv contains molecules with their corresponding quantitative hydration free energy. HIV contains molecules labeled by their HIV active labels, with predefined scaffold data split.
val-
83-us-west-
amnawssss
Datasets
PCQM4M-LSC
QM9
QM8
QM7
FreeSolv
HIV
#molecules
3803453
133885
21786
6834
642
41127
#tasks
12
16
Task Type
Regression
Classification
Metric
MAE
Multi-MAE
MAE
RMSE
ROC-AUC
Table 5: Statistics of Datasets.
A.2 EXPERIMENT DETAILS
A.2.1 KNOWLEDGE DISTILLATION ON PCQM4M-LSC
All parameters involved are listed in Table During experiments we find by adding an initial stage when there is task-related supervision, the model can have faster convergence. Parameters of this initial stage 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 exper- iments we discover that results can benefit from changing the model dropout value this finetune training process. All related parameters are listed in Table and g
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.
13

---

Under review as a conference paper at ICLR 2022
Parameter
Value
#Transformer Layers
#Biased Transformer Layers
3
Hidden Dimension d
512
Key Dimension of Attention Block da Ford-Forward Network Dimension
32
2048
#Attention Heads h
16
Dropout Value
0.1
Feature Distillation GT-ST Layer Pairs Attention Distillation GT-ST Layer Pairs
{(93),, (10,4), (11,5)
(993), (10,4), (11,5)
Task Loss Weight
1.0
Feature Distillation Loss Weight a
0.5
Attention Distillation Loss Weight
2.0
Initial Stage Epochs
5
Initial Stage Task Loss Weight
0.0
Initial Stage Feature Distillation Loss Weight Initial Stage Attention Distillation Loss Weight |
1.0
4.0
Optimizer
AdamW W
Max Epochs
200
Batch Size
256
Learning Rate
le-4
Weight Decay
1e-5
of Adam W (B1, (2)
1e-8
of Adam
(0.9,0.999)
Table 6: Parameter settings of knowledge distillation experiment on PCQM4M-LSC.
Parameter
Value
Parameter
Value
Optimizer
AdamW
Optimizer
AdamW
Max Epochs
100
Max Epochs
100
Batch Size
80
Batch Size
Dropout Value
0.0
Dropout Value
0.1
Learning Rate
le-4
Learning Rate
lc-4
Weight Decay
le-5
Weight Decay
le-5
Adam
1e-8
c of Adam W (31, B2) of AdamW
1e-8
(31, B2) of AdamW
(0.9,0.999)
(0.9,0.999)
Table 7: Parameter settings of molecular property prediction experiments on QM9, QM8 and QM7.
Table 8: Parameter settings molecular
property prediction experiments on FreeSolv and HIV.
Task
MAE (Standardized)
MAE
0.2319:t6.21ee-3
0.
alpha
0.228-3-
0.2563-19.893--3
homo
0.1452
2.390e-3+1.9122-5
lumo
0.05173=7.760--4
14226--3+3.640e-5
gap
0.0680
3.2320-3:5.544e-5
r2
0.06301 10:00-11.000 -.061ee3
15.73+0.2261
zpve
3.63463
4.5.8e-45..400-55
u0
6.300e-342.160c-4
0,253418..55o--
u298
6.233e-3+1.886e-4
0.249747.554c-3
h298
6.402e-315.57c-4 0.256412.2666--2
g298
6.567e-342.99e-4
12.3149
0.02820=7.118-44
0.1146+2..822-3
Table 9: Performances of ST-KD on QM9, reported by seperate tasks.
14


Reviews:
[{"review_id": "k8S-xbmmbKe", "reviewer": "Reviewer_PfwA", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The knowledge distillation framework is quite interesting, where the author carefully provide the way to perform it Fig2 and Sec 3.2. \n\nMy major concern about this work is the performance of the SMILES-based transformer is not worse than the graph-based model. In fact, in many cases SMILES-based transformers outperform graph-based models (or at least comparable performances). Why is there a need to perform knowledge distillation to learn a SMILES Transformer (student) from graph Transformer (teacher) ? \n\n\nThe SMILES-based model (e.g. ST-BASE) used as baseline in this paper seem not a competitive SMILES-based model. Other competitive SMILES based models should be included in the comparison. For example, https://github.com/pschwllr/MolecularTransformer\nwhich has a different way of tokenizing and utilizing SMILES-based input. The SMILES-Transformer's performance can be largely boosted by simple augmentation trick, e,g using random-smiles as augmentation when training the model. The performance of this model may outperform graph-based models. Can you please include the comparison against this work?\n\n\nAdditionally, if we refer to table 2, it seems that there is not a clear improvement of graph-based models over ST-BASE. \n\nCan the author provide intuition of why the proposed methods ST-KD can outperform both graph and SMILES based methods in table 2, when it is trained to minimize the difference between graph transformer and SMILES transformer (by eq 16)? Should the teacher ( the graph-based methods)  be the upper bond? \n\n", "summary_of_the_paper": "The authors provide a knowledge distillation framework to transformer knowledge from graph transformer (teacher) to SMILES transformer (student), where the  SMILES transformer (student) is forced to mimic graph transformer (teacher) 's desired property. ", "summary_of_the_review": "The authors provide an interesting knowledge distillation framework to learn molecule embedding, however, the reviewer is not fully onboard with the motivation of doing the knowledge distillation (e.g. learning SMILES based transformer from graph-based transformer), since smiles-based transformer has a competitive performance. \n\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "PgG_6XK3X0x", "reviewer": "Reviewer_Maqb", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I am uncertain about the efficiency claim in this paper. According to my experience, SMILES-to-graph preprocessing is not that time-consuming (in my case probably 10 mins for 10M conversion in 1 process), and this is a preprocessing step, meaning that it only executes once for the future repeated use. Furthermore, a lot of molecule libraries are stored in .sdf format that we only need to read without conversion. I thus feel the motivation is exaggerated while I am also glad to listen to other reviewers.\n\nBesides, it sacrifices the training efficiency to trade the inference one. According to Table 1, compared with GNNs, ST-KD is 3-10 times larger (so what is the training cost which is not reported?) considering training is usually much more expensive than inference. Since graph transformer is also used in training, SMILES-to-graph conversion exists thus the comparison would be training complexity.\n\nI am also concerned with the tokenization step in section 3.2 that authors convert SMILES into a sequence composed of atom characters and bond characters. For instance, cc=cc will be converted to cccc= if I understand correctly. If it is the case, there will be two types of mistakes: 1) different molecules will produce the same token sequences as both cc=cc and c=ccc will be cccc=; 2) the same molecules will produce different token sequences as c=occ and ccc=o will be cocc= and ccco=, respectively. I do not judge on the second mistake since this is essentially the drawback of SMILES representations which is not permutation invariant, but I am worried about the first one.", "summary_of_the_paper": "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.", "summary_of_the_review": "I am concerned about the over-claim of efficiency and the correctness of tokenization.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "PSZx0GDX3rZ", "reviewer": "Reviewer_xrsP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1.\tThe core ideas of this paper are interesting and well-implemented. To my knowledge, end-to-end SMILES-based molecular fingerprint generators are discussed by many prior works, but none of them have proposed to improve the model with knowledge from existing state-of-the-art graph models. Followed by the key insight, the authors have proposed effective methods to distill knowledge from graph Transformers to SMILES Transformers. \n2.\tConvincing experimental results. Results on PCQM4M and ablation studies have corroborated the effectiveness of distillation methods. The model also achieves outstanding performances on downstream tasks like QM9, where the 0.62 multi-mae surpasses all existing models I know by a large margin, including those with molecular conformation provided.\n3.\tClear writing and illustrations. The authors have stated their motivations and implement details clearly with well-drawn figures. I find no overstatements in this paper.\n\nWeaknesses:\n1.\tMore experimental datasets would be welcome. The authors have performed molecular property prediction experiments on QM datasets and FreeSolv. More experiments would be welcome on more datasets, such as in MoleculeNet and OGB, to further evaluate the performances of ST-KD on different downstream molecular tasks.\n2.\tMore experiments on parameter analysis would be good. The authors run the knowledge distillation experiments on PCQM4M with a fixed network setting. Additional experiments can be done to explore the effects of different network structure and training hyperparameters.\n", "summary_of_the_paper": "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.", "summary_of_the_review": "This paper is based on an interesting idea of building a fast and high-quality SMILES-based molecular fingerprint generator with knowledge distillation from state-of-the-art graph-based models. The authors have described the proposed methods clearly, with their claims being supported by experimental results. Although more experiments can be performed to give a thorough evaluation to the model’s performance, I believe the paper should be recommended due to its novel idea and convincing results.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "NWCu4qf8j-L", "reviewer": "Reviewer_2H4F", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n(1) Knowledge distillation between graph transformer and SMILES transformer is proposed for molecular representation learning. \n(2) Feature distillation and attention transfer losses are used for knowledge transfer.\n(3) ST-KD achieves strong performance, similar to graph-based models, on molecular property prediction tasks, with a fast inference time.\n\nWeaknesses:\n(1)\tRecent SMILES transformer models trained on large data have shown performance similar to state-of-the-art graph-based models without any knowledge distillation, see e.g. Ross et al 2021 https://arxiv.org/abs/2106.09553. Such works should be discussed in comparison to ST-KD.  \n(2)\tST-KD performs slightly worse than Graphformer on PCQM4M-LSC as in Table 1. At the same time, the comparisons of ST-KD with Graphformer on QM* and Freesolv are not presented in Table 2. Therefore, it is not evident to what extent knowledge transfer is successful on QM* and Freesolv.\n(3)\t It is not clear to what extent ST-KD performance is generalizable beyond the tasks discussed in the paper – for example, Moleculenet classification tasks or on bigger molecules. \n(4)\tGraph-to-text knowledge distillation has been investigated before, e.g. https://aclanthology.org/2020.emnlp-main.551.pdf. Such works should be at least cited in this work. \n(5) Inference time comparison is not shown for datasets other than PCQM4M-LS.\n", "summary_of_the_paper": "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).\n\nMain contributions are:\n\n(1)  An end- to-end SMILES Transformer for molecular representation learning is proposed that is boosted by graph transformer to smiles transformer knowledge distillation.\n(2) Biased transformer layers are designed to achieve knowledge transfer; learnable SMILES token embedding are used to encode structural information. \n(3) ST-KD shows competitive results on PCQM4M-LSC and QM9, with fast inference compared to existing graph models. \n(4)", "summary_of_the_review": "The work is interesting, as it explores the idea of knowledge distillation from a Graph transformer to a SMILES transformer for performant molecular prediction with fast inference. However, the idea of cross-modal knowledge distillation is not novel. While it is well understood that the SMILES transformer after KD provides the advantage of fast inference, it is not clear if knowledge distillation is the way to get to SOTA performance with fast inference, as recent pre-trained SMILES-based efficient transformer models are providing competitive performance.  Also, the work does not provide any strong evidence of the  generalizability of the proposed approach.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper  builds  fast and high-quality SMILES-based molecular embeddings  by distilling  state-of-the-art graph-based models teachers.
This has the advantage of speeding inference time w.rt to graph based methods. 

The reviews were split regarding the motivation of the work, in the sense of why not train directly on SMILES instead of distilling graph based methods that are in some tasks behind SMILES transformer. Authors provided clarifications in the rebuttal showing that on Knowledge distillation of graph models  surpasses  SMILES only model training. 

I think given the experimental nature of the paper the main motivation of the paper should be better clarified and supported with more experimentation and downstream tasks.
Decision: Reject