Substructure-Atom Cross Attention for Molecular Representation Learning
Abstract
Designing a neural network architecture for molecular representation is crucial for
AI-driven drug discovery and molecule design.
In this work, we propose a newframework for molecular representation learning.
Our contribution is threefold: (a)demonstrating the usefulness of incorporating substructures to node-wise featuresfrom molecules, (b) designing two branch networks consisting of a transformer anda graph neural network so that the networks fused with asymmetric attention, and(c) not requiring heuristic features and computationally-expensive information frommolecules.
Using 1.8 million molecules collected from ChEMBL and PubChemdatabase, we pretrain our network to learn a general representation of moleculeswith minimal supervision.
The experimental results show that our pretrainednetwork achieves competitive performance on 11 downstream tasks for molecularproperty prediction.
Introduction
Predicting properties of molecules is one of the fundamental concerns in various fields.
For instance,researchers apply deep neural networks (DNNs) to replace expensive real-world experiments tomeasure the molecular properties of a drug candidate, e.g., the capability of permeating the bloodbrain barrier, solubility, and affinity.
Such an attempt significantly reduces wet-lab experimentationthat often takes more than ten years and costs $1 million [1, 2].
Among the DNN architectures, graph neural networks (GNNs) and Transformers are widely adoptedto recognize graph structure of molecules.
GNN architectures for molecular representation learninginclude message-passing neural network (MPNN) and directed MPNN [3, 4], where they investigatehow to obtain effective node, edge, and graph representation.
GNNs are powerful in capturing localinformation of a node, but may lack the ability to encode information from far-away nodes due toover-smoothing and over-squashing issues [5, 6].
On the other hand, Transformer-based architecturessuch as MAT [7] and Graphormer [8] augment the self-attention layer of a vanilla Transformer usinghigh-order graph connectivity information.
Transformers can encode global information as theyconsider attention between every pair of nodes from the first layer.
To guide a structural bias in theattention mechanism, previous work relies on heuristic features such as the shortest path between twonodes since the naive Transformers cannot recognize the graph structure.
From the understanding of chemical structure, it is known that meaningful substructures can be foundacross different molecules, also known as motif or fragments [9].
For example, carbon rings and
NO2 groups are typical substructures contributed to mutagenicity [10] showing that proper usageof substructures can help a property prediction.
Molecular substructures are often represented as
[Figure]
molecular fingerprints or molecular fragmentation.
Molecular fingerprints such as MACCS (Molecular ACCess System) keys [11] and Extended-Connectivity Fingerprints (ECFPs)
[12] represent amolecule into a fixed binary vector where each bit indicates the presence of a certain motif in themolecule.
With a predefined fragmentation dictionary, such as BRICS [13] or tree decomposition [14],a molecule can be decomposed into distinct partitions.
Interestingly, machine learning algorithmsthat utilize molecular substructures still show competitive performance on some datasets to deeplearning models [15, 7].
We propose a fusion architecture between a GNN and Transformer to incorporate molecular graphinformation and molecular substructures.
Molecular substructures and graph are encoded through
Transformer and GNN, respectively.
The Transformer is designed to recognize the molecularsubstructures, where the substructure information is mixed through self-attention to obtain betterrepresentations.
With the Transformer only architecture, however, local information of molecules,such as atoms, bonds, and connectivity, can be lost from the structures.
For example, the twomolecules shown in Figure 1b share the same representation with MACCS keys while havingdifferent structures.
To overcome, we use a separate GNN branch for preserving local information.
Inour model, we inject the GNN feature into the intermediate Transformer layers through the fusingnetwork.
In this way, substructures and local node information are interactively fused, producing afinal representation for molecular graphs.
We name our network as Substructure-Atom Cross Attention (SACA) as it uses substructure aswell as atom information in molecules and fuses them through cross-attention.
The architecturalchoices allows us to avoid heuristic features such as the shortest-path or 3-dimensional distancein attention layers, which are required in many existing Transformer architectures for molecularrepresentation learning [7, 8].
Furthermore, our model reduces the complexity for attention calculationover the node-level Transformer models from O ( N 2 ) to O ( N ) , where N is the number of atoms.
To demonstrate the empirical effectiveness of the proposed network and see the ability to capturethe general representation of molecules, we evaluate our model on 11 downstream tasks from
MoleculeNet [16].
Our approach achieves the competitive performance on 11 downstream tasks.
In what follows, we summarize the key contributions and benefits.
• We propose a novel network that combines the information from substructures and node
features in a molecule.
Our model combines the advantages of both Transformer and GNN
architectures to represent the information given to each architecture.
• We show the effectiveness of our model for molecular representation learning.
Our model
achieves competitive performance upon strong baseline models on 11 molecular property
tasks.• Our model does not require computationally-expensive heuristic information of molecular
graphs.
Related Work2.1 Architectures for Molecular representation learning
Graph neural networks (GNNs)
Most common architecture for molecular representation learningis the GNNs since molecules can be naturally represented as a graph structure; a node as an atom,an edge as a connection.
Researchers have actively investigated variations of GNN architectures[3, 4, 17, 18], for molecules.
For example, MPNN [3] generalizes the message passing frameworksand explores some variants that predict molecular properties.
Directed MPNN (DMPNN) [4] proposedto replace the node-based message by edge-based messages to avoid unnecessary message loops.
Next, communicative MPNN (CMPNN) [18] improved DMPNN by additionally considering thenode-edge interaction during the message passing phase.
AttentiveFP [17] extends the graph attentionmechanism to allow for nonlocal effects at the intramolecular level.
Despite the advance of GNN architectures for molecular representation learning, there are knownproblems in GNN, such as over-smoothing and over-squashing problems [5, 6], which means thenode representations become too similar, and the information from far nodes does not propagate wellas the number of neighbors increases exponentially.
Furthermore, [19] analyzes expressive powersof standard GNNs following neighborhood aggregation scheme and shows that they are boundedto Weisfeiler-Lehman test (WL-test) [20].
Therefore, GNNs with standard message passing cannotlearn to discern a simple substructure such as cycles.
For example, the two molecules in 1a cannotbe distinguished by WL-test, hence standard GNNs cannot distinguish these two molecules [21].
Asolution to this limited representation power of GNNs is to directly incorporate important substructuresin the representation learning framework.
Transformers With recent advance of Transformer architectures and their promising performancein various domains, including NLP and computer vision [22, 23], Transformer-based architecturesfor molecular representation learning [7, 8, 24, 25] have been developed.
Transformer architecturecalculates pair-wise attention between every node from the first layer.
Therefore, it can effectivelycapture the global information of a graph.
However, a vanilla Transformer architecture [26] isnot directly applicable for molecular graph representation because it cannot incorporate structuralinformation such as the edge and connectivity in graphs.
To bridge the gap, advanced Transformerarchitectures alter the self-attention layer [7, 8, 25] or incorporate message passing networks into
Transformer architectures for input feature [24].
Specifically, MAT [7] uses adjacency and distancematrix of atoms to augment the self-attention layer.
GROVER [24] runs Dynamic Message Passing
Network (dyMPN) over the input node and edge features to extract queries, keys and values forself-attention layers.
CoMPT [27] uses the shortest path information between two nodes, and
Graphormer [8] encodes node’s degree, edges and the shortest path with edges between two nodesfor molecular data.
Although these graph-specific features such as the shortest path or 3D distanceare found to be useful in graph representation, these features can be heuristic and impose excessivecomputational overhead, which limits the model’s applicability to large molecules such as proteins.
Our proposed model does not require any computationally expensive heuristic features such as 3Dinformation or shortest path for preprocessing or computation in attention layer.
2.2 Molecular Substructure
Substructure information extracted from molecules has been widely used in molecular generation,property prediction, and virtual screening [28, 29, 30, 14].
ECFPs [12] encode existing substructureswithin a circular distance from each atom in a molecule.
PMTNN [31] is a multi-task network thattakes ECFPs as an input to predict molecular properties on various datasets.
MACCS keys [11] extractmolecules substructure depending on the presence of pre-defined functional groups.
One example ofthe usage of MACCS keys is to encode known ligands of each protein, which leads to an improvementof prediction performance in protein-ligand interaction [32].
Molecular fragmentation method suchas BRICS fragmentation [13, 33] or tree decomposition [14] is to decompose molecules in nonoverlapping partitions with pre-defined rules.
BRICS fragmentation divides molecules followingchemical reaction based rules.
Tree decomposition proposed in [14] also extracts junction tree by contracting certain edges.
One node in the junction tree represents a substructure of originalmolecules.
Our model receives molecular substructures as input to supplement GNN’s expressivityand it can flexibly encode any substructure vocabulary.
Our model combines Transformer and GNN, with molecular substructures and molecular graph asinputs to each network.
There are existing studies that have also considered combining GNN and
Transformer architecture [34, 35, 36] to learn graph representations.
Specifically, DMP [34] utilizes
GNN and Transformer to encode graphs and SMILES representation of molecules and train bothbranches using a consistency loss to match the two outputs of input molecules.
PoseGTAC [35] and
GraphFormers [36] combine GNN and Transformer layer alternatively to enlarge the receptive field orto mix the output of Transformer.
Our model is the first to encode substructures through Transformerand inject atom features through a separate GNN.
In this way, we preserve both substructures andlocal atom features of molecules.
Model
[Figure]
The final CLS token is used to predict the molecular properties. In this section, we explain the architectureof our model.
Overall architecture Figure 2 shows theoverall architecture of our model.
Our network consists of two branches: (i) a Transformer branch that uses the molecular substructures as input and (ii) a GNN branchthat uses a molecular graph as input.
Thetwo branches have different roles.
First,the Transformer branch is intended to capture global information of molecules.
Itreceives the molecular substructures thathave important role in molecular properties, but cannot be easily captured by GNNs, and learn the overall representationof molecules.
On the other hand, the GNNbranch is intended to capture local node information of molecules.
The two differentlevels of information are mixed through thefusing network in the Transformer branch.
Transformer branch Our Transformerbranch is to incorporate both molecular substructure information and local node features.
The input token for Transformer isthe substructure embeddings of molecules.
Predefined substructures are first detectedand then projected into separate embedding vectors.
For example, in Figure 2, substructures such as N-Heterocycle, carbon-oxygen bond and methyl group are detected and embedded into learnable embedding vectors.
To identify substructure from the input molecules, we use MACCS keys [11], whichindicates the presence of motifs in a molecule.
Note that our architecture is not limited to certainmolecular substructures, but it can flexibly receive any substructure vocabulary.
The embeddings ofsubstructures are mixed together and refined as they are passed through the self-attention module.
The substructure embeddings after self-attention layer are fused with node embeddings from aseparate GNN branch.
The fusing network computes cross-attention between substructures and nodeswhere substructures are used as query and nodes are used as key and value.
The detailed computation of the fusing network is shown in Figure 3.
In the fusing network, the cross-attention between eachpair of substructure embedding and node embedding is computed.
To be specific, for a given molecule having n atoms and m extracted substructures, we have substructure embeddings E s ∈ R m × d and node embeddings E n ∈ R n × d where d is the embeddingdimension.
Then, the cross-attention is computed as follows:
[Equation]
where W Q , W K , W V ∈ R
d × d k are learnable parameters.
The cross-attention module outputs E ′ s ∈
R m × d k .
Through this fusing network, the substructure embeddings aggregate the local informationfrom node embeddings.
Structurally important nodes are aggregated with more weights.
Insteadof designing heuristic weights on the nodes, our model can learn to select structurally importantnodes related to graph-level property by cross-attention.
Additionally, as the attention is computedbetween substructures and nodes, the space and time complexity of the self and cross attention mapof our model is linear to the number of atoms, i.e., O ( N ) , whereas other transformer architecturesfor molecular graphs have quadratic complexity.
[Figure]
𝑚×𝑑 key and value.
The self-attention between substructures and fusing network with node embeddings are repeatediteratively N times.
Before making the final prediction, we add M self-attention layers at theend of the network for making refinement onthe substructures.
We add a residual connectionfrom the input tokens to every output of the fusing network.
This ensures the input structuralinformation last throughout the entire network.
Furthermore, we use a CLS token similar tospecial classification token in BERT [22] to aggregate the global representation of molecules.
The CLS token is shown in Figure 2.
It is alearnable latent vector and passed to the Transformer branch attached to the substructure tokens asinput.
A CLS token has been used in many Transformer-based architectures [7, 8] for molecularrepresentation learning.
We attach an MLP head to the final CLS token to perform graph-levelprediction tasks.
GNN branch The GNN branch is used to extract local node features from molecular graphs.
For
GNN architecture, we use GIN [19] with jumping knowledge [37].
The computed node features areinjected into the Transformer through the fusing network with the same hierarchy.
For example, 0-hopnode representations are injected into the first cross-attention module and 1-hop representations forthe second cross-attention module.
This allows the model to encode local node features progressivelyfrom the shallow to deeper layers.
The computation of node representation through GNNs and injection to the Transformer allow usto take advantage of both GNNs and Transformer architectures.
Substructures that are hard to becaptured by GNNs, but essential to molecular properties, are first detected and encoded through
Transformer.
Meanwhile, local node information that can be lost in using substructures alone iseffectively captured by GNNs and fused with substructures.
Additionally, our architecture doesnot require computationally expensive high-order graph-level information.
As Transformer cannotnaturally incorporate a graph’s edge connectivity information, existing work [24, 8, 27] mainlyfocuses on how to add structural bias such as the shortest path or 3D distance between two nodes intothe Transformer self-attention computation.
However, these structural biases require computationallyexpensive preprocessing of the molecular datasets.
Our network that utilizes GNN as a separatebranch can avoid these limitations.
Experiment
4.1 Experimental Setting
Pretraining We pretrain our network to obtain molecular representation transferable to variousmolecular datasets and tasks.
For pretraining, we extracted 200 real-valued descriptors of physicochemical properties from the pretraining datasets using RDKit [38] and train our network to predictthese properties.
As the 200 molecular descriptors include a diverse set of molecular properties, themodel can learn a representation of molecules that can be used for various downstream tasks.
Dataset We collected 1,858,081 number of unlabeled molecules from ChEMBL and PubChemdatabases [39, 40].
ChEMBL and PubChem are large-scale databases that include a variety ofchemical and physical properties, and biological activities of molecules.
To obtain molecularsubstructures, we utilize MACCS key [11], a 166 dimensional vector that indicates a presence ofcertain substructure in molecules, and extract this for every molecule using RDKit [38].
We use OGBpackage [15] to convert SMILES [41], a text-representation for molecules, to molecular graphs.
Implementation details We set M = 4 and N = 3 , where M and N are defined in section 3.
Weset 768-dimensional hidden units and 16 attention heads.
When pretraining, we used the AdamWoptimizer [42] with a learning rate of 1e-4.
We divide the pretraining dataset into a 9:1 ratio and usethem for training and validation sets.
The model is trained for 10 epochs, and the model with the bestvalidation loss is used for downstream tasks.
Further details for pretraining setting is presented in
Appendix D.
4.2 Downstream tasks
Tasks We evaluate the performance on six classification tasks (BBBP, SIDER, ClinTox, BACE,
Tox21, and ToxCast) and five regression tasks (FreeSolv, ESOL, Lipo, QM7 and QM8) from MoleculeNet [16].
The dataset statistics are summarized in Table 4.
Each task is related to molecular propertyfrom low-level, for example, water solubility in ESOL to high-level, possibility of blood-brain barrierpenetraion in BBBP.
Further details about the downstream tasks are available in Appendix B.
Experimental setting To evaluate each model, we use 3 different scaffold splits [16] following [24].
Scaffold split divides structurally different molecules into different subsets and provides morechallenging and realistic test environment.
From the pretrained model, we replace the last MLPlayer of the network with the task-specific MLP heads.
For each downstream dataset, we train ourmodel for 100 epochs and report the test score corresponding to the best validation epoch.
We tunehyperparameters with Bayesian optimization search with a budget of 100 for learning rate, dropout,weight decay and the number of last prediction heads.
The hyperparameter search range is providedin Appendix D.
We compare the performance of our model on the downstream tasks with several GNN and Transformer based state-of-the-arts approaches for molecule representation learning.
TF-Robust [31] is a
DNN-based model that takes molecular fingerprints.
GNN-based models include GraphConv [43],
Weave [44] and SchNet [45] which are 3 graph convolutional networks, MPNN [3], DMPNN [4],
MGCN [46] and CMPNN [18] which are GNN models considering the edge features during messagepassing.
AttentiveFP [17] is an extension of graph attention network for molecule representation.
Transformer-based models include GROVER [24], MAT [7], Graphormer [8] and CoMPT [27].
Among the baselines, N-GRAM [47], [48], GraphLoG [49], MAT, GROVER, Graphormer, GEM [50]and MPG [51] are models that use pretraining strategies.
We report GROVER base model for afair comparison in terms of the number of parameters.
To reproduce the results for models usingpretraining strategies, we use the pretrained model made available by the authors.
We reproduce theresults of MPG due to the different splits used in [51] 1 .
We also evaluate our model with the samedata split used for MPG and the results can be found in the appendix.
[Table]
7 MGCN [46] - .850 (.064) .552 (.018) .634 (.042) .734 (.030) .707 (.016) .663 (.009) 17.0 AttentiveFP [17] - .908 (.050) .605 (.060) .933 (.020) .863 (.015) .807 (.020) .579 (.001) 10.0 CMPNN [18] - .
[Table] [Table]
[Table] [Table]
CoMPT [27]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
010) 2.3 Graphormer [8] ✓ .938 (.032) .625 (.009) .913 (.056) .848 (.023) .801 (.013) .718 (.007) 6.7 GraphLoG [49] ✓ .913 (.024) .595 (.039) - .845 (.012) .773 (.010) .677 (.008) 13.0 MPG* [51] ✓ .922 (.039) .628 (.014) - .
[Table] [Table]
[Table] [Table]
[Table] [Table]
GEM
[Table]
50] ✓
[Table] [Table]
[Table] [Table]
[Table] [Table]
Ours
[Table] [Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
[44] - 2.398 (.250) 1.158 (.055) .813 (.042) 94.7 (2.7) .022 (.001) 11.8 GraphConv [43] - 2.900 (.135) 1.068 (.050) .712 (.049) 118.9 (20.2) .
(.001) 12.2 SchNet [45] - 3.215 (.755) 1.045 (.064) .909 (.098) 74.2 (6.0) .020 (.002)
11.4 MPNN [3] - 2.185 (.952) 1.167 (.430) .672 (.051) 113.0 (17.2) .015 (.002) 10.0 DMPNN [4] - 2.177 (.914) .980 (.258) .
(.002) 7.8 MGCN [46] - 3.349 (.097) 1.266 (.147) 1.113 (.041) 77.6 (4.7) .022 (.002)
13.6 AttentiveFP [17] - 2.030 (.420) .853 (.060) .650 (.030) 126.7 (4.0) .
(.001) 8.8 CMPNN [18] - 2.254 (.356) .841 (.090) .606 (.040)
(.001) 4.8 CoMPT [27] - 2.125 (.590) .898 (.053) .
N-GRAM [47] ✓ 2.512 (.190) 1.100 (.160) .876 (.033)
(.003) 14.0 MAT [7] ✓ 2.116 (.152) .833 (.122) .668 (.025) 93.6 (13.8) .
(.002) 6.6 GROVER [24] ✓ 1.592 (
[Table]
Graphormer [8]
[Table]
Results Table 1 shows the overall results of the baselines and our model on 11 MoleculeNet datasets.
Our model achieved the best performance on five downstream tasks: ClinTox, Tox21, ToxCast, ESOLand QM8, and the second best performance on five downstream tasks: SIDER, BACE, FreeSolv, Lipoand QM7.
We also computed the average rank for the classification and regression tasks, separately.
Our model achieved the best average rank among all compared models.
The result shows that ourmodel generalizes well across different downstream tasks, which means local information aggregatedfrom GNN can propagate globally as interacting with structural information in Transformer.
We findthat some substructures weight high attention to nodes consisting of the substructures, and CLS tokenused for prediction also focuses on the substructures from attention scores in self-attention.
[Figure]
[Figure] [Figure]
[Figure]
[Figure]
[Figure] [Figure]
[Figure]
4.3 Ablation study
We report ablation study to justify each component and flexibility of our model architecture.
For theablation study, we reduce the hyperparameter search space to only 6 learning rates {1e-3, 5e-4, 1e-4,5e-5, 1e-5, 5e-6} to facilitate the comparison between different models.
We report the test scoresbased on the hyperparameters from the best validation scores.
Ablation on model components Figure 4 shows the performance comparison between four variations of our model on four downstream datasets.
We first verify the performance of our model without
GNN branch.
To do that, we replace all cross-attention layers with self-attention layers and exclude
GNN branch.
Begin-Concat and End-Concat examine different ways of combining substructure andlocal features.
Begin-Concat runs a vanilla-transformer encoder on top of a concatenated atom andsubstructure embeddings.
End-Concat runs a vanilla-transformer encoder on top of atom embeddingsand concatenates the substructure feature at the end to make the final prediction.
Random Embeddingexamines the effect of substructural embeddings by replacing them with random learnable embeddings.
Figure 4 shows our model outperforms other variations on all datasets, which justifies thenecessity of each component.
4.4 Analysis
Cl Cl
Cl Cl Cl
Cl
Cl Cl
C C s ub s t r u c t u r e t
ok e n s
CLS
Atom tokens
C
C (a) Ring-Chain-Ring bond  0.
⋯
Cl Cl
Cl Cl Cl
Cl
Cl Cl Cl Cl Cl
Cl
Cl Cl Cl
Cl
Atom tokens CLS s ub s t r u c t u r e t ok e n s
Cl Cl
Cl (b) Chlorine atom
More interestingly, often a substructure givesmore attention on the nodes that consists thesubstructure itself.
Despite of not having anystructural information between input substructures and nodes as input, our model can identifystructurally related nodes in the cross attentionlayer, showcasing the ability of understanding molecular structure.
Conclusion
In this paper, we propose a novel framework that incorporates Transformer and GNN architecture formolecular representation learning.
Our model takes advantages of the two architectures to aggregatesubstructure and local information.
With the cross attention mechanism in fusing network, our modelcould achieve state-of-the-art performance on various molecular property prediction benchmarks.
Overall, our work highlights the effectiveness of SACA for molecular representation learning.
A Node and Edge Features
In this section, we present the node and edge features of molecules used for GNN branch.
We used
OGB package [15] to convert SMILES strings [41] to molecular graphs.
The molecular graphs areencoded through GIN [19] and injected into the Transformer branch by the cross-attention.
Themolecular graphs have the following node and edge features.
Node features.
Each node has the following 9 dimensional features as shown in Table 2.
[Table]
Edge features.
Each edge has the following 3 dimensional features as shown in Table
[Table] [Table]
B Details of Downstream Datasets
We used 11 binary graph classification and regression datasets: BBBP, SIDER, ClinTox, BACE,
Tox21, ToxCast, FreeSolv, ESOL, Lipophilicity, QM7 and QM8 from Moleculenet [16].
The statisticsof each dataset is shown in Table 4.
The details of each dataset are shown in Table 5.
Through thevarious datasets, we can test the generalization ability of our pretrained model.
[Table] [Table]
C Molecular Structural Keys
In this section, we explain the details of molecular substructures that we used for our model, i.e.,
Molecular ACCess System (MACCS) keys [11].
We extract MACCS keys for each molecule using
RDKit package [38].
MACCS keys have 166-dimensional features where each binary label indicatesthe presence of a particular substructure in the given molecule.
For example, the 139 th index of
MACCS keys indicates the presence of the Hydroxy group (-OH), and the 162 nd index indicates thepresence of the aromatic ring in a molecule.
The full list of 166 MACCS keys can be found in thedocument 2 .
[Figure]
Figure 6 shows examples of how MACCS keys encode substructure information into a binary bit string.
The canonical SMILES representations of the left and right molecules are C1=
CC=C(C
=
C1)O and
Cc1cc(Cl)cc(Cl)c1CNC(=O)c1cccs1 , respectively.
Table 6 shows the MACCS keys statistics for pretraining and downstream datasets.
It includes theaverage number of MACCS keys that molecules in each dataset contain and the standard deviation.
Also, the minimum and the maximum number of MACCS keys for each dataset are reported. [Table]
D Experimental Setting
In this section [Table] we explain further details of our experimental setting.
Pretraining.
The default hyperparameters for pretraining are listed in Table 7.
All attention layershave 16 attention heads and 768 hidden dimension.
Our model has 41M parameters.
[Table]
Hyperparameter Search Range.
For each downstream task, we search for the best hyperparameter
[Table]
combinations.
We perform the Bayesian optimization over the validation set and use the hyperparameters for the best validation score to report the test score.
The hyperparameter range that we searchedover is shown in Table 8. [Table]
Pretraining Task.
To pretrain our network, we extract the 200 real-valued descriptors for each [Table] molecule using RDKit package [38].
This task is proposed by [52].
Through this task
[Table]
we can make [Table] the model to learn the physicochemical properties of the input molecules. [Table]
E Parameters and space complexity of different models
We present the number of parameters of transformer-based molecular representation learning modelsin Table 9.
Except for CoMPT, our model has a comparable or lower number of parameters thanother transformer-based models.
We also present the space complexity in Table 9.
The space complexity for our attention computationis linear to the number of atoms (i.e., O ( n ) where n is the number of atoms).
On the other hand, thespace complexity of other transformer-based molecule representation learning models is quadratic tothe number of atoms (i.e., O ( n 2 ) ).
[Table]
F Additional Ablation StudyDifferent GNN
branch Our model can flexibly utilize other GNN architectures as our GNNbranch.
We test the changes in performance when our model design is applied with different
GNN architectures.
Figure 7 shows the comparison between commonly-used GNNs, i.e., Graph
Convolutional Network (GCN), Graph Attention Network (GAT) [53] and GIN, and our models withthe GNN branch switched to each corresponding GNN on ToxCast dataset.
There is a significantperformance improvement when our model is adopted to each GNN model.
Effectiveness of substructures Figure 8 shows the t-SNE embeddings of molecules with differentsubstructures.
Our model discriminates molecules with similar substructures (aromatic rings withone or two different atoms) better than GIN.
[Figure]
[Figure]
Ours GIN
Ours SGCN
QM7 (↓)
Change GNN branch to 3D-aware GNN To observe the effectof the choice of GNN architecture for GNN branch, we conductadditional experiments with 3D-aware GNN.
We changed our GNNbranch from GIN to SGCN [54] that utilizes 3d coordinates of inputmolecules.
Please note that the two models are not pre-trained andinstead trained on the downstream dataset from scratch.
Figure 9shows that the model with 3D-aware GNN shows better performancethan our original model on QM7 dataset whose task is closely relatedto 3D information.
The result shows the flexibility of our frameworkthat can be further tuned by using task-appropriate GNN architecture.
Other substructures as input tokens for Transformer Despiteour model utilizes MACCS keys for the input of Transformer, ourmodel can flexibly receive any substructure vocabulary.
We conductexperiments using other substructures, that are ECFP fingerprintsand Tree decomposition.
Table 10 shows the results.
It shows using
MACCS keys achieves better performance than ECFP fingerprintin most of the downstream tasks.
We speculate that MACCS keysinclude predefined functional groups such as the Hydroxyl group (-OH) whereas ECFP encodes localsubstructure around an atom (i.e., a certain radius neighborhood of an atom).
In this aspect, ECFPfingerprint is similar to how GNN encodes node information and as we already utilize GNN to encodelocal information, ECFP would not bring new information about molecules.
[Table]
0.0140 ECFP0. 0.0152 ECFP0.
0.0151 Tree Decomposition 0.
Computation time for graph features Figure 10 shows the time required in millisecond to computethe MACCS keys, shortest path and 3D distance for each molecule with different number of atoms.
Asshown in the figure, the time required for 3D distance and the shortest path increases dramatically asthe number of atoms increase, which make these features impossible to be applied in large molecules.
However, identifying the MACCS key substructures does not depend on the input molecule’s numberof atoms.
G Performance on MPG split
As MPG uses different splits to GROVER on the paper, we reproduce the result of our model on thesame split that MPG used for fair comparison.
As shown in Table 11, our model outperforms MPGon most of the downstream datasets.
[Table]
H
Others
URLs for pretrained model for baselines.
We use the following sources for pretrained models toreproduce the results.
Graphormer: https://github.com/microsoft/Graphormer
MAT: https://github.com/ardigen/MAT
GraphLoG: https://github.com/DeepGraphLearning/GraphLoG
[Table]
GEM: https://github.com/PaddlePaddle/PaddleHelix/tree/dev/apps/pretrained_compound/ChemRL/GEM