Atmol: Attention-Aware Encoded Graph Contrastive Learning for Molecule Property Prediction

Abstract

Graph neural networks have demonstrated significant successes for molecular property prediction. However, previous designs of graph neural networks are frequently restricted from insufficient labeled data and weak generalization capability. In this work, we present AtMol, a multi-structure self-supervised contrastive learning framework for molecular graph representation on a large amount of unlabeled data that incorporates self-attention into graph neural networks. After pre-training on 10 million unlabeled molecules , and then fine-tuning on multiple types of downstream tasks, outstanding results revealed the capacity of the model to identify chemically reasonable molecular similarities and the interpretability of AtMol via visualization methods. On most benchmarks, AtMol rivals or surpasses supervised learning methods with sophisticated feature engineering. Compared to previous CL framework, AtMol shows an average 4.9% gain of ROC-AUC on 7 classification benchmarks and a 7.2% decrease of error on 6 regression benchmarks.  Extensive experiments indicate that our contrastive learning framework significantly boost GNN performance on a variety of molecular property benchmarks.

Keywords: Molecular Machine Learning, Contrastive Learning, Graph neural networks, Augmentation Methods, molecular property prediction