Mole-BERT: Rethinking Pre-training Graph Neural Networks for Molecules

Abstract

Recent years have witnessed the prosperity of pre-training graph neural networks (GNNs) for molecules. Typically, atom types as node attributes are randomly masked and GNNs are then trained to predict masked types as in AttrMask \citep{hu2020strategies}, following the Masked Language Modeling (MLM) task of BERT~\citep{devlin2019bert}. However, unlike MLM where the vocabulary is large, the AttrMask pre-training does not learn informative molecular representations due to small and unbalanced atom `vocabulary'. To amend this problem, we propose a variant of VQ-VAE~\citep{van2017neural} as a context-aware tokenizer to encode atom attributes into chemically meaningful discrete codes. This 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 mask some discrete codes randomly and then pre-train GNNs to predict them. MAM also mitigates another issue of AttrMask, namely the negative transfer. It can be easily combined with various pre-training tasks to improve their performance. Furthermore, we propose triplet masked contrastive learning (TMCL) for graph-level pre-training to model the heterogeneous semantic similarity between molecules for effective molecule retrieval. MAM and TMCL constitute a novel pre-training framework, Mole-BERT, which can match or outperform state-of-the-art methods in a fully data-driven manner. We release the code at \textcolor{magenta}{\url{https://github.com/junxia97/Mole-BERT}}.

MPEK: a multi-task learning based on pre-trained language model for predicting enzymatic reaction kinetic parameters

“…We experimentally compared two frequently used protein pre-training models, ProtT5 17 and ESM-2 18 , to extract enzyme sequence features, and two small-molecule pre-training models, MolCLR 19 and Mole-BERT 20 , to extract substrate structural features. Four feature-extraction combinations were generated: ProtT5+Mole-BERT, ProtT5+MolCLR, ESM-2+Mole-BERT, and ESM-2+MolCLR.…”

“…However, the current predictive models still have much room for improvement due to the limitations of the small number of datasets, insufficient types of enzymes, and incomplete structural characterization. Based on the ongoing development of pre-trained models for proteins and small molecules [17][18][19][20][21] , molecular characterization using large-scale pre-trained models is now feasible, and such models effectively compensate for insufficient molecular characterization and small sample sizes [22][23][24] .…”

Mol-AE: Auto-Encoder Based Molecular Representation Learning With 3D Cloze Test Objective

“…As shown in Figure 4, a 3D molecule contains two types of information: the 3D structure and the atom type. Since atom type modeling is a well-defined problem and can be easily achieved by atomic MLM objective (Xia et al, 2022; Wang et al, 2019), thus, we mainly focus on how to model 3D structure and efficiently address the aforehead mentioned problems in EnCD .…”

“…Early approaches to molecular representation learning predominantly focused on 1D SMILES (Wang et al, 2019; Chithrananda et al, 2020; Guo et al, 2021; Honda et al, 2019) and 2D graphs (Li et al, 2021; Lu et al, 2021; Fang et al, 2022b; Xia et al, 2022). Recently, there has been a growing interest in 3D molecular data, which could provide a more comprehensive reflection of physical properties, including information not captured by 1D and 2D data, such as conformation details.…”

“…Pre-training based molecular representation learning has shown remarkable performance across various molecular understanding tasks, such as drug discovery (Pinzi & Rastelli, 2019;Adelusi et al, 2022), molecular property prediction (Luo et al, 2022;Liu et al, 2022b;Zhou et al, 2023;Yu et al, 2023) and reaction prediction (Gastegger et al, 2021;Schwaller et al, 2021). Early approaches tend to model 1D SMILES (Wang et al, 2019;Guo et al, 2021;Honda et al, 2019) or 2D graphs (Li et al, 2021;Lu et al, 2021;Fang et al, 2022b;Xia et al, 2022). More recently, there has been a growing interest in 3D molecular data, with its inclusion of 3D structure information providing more comprehensive information of molecules.…”