M OLE -BERT: R ETHINKING P RE - TRAINING G RAPH N EURAL N ETWORKS FOR M OLECULES
Anonymous authors
Paper under double-blind review
A BSTRACT
Recent years have witnessed the prosperity of pre-training graph neural networks (GNNs) for molecules.
Typically, following the Masked Language Modeling (MLM) task of BERT (Devlin et al., 2019), Hu et al.
(2020) first randomly mask the atom types and then pre-train the GNNs to predict them.
However, unlike MLM, this pre-training task named AttrMask is too simple to learn informative molecular representations due to the extremely small and unbalanced atom ‘vocabulary’.
As a remedy, we adopt the encoder of a variant of VQ-VAE (Van Den Oord et al., 2017) as a context-aware tokenizer to encode atoms as meaningful discrete values, which can enlarge the atom vocabulary size and mitigate the quantitative divergence between dominant (e.g., carbons) and rare atoms (e.g., phosphorus).
With the enlarged atom ‘vocabulary’, we propose a novel node-level pre-training task, dubbed Masked Atoms Modeling ( MAM ), to randomly mask the discrete values and pre-train GNNs to predict them.
MAM mitigates the negative transfer issue of AttrMask and can be combined with various pre-training tasks to advance their performance.
Furthermore, for graph-level pre-training, we propose triplet masked contrastive learning ( TMCL ) to model varying degrees of semantic similarity between molecules, which is especially effective for molecule retrieval.
MAM and
TMCL constitute a novel pre-training framework, Mole-BERT , which can match or outperform state-of-the-art methods that require expensive domain knowledge as guidance.
The codes, the tokenizer, and the pre-trained models will be released.
I NTRODUCTION
Pre-training Language Models (PLMs) have revolutionized the landscape of Natural Language
Processing (NLP).
The representative one is BERT (Devlin et al., 2019), whose Masked Language
Modeling (MLM) pre-training task first randomly masks some proportions of tokens within a text, and then recovers the masked tokens based on the encoding results of the corrupted text.
Although BERT also includes the pre-training task of next sentence prediction, MLM is verified as the only success recipe for BERT (Liu et al., 2019; Raffel et al., 2020).
Inspired by this, MLM-style pre-training task has been extended to many other domains (Hu et al., 2020; Dosovitskiy et al., 2021;
He et al., 2022).
Molecules can be represented naturally as graphs with their atoms as nodes and chemical bonds as edges.
Hence, Graph Neural Networks (GNNs) can be utilized to process molecular graph data.
To exploit the abundant unlabeled molecules, tremendous efforts have been devoted to pre-training
GNNs for molecular graph representations (Xia et al., 2022b).
The pioneering work (Hu et al., 2020) on this topic first pre-trains GNNs with a MLM-style pre-training task (AttrMask) on large-scale unlabeled molecular graph datasets.
Specifically, they randomly mask some proportions of atoms and then pre-train the models to predict them.
AttrMask has emerged as a fundamental pre-training task and many subsequent works adopt it as a sub-task for pre-training (Zhang et al., 2021; Li et al., 2021a).
During the tuning stage, researchers replace the top layer of the pre-trained models with a task-specific sub-network and train the new model with the labeled molecules of the downstream tasks.
However,
Hu et al.
(2020) observe that pre-training only with AttrMask (node-level pre-training task) will incur the negative transfer issue (i.e., pre-trained models fall behind no pre-trained models) sometimes.
Intuitively, they think this phenomenon can be attributed to the lack of graph-level pre-training tasks and thus introduce supervised graph-level pre-training strategies, which are impractical because the labels are often expensive or unavailable.
Additionally, some supervised pre-training tasks unrelated to the downstream task of interest can even degrade the downstream performance (Hu et al., 2020).
[Figure]
In this paper, we provide a second voice on this predominant belief and aim to explain the negative transfer in molecular graph pre-training.
Firstly, as can be observed in Figure 1(a), the pre-training accuracy of AttrMask converges to ∼ 96% quickly, which indicates AttrMask task (118-way classification, 118 is the number of common chemical elements in nature) is extremely simple for the small atom vocabulary size.
The atom vocabulary is the set of unique atom types of common chemical elements in nature.
In contrast, the training accuracy of the MLM task ( ∼ 30k-way classification) in BERT only grows to 70% and hardly converges for the large text vocabulary size ( ∼ 30k tokens)
(Kosec et al., 2021).
Text vocabulary is the set of unique tokens in the corpus.
Secondly, the quantitative divergence between different atoms is extremely significant (see Figure 1(b)), which will bias the models’ prediction toward dominant atoms (e.g., carbons) and lead to fast convergence.
Previous works (Clark et al., 2020; Robinson et al., 2021) have revealed that simple pre-training tasks will capture less transferable knowledge and impair the generalization or adaptation to novel tasks.
Tokenization is the first step in any NLP pipeline, which separates a piece of text into smaller units called tokens (Sennrich et al., 2015; Wu et al., 2016).
Hence, pre-training language models include two stages: the first stage is tokenizer training and the second stage is language models pre-training.
However, for GNNs pre-training, previous works adopt the atoms’ types as tokens, which will result in a small-size and unbalanced atom vocabulary.
We argue that atoms with different contexts should be tokenized into different discrete values even if they belong to the same type.
For example, Aldehyde carbons and ester carbons indicate different properties of molecules even if both of them are carbons.
Hence, we introduce a context-aware tokenizer to encode atoms to meaningful discrete values.
Specifically, these discrete values are the latent codes of a variant of graph VQ-VAE (Van Den Oord et al., 2017).
The tokenizer is context-aware because the encoder of graph VQ-VAE is a GNN model.
In this way, we can categorize the dominant atoms (e.g., carbons) into several chemically meaningful sub-classes (e.g., aldehyde carbons, ester carbons, etc.)
considering the atoms’ context, which will enlarge the atom vocabulary size and mitigate the quantitative divergence between dominant and rare atoms.
To support the above claims, we provide the t-SNE visualization of carbon representations learned by the proposed tokenizer in Figure 1(c).
As can be observed, the representations of carbons are clustered based on types of functional groups, which indicates our tokenizer can encode atoms to chemically meaningful values.
With the new vocabulary, we propose a node-level pre-training task, dubbed Masked Atoms Modeling (MAM), to randomly mask the discrete values and pre-train
GNNs to predict them.
Alternatively, we can also regress the output atoms embeddings of VQ-VAE as the pre-training objective, which we study in Appendix D. Additionally, we can also tokenize the chemical bonds to boost the performance, which we elaborate in Appendix G. For molecular graph-level pre-training, graph contrastive learning (You et al., 2020) is a feasible pre-training strategy.
However, contrastive approaches push different molecules away equally regardless of their varying degrees of similarities.
To remedy this deficiency, we propose triplet masked contrastive learning (TMCL), which mimics the varying degrees of molecular similarities with different masking ratios.
We highlight the following contributions: (i)
We find the negative transfer issue of AttrMask can be attributed to the extremely small and unbalanced atom vocabulary.
(ii)
As a remedy, we contribute a context-aware tokenizer for molecular graphs using a variant of VQ-VAE.
Also, the tokenizer can be re-used as an off-the-shelf tool for subsequent works like the NLP community.
(iii)
With the new vocabulary, we propose a tailored pre-training task, MAM, to alleviate the negative transfer issue.
MAM serves as a fundamental pre-training task and can be combined with various pre-training tasks to advance their performance.
(iv) We propose a novel graph-level pre-training task, TMCL, to consider varying degrees of similarities between molecules, which is especially effective for molecule retrieval.
(v) We combine MAM and TMCL as a joint pre-training framework (Mole-BERT), which matches or outperforms state-of-the-art models that require expensive domain knowledge as guidance.
R ELATED W ORK 2. P RE - TRAINING ON M OLECULES
Recent years have witnessed the proliferation of Neural Networks in molecular representation learning.
While effective and prevalent, they require expensive annotations and barely generalize to unseen molecules, which poses a hurdle to practical applications.
To remedy these deficiencies, tremendous efforts have been devoted to pre-training on molecules.
Initially, one line of these works (Wang et al., 2019; Chithrananda et al., 2020) adopts MLM-style pre-training strategy on molecular SMILES (Weininger et al., 1989) strings.
Subsequently, recent works follow the contrastive paradigm (Zhu et al., 2021b;a; Qiu et al., 2020; Thakoor et al., 2022).
For molecular pre-training, GraphCL (You et al., 2020) and its variants (You et al., 2021;
Suresh et al., 2021; Xia et al., 2022a; Wang et al., 2022; Sun et al., 2021; Fang et al., 2022b) embed augmented versions of the anchor molecular graph close to each other and push the embeddings of other molecules apart.
Additionally, DGI (Velickovic et al., 2019) and InfoGraph (Sun et al., 2020a) is proposed to obtain expressive representations for graphs or nodes via maximizing the mutual information between graph-level representations and substructure-level representations of different granularity.
The other line of work adopts generative or predictive pretext tasks.
Typically, GPT-GNN (Hu & others., 2020) introduces an attributed graph generation task to pre-train GNNs so that they can capture the structural and semantic properties of the graph.
For molecular graph pre-training, Hu et al.
(2020) and Li et al.
(2021b) conduct attribute and structure prediction at the level of individual nodes as well as entire graphs.
To capture the rich information in molecular graph motifs, GROVER (Rong et al., 2020a) and MGSSL (Zhang et al., 2021) propose to predict or generate the motifs.
Considering that 3D geometric information also plays a vital role in predicting molecular properties, several recent works (Liu et al., 2022a; St ¨ ark et al., 2021; Fang et al., 2022a; Zhu et al., 2022) pre-train the GNN encoders on molecular datasets with 3D geometric information.
We recommend readers refer to several recent reviews (Liu et al., 2022b; Xie et al., 2022) for more relevant literatures.
Many above-mentioned works adopt AttrMask (Hu et al., 2020) as a fundamental pre-training sub-task.
However, AttrMask will incur the negative transfer issue sometimes.
We explain this phenomenon and contribute a novel pre-training strategy to remedy this deficiency.
2. MLM- STYLE P RE - TRAINING S TRATEGIES
The masked language modeling (MLM) task proposed in BERT Devlin et al., 2019 has emerged as one of the most popular and successful pre-training tasks.
Empirically, RoBERTa (Liu et al., 2019) finds that MLM is the only success recipe of BERT and discards the sentence-level task in
BERT.
Also, BART (Lewis et al., 2019) and T5 (Raffel et al., 2020) both observed that MLM is often the most effective task.
Alternatively, various MLM variants improve MLM with dynamic masking strategy (Liu et al., 2019) and blockwise masking strategy (Joshi et al., 2020).
Additionally,
XLNet (Yang et al., 2019) proposes permutation language modeling that conducts MLM in an autoregressive manner.
UniLM (Bao et al., 2020) uses pseudo MLM which unifies autoregressive and MLM tasks.
Inspired by these advances, for vision transformer pre-training, the model receives incomplete images with a large portion of the patches removed and learns to reconstruct the missing contents on low-level image pixels (He et al., 2022; Li et al., 2021c; Dosovitskiy et al., 2021; Chen et al., 2020), high-level semantics (Bao et al., 2021) or hand-crafted feature descriptors (Wei et al., 2022).
Different from the above-mentioned works, our MAM adopts the encoder of VQ-VAE as the atoms tokenizer, which discretizes a continuous semantic space to discrete codes in a context-aware manner.
Moreover, we observe that atoms of different types might be allocated with the same token id with vanilla VQ-VAE.
As a remedy, we introduce a group VQ-VAE to address this issue.
P RELIMINARY
Graph Neural Networks (GNNs) are the dominant tools for modeling graph data (Kipf & Welling, 2016;
Velickovic et al., 2018; Hamilton et al., 2017; Xu et al., 2019).
The structure of graph data guides the aggregation of local neighborhood information and leads to a more contextual representation for each node.
Also, we can adopt a graph pooling operation (Mesquita et al., 2020) to get the representation for the whole graph.
Let G =
( V , E ) denotes a graph with node attributes x v for v ∈ V and edge attributes e uv for ( u, v ) ∈ E .
Formally, supposing that h ( l ) v is the representation of node v at the l -th GNN layer and N ( v ) are all the neighbor nodes of node v , the update procedure from the ( l − 1) -th layer to the l -th layer is:
[Equation]
where e uv denotes the edge between node u and v .
AGGREGATE ( · ) is the aggregation function (e.g., mean operator) of the neighborhood information.
COMBINE ( · ) combines the information of neighbours and node v (e.g., concatenation operator).
After L iterations of message passing, the hidden states h
( L ) v in the last iteration are the embeddings of v .
Finally, we adopt a READOUT ( · ) operation (e.g., averaging, sum or graph pooling) to get the representation h G for the whole graph G :
[Equation]
P ROPOSED P RE - TRAINING F RAMEWORK : M OLE -BERT
In this section, we will introduce the proposed pre-training framework Mole-BERT, which includes a node-level per-training approach (MAM) and a graph-level pre-training approach (TMCL).
[Figure]
4. M ASKED A TOMS M ODELING (MAM)
Similar to NLP, we represent the atoms as discrete tokens using group VQ-VAE shown in Figure 2.
Formally, the atoms V = { v 1 , v 2 , · · · , v n } of a molecule graph G is tokenized to z = { z 1 , z 2 , · · · , z n } ∈ A n , where the atom vocabulary A contain |A| ( |A| = 512 ) discrete codes.
Firstly, the GNN encoder of group VQ-VAE encodes the atoms to atoms embeddings.
Next, the vector quantizer (VQ) looks up the nearest neighbor in the codebook for each atom embedding h i .
Let { e 1 , e 2 , · · · , e |A| } denote the codebook embeddings.
The quantized code of the i -th atom is:
[Equation]
After quantizing the atoms to discrete tokens, we feed the corresponding codebook embeddings { e z 1 , e z 2 , · · · , e z n } to the decoder to reconstruct the input molecule graph.
Kindly note that the vector quantization process is non-differentiable.
In order to train the encoder, the gradient is approximated like the straight-through estimator (Bengio et al., 2013) and copied from the decoder to the encoder.
Intuitively, the quantizer looks up the nearest code for each encoder output, so the gradients of codebook embeddings can be utilized to train the encoder approximately.
With the input attributes v i and reconstructed attributes (cid:98) v i , the training loss of the tokenizer for the molecule graph G is,
[Equation]
where the first term is a reconstruction loss with the scaled cosine error ( γ ≥ 1 ), the second term is a
VQ loss aiming to update the codebook and the third term is a commitment loss which encourages the output of the encoder to stay close to the chosen codebook embedding.
sg [ · ] denotes stop-gradient, β is a hyper-parameter set to 0.25 in our experiments.
This idea is inspired by DALL-E (Ramesh et al., 2021) which uses discrete VAE as a tokenizer for text-to-image generation.
We show the superiority of our tokenizer over discrete VAE empirically in Appendix I and clarify the difference here: (1) Our tokenizer tokenizes atoms to discrete codes in a context-aware manner; (2) We use the tokens for pre-training instead of autoregressive generation; (3) We observe that atoms of different types might be allocated with the same token id with vanilla VQ-VAE.
As a remedy, we introduce a group VQ-VAE to address this issue.
Specifically, we divide the codebook embeddings into several groups, each of which corresponds to specific atom types.
For example, the quantized codes of carbon, nitrogen and oxygen are restricted to [1 , 128] , [129 , 256] and [257 , 384] , respectively.
The left rare atoms are restricted to [385 , 512] because they are less likely to conflict with each other.
Additionally, we will release the tokenizer as an off-the-shelf tool for better or larger-scale molecular graph pre-training, just like what WordPiece (Wu et al., 2016) and BPE (Sennrich et al., 2015) have done for the NLP community.
With the new tokenizer, we propose MAM to pre-train GNNs.
Given an input molecule graph G , we randomly mask its 15% atoms’ tokens and pre-train GNNs to predict them.
We study the influence of the masking ratios in Appendix C. We term the masked atoms’ index set as M and the masked molecular graph as G M .
For each masked atom i ∈ M , a softmax classifier is adopted to predict the discrete values over the vocabulary A .
The pre-training loss of MAM is:
[Equation]
where D denotes the datasets and z i is the token of the atom v i .
Theoretically, we provide an explanation of MAM from the perspective of denoising variational autoencoder (Kingma & Welling, 2014).
Given that (cid:101) v i is the masked atom, The evidence lower bound (ELBO) of the log-likelihood is, (cid:88)
G∈D (cid:88)
i
[Equation]
where q ϕ ( z | v ) is the encoder that outputs atoms tokens, p ψ ( v | z ) is the decoder that reconstructs the molecules given the tokens.
Additionally, p θ ( z | (cid:101) v ) predicts the masked atoms’ tokens given the corrupted molecules, which is the MAM pre-training task indeed.
In the stage of tokenizer training, we minimize the graph reconstruction loss term.
In the pre-training stage, the tokenizer q ϕ and decoder p ψ are fixed.
Hence, we can simplify q ϕ ( z | v ) to a one-point distribution with the most likely atom tokens
ˆ z i = arg max z q ϕ ( z | v i ) .
Then the right side of Eq.
(6) is,
[Equation]
which means that the joint objectives of tokenizer pre-training and MAM are maximizing the ELBO of the log-likelihood.
Additionally, we study the time complexity of MAM in Appendix H.
4. G RAPH - LEVEL T ASK : T RIPLET M ASKED C ONTRASTIVE L EARNING (TMCL)
Although MAM can mitigate the negative transfer issue, we observe that it fails to capture molecular graph-level semantics.
Specifically, we first calculate the widely-used Tanimoto coefficient (Bajusz et al., 2015) of the extended connectivity fingerprints (ECFP) (Rogers & Hahn, 2010) between two molecules as their chemical similarity and pick the molecule pairs with top 15% similarity as ‘similar’ ones.
All the molecule pairs in the datasets are ‘random’ pairs.
Figure 4(a) shows significant inconsistency between the learned representations (by MAM) and
[Figure]
ECFP, which will impair molecule retrieval using MAM (See Section 5.3) because two random molecules may have high similarity scores (lack of uniformity), while closely related molecules may have more different representations (lack of alignment)
(Wang & Isola, 2020).
On the other hand, existing graph-level pre-training tasks often follow a supervised paradigm (Hu et al., 2020), where the labels for molecules are expensive for the laborious wet-lab experiments.
[Figure]
Cosine similarity measures the similarity between the learned representations while ‘Random’ and ‘Similar’ are defined with the chemical similarity between molecular fingerprints (ECFP).
The histogram of GraphCL (You et al., 2020) can be seen in Appendix K
Graph contrastive learning (You et al., 2020) is a possible remedy for the above issues.
For each molecule (anchor), they maximize the agreement between paired molecular graph augmentations (positive pairs) and push away other molecules in the batch as negative pairs (dissimilar molecules) indiscriminately.
However, we argue that there are varying degrees of semantic similarity between molecules.
For example, the similarity between formic acid (negative) and acetic acid (anchor) should be more significant than the one between ethanol (negative) and acetic acid (anchor).
Previous contrastive approaches fail to model such partial order of semantic similarity.
We introduce triplet masked contrastive learning, dubbed TMCL, to mitigate this crucial defect.
More specifically, for each molecular graph G , we first generate its augmented version G M 1 with masked atoms index M 1 and a smaller masking ratio (e.g., 15%).
Then we enlarge the masking ratio (e.g., 30%) and obtain the other augmented version G M 2 with atoms index M 2 .
Now, we constitute a triplet ( G , G M 1 , G M 2 ) with latent relation among them, i.e., G M 1 is more similar to G than G M 2 to
G .
Although index M 1 may lie in some crucial atom sets (e.g., functional group) that are indicative of molecular properties sometimes, the latent relation holds true in most cases and pre-training with abundant data will help alleviate this issue.
Given that h G , h M 1 , h M 2 are the graph-level representations for G , G M 1 and G M 2 , respectively, we can formulate the triplet loss as,
[Equation]
where sim ( · , · ) denotes the cosine similarity and m is the margin.
Now, we can combine the triplet loss L tri with commonly-used contrastive loss L con as graph-level pre-training objective L TMCL ,
[Equation]
B is the sampled batch including G , and τ is the temperature hyper-parameter.
Finally, MAM and
TMCL constitute a unified pre-training framework, Mole-BERT (See Figure 3), whose hybrid loss is,
[Equation]
[Equation]
E XPERIMENTS 5. D ATASETS
For the pre-training stage, we use 2 million molecules sampled from the ZINC15 database (Sterling & Irwin, 2015) following previous works (Hu et al., 2020).
The main downstream task is molecular property prediction, where we adopt the widely-used 8 binary classification datasets contained in MoleculeNet (Wu et al., 2018).
Kindly note that we use scaffold splitting (Ramsundar et al., 2019), which splits the molecules according to their structures to mimic the real-world use cases.
Additionally, we validate the effectiveness of Mole-BERT on a broader range of downstream tasks and datasets (See Section 5.3).
The detailed information of all the datasets can be seen in Appendix A.
5. E XPERIMENTS C ONFIGURATION
We use a 5-layer Graph Isomorphism Networks (GINs) whose hidden dimension is 300 (Xu et al., 2019) as the backbone architecture, which is one of the state-of-the-art GNNs for graph-level tasks.
We adopt mean pooling as the readout function.
During the pre-training stage, GNNs are pre-trained for 100 epochs with batch-size as 256 and the learning rate as 0.001.
During the fine-tuning stage, we train for 100 epochs with batch-size as 32 and report the test score with the best cross-validation performance.
The split for train/validation/test sets is 80% : 10% : 10%.
The hyper-parameter λ is picked from { 0 . 1 , 0 . 3 , 0 . 5 } with the validation set.
All the experiments are conducted on Tesla
V100 GPUs.
Additionally, considering that previous works adopt different evaluation protocols, we reproduce all the results with the same protocol (10 random seeds + scaffold splitting) as the pioneering work (Hu et al., 2020) rigorously for fairness.
More details can be found in Appendix B. bold ∗ indicates our methods outperform baselines with ≥ 0.5/ ≥ 2 std respectively. pre-train’ denotes training from scratch
[Table]
5. R ESULTS AND A NALYSIS
Primary Results and Analysis.
We document the main results of molecular property prediction in
Table 1 and Table 2.
Our systematic study suggests the following trends:
Observation I: The pre-training task of AttrMask incurs negative transfer issue on some datasets (HIV and BBBP).
In contrast, MAM achieves consistent and notable improvements over AttrMask and ‘No pretrain’, although MAM pre-trains GNNs only with a node-level task.
This observation verifies that only node-level pre-training tasks can also mitigate the negative transfer, which overturns the previous belief that pre-training GNNs at the level of individual nodes may give limited improvements.
The reasons why AttrMask fails lie in the extremely small and unbalanced atoms vocabulary.
Observation II: Mole-BERT can achieve competitive or better performance than previous pretraining strategies under the same experimental protocols.
More specifically, Mole-BERT outperforms ‘No pretrain’ model by 6.36% and the current state-of-the-art method GraphMVP by nearly 1.34%, though GraphMVP pre-trains GNNs on another molecular dataset with 3D geometry.
To further support the above claims, we plot the training and testing accuracy curves in Appendix J.
Observation III: As demonstrated in Table 2, MAM can serve as a fundamental pre-training sub-task like AttrMask.
Furthermore, MAM shows significant superiority over AttrMask when they serve as sub-tasks for multi-task GNNs pre-training.
Observation IV (Ablation Study): We substitute or remove some components of the proposed approaches to study their effectiveness.
As can be observed in Table 1, the group VQ-VAE is superior to vanilla VQ-VAE in MAM because it prevents the atoms of different types be allocated with the same token id. Additionally, we observe a notable performance drop when we remove the triplet loss
L tri or contrastive loss L con in TMCL, which indicates both of them are necessary and effective.
-training (Hu et al., 2020) comes from a preprocessed ChEMBL dataset (Gaulton et al., 2012a) with some labels from biochemical assays
[Table]
Influence of GNNs Backbone.
As shown in Table
[Table]
we verify that Mole-BERT is agnostic to the
GNN architectures by trying four popular GNN models including GIN (Xu et al., 2019), GCN (Kipf.
& Welling, 2017), R-GCN (Schlichtkrull et al., 2018) and GraphSAGE (Hamilton et al., 2017).
As can be observed, Mole-BERT achieves consistent and notable improvements over training from scratch with various GNNs.
Additionally, pre-training with GIN achieves the most significant gains.
[Table]
Broader Range of Downstream Tasks.
We report the performance in regressive property prediction and Drug-target affinity (DTA) tasks in Table 4.
DTA is a crucial task in drug discovery, where we aim to predict the affinity scores between the molecular drugs and protein targets.
We follow the settings of a recent work (Nguyen et al., 2021) on DTA which models the molecular graphs with
GNN and target protein (as an amino-acid sequence) with a convolution neural network (CNN).
We substitute the GNN in their approach with pre-trained GNNs.
The superior performance indicates that Mole-BERT can work well in a broader range of downstream tasks.
[Table]
[Table]
Influence of the Vocabulary Size.
In our previous experiments, we set the vocabulary size as 512.
However, we cannot decide the optimal vocabulary size for MAM because the exact atoms sub-types are unavailable.
In Table 5, we study the vocabulary size ranging from 128 to 2,048, from which we can observe: (1) Even if when we set the vocabulary size as 128, near 119 of AttrMask, MAM can outperform AttrMask, which indicates that the tokens derived by VQ-VAE are context-aware and is superior to pure atoms’ identities; (2) Vocabulary size also influence on MAM’s performance.
Although the vocabulary size of 1, 024 outperforms 512, the superiority is not significant.
Hence, we set 512 as the default vocabulary size considering the computational budget.
[Figure]
Molecule Retrieval.
For more comprehensive evaluations, we first extract the representation for a query molecule.
And then, we calculate its cosine similarities with all reference molecules in
ToxCast dataset.
We demonstrate 5 molecules that are most similar to the query molecule with the cosine similarities in Figure 5.
As can be observed, the representation similarities of Mole-BERT are approximately aligned with the fingerprint similarities, which indicates that Mole-BERT learns chemically meaningful representations.
Moreover, the representations extracted from MAM fail to model the varying degree of similarities between molecules.
Hence, graph-level tasks like TMCL are necessary and effective for molecule retrieval.
More ablations and results can be seen in Appendix E.
C ONCLUSIONS AND F UTURE W ORKS .
In this paper, we find the negative transfer issue of AttrMask can be attributed to the extremely small and unbalanced atom vocabulary.
As a remedy, we contribute a context-aware tokenizer with group
VQ-VAE.
With the new vocabulary, we propose a more suitable pre-training task, MAM, to mitigate the negative transfer issue of AttrMask.
Additionally, we develop triplet masked contrastive learning (TMCL) to model the varying degree of molecular similarities.
MAM and TMCL constitute a joint pre-training framework (Mole-BERT), which achieves superior performance over state-of-the-art methods while not requiring any domain knowledge.
For the future, it remains to be explored whether the proposed pre-training strategies can be applied to SMILES (Weininger et al., 1989) sequencebased molecular pre-training, where small and unbalanced vocabulary also impairs performance
A D ATASETS
In this section, we provide detailed information of the datasets used for downstream tasks.
Molecular Property: Pharmacology
The Blood-Brain Barrier Penetration (BBBP) (Martins et al., 2012) dataset measures whether a molecule will penetrate the central nervous system.
All three datasets, Tox21 (tox, 2017), ToxCast (Wu et al., 2018), and ClinTox (Gayvert et al., 2016) are related to the toxicity of molecular compounds.
The Side Effect Resource (SIDER) (Kuhn et al., 2015) dataset stores the adverse drug reactions on a marketed drug database.
Molecular Property: Physical Chemistry Dataset proposed in (Delaney, 2004) measures aqueous solubility of the molecular compounds.
Lipophilicity (Lipo) dataset is a subset of ChEMBL (Gaulton et al., 2012b) measuring the molecule octanol/water distribution coefficient.
CEP dataset is a subset of the Havard Clean Energy Project (CEP) (Hachmann et al., 2011), which estimates the organic photovoltaic efficiency.
Molecular Property: Biophysics Maximum Unbiased Validation (MUV) (Rohrer & Baumann, 2009) is another sub-database from PCBA, and is obtained by applying a refined nearest neighbor analysis.
HIV is from the Drug Therapeutics Program (DTP)
AIDS Antiviral Screen (Zaharevitz, 2015), and it aims at predicting inhibit HIV replication.
BACE measures the binding results for a set of inhibitors of β -secretase 1 (BACE-1) and is gathered in MoleculeNet (Wu et al., 2018).
Malaria (Gamo et al., 2010) measures the drug efficacy against the parasite that causes malaria.
Drug-Target Affinity Davis (Davis et al., 2011) measures the binding affinities between kinase inhibitors and kinases, scored by the K d value (kinase dissociation constant).
KIBA (Tang et al., 2014) contains binding affinities for kinase inhibitors from different sources, including K i , K d and
IC 50 .
KIBA scores ( ¨Ozt ¨ urk et al., 2018) are constructed to optimize the consistency among these values.
[Table]
Input graph representation.
For simplicity
[Table]
we use a minimal set of node and bond features that unambiguously describe the two-dimensional structure of molecules following previous works (Hu et al., 2020).
We use RDKit (Landrum et al., 2013) to obtain these features.
• Node features: – Atom number: 1 ∼– Chirality tag: { unspecified , tetrahedralcw , tetrahedralccw , other } • Edge features: – Bond type: { single , double , triple , aromatic } – Bond direction: {− , endupright , enddownright }
B M ORE E XPERIMENTAL D ETAILS
As we describe in the main text, we use the GIN architecture as the main encoder.
To incorporate edge features, following previous works (Hu et al., 2020), we make some minor modifications to include bond features.
Specifically, the raw node features and edge features are both 2-dimensional categorical vectors (see Appendix C for details), denoted as ( i v, 1 , i v, 2 ) and ( j e, 1 , j e, 2 ) for node v and edge e , respectively.
Note that we also introduce unique categories to indicate masked nodes/edges as well as self-loop edges.
As input features to GNNs, we first embed the categorical vectors by
[Equation]
where EmbNode 1 ( · ) , EmbNode 2 ( · ) , EmbEdge ( k ) 1 ( · ) , and EmbNode ( k ) 1 ( · ) represent embedding operations that map integer indices to d -dimensional real vectors, and k represents the index of GNN layers.
At the k -th layer, GNNs update node representations by
[Equation]
where N ( v ) is a set of nodes adjacent to v , and e = ( v, v ) represents the self-loop edge.
Note that for the final layer, i.e., k = K , we removed the ReLU from Eq.
(11) so that h ( k ) v can take negative values.
This is crucial for pre-training methods based on the dot product, e.g., Context Prediction and
Edge Prediction, as otherwise, the dot product between two vectors would be always positive.
The graph-level representation h G is obtained by averaging the node embeddings at the final layer, i.e.,
[Equation]
The label prediction is made by a linear model on top of h G .
In our experiments, we set the embedding dimension d to 300.
For MLPs in Eq.
(11), we use the ReLU activation with 600 hidden units.
We apply batch normalization right before the ReLU in Eq.
(11) and apply dropout to h ( k ) v at all the layers except the input layer.
For tokenizer training, we adopt the above 5-layer GINs as the encoder and the decoder, which is trained for 60 epochs on the 2 million unlabeled molecules sampled from the ZINC15 database with batch size as 256 and the learning rate as 0.001.
For TMCL, we set the masking ratios as 0.15 and 0.30, respectively.
The temperature parameter τ is set to 0.1 and the margin m is set to 0.
C T HE I NFLUENCE OF THE M ASKING R ATIO
In this section, we study the influence of the masking ratio of MAM and Mole-BERT.
As can be observed in Table 7, the performance increase when the masking ratio varies from 0.10 to 0.20 while witnessing a notable drop when the masking ratio varies from 0.20 to 0.30, which indicates over-large masking ratio will impair the GNNs pre-training.
Additionally, as shown in Table 8, the masking ratio pairs of two branches in Mole-BERT matter for pre-training.
Specifically, when the masking ratio pairs are with larger differences, the performance of Mole-BERT will be pronounced.
Also, the performance of Mole-BERT drops sharply when the masking ratio pairs are both larger than 0.25.
[Table] [Table]
D R EGRESS THE F EATURES THAT ARE U SED TO I NFER THE A TOM L ABELS .
In this section, we aim to investigate a possible alternative to MAM.
Specifically, we try to regress the features that are used to infer the atom labels.
Namely, the new pre-training objective is to regress the output of the group VQ-VAE (denoted as ‘MAM (regress)’).
As shown in Table 9, MAM or MoleBERT with regressive pre-training objective consistently falls behind those with one-hot classification objective.
This regressive objective may make the model over-fit to the noise or perturbations in the continuous values and hinders the models to learn high-level semantics.
[Table]
E M ORE RESULTS OF M OLECULE R ETRIEVAL
[Figure]
[Table]
We show more results of molecule retrieval in Figure 6, where ‘Mole-BERT w/o L tri ’ means that we remove the triplet loss of Mole-BERT.
As can be observed, the retrieval results of MAM are unsatisfactory.
In contrast, both ‘Mole-BERT w/o L tri ’ and ‘Mole-BERT’ can find chemically similar molecules in the database.
However, ‘Mole-BERT w/o L tri ’ fails to sort the molecules in a chemically meaningful order.
Hence, the triplet loss is necessary and effective for Mole-BERT.
F I MPLEMENTATION D ETAILS OF B ASELINES
Considering that previous works adopt different evaluation protocols, we reproduce all the results with the same protocol as the pioneering work (Hu et al., 2020) rigorously for fairness.
Specifically, we fine-tune the respective publicly available pre-trained models with 10 random seeds (0-9) and scaffold splitting following the pioneering work (Hu et al., 2020).
We use a batch size of 32 and a dropout rate of 0.5.
We train models on each dataset for 100 epochs and the datasets with multiple prediction tasks are fit jointly.
Additionally, we evaluate test performance on downstream tasks using
ROC-AUC with the validation early stopping protocol, i.e., test ROC-AUC at the best validation epoch is reported.
For datasets with multiple prediction tasks, we take the average ROC-AUC across all their tasks.
For the recent state-of-the-art method GraphMVP (Liu et al., 2022a), we adopt the contrastive variant (GraphMVP-C).
Specifically, we pre-train the model from scratch with the default settings in their paper and fine-tune the pre-trained model with the above-mentioned protocols.
G T OKENIZATION FOR THE C HEMICAL B ONDS
In addition to masking the atoms, Hu et al.
(2020) also mask the chemical bond type and pre-train
GNNs to predict them (4-way classification task).
Similarly, this pre-training task is too simple for its small vocabulary ( { single , double , triple , aromatic } ).
Additionally, the vocabulary is extremely unbalanced because the type of single bond is more common than other types.
As a remedy, we mask the bond type and pre-train GNNs to predict the tokens of the two atoms that constitute the bond.
Namely, we regard the pre-training task as a multi-label classification task, where two heads is adopted to predict the two tokens of the connected atom pairs, respectively.
The results shown in Table 10 indicate that pre-training with both atoms and bonds attributes masking tasks is not superior to pre-training with only atom attributes masking.
However, after we tokenize the chemical bonds, we achieve significant performance gains, which verifies that tokenization for the chemical bonds is necessary and effective.
[Table]
H T IME C OMPLEXITY OF MAM
Similar to AttrMask, the time complexity of MAM is linear with respect to the number of edges/nodes as it only involves sampling nodes/edges to be masked.
Hence, it enjoys the same computational complexity as supervised training.
The additional computational overhead originates from the memory consumption of the tokenizer.
The pre-training of both MAM is nearly 24 hours, which is more time-consuming than AttrMask but the divergence is not significant.
[Equation]
[Table]
Both discrete VAE (Ramesh et al., 2021) and our group VQ-VAE can tokenize the atoms to compact codes.
In Table 11, we compare their performance, from which we can observe that our group
VQ-VAE outperforms discrete VAE by a significant margin.
The reason is that group VQ-VAE can tokenize the atoms in a context-aware and semantics-aware manner.
Also, group VQ-VAE can prevent the atoms of different types from being allocated with the same token id.
J T RAINING AND V ALIDATION C URVES
We plot the training and validation curves in Figure 7, from which we can observe that the pre-trained models outperform training from scratch by significant margins.
Additionally, for small-scale datasets such as Bace, training from scratch tends to overfit the training data of downstream tasks.
In contrast, the pre-trained models can mitigate the over-fitting issue.
K G RAPH CL V . S . TMCL
GraphCL (You et al., 2020) is a representative contrastive pre-training strategy for molecular graphs.
In this section, we compare GraphCL and TMCL as graph-level pre-training tasks.
The results can
[Figure]
[Figure] be seen in Figure 8 and Table 12, from which we can observe that TMCL consistently outperforms
GraphCL by considering the order of similarity between molecules.
[Table]
L A BLATIONS FOR G ROUP VQ-VAE.
In this section
[Figure]
we substitute group VQ-VAE with vanilla VQ-VAE or discrete VAE
(Ramesh et al., 2021)
and show the t-SNE visualization of the carbon representations learned by them.
As can be observed in Figure
[Figure]
although vanilla VQ-VAE and discrete VAE can distinguish the carbons well based on the types of functional groups that carbons belong to, they cannot prevent some atoms of different types from being clustered in the same region.
In contrast, our group VQ-VAE can alleviate this issue.
M D ETAILED I NFORMATION OF THE S TRUCTURAL A BBREVIATIONS .
Table 13 shows the seven local structure categories of the carbons visualized in the main text. [Figure]
[Table]